# Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency

arXiv: https://arxiv.org/abs/2505.13499 · HTML: https://mlsanity.com/paper/2505.13499/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 62 of its 94 references that are in the map.
- First posted on arXiv 16 May 2025.

## Builds on (most cited first, 30 of 62)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space](https://mlsanity.com/paper/1706.02413/) (2017)
- [wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations](https://mlsanity.com/paper/2006.11477/) (2020)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Dynamic Graph CNN for Learning on Point Clouds](https://mlsanity.com/paper/1801.07829/) (2018)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [LaMDA: Language Models for Dialog Applications](https://mlsanity.com/paper/2201.08239/) (2022)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Is BERT Really Robust? A Strong Baseline for Natural Language Attack on Text Classification and Entailment](https://mlsanity.com/paper/1907.11932/) (2019)
- [Revisiting Point Cloud Classification: A New Benchmark Dataset and Classification Model on Real-World Data](https://mlsanity.com/paper/1908.04616/) (2019)
- [Harnessing the Power of LLMs in Practice: A Survey on ChatGPT and Beyond](https://mlsanity.com/paper/2304.13712/) (2023)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Stable Architectures for Deep Neural Networks](https://mlsanity.com/paper/1705.03341/) (2017)
- [Data-driven Distributionally Robust Optimization Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations](https://mlsanity.com/paper/1505.05116/) (2015)
- [Transformer for Partial Differential Equations' Operator Learning](https://mlsanity.com/paper/2205.13671/) (2022)
- [The GEM Benchmark: Natural Language Generation, its Evaluation and Metrics](https://mlsanity.com/paper/2102.01672/) (2021)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)

## Built on by (most cited first, 8 of 8)

- [Stability of Transformers under Layer Normalization](https://mlsanity.com/paper/2510.09904/) (2025)
- [Zero-Shot Transferable Solution Method for Parametric Optimal Control Problems](https://mlsanity.com/paper/2509.18404/) (2025)
- [Sparse Transformer Architectures via Regularized Wasserstein Proximal Operator with L1 Prior](https://mlsanity.com/paper/2510.16356/) (2025)
- [Dynamical Implicit Neural Representations](https://mlsanity.com/paper/2511.21787/) (2025)
- [State-space models through the lens of ensemble control](https://mlsanity.com/paper/2603.13587/) (2026)
- [Transformer-like Inference from Optimal Control](https://mlsanity.com/paper/2605.15608/) (2026)
- [End-to-End Learning of Safe Optimal Feedback Control in High Dimensions with Control Barrier Function Layers](https://mlsanity.com/paper/2607.20674/) (2026)
- [Generalization Bounds on Optimal Control for Transformer Training and Wasserstein Distributional Robustness](https://mlsanity.com/paper/2607.27975/) (2026)

## Abstract (arXiv)

We study Transformers through the perspective of optimal control theory, using tools from continuous-time formulations to derive actionable insights into training and architecture design. This framework improves the performance of existing Transformer models while providing desirable theoretical guarantees, including generalization and robustness. Our framework is designed to be plug-and-play, enabling seamless integration with established Transformer models and requiring only slight changes to the implementation. We conduct seven extensive experiments on tasks motivated by text generation, sentiment analysis, image classification, and point cloud classification. Experimental results show that the framework improves the test performance of the baselines, while being more parameter-efficient. On character-level text generation with nanoGPT, our framework achieves a 46% reduction in final test loss while using 42% fewer parameters. On GPT-2, our framework achieves a 9.3% reduction in final test loss, demonstrating scalability to larger models. To the best of our knowledge, this is the first work that applies optimal control theory to both the training and architecture of Transformers. It offers a new foundation for systematic, theory-driven improvements and moves beyond costly trial-and-error approaches.