Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency

Kelvin Kan, Xingjian Li, Benjamin J. Zhang, Tuhin Sahai, Stanley Osher, Markos A. Katsoulakis

TL;DR

This paper frames Transformer training as a continuous-time optimal-control problem, introducing OT-Transformer—a plug-and-play method that adds an optimal transport regularization to the training objective. The authors prove theoretical properties including well-posedness, stable forward propagation, and distributional robustness, yielding non-asymptotic generalization guarantees. Empirically, OT-Transformer delivers consistent improvements in test performance and parameter efficiency across seven diverse tasks, from point-cloud and image classification to sentiment analysis and large-scale text generation, while also showing enhanced robustness to input perturbations. This work provides a theory-driven foundation for systematic improvements in Transformer architectures and training, with practical implications for efficiency and scalability.

Abstract

We study Transformers through the perspective of optimal control theory, using tools from continuous-time formulations to derive actionable insights into training and architecture design. This framework improves the performance of existing Transformer models while providing desirable theoretical guarantees, including generalization and robustness. Our framework is designed to be plug-and-play, enabling seamless integration with established Transformer models and requiring only slight changes to the implementation. We conduct seven extensive experiments on tasks motivated by text generation, sentiment analysis, image classification, and point cloud classification. Experimental results show that the framework improves the test performance of the baselines, while being more parameter-efficient. On character-level text generation with nanoGPT, our framework achieves a 46% reduction in final test loss while using 42% fewer parameters. On GPT-2, our framework achieves a 9.3% reduction in final test loss, demonstrating scalability to larger models. To the best of our knowledge, this is the first work that applies optimal control theory to both the training and architecture of Transformers. It offers a new foundation for systematic, theory-driven improvements and moves beyond costly trial-and-error approaches.

Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency

TL;DR

This paper frames Transformer training as a continuous-time optimal-control problem, introducing OT-Transformer—a plug-and-play method that adds an optimal transport regularization to the training objective. The authors prove theoretical properties including well-posedness, stable forward propagation, and distributional robustness, yielding non-asymptotic generalization guarantees. Empirically, OT-Transformer delivers consistent improvements in test performance and parameter efficiency across seven diverse tasks, from point-cloud and image classification to sentiment analysis and large-scale text generation, while also showing enhanced robustness to input perturbations. This work provides a theory-driven foundation for systematic improvements in Transformer architectures and training, with practical implications for efficiency and scalability.

Abstract

We study Transformers through the perspective of optimal control theory, using tools from continuous-time formulations to derive actionable insights into training and architecture design. This framework improves the performance of existing Transformer models while providing desirable theoretical guarantees, including generalization and robustness. Our framework is designed to be plug-and-play, enabling seamless integration with established Transformer models and requiring only slight changes to the implementation. We conduct seven extensive experiments on tasks motivated by text generation, sentiment analysis, image classification, and point cloud classification. Experimental results show that the framework improves the test performance of the baselines, while being more parameter-efficient. On character-level text generation with nanoGPT, our framework achieves a 46% reduction in final test loss while using 42% fewer parameters. On GPT-2, our framework achieves a 9.3% reduction in final test loss, demonstrating scalability to larger models. To the best of our knowledge, this is the first work that applies optimal control theory to both the training and architecture of Transformers. It offers a new foundation for systematic, theory-driven improvements and moves beyond costly trial-and-error approaches.
Paper Structure (64 sections, 10 theorems, 73 equations, 9 figures, 13 tables, 1 algorithm)

This paper contains 64 sections, 10 theorems, 73 equations, 9 figures, 13 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Notation.
  4. Transformers.
  5. Continuous-time Neural Networks.
  6. OT-Transformers
  7. Model Formulation.
  8. Plug-and-Play Formulation.
  9. Theoretical Analysis
  10. Assumptions and Practical Relevance.
  11. Non-parametric Formulation.
  12. Well-posedness of Training Problem.
  13. Stable Forward Propagation.
  14. Generalization and Distributional Robustness.
  15. Experimental Results
  16. ...and 49 more sections

Key Result

Theorem 1

For input-target output pairs $({\bf x}_1, {\bf y}_1)$ and $({\bf x}_2, {\bf y}_2)$, the corresponding output of the optimally trained model $\tilde{ {\bf y}}_1$ and $\tilde{ {\bf y}}_2$ satisfy where $C,C'>0$ are constants that can be controlled by adjusting the strength of the regularization. In the absence of the regularization, these constants become unbounded, and the model can exhibit uns

Figures (9)

  • Figure 1: Schematic Illustration of our Plug-and-Play Model \ref{['eq:training_obj']}. Given an existing Transformer, we construct a continuous-time formulation by using the Transformer to parametrize the velocity field of a dynamical system. Optimal control theory informs the learning of the velocity field by imposing regularity, in particular, through the use of optimal transport regularization. Empirical performance improvements (Section \ref{['sec:experiment']}) are consistent with our theoretical guarantees (Section \ref{['sec:theory']}).
  • Figure 2: Accuracy and data-fitting loss for the point cloud experiment (averaged over five trials)
  • Figure 3: Accuracy and data-fitting loss for the MNIST image classification experiment (averaged over five trials)
  • Figure 4: Accuracy and data-fitting loss for the cats and dogs image classification experiment
  • Figure 5: Accuracy and data-fitting loss for the sentiment analysis experiment
  • ...and 4 more figures

Theorems & Definitions (16)

  • Theorem : Stable Forward Propagation
  • Theorem 1
  • Theorem 2
  • Theorem 3
  • Lemma 1
  • Theorem 3
  • proof
  • Proposition 1
  • proof
  • Proposition 2
  • ...and 6 more