Research map / Architectures and efficiency / RNNs and normalization
Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes OT-Transformer, a plug-and-play continuous-time reformulation of pretrained-architecture Transformers (single ODE over composed blocks) with an optimal-transport kinetic-energy regularizer, showing better test performance with fewer parameters across seven experiments and proving stable forward propagation and distributional robustness.
Problem
Transformer architectures are chosen by trial and error, and training lacks theoretical guarantees on stability, robustness, and generalization; prior continuous-time Transformer formulations did not improve performance.
Approach
Treat the full stack of Transformer blocks as one dynamical system whose velocity field is the Transformer, and add an optimal-transport (kinetic energy) regularizer on the velocity; optimal control/HJB regularity then yields Lipschitz stable forward propagation and Wasserstein distributional robustness.
Key results
- On char-level nanoGPT Shakespeare, OT-Transformer reaches test loss 1.44 with 6.16M params vs baseline 2.68 with 10.65M params.
- On GPT-2 OpenWebText (9B tokens), test loss 2.91 vs baseline 3.21 at equal 123.7M params.
- On MNIST, OT-Transformer reaches 97.1% accuracy with 18K params vs baseline 93.0% with 93K params.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply OT regularization and continuous-time formulation | Set Transformer · ModelNet40 point cloud classification | accuracy: 87.4% ± 0.45% -> 89.9% ± 0.42% | Helped |
| Apply OT regularization and continuous-time formulation | ViT · MNIST image classification | accuracy: 93.0% ± 0.69% -> 97.1% ± 0.16% | Helped |
| Apply OT regularization and continuous-time formulation | ViT · Cats and Dogs image classification | accuracy: 77.6% ± 0.86% -> 79.0% ± 0.31% | Helped |
| Apply OT regularization and continuous-time formulation | Encoder-only Transformer · IMDb sentiment analysis | accuracy: 83.9% ± 0.26% -> 84.6% ± 0.55% | Helped |
| Apply OT regularization and continuous-time formulation | nanoGPT · Shakespeare character-level | test loss: 2.68 ± 0.006 -> 1.44 ± 0.005 | Helped |
| Apply OT regularization and continuous-time formulation | GPT-2 · Shakespeare word-level | test loss: 5.18 ± 0.032 -> 4.96 ± 0.012 | Helped |
| Apply OT regularization and continuous-time formulation | GPT-2 · OpenWebText | test loss: 3.21 -> 2.91 | Helped |
| Add OT regularization (ablation) | OT-Transformer · MNIST image classification | accuracy: 96.8% ± 0.23% -> 97.1% ± 0.16% | Helped |
| Add OT regularization (ablation) | OT-Transformer · Cats and Dogs image classification | accuracy: 78.2% ± 0.39% -> 79.0% ± 0.31% | Helped |
| Add OT regularization (ablation) | OT-Transformer · IMDb sentiment analysis | accuracy: 82.7% ± 0.38% -> 84.6% ± 0.55% | Helped |
| Vary OT regularization strength λ (ablation) | nanoGPT · Shakespeare character-level | test loss: 2.68 ± 0.019; 2.29 ± 0.017; 2.05 ± 0.009; 1.52 ± 0.002; 1.44 ± 0.004; 1.48 ± 0.003 | Mixed |
| Vary integration time step T/Δt (ablation) | nanoGPT · Shakespeare character-level | test loss: 1.46 ± 0.003; 1.43 ± 0.002; 1.44 ± 0.004; 1.46 ± 0.002; 1.49 ± 0.008 | No clear effect |
Surprising findings
- The prior continuous-time Transformer formulation (block-wise ODEs) failed to improve over the vanilla Transformer and the OT regularization could even degrade it.
- Unregularized OT-Transformer underperforms the baseline on sentiment analysis (82.7% vs 83.9%).
Limitations
- Computational overhead from numerical integration makes training roughly 2.7x slower per iteration than baseline on nanoGPT.
- Theory assumes convex loss in first argument, which holds for encoder-only/decoder-only but not generally for encoder-decoder models.
- Only a first step; layer normalization, attention mechanism, and integration schemes are not analyzed.
Experimental setup
Datasets: Shakespeare OpenWebText ModelNet40 MNIST Cats and Dogs IMDb
Benchmarks: Shakespeare char-level Shakespeare word-level OpenWebText ModelNet40 MNIST Cats and Dogs IMDb
Metrics: test loss test accuracy parameter count accuracy drop under noise/dropout
Baselines: nanoGPT baseline GPT-2 baseline Set Transformer ViT N-ODE Transformer Unregularized OT-Transformer
Builds on: Neural ODE Benamou-Brenier formulation Transformer (Vaswani et al.) N-ODE Transformer
Closest papers by meaning
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.89)
- Linear attention is (maybe) all you need (to understand Transformer optimization) (ICLR 2024 · similarity 0.88)
- Peri-LN: Revisiting Normalization Layer in the Transformer Architecture (ICML 2025 · similarity 0.87)
- PIDformer: Transformer Meets Control Theory (ICML 2024 · similarity 0.87)
- MambaPEFT: Exploring Parameter-Efficient Fine-Tuning for Mamba (ICLR 2025 · similarity 0.86)
- Dynamical Properties of Tokens in Self-Attention and Effects of Positional Encoding (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.