Generalization Bounds on Optimal Control for Transformer Training and Wasserstein Distributional Robustness

• URL: http://arxiv.org/abs/2607.27975v1
• Date: Thu, 30 Jul 2026 10:22:38 GMT
• Title: Generalization Bounds on Optimal Control for Transformer Training and Wasserstein Distributional Robustness
• Abstract summary: 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.
• Score: 8.564319625930894
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Reachability and asymptotics of Gaussian Transformer dynamics [2.669398324965473]
  We formulate data propagation through the Transformer, the machine learning architecture powering large language models.<n>We prove that Gaussian distributions remain exactly Gaussian along the induced flow.<n>We reformulates the expressive capacity of Transformers as a reachability problem for prescribed Gaussian moments.
  arXiv  Detail & Related papers  (2026-05-29T07:51:03Z)
• An Optimal Control Approach To Transformer Training [7.136933021609078]
  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.
  arXiv  Detail & Related papers  (2026-03-10T12:17:48Z)
• Stability and Generalization of Push-Sum Based Decentralized Optimization over Directed Graphs [55.77845440440496]
  Push-based decentralized communication enables optimization over communication networks, where information exchange may be asymmetric.<n>We develop a unified uniform-stability framework for the Gradient Push (SGP) algorithm.<n>A key technical ingredient is an imbalance-aware generalization bound through two quantities.
  arXiv  Detail & Related papers  (2026-02-24T05:32:03Z)
• Variational Entropic Optimal Transport [67.76725267984578]
  We propose Variational Entropic Optimal Transport (VarEOT) for domain translation problems.<n>VarEOT is based on an exact variational reformulation of the log-partition $log mathbbE[exp(cdot)$ as a tractable generalization over an auxiliary positive normalizer.<n> Experiments on synthetic data and unpaired image-to-image translation demonstrate competitive or improved translation quality.
  arXiv  Detail & Related papers  (2026-02-02T15:48:44Z)
• Stochastic Control Methods for Optimization [0.0]
  In the Euclidean setting, we analyze the problem of regularized control problems.<n>For global measures, we formulate a regularized mean-field problem characterized by a master-field problem.
  arXiv  Detail & Related papers  (2026-01-03T17:55:26Z)
• From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers [67.02076505996284]
  We study how the choice of pretraining data distribution steers a shallow transformer toward one behavior or the other.<n>Our results shed light on the algorithmic biases of pretrained transformers and offer conceptual guidelines for data-driven control of their learned behaviors.
  arXiv  Detail & Related papers  (2025-12-21T08:10:26Z)
• Closed-Loop Transformers: Autoregressive Modeling as Iterative Latent Equilibrium [0.6820746164515952]
  We introduce the closed-loop prediction principle, which requires that models iteratively refine latent representations until reaching a self-consistent equilibrium.<n>We instantiate this principle as Equilibrium Transformers, which augment standard transformer layers with an Equilibrium Refinement Module.<n>Preliminary experiments on the binary parity task demonstrate +3.28% average improvement on challenging sequences, with gains reaching +8.07% where standard transformers approach random performance.
  arXiv  Detail & Related papers  (2025-11-26T20:02:59Z)
• Transformers Are Universally Consistent [14.904264782690639]
  We show that Transformers equipped with softmax-based nonlinear attention are uniformly consistent when tasked with executing Least Squares regression.<n>We derive upper bounds on the empirical error which, in the regime, decay at a provable rate of $mathcalO(t-1/2d)$, where $t$ denotes the number of input tokens and $d$ the embedding dimensionality.
  arXiv  Detail & Related papers  (2025-05-30T12:39:26Z)
• Unveiling the Statistical Foundations of Chain-of-Thought Prompting Methods [59.779795063072655]
  Chain-of-Thought (CoT) prompting and its variants have gained popularity as effective methods for solving multi-step reasoning problems. We analyze CoT prompting from a statistical estimation perspective, providing a comprehensive characterization of its sample complexity.
  arXiv  Detail & Related papers  (2024-08-25T04:07:18Z)
• Efficient CDF Approximations for Normalizing Flows [64.60846767084877]
  We build upon the diffeomorphic properties of normalizing flows to estimate the cumulative distribution function (CDF) over a closed region. Our experiments on popular flow architectures and UCI datasets show a marked improvement in sample efficiency as compared to traditional estimators.
  arXiv  Detail & Related papers  (2022-02-23T06:11:49Z)
• 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)
• Learning Likelihoods with Conditional Normalizing Flows [54.60456010771409]
  Conditional normalizing flows (CNFs) are efficient in sampling and inference. We present a study of CNFs where the base density to output space mapping is conditioned on an input x, to model conditional densities p(y|x)
  arXiv  Detail & Related papers  (2019-11-29T19:17:58Z)

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.