Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2025
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866912666971275264 |
|---|---|
| author | Kan, Kelvin Li, Xingjian Zhang, Benjamin J. Sahai, Tuhin Osher, Stanley Katsoulakis, Markos A. |
| author_facet | Kan, Kelvin Li, Xingjian Zhang, Benjamin J. Sahai, Tuhin Osher, Stanley Katsoulakis, Markos A. |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2505_13499 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency Kan, Kelvin Li, Xingjian Zhang, Benjamin J. Sahai, Tuhin Osher, Stanley Katsoulakis, Markos A. Machine Learning Artificial Intelligence Optimization and Control 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. |
| title | Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency |
| topic | Machine Learning Artificial Intelligence Optimization and Control |
| url | https://arxiv.org/abs/2505.13499 |