An Optimal Control Approach To Transformer Training

• URL: http://arxiv.org/abs/2603.09571v1
• Date: Tue, 10 Mar 2026 12:17:48 GMT
• Title: An Optimal Control Approach To Transformer Training
• Abstract summary: We develop a rigorous optimal control-theoretic approach to Transformer training that respects key structural constraints.<n>We show that lifting it to probability measures produces a fully-observed Markov decision process (MDP)<n>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.
• Score: 7.136933021609078
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Generalization Bounds on Optimal Control for Transformer Training and Wasserstein Distributional Robustness [8.564319625930894]
  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.<n>We show that the same machinery yields a distributionally robust control formulation of the training problem, connecting Transformer generalization to Wasserstein distributionally robust optimization.
  arXiv  Detail & Related papers  (2026-07-30T10:22:38Z)
• Actor-Critic Learning for Extended Mean Field Control with Deterministic Policies [11.230955775247997]
  This paper develops a model-free reinforcement learning framework for continuous-time extended mean field control problems.<n>We adopt deterministic feedback policies, under which the state-action distribution is induced directly as a push-forward of the state law.<n>Experiments on Cucker--Smale consensus control and optimal liquidation with trade crowding demonstrate the efficiency, stability, and robustness of the proposed method.
  arXiv  Detail & Related papers  (2026-07-13T02:10:07Z)
• From Noise to Control: Parameterized Diffusion Policies [54.941343992110426]
  We propose a framework for learning diffusion policies conditioned on continuous parameters embedded in a learned behavior manifold.<n>Our approach enables smooth distances between known strategies and efficient adaptation to novel constraints without updating policy weights.
  arXiv  Detail & Related papers  (2026-05-29T20:21:50Z)
• Guided Streaming Stochastic Interpolant Policy [8.673777984856075]
  Inference-time guidance is essential for steering generative robot policies toward dynamic objectives without retraining.<n>In this work, we derive the optimal guidance term for dynamic, unstructured Interpolants (SI)<n>We apply this framework to real-time control through the Streaming Interpolant Policy (SSIP)
  arXiv  Detail & Related papers  (2026-05-11T06:24:21Z)
• Score-Guided Proximal Projection: A Unified Geometric Framework for Rectified Flow Editing [1.0312968200748118]
  Rectified Flow models achieve state-of-the-art generation quality, yet controlling them for precise tasks remains a challenge.<n>Current approaches bifurcate into inversion-based guidance, which suffers from "geometric locking"<n>We propose Score-Guided Proximal Projection, a unified framework that bridges the gap between deterministic optimization and contraction sampling.
  arXiv  Detail & Related papers  (2026-03-05T23:44:45Z)
• VESPO: Variational Sequence-Level Soft Policy Optimization for Stable Off-Policy LLM Training [18.849117699859622]
  Training stability is a central challenge in reinforcement learning for large language models.<n>We propose Variational sEquence-level Soft Policy Optimization (VESPO)<n> Experiments on mathematical reasoning benchmarks show that VESPO maintains stable training under staleness ratios up to 64x and fully asynchronous execution.
  arXiv  Detail & Related papers  (2026-02-11T09:48:08Z)
• Iterative Refinement of Flow Policies in Probability Space for Online Reinforcement Learning [56.47948583452555]
  We introduce the Stepwise Flow Policy (SWFP) framework, founded on the key insight that discretizing the flow matching inference process via a fixed-step Euler scheme aligns it with the variational Jordan-Kinderlehrer-Otto principle from optimal transport.<n>SWFP decomposes the global flow into a sequence of small, incremental transformations between proximate distributions.<n>This decomposition yields an efficient algorithm that fine-tunes pre-trained flows via a cascade of small flow blocks, offering significant advantages.
  arXiv  Detail & Related papers  (2025-10-17T07:43:51Z)
• Distributionally Robust Control with End-to-End Statistically Guaranteed Metric Learning [5.309590159815129]
  We propose a novel end-to-end finite-horizon Wasserstein DRC framework.<n>It integrates the learning of anisotropic Wasserstein metrics with downstream control tasks in a closed-loop manner.<n>We show that the proposed framework achieves superior closed-loop performance and robustness compared with state-of-the-art methods.
  arXiv  Detail & Related papers  (2025-10-11T13:40:49Z)
• Stabilizing Policy Gradients for Sample-Efficient Reinforcement Learning in LLM Reasoning [77.92320830700797]
  Reinforcement Learning has played a central role in enabling reasoning capabilities of Large Language Models.<n>We propose a tractable computational framework that tracks and leverages curvature information during policy updates.<n>The algorithm, Curvature-Aware Policy Optimization (CAPO), identifies samples that contribute to unstable updates and masks them out.
  arXiv  Detail & Related papers  (2025-10-01T12:29:32Z)
• ACPO: Adaptive Curriculum Policy Optimization for Aligning Vision-Language Models in Complex Reasoning [17.928214942495412]
  ACPO employs a dynamic curriculum that orchestrates a principled transition from a stable, near on-policy exploration phase to an efficient, off-policy exploitation phase.<n>We conduct extensive experiments on a suite of challenging multimodal reasoning benchmarks, including MathVista, LogicVista, and MMMU-Pro.<n>Results demonstrate that ACPO consistently outperforms strong baselines such as DAPO and PAPO, achieving state-of-the-art performance, accelerated convergence, and superior training stability.
  arXiv  Detail & Related papers  (2025-10-01T09:11:27Z)
• A learning-based approach to stochastic optimal control under reach-avoid constraint [8.354731976915588]
  We develop a model-free approach to optimally control, Markovian systems subject to a reach-avoid constraint.<n>We prove that under suitable assumptions, the policy parameters converge to the optimal parameters, while ensuring that the system trajectories satisfy the reach-avoid constraint with high probability.
  arXiv  Detail & Related papers  (2024-12-21T10:07:40Z)
• Formal Controller Synthesis for Markov Jump Linear Systems with Uncertain Dynamics [64.72260320446158]
  We propose a method for synthesising controllers for Markov jump linear systems. Our method is based on a finite-state abstraction that captures both the discrete (mode-jumping) and continuous (stochastic linear) behaviour of the MJLS. We apply our method to multiple realistic benchmark problems, in particular, a temperature control and an aerial vehicle delivery problem.
  arXiv  Detail & Related papers  (2022-12-01T17:36:30Z)
• Robust Value Iteration for Continuous Control Tasks [99.00362538261972]
  When transferring a control policy from simulation to a physical system, the policy needs to be robust to variations in the dynamics to perform well. We present Robust Fitted Value Iteration, which uses dynamic programming to compute the optimal value function on the compact state domain. We show that robust value is more robust compared to deep reinforcement learning algorithm and the non-robust version of the algorithm.
  arXiv  Detail & Related papers  (2021-05-25T19:48:35Z)
• Improper Learning with Gradient-based Policy Optimization [62.50997487685586]
  We consider an improper reinforcement learning setting where the learner is given M base controllers for an unknown Markov Decision Process. We propose a gradient-based approach that operates over a class of improper mixtures of the controllers.
  arXiv  Detail & Related papers  (2021-02-16T14:53:55Z)
• Gaussian Process-based Min-norm Stabilizing Controller for Control-Affine Systems with Uncertain Input Effects and Dynamics [90.81186513537777]
  We propose a novel compound kernel that captures the control-affine nature of the problem. We show that this resulting optimization problem is convex, and we call it Gaussian Process-based Control Lyapunov Function Second-Order Cone Program (GP-CLF-SOCP)
  arXiv  Detail & Related papers  (2020-11-14T01:27:32Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.