OT-Transformer: A Continuous-time Transformer Architecture with Optimal Transport Regularization
Kelvin Kan, Xingjian Li, Stanley Osher
TL;DR
The paper tackles instability and poor generalization in continuous-time transformer models by introducing OT-Transformer, a continuous-depth transformer whose hidden states evolve under a single dynamical system. It regularizes training with an optimal transport term that penalizes the velocity of the state trajectory, and it provides theoretical support via optimal-control and HJB analysis showing well-posedness. Empirically, the method yields improved accuracy and stability across point-cloud, image, and sentiment tasks while achieving parameter efficiency and compatibility with existing transformer architectures. The approach demonstrates the practical viability and advantages of combining continuous-time modeling with transport-based regularization for robust, scalable transformer designs.
Abstract
Transformers have achieved state-of-the-art performance in numerous tasks. In this paper, we propose a continuous-time formulation of transformers. Specifically, we consider a dynamical system whose governing equation is parametrized by transformer blocks. We leverage optimal transport theory to regularize the training problem, which enhances stability in training and improves generalization of the resulting model. Moreover, we demonstrate in theory that this regularization is necessary as it promotes uniqueness and regularity of solutions. Our model is flexible in that almost any existing transformer architectures can be adopted to construct the dynamical system with only slight modifications to the existing code. We perform extensive numerical experiments on tasks motivated by natural language processing, image classification, and point cloud classification. Our experimental results show that the proposed method improves the performance of its discrete counterpart and outperforms relevant comparing models.