An Optimal Control Approach To Transformer Training
Authors: Kağan Akman, Naci Saldı, Serdar Yüksel
Organizations: Department of Mathematics, Bilkent University, Ankara, 06800, Turkey · Department of Mathematics and Statistics, Queen’s University, Kingston K7L 3N6, ON, Canada
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.
Figures & tables
Figure 1: Equivalence Properties of Control Policies for Transformer Design.
Figure 2: Four-level hierarchy of dynamics.
Figure 3: Decompositional illustration of the flow ( 1 ).
Figure 4: Flow of the ensemble of empirical measures on particle via the map Φ .
Figure 5: Comparison of information structures for open-loop and closed-loop controls. In the open-loop case, Ut∗ ’s are fixed before and do not depend on the current measure unlike the closed-loop case.
Qn(p,x)=(p,argmins∈Sn∥x−s∥)
Algorithm 1 Triply Quantized Dynamic Programming
Figure 6: Training and test errors as a function of action level.
Number of actions
Train error
Train improv. (%)
Test error
Test improv. (%)
10
0.01561
0.0
0.01117
0.0
20
0.00714
54.2
0.00737
34.0
30
0.00650
58.3
0.01171
-4.8
40
0.00650
58.3
0.01171
-4.8
50
0.00620
60.3
0.00816
27.0
60
0.00620
60.3
0.00816
27.0
Table 1: Relative improvement compared to number of actions.
Action level
Training error
Test error
Time (min:sec)
10
0.01561
0.01117
0:06
20
0.00714
0.00737
0:17
30
0.00650
0.01171
0:35
40
0.00650
0.01171
1:03
50
0.00620
0.00816
1:40
60
0.00620
0.00816
2:17
Table 2: Performance versus action level with runtime.
We derive finite-sample generalization bounds for Transformers trained with dynamic programming recursions. Building on the doubly lifted, measure-valued formulation of Transformer dynamics, we view data sets as probability laws on pairs of empirical input-output measures, allowing us to interpret the training problem as a finite-horizon Markovian control problem. We then analyze a quantized model, derived by quantizing the state, action, and measure-state spaces, and derive explicit finite-sample generalization bounds using concentration inequalities for empirical laws on finite metric spaces together with a Lipschitz stability estimate for the value function. These bounds are transferred to the base model at the cost of an explicit approximation error. Finally, we show that the same machinery yields a distributionally robust control formulation of the training problem, connecting Transformer generalization to Wasserstein distributionally robust optimization.
Kağan Akman, Naci Saldi, Serdar Yüksel
Department of Mathematics, Bilkent University, 06800, Ankara, Turkey · Department of Mathematics and Statistics, Queen’s University, Kingston, K7L 3N8, Ontario, Canada
Transformers can implement reinforcement-learning algorithms in context, yet how off-policy control emerges within Transformers remains poorly understood. We study Q-learning, a foundational model-free control algorithm that combines off-policy bootstrapping with greedy policy improvement. We prove that a two-layer attention-only Transformer can implement a finite-temperature form of Q-learning, where the temperature τ>0 controls the approximation to greedy action selection, recovering window Q-learning as τ→0. We train the structured Transformer on randomly generated MDP tasks using output-only supervision. Experiments show that the theoretically predicted parameter structure emerges during training, including the internal greedy action-selection mechanism, and that the trained Transformer achieves performance comparable to Q-learning on unseen tasks. Our theoretical analysis further shows that the equivalent realizations form a solution manifold and that, under suitable task-diversity conditions, the training iterates converge locally to this manifold at an exponential rate. For the random-task model analyzed here, the local worst-case number of gradient steps needed to reduce the distance to the manifold by a factor ε is Θ(τ−5log(1/ε)), quantifying the increasing optimization cost as the learned update approaches window Q-learning.
Yingjun Mo, Yajie Zhao, Fang Kong
Southern University of Science and Technology (SUSTech) · University of Electronic Science and Technology of China (UESTC)
Decoder-only transformers compute the conditional probability of the next token from a sequence of past observations. This paper derives, from first principles, inference architectures that solve the same prediction problem - and in doing so, recovers transformer-like layer operations as a consequence of optimal control theory. The framework is developed for two model classes: a nonlinear model of discrete-valued processes, directly motivated by the transformer, and a linear Gaussian model as a tractable baseline. For both model classes, the prediction objective is reformulated as an optimal control problem whose solution yields an explicit inference algorithm, the dual filter, with a layer structure that mirrors the layer structure of a decoder-only transformer. Numerical experiments provide a comparison of the optimal control to attention weights from a trained transformer. These experiments reveal that when the embedding dimension is insufficient, the transformer implicitly exploits non-Markovian structure.
Aditya Kudre, Heng-Sheng Chang, Prashant G. Mehta
Coordinated Science Laboratory · Electrical and Computer Engineering · University of Illinois Urbana-Champaign +1