Search papers, labs, and topics across Lattice.
This paper formulates Transformer training as an optimal control problem, modeling the architecture as a discrete-time controlled particle system with McKean-Vlasov dynamics. By lifting the dynamics to probability measures, the problem is transformed into a fully-observed Markov decision process (MDP) that incorporates positional encodings. The authors propose a triply quantized training procedure for the lifted MDP and demonstrate that optimal policies for the quantized model are near-optimal for the original training problem, offering a globally optimal, gradient-free alternative.
Ditch finicky gradient descent: this paper recasts Transformer training as an optimal control problem, guaranteeing global optimality and robustness.
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.
