Generalization Bounds on Optimal Control for Transformer Training and Wasserstein Distributional Robustness
arXiv:2607.27975
The paper derives finite‑sample generalization bounds for Transformer models by formulating their training as a finite‑horizon Markovian control problem, using quantized approximations and concentration inequalities, and connects the analysis to Wasserstein distributionally robust optimization.
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.
25 pages