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

In this paper, we establish a collection of covering number bounds for linear function classes under various norm constraints on the inputs and matrices. We then combine these results with existing covering number bounds to derive improved estimates and, based on these estimates, develop generalization error bounds for single-layer Transformers. The resulting generalization bounds improve upon several existing results in the literature and, in particular, are independent of the input sequence length. Moreover, our generalization error bound decays at the rate $O(1/\sqrt{n})$, where $n$ denotes the sample size, thereby improving upon existing bounds that scale as $O((\log n)/\sqrt{n})$. Furthermore, our covering number analysis explicitly incorporates rank constraints on the underlying matrix classes, allowing us to characterize how low-rank structures affect the metric entropy and, consequently, the resulting generalization bounds for Transformer architectures.

In recent years, models based on the Transformer architecture have seen widespread applications and have become one of the core tools in the field of deep learning. Numerous successful techniques, such as parameter-efficient fine-tuning and efficient scaling, have been proposed surrounding their applications to further enhance performance. However, the success of these strategies has always lacked the support of rigorous mathematical theory. To study the underlying mechanisms behind Transformers and related techniques, we first propose a Transformer learning framework motivated by distribution regression, with distributions being inputs, connect a two-stage sampling process with natural language processing, and present a mathematical formulation of the attention mechanism called attention operator. We demonstrate that by the attention operator, Transformers can compress distributions into function representations without loss of information. Moreover, with the advantages of our novel attention operator, Transformers exhibit a stronger capability to learn functionals with more complex structures than convolutional neural networks and fully connected networks. Finally, we obtain a generalization bound within the distribution regression framework. Through the aforementioned theoretical results, we further discuss some successful techniques emerging with large language models (LLMs), such as prompt tuning, parameter-efficient fine-tuning, and efficient scaling. We also provide theoretical insights behind these techniques within our novel analysis framework.

We introduce a constrained optimization framework for training transformers that behave like optimization descent algorithms. Specifically, we enforce layerwise descent constraints on the objective function and replace standard empirical risk minimization (ERM) with a primal-dual training scheme. This approach yields models whose intermediate representations decrease the loss monotonically in expectation across layers. We apply our method to both unrolled transformer architectures and conventional pretrained transformers on tasks of video denoising and text classification. Across these settings, we observe constrained transformers achieve stronger robustness to perturbations and maintain higher out-of-distribution generalization, while preserving in-distribution performance.