🤖 AI Summary
This work addresses the weak generalization, poor robustness, and low parameter efficiency of Transformers by systematically introducing optimal control theory into their modeling and training—first of its kind. We formulate a continuous-time dynamical framework that integrates variational inference, dynamics-based regularization, and lightweight controller embedding, enabling theoretically grounded training optimization and architecture design. Our approach departs from conventional black-box hyperparameter tuning, offering an interpretable and analyzable modeling paradigm. Empirical evaluation demonstrates a 46% reduction in training loss with 42% fewer parameters on nanoGPT; a 5.6% loss reduction on GPT-2; and consistent performance gains across diverse tasks—including text generation, sentiment analysis, image classification, and point cloud classification—validating its universality, strong generalization, and robustness.
📝 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 5.6% 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.