# An Optimal Control Approach To Transformer Training

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 45 references that are in the map.
- First posted on arXiv 10 Mar 2026.

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Gradient Descent Provably Optimizes Over-parameterized Neural Networks](https://mlsanity.com/paper/1810.02054/) (2018)
- [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)
- [Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers](https://mlsanity.com/paper/1811.04918/) (2018)
- [Information-theoretic analysis of generalization capability of learning algorithms](https://mlsanity.com/paper/1705.07809/) (2017)
- [Deep Neural Networks Motivated by Partial Differential Equations](https://mlsanity.com/paper/1804.04272/) (2018)
- [A Convergence Analysis of Gradient Descent for Deep Linear Neural Networks](https://mlsanity.com/paper/1810.02281/) (2018)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [The emergence of clusters in self-attention dynamics](https://mlsanity.com/paper/2305.05465/) (2023)
- [Maximum Principle Based Algorithms for Deep Learning](https://mlsanity.com/paper/1710.09513/) (2017)
- [Formal Algorithms for Transformers](https://mlsanity.com/paper/2207.09238/) (2022)
- [A Mean-Field Optimal Control Formulation of Deep Learning](https://mlsanity.com/paper/1807.01083/) (2018)
- [Model-Free Mean-Field Reinforcement Learning: Mean-Field MDP and Mean-Field Q-Learning](https://mlsanity.com/paper/1910.12802/) (2019)
- [Neural ODE control for classification, approximation and transport](https://mlsanity.com/paper/2104.05278/) (2021)
- [A Unified Perspective on the Dynamics of Deep Transformers](https://mlsanity.com/paper/2501.18322/) (2025)
- [Generalization Bounds via Convex Analysis](https://mlsanity.com/paper/2202.04985/) (2022)
- [Mean-field Markov decision processes with common noise and open-loop controls](https://mlsanity.com/paper/1912.07883/) (2019)
- [Dynamic programming for optimal control of stochastic McKean-Vlasov dynamics](https://mlsanity.com/paper/1604.04057/) (2016)
- [Mean Field Markov Decision Processes](https://mlsanity.com/paper/2106.08755/) (2021)
- [Discrete time McKean-Vlasov control problem: a dynamic programming approach](https://mlsanity.com/paper/1511.09273/) (2015)
- [Randomized dynamic programming principle and Feynman-Kac representation for optimal control of McKean-Vlasov dynamics](https://mlsanity.com/paper/1606.08204/) (2016)
- [Near Optimality of Discrete-Time Approximations for Controlled McKean-Vlasov and Large Interacting Particle Diffusions](https://mlsanity.com/paper/2510.21208/) (2025)
- [Optimality of Symmetric Independent Policies under Decentralized Mean-Field Information Sharing for Stochastic Teams and Equivalence with McKean-Vlasov Control of a Representative Agent](https://mlsanity.com/paper/2404.04957/) (2024)

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

- [Reachability and asymptotics of Gaussian Transformer dynamics](https://mlsanity.com/paper/2606.07600/) (2026)
- [Transformer-like Inference from Optimal Control](https://mlsanity.com/paper/2605.15608/) (2026)
- [Generalization Bounds on Optimal Control for Transformer Training and Wasserstein Distributional Robustness](https://mlsanity.com/paper/2607.27975/) (2026)

## Abstract (arXiv)

In this paper, we develop a rigorous optimal control-theoretic approach to Transformer training that respects key structural constraints such as (i) realized-input-independence during execution, (ii) the ensemble control nature of the problem, and (iii) positional dependence. We model the Transformer architecture as a discrete-time controlled particle system with shared actions, exhibiting noise-free McKean-Vlasov dynamics. While the resulting dynamics is not Markovian, we show that lifting it to probability measures produces a fully-observed Markov decision process (MDP). Positional encodings are incorporated into the state space to preserve the sequence order under lifting. Using the dynamic programming principle, we establish the existence of globally optimal policies under mild assumptions of compactness. We further prove that closed-loop policies in the lifted is equivalent to an initial-distribution dependent open-loop policy, which are realized-input-independent and compatible with standard Transformer training. To train a Transformer, we propose a triply quantized training procedure for the lifted MDP by quantizing the state space, the space of probability measures, and the action space, and show that any optimal policy for the triply quantized model is near-optimal for the original training problem. Finally, we establish stability and empirical consistency properties of the lifted model by showing that the value function is continuous with respect to the perturbations of the initial empirical measures and convergence of policies as the data size increases. This approach provides a globally optimal and robust alternative to gradient-based training without requiring smoothness or convexity.