Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency
Authors: Kelvin Kan
Xingjian Li
Benjamin Zhang
Tuhin Sahai
Stanley Osher
Markos Katsoulakis
Venue: NeurIPS 2025 poster ()
Keywords: Transformer
LLM
Uncertainty Quantification
Optimal Control Theory
Optimal Transport
Primary area: deep_learning

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
Kelvin Kan
Xingjian Li Oden Institute
Department of Mathematics
UCLA kelvin. kan&math. wwla.lllaee..o edu
University of Texas at Austin
xingj ian. 1ellli@aaa li@austin... wtw.as.eeaa..mmmmmmmmmmm edu
Benjamin J. Zhang
Tuhin Sahai
Stanley Osher
School of Data Science and Society
SRI International
Department of Mathematics
UNC Chapel Hill
tuhin sahaissaaiiii @sri. com
UCLA sj sjo0math. ucla. edu
bjz@unc.ed edu
Markos A. Katsoulakis
Department of Mathematics and Statistics
University of Massachusetts Amherst markos@umass, edu
Abstract
We study Transformers through the perspective of optimal control theory, using tools from continuous-timg 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 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-lewe text generation with nanoGPT, our framework achieves a 46% reduction in final test loss while using 42% fewer parameters. On GPT-2, framework achieves a 9.3% reduction in final test loss, demonstrating scalability 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 foundation for systematic, theory-driven improvements and moves beyond costly trial-and-error approaches.
Introduction
Transformers have achieved state-of-the-art performance in various applications, including natural
language processing 669 83], computer vision (22), program synthesis [14], computational biol-
ogy [43], speech processing reinforcement learning [3 57], operator learning [56] 96) and
climate modeling
B1 60 611.
The popularity of Transformers has led to myriad architectural variants 17 19 87 each developed to
exhibit certain advantages. Practitioners, however, often discover effective Transformer architectures through trial-and-error approach. The objective function used to train Transformers admits natural
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

formulation as optimal control problem, with the loss serving as a terminal cost and model depth corresponding time. By examining the optimality conditions of this formulation, optimal transport (OT) emerges as principled regularizer that ensures well-posedness. This optimal control framework informs both the structure of the loss through regularization - and the design of the final layer, grounding these architectural choices in control-theoretic principles rather than heuristics.
Guided by this framework, we propose OT-rransformer, plug-and-play model grounded in optimal control theory. Our model is flexible and straightforward to implement in the sense that one can directly insert a predefined Transformer architecture into the OT-Transformer model. This requires only slight modifications to existing code and allows seamless integration with established models. An illustration of our framework is given in Figure
Ourgus f(Xg)
Ourpot X(7)
Optimal Control Thenry
Theoretical Guarantees
Transforme Block Io
Contimuous time Formudation
Robsuttess
Dynonical System
Geeeral Ination
x() (()(1), (). where 10 ....0 ff.
Thoeey euy Empiriod results.
Transformer Block f
Optimal Transport Regulatizztionn (1)(((()))))
Empirical Benefits
Transformer Block
Improved
Perfortannce
Inpul
Inpul X(0) Xo
Efficieney
OT Transformer
Figure 1: Schematic Illustration of our Plug-and-Play Model (8). Given an existing Transformer, construct continuous formulation by using the Transformer to parametrize the velocity field of dynamical system. Optimal control theory informs the learning of the velocity field by imposing regularity, in particular, through the of optimal transport regularization. Empirical performance improvements (Section 5) are consistent with our theoretical guarantees (Section
Empirically, model consistently improves the test performance of the original models across seven diverse experiments and demonstrates the following benefits.
Efficiency
Impementation Efficiency: Our optimal control framework can be directly applied to existing models, improving their performance while bypassing the costly and time- consuming and--ord approach to manual parameter tuning.
Parameter Efficiency: In our comprehensive experiments across various applications, we show that our framework can improve the performance of the original model with reduced parameter count. This also improves memory efficiency during inference.
A subset of our experimental results is reported in Table | We show that for variety of text generation tasks with nanoGPT and GPT-2, our Tranrsosfrrmer consistently improves the test performance using models of the same or smaller size.
Table 1: Results for text generation experiments in Section 5.4
Due to space constraints, additional L.M metrics provided in Appendix [I. method consistently outperforms the baseline across all reported metrics.
Experiment
Method
Para. Count
Test Loss
Baseline
10.65M
2.68 + 0.006
nanoGPT on Shakespeare (Char.-level)
OT-Trans. (Ours)
6.16M
1.44 + 0.005
Baseline
123.7N
5.18 + 0.032 4.96 + 0.012
GPT-2 on Shakespeare (Word-level)
OT-Trans. (Ours)
123.7M
Bascline
123.7M
3.21
GPT-2 on Open WebTet (9B tokens)
OT-Trans (Ours)
123.7M
2.91

---

The empirical findings are further supported by theoretical analysis in Section 4 Our main theorem, Theorem 2 which we present simplified version below, shows that Transsoorrerr exhibits stable forward propagation 39 73
Theorem (Stable Forward Propagation). For input-target output pairs (x1,y1) and (x2,82) the corresponding output of the optimally trained model y1 and Y2 satisfy
1 y21 - C||x] - x21l C'lly - -ylll
(1)
where C,C O 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 unstable forward propagation.
This stability induces highly favorable properties in practice, highlighted as follows.
Robustness
Robustness to Input Perturbations: On data level, controls the extent to which input perturbations are amplified in the output. As a result, this enhances robustmess against input perturbations, such as those caused by noise adversarial attacks.
Distributional Robusiness: On a distribution level. (I) guarantees that distributional perturbations in the input space not disproportionatell amplified at the output level, which we show in Theorem 3 Distributional robustness induces robustness to training data perturbation. In particular, it ensures that model trained perturbed data maintains good predictive performance on the original data, with the prediction error controlled by the degree of the perturbation.
Generalization
In-Distribution Generalization: Stable forward propagation improves generalization to test inputs resembling the training data, as ( ensures that similar inputs are mapped to similar outputs. Also, distributional robustness ensures that model trained on the training distribution (sampled from the true distribution) can generalize well to data drawn from the true distribution. Combined with distributionally robust optimization (DRO), this robustness result produces non-asymptotic generalization bounds for the learned Transformer model 24 75 76
Out-of- Distriibution Generalization: Stable forward propagation combined with DRO also provides bounds expected test loss for distributions within Wasserstein ball around the empirical training distribution, offering principled insight into the model's performance
samples.
We summarize our contributions as follows:
We analyze the training and architecture of Transformers using optimal control theory. To the best of our knowledge, this is the first such analysis of Transformers from this perspective.
Based on optimal control methods, propose plug-and-play model called OT-Transformer ( ( which formulates Transformers in a continuous-time setting and incorporates an OT regularization into the training objective. Our framework is versatile, allowing for easy adaptation of existing Transformer architectures. We remark that previous attempts to apply OT to the design of Transformer models remain underexplored and have achieved limited success
5
Our theoretical analysis demonstrates that the OT-Transformer model confers highly beneficial theoretical properties. The most notable property its stable forward propagation, which induces generalization and robustness of the learned model. These properties further supported by our experimental results. We emphasize that our use of optimal control theory represents only initial step in analyzing Transformers, and our framework provides new foundation for future theoretical extensions to many other aspects of these models.
We conduct seven extensive experiments across different fields and applications. Our frame- work improves the performance of the original architectre. In particular, yields better test performance and is better at avoiding overfitting while using reduced number of parameters.
3

---

2 Background
This section provides background on Transformers and continuous-time neural networks.
Notation. We use bold upperease letters (e.g., X) to denote matrices and bold lowercase letters (e.g., x) to denote vectors. We also use x, (resp. x1.)) to represent the jth column of X (resp. X,).
Transformers. Let {zi)i=1, where 2; RRy, be a sequence of n vectors. In the language of Transformers, each vector Z; is referred to as a token. Transformers are mappings which take in inputs Z [ E Rdyxn, with the specific form of the output depending on the downstream application. Each token is first emhedded into IRd through the mapping
Xu.j (z); ()), for = 1,2,.....
where Xo, E Rd The input embedding 9I. parametrized by weights Y. embeds each token into a d-dimensional space and incorporates positional encoding into each token. Then, it is processed through series of Transformer blocks, where the output of each block serves the input to the next. At each step, the model sequentially applies the operation X = fi+1 (X;) where fi+1 : Rd Rd given by
H W V softmax (K)X,) X,) TQ/
W x,
(3)
VT
h=1
X U.j + 91 J 0.),
(4)
for 1,2,..., and 0, 1, .... D where D is the total number of Transformer blocks. Here, each summand of (3) are called self-attention heads where is the number of heads, Q K. V Rexd are known as query, key, and value matrices, and W E Baxk are weight matrices. In 4 fully connected layer gs 12d Fd parametrized by weights 0,, is applied individually to each of the n tokens. This layer is common for all embedded tokens but changes for each i. The first equation 33 is known as multihead self-attention layers and the key feature of Transformer architectures. A detailed derivation of (3) is given in Appendix [ This self-attention mechanism enables the model to focus on the most relevant parts of an input sequence of tokens, adapting dynamically the context. The model can flexibly capture complex, long-distance relationships in sequential data. Such features make Transformers particularly powerful for tasks such as language processing and image recognition. In addition, the self-attention mechanism can be implemented efficiently as the manner in which the matrices are applied to process input data can be done in parallel, rendering them particularly effective for handling long sequences of tokens, i.c., when is large. In encoder-only and encoder-decoder setups, the series of Transformer blocks is referred to as the encoder; in decoder-only setups, it is referred as the decoder.
For sequence generation tasks, the final output XD is then either passed a decoder (in the incoder-decoder setup), which consists of another series of Transformer blocks and then a multilayer perceptron (MLP). More commonly, the output feeds directly to an MLP (in the encoder-only or decoder-only setup) in downstream tasks including sequence generation, classification, or regression. Irrespective of the setup, the Transformer output ý given by
ý go (X 70),
(5)
where yo is either the composition of decoder and an MLP or just an MLP, parametrized by Yo-
Continuous-time Neural Networks. Continuous-tinn time models use neural networks to define the dynamics of the model's hidden states. In particular, the hidden state x(t) evolves according to
dx(() INN(x(1),
dt
where E [0. 7]. The model takes x(0) input and produces x(T) as output, with the velocity field parametrized by the neural network fNN. Various architectures have been proposed and successfully
'Layer normalization is commonly applied in each Transformer block 992 We omit layer normalization in this exposition for brevity, but it is included in experiments.

---

applied in [15
26 40
04
95, 99] A notable and relevant advantage of continuous- archi- 84,
tectures their parameter efficiency, as the continuous formulation allows them model complex transformations over time with fewer parameters compared to discrete models. However, the use of Transformers a continuous-time setting remains largely unexplored, with existing approaches showing limited success.
3 OT-Transformers
In this section, introduce OT-Trassformer, model that can be flexibly combined with existing Transformers. The model inserts a given Transformer into a continuous-time architecture and incorporates optimal transport regulaaaaaii into the training objectíve, which are motivated from the theoretical analysis using optimal control theory in Section
Model Formulation. Given an input sequence Z = [Z], 22, Zn| of length n, we first apply (2) to obtain the embedded input Xo E Bdx=, The dynamics of the hidden state are then governed by the dynamical system
dx (t) - f(X()),,; for iE [0,T], with X(0) = Xo: di
where f is the composition of a sequence of Transformer blocks defined in and ( that is, f f= SD .... o fI, and 0 collectively denotes their trainable parameters O,, K V Q and w for
all h and i. Finally, we obtain the Transformer output y by applying (5) to the terminal state (T).
Plug-and-Play Formulation.
We formulate the discretized training problem
Ade M-1
8,7 min E(xo,r)
+ G(XM,Ni7)
Σ
llf( m, tm 0)12 }
(8)
2
m==
subject to X Xm At. f(X km; O),
m 0, 1,......,A 1
Here, we adopt a discretize-then-optiine of the continuous dynamics in ( obtained by splitting the time interval 00. T into M uniform steps with step size Al = T/M and Lm = mAL. The optimization problem ( the continuous-time training problem (10) over parametrized family of Transformer models. The expectation is taken over the embedded input-output pairs (X..).. ||- || denotes the Frobenius norm, 7 collectively denotes the weights of the input embedding 7 and output layer Yo. The loss function G measures the discrepancy between the target output and model output yX Y) in (5). For instance, in classification [22] and sequence generation (83] tasks, one commonly uses the softmax loss. The second term is an OT regularization penalizing the squared norm of the velocity (the right hand side of the dynamical system (7) at every time step. enhances the regularity of the hidden state dynamics (7): see Section 4 We highlight that our model is plug-and-pll aand-lay and easy implement in the sense that it can be flexibly applied almost any existing Transformer architectures. In particular, it can directly reuse the architecture of an existing Transformer's 's input embedding, encoder/decoder and output layers and use its Transformer blocks fi's to construct the dynamical system (9). In practice, the OT regularization term is computed efficiently and incurs negligible overhead, as it is calculated alongside each X,, in 9 A schematic illustration of our model is shown in Figure [ The regularization parameter balances the effects of the two terms. Our model generalizes the vanilla formulation of Transformer blocks to continuous-tima Specifically, when T = and M 1, our model is consistent with the original discrete Transformer formulation.
approach 66 651 where p represents the discretized form
( is the discretization of
On one hand, the training formulation is derived from the optimal transport problem arising from the Banamou--Beniee formulation
6 as shown in Appendix E] The derivation reveals that for an appropriate choice of > 0. the solution to the training problem corresponds to the optimal solution of the Benamou-Brente problem. On the other hand, we will demonstrate in the next section that the training problem can be understood through optimal control theory. And our extensive theoretical analysis shows that our model confers highly advantageous properties.
5

---

4 Theoretical Analysis
We provide theoretical analysis demonstrating the benefits of our model. The theoretical benefits presented in this section and the empirical evidence in Section SectioonSelln eluucidate and support each other.
Assumptions and Practical Relevance. Our assumptions specified in Appendix F.1 are realistic in that they hold for encoder-only and decoder-only Transformer models in the continuous-time settings. They are also general and hold for other continuous tieeuo models, including Neural ODEs [[5]. These settings cover many leading models, including Vision Transformers [22]. and language
models including BERT [2]] the GPT series [10] 69, [70] of OpenAl, PaLM (1), GLaM (23)
and LaMDA 79 of Google, OPT (100) of Meta AI, and Granite of IBM (38 among others.
One of our key assumptions deriving the theoretical advantages is for the loss function G to be convex with respect to its first argument. Surprisingly, this holds for encoder-only and decoder-only Transformer architectures; see Appendix [F.T] Generally, the convexity does not hold for encoder- decoder models. We remark that this distinction is not just theoretical, it is consistent with empirical findings from the literature. In particular, decoder-only models can match or even surpass the performance of encoder-decoder models for various large-scale language modeling tasks (28 88 94],
despite having a simpler architecture. Moreover, current industry trends [] 11 80 81 also suggest that decoder-only Transformers are replacing encoder-decodes models for LLMs. Meanwhile, in areas such as image classification or time series forecasting, encoder-only models remain a popular option, as pointed out in [41 49 197. The alignment between our theoretical analysis and reported empirical performance further shows that our assumptions are relevant in practice. Our analysis offers a possible explanation for the performance gap between prooder-only/ddccoer-on. models and encoder-decoder models, a distinction that arises from our convexity assumption. For more details, see Appendices and F.1
Non-parametrie Formulation.
We consider a continuous-time time non- parametric formulation of (8)
A min x) ((X(T) + 2 [" (()x di } subject to 7
(10)
where the velocity f is optimized over some admissible class of functions instead of model weights. The analysis of the non- non-parametric formulation justified by the universal approximation property of Transformers [98 Theorem 3], which ensures that, with enough model complexity, they can represent the optimal velocity. A central component of our approach is that the non-parametric formulation can be equivalently cast as mean-field control problem, enabling theoretical analysis via optimal control theory. the following, we highlight key theorems that show the benefits of OT- Transformer, including well-posedness of the training objective, regularity of the learned Transformer, and non-asymptotic generalization bounds. For detailed results and proofs, see Appendices F and H
Well-posedness of Training Problem. We first show that the training problem is ill- ill-posed without regularization and that the OT regularization renders the problem well-posed.
Theorem 1. There exists a unique solution to the optimization problem (10) if and only if > O.
The ill poseddess of (10p that arises when X = 0 is because there exists infinitely many velocities that are optimal. Among the infinitely many optimal velocities for the unregularized problem, some are highly irregular have arbitrarily large magnitudes. This phenomenon has been noted in flow-based generative modeling [6 64). In particular, highly irregular velocities fields can produce winding hidden state trajectories which can pose challenges in numerical integration and result in numerical instability during training, as demonstrated by our experiments.
Our analysis is sufficiently general to apply other continuous-tin models beyond the Transformer architecture. In particular, Theorem provides insight into the training of any neural network model under continuous-tin time formulation. In the absence of regularization, such training is, in general, ill-posed. OT regularization is only one strategy that can ensure well poseedess,,,ss, and optimal control methods can be used to construct other regularizations that ensures well-posedness (99].
detailed proof is provided in Appendix F.4 The core idea to inspect the optimality conditions for the training problem (10), which is Jamolto Bellman (HJB) partial differential equation

---

(PDE) coupled with a continuity equation. When À = O, the HJB PDE is not well-defined. On the other hand, if and only the HJB PDE is well-defined, and the well poseednesss of the resulting PDE system is, in turn, equivalent to the well-posedness of the training problem 7 51
Stable Forward Propagation. Furthermore, we show that À > 0 not only renders the optimal velocity unique, but also guaranteed that is highly regular; see Appendices [F5 and F6 for details. Leveraging this regularity, we prove our main result: stable forward propagation.
Theorem 2. Assume that X> TI?, where T is the terminal time and T. is the 2-operator norm of the output layer weights (5) Under the assumptions in Appendix [.1] for embedded input-outpi pairs (1), and (X2 82), the corresponding model outputs yı and ý2 satisfy
1 - velle S (1 ) -
TL? -1
TI²
Xx(O))l -
ly --2lll))) (1I)
This shows that the model, as a function of the embedded input and target output data, is Lipschitz continuous. In particular, any perturbation in the input leads to a uniformly and proportionally bounded change the model output. This theorem also informs the selection of the hyperparameter A >> 0. On one hand, it should not be too large that the model outputs are insensitive to the input, On the other hand, it should be large enough to satisfy À > TI.². In practice, the condition À > TT,2 can be enforced, and the bound (11) can be controlled through increasing À or imposing deight-decay regularization on the output layer weights
5
which reduces its operator norn L.
The proof of Theorem 2 is given in Appendix F7
The core mathematical novelty of our paper is the use of regularity theory of HJB PDEs to prove stable forward propagation of the Transformer map in Theorem which yields the distributional robustness result in Theorem 3] see Appendices F.5 and F.6 Crucially, we also use the regularity assumptions on G. such its convexity.
Generalization and Distributional Robustness. Stable forward propagation informs out-of- distribution performance, generalization bounds, and the robustness of learned Transformers to adversarial attacks. For each target output y, we study the regularity of the flow map of the learned dynamical system that evolves inputs X(0) to predictions y. In this setting, Theorem 2 the flow map is Lipschitz with constant I.(1 -X--TT?2-- This result shows that while the map
implies that
not explicitly enforced to be Lipschitz during training, the OT regularization nevertheless induces this property. The Lipschitzness of the learned map implies distributional robustness, which we state formally as follows.
Theorem 3. Denote the pushforward operator under the trained Transformer to be T; and ()) = inf in Fraxze ||x x'Il to be the p- p-Wasserstein distance. Under the as-
sumptions of Theorem 2 there exists constant ((p) > that depends on p such that for any p 2 and any distributions and v,
TL WT Tu²) < C(p) (1 (1- X Wp(m,
(12)
This is classic optimal transport result 886 Chapter 6]. which holds under the Lipschitzness of the trained model proved in Theorem 2 See Appendix[H.] for a detailed proof. The bound (12) quantifies how perturbations to the input distributions y and μ propagate to the model output distributions. The perturbations can be due to noise, new data, or adversarial modifications. Here, Wasserstein distances crucial, as they able to measure discrepancies between empirical distributions, which is not possible for probability divergences (e.g., Kullback-Leibler or f-divergences). Crucially, 12) also depends on the convexity assumption G and the magnitude of OT regularization.
The bounds (TT) and (12) induce highly desirable practical generalization. Due to space constraints, we refer the reader to the discussion in Section Explicit formulas and detailed derivations are provided in Appendix H
paperttis--naeeee robustness and
'The squared Frobenius norm of the weight matrix equals to the sum of its squared singular values, while the operator norm of the matrix is its largest singular value. Therefore, applying weight decay regularization can decrease its operator norm L.
7

---

Experimental Results
We demonstrate the effectiveness of Or-aransformers through seven experiments spanning diverse applications, including point cloud classification, image classification, sentiment analysis, and text generation. Our model performs competitively and generalizes well across all tasks.
For each task, we use commonly adopted Transformer models as baselines, including encoder- only, decoder-only, and encoder-decoder architectures. This ensures comprehensive evaluation. Moreover, while our theoretical analysis provides guarantees for encoder-only and decoder-only architectures, the empirical results demonstrate that our model also performs well with encoder- decoder architectures. We base our experiments on the setups from [48] [7],, building on their code and closely following their experimental protocols. Hyperparameters. including model architectures, number of training epochs, learning rates, and layer normalization, closely follow the original setups.
For the OT-Transformer, we employ the same architectures as the baselines but with reduced hidden dimensions, number of attention heads, and/or number of Transformer blocks. Through this, we demonstrate that OT-Transformers can achieve better performance across various tasks while having reduced model sizes. To demonstrate the effectiveness of the OT regularization, we also perform the experiments with À 0 in (8) effectively creating an unregularized model. We label this model unregularized OT- Trannsssfrrerr in the reported results.
See Appendix [ for experimental details. The source code is publicly available at htttp:///gtthi
com/ Ben/Kelvinn weweevvinnannOO-- Coaoo--raaffrrrrn
5.1 Point Cloud Classification
We the ModelNet 40 dataset 990,, which is among the most widely used benchmark for point cloud classification 82). The dataset contains roughly 10,000 Computer-Aided Design (CAD) models that are categorized into 40 distinct classes, including common objects such as airplanes, cars, and furniture. We experiment with the Set Transformer model [53]. which notably has an encoder-decoder architecture.
5.2 Image Classification
To further demonstrate the applicability of our proposed method, we perform experiments on imaging tasks. We use the Vision Transformer (ViT) [22]. which an encoder-only model. Since then, the model and its variants have achieved state-of-the-arr performance in computer vision tasks [7] 91]] The key feature of VITs is that they divide an image into fixed-size patches, which are treated as sequences of data. ViTs then apply self-attention mechanisms to capture relationships between these patches, enabling to learn complex structures across the entire image. We perform two image classification experiments following the same setup as in (74)
MNIST Classification. We first conduct a small-scale image classification experiment with the MNIST dataset [52). The dataset consists of hand- written digit images, with 50,000 images used for training and 10,000 images reserved for testing. Each image is of size 28 by 28.
Cats and Dogs Classification. We perform experiments on binary cats and dogs image classifi- cation task, following [74]. The dataset contains 25, 000 training samples and 12, 500 test samples, each image in the dataset an RGB image of size 460 320.
5.3 Sentiment Analysis
We perform sentiment analysis on the IMDb movie review dataset 59 which aims to predict whether each movie review is positive or negative. The dataset is balanced and contains a total of 50, 000 different reviews. The model used in the experiment is an encoder-only Transformer [74]
5.4
Text Generation
To further demonstrate the capabilities of the OT-Transformer, we conduct experiments on text generation. We use nanoGPT [48 and GPT-2 669) both of which are decoder-only models with

---

Table 2: Results for experiments from Sections 5.1 to 5.3
Experiment
Method
Para. Count
Test Accuracy
Bascline
0.86M
87.4% + 0.45%
Point Cloud Classification
OT-Trans. (Ours)
0.65M
89.9% 0.42%
Baseline
93K
93.0% + 0.69%
Image Classification (MNIST)
Unreg. OT-Trans.
18K
96.8% + 0.23%
OT-Trans. (Ours)
18K
97.1% + 0.16%
Baseline
1.77M
77.6% + 0.86%
Image Classification (Cats Dogs)
Unreg. OT-Trans.
1.48M
78.2% 0.39%
OT-Trans. (Ours)
1.48M
79.0% + 0.31%
Baseline
4.74M
83.9% + 0.26%
Sentiment Analysis
Unreg. OT-Trans.
2.37M
82.7% + 0.38%
OT-Trans. (Ours)
2.37M
84.6% + 0.55%
10.7 million and 124 million parameters, respectively. We conduct three different text generation experiments using different data. The goal to evaluate the performance of OT "T-Transformer on text generation tasks and assess its scalability to large models with over 100 million parameters. The results are reported in Table of Section
Shakespeare Dataset with nanoGPT. We first conduct experiments using the smaller-sized nanoGPT architecture on the benchmark Shakespeare dataset. The source text is taken from Shake- speare's is works, and the goal to make predictions at the character level based on input sequences.
Shakespeare Dataset with GPT-2. We next perform in-depth experimentation the Shakespeare dataset using the much larger GPT-2 architecture, which contains over 100 million trainable parame- ters. Note that for this experiment, token prediction is performed at the word level, making the task more difficult compared to the previous example.
OpenWeb WebTeet Dataset with GPT-2. Lastly, to demonstrate the applicability of our model to large-scale problems, we conduct experiments using the GPT-2 architecture baseline the OpenWebTTe WebText dataset. This dataset, originally curated in (36) from Reddit posts, includes a training set of approximately 9 billion tokens and validation of around 4 million tokens. This experiment is large-scale in both model and dataset size, representing a realistic application setting.
5.5 Summary of Numerical Results
Compiled Results. We present the compiled numerical results in Tables Tand] 2 and summarize the main findings follows. First, our proposed model demonstrates competitive performance across all seven test examples, ranging from small-scale to large-scale experiments, highlighting its ability consistently improve performance
baseline models. Second, OT--Transfrrmer outperforms baseline models across various examples while using significantly smaller models, showcasing its parameter efficiency. Finally, Or-Transformer avoids overfitting more effectively, resulting in improved generalization and lower test loss compared to the baseline; see Figures 5 and g in Appendix
Note that while our experiments focus on generalization metrics like test loss and accuracy, strong performance in these settings also reflects robustness. In real-world scenarios, test data often differ from training data due to noise, sampling variability, or distribution mismatch. The consistent results across different experiments support theory Theorems and 3 robustness to input, distributional and training data perturbations.
Robustness Tests. To further validate our stability theory, we conduct four additional experiments evaluating the performance of TTTTraasforrrr on test data corrupted by varying levels of random noise. The noise was absent from the training data. This setup assesses the model's forward stability, distributional robustness and out-of-distributio generalization.

---

Table 3: Robustness tests under varying levels of noise.
Experiment
Metric/rephae rate
0.0
0.005
0.01
0,05
0.1
Loss (#istd) Base.
2.68=0004
2.78=0004
2.88
3.60
4.60
NanoGPT
Loss (#stdd Ours
1.4
1.49:0004
1.55
1.95
2.42:0002
replace)
Drop (L) Base.
0.10
0.20
0.97
1.92
Drop (U) Ours
0.05
0.11
0.51
0.98
Experiment
Metric/ /rrop rate
0.0
0.01
0.05
0.1
0.2
0,5
Acc. (#hstd) Base.
86.6%
86. 60.- aes
85, 896-0
84.3% ono
76.9%-2000e
34.5%
Point cloud (dropout)
(dhttd) Ours
89.3% 19358
0..5%
83.9% 10.00%
55.4%
Drop (1) Base.
0.0%
0.8%
2.3%
9.7%
52.1%
Drop (L) Ours
0.0%
0.5%
1.7%
5.4%
33.9%
Experiment
Metric/ /noise
0.01
0,05
Acc. (dstd) Base.
92.97% sours
=2.999%-=000
92.76% 20 128
91. 91.70%=corre
80.6-458:1..55
MNIST (Ganss, noise)
(d:std)
97.05 % =0.5%
97.05 -0.16%
96.99%
96.89 #
96.39%
90.10%
Drop (1) Base.
-0.02%
0.01 %
0.21%
1.27%
12.33%
Drop (4) Ours
0.00%
0,06%
0.16%
0.66%
6.95%
Acc. (:sttd Base.
92.97% Loars
92. 1999-1
92.9 Lusas
92. 909 1.
92.644 _0977
90.02% Luuee
MNIST (Uni, noise)
Acc. (#btth) Ours
97.05% 20155
97.03 % 20.14
97.00%
96.97% 8..4%
96.79 % 20.6%
95.57 %6 10.11
Drop (1) Base.
-0.02%
0.33%
2.95%
Drop (1) Ours
0.02% %
0.05%
0.08%
0.26 %
1.48%
In the first experiment, we evaluate the NanoGPT model, where each character in the test data is randomly replaced at specified rate, following the setup in [42]. In the second experiment, we perform point cloud classification under point dropout, where a fraction of points in each test sample removed randomly, following [68
89 The third and fourth experiments focus on MNIST dataset with Gaussian and uniform noise added to the test images, respectively. The setups follow [37]
The results reported in Table Our model consistently outperforms the baselines over all these tests across all noise levels. Importantly, the performance gap widens at higher noise levels. These additional results again corroborate our theory and other empirical findings. They also further highlights the attractiveness of our model in practice.
We also provide an empirical sensitivity study on the model's hyperparameters. The results show that our model again consistently outperforms the baseline across various numbers of integration steps and over two orders of magnitude of the regularization strength x. Due to space constraints; detailed results are provided in Appendix C
Overall, the OT- Tr-nnssffrrer results provide evidence - backed up by our theory in Section 4 that optimal transport, acts as a unifying regularization principle across text, image, and 3D modalities. For more details, we direct readers to Appendix
6 Discussion and Summary
In this work, we analyze the Transformer architecture and training through a proposed optimal control framework. Based on this, we propose Transoosf,r,,,, a plug-and-play model which can be flexibly applied to established Transformers with minimal code changes. OT Tranffoffrrerr improves the performance of existing models while conferring strong theoretical guarantees. These include generalization and robustness, established through our optimal control-based analysis. We highlight that key theoretical results rely on the novel application of the regularity theory of HJB PDEs to prove stable forward propagation and distributional robustness of the learned Transformer. This further proves non-asynn non-asymptotic generalization bounds through DRO. These theoretical results are supported by extensive experiments that demonstrate the effecttivnness of the optimal control framework and further show that the resulting -T-aranSformer model improves parameter efficiency. Overall, our framework provides a foundation for systematic and theory-drived improvements for Transformers.
We emphasize that this work represents only first step toward building control- -hhooreicc foundation for designing and analyzing Transformer architectures and training. Many additional components of Transformers - such as layer normalization, attention mechanism, and others - present promising directions for future investigation through optimal control methods. In addition, exploring alternative numerical integration schemes may offer a path toward improving training efficiency; we leave this for future work.
10

---

Acknowledgements
The authors would like to thank Lars Ruthotto (Emory University) and Tingwei Meng (Amazon Robotics) for their valuable advice and insightful discussion.
The authors would also like to thank the four anonymous reviewers for their thorough review and constructive suggestions.
K. Kan and S. Osher were partially funded by the U.S. Department of Energy (DOE), Office of Science (SC), Advanced Scientific Computing Research program under award B&R# KJ0401010, FWP# CC147.
X. Li was partially funded by NSF 2339678 and 2321040.
B. Zhang and M. Katsoulakis were partially funded by AFOSR grant FA9550-21-1-0354
This material is based upon work of M. Katsoulakis and T. Sahai supported by the Defense Advanced Research Projects Agency (DARPA) under Agreement No. HR001 125901 12. Approved for public release; distribution is unlimited.
Osher was partially funded by DARPA under grant HR001 1259007, NSF under grants 1554564 and 220827, AFOSR under MURI grant N00014-20-1-2777, and ARO under grant 1911NF-24--00157
References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.0877 2023.
[2] Daniel Owusu Adu and Bahman Gharesifard. Approximate controllability of continuity equation of transformers. IEEE Control Systems Letters, 2024.
[3] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008.
[4] Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for self-supervised learning of speech representations. Advances in neural information processing systems, 33:1 13:1244912266 2020.
[5] Aaron Baier-Reinia and Hans De Sterck. N-ODE transformer: A depth-adaptive variant the transformer using neural ordinary differential equations. arXiv preprint arXiv:2010.11358. 2020.
[6] Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the monge-kantorovice mass transfer problem. Numerische Mathematik, 84(3)3375-333, 2000.
Alain Bensoussan, Jens Frehse, Phillip Yam, et al. Mean field games and mean field rype control theory, volume 101. Springer, 2013.
[8] Shiba Biswal, Karthik Elamvazhuthi, and Rishi Sonthalia. Identification of mean-field dynam- ies using transformers. arXiv preprint arXiv:2410. 16295, 2024.
[9] Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics, 44(4):375-41) 1991.
[10] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877- 1901, 2020.
[11] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advonces Neural Information Processing Systems, 33:1877- 1901, 2020.
11

---

[12] Giuseppe Bruno, Federico Pasqualotto, and Andrea Agazzi. Emergence of meta-stable cluster- ing in meean-field transformer models. arXiv preprint arXiv:2410.25222,, 2024.
[13] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 34:15084-10097, 2021.
[14] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint
2021.
[15] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[16] Yali Cheng, Ruichen Li, Juan Cheng, and Menglin Kong. Rktrans: Transformer with improved residual connection units for power forecasting. In 2024 2nd International Conference Pattern Recognition, Machine Vision and Intelligent Algorithms (PRMVIA), pages 54-58. IEEE, 2024.
[17] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[18] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. PaLM: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1 113, 2023.
[19] Zihang Dai, Zhilin Yang. Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint htt1:/9911.28600 2019.
[20] Etienne Denoual and Yves Lepage. BLEU in characters: towards automatic evaluation in languages without word delimiters. In Companion Volume to the Proceedings of Conference including Posters/Demos and tutorial abstracts, 2005.
[21] Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers), pages 4171-4186, 2019.
[22] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold. Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale. In International Conference on Learning Representations, 2021.
[23] Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. GLaM: Efficient scaling of language models with mixture-of-experts. In International conference on machine learning. pages 5547-5569. PMLR, 2022.
[24] Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the Wasserstein metric: Performance guarantees and tractable reformulations. Mathe- mutical Programming, 171(1):115-166 2018
[25] Lawrence C Evans. Partial Differential Equations, volume 19. American Mathematical Soc., 2010.
[26] Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ODE: the world of Jacobian and kinetic regularization. In International Conference on Machine Learning, pages 3154-3164. PMLR, 2020.
[27] Nicolas Fournier and Arnaud Guillin. On the rate convergence in Wasserstein distance of the empirical measure. Probability Theory and Related Fields, 162(3-4):707-738, 2015.
12

---

[28] Zihao Fu, Wai Lam, Qian Yu, Anthony Man-Cho So, Shengding Hu, Zhiyuan Liu, and Nigel Collier. Decoder-only or encoder-decoder" interpreting language model as a regularized encoder-decoder. arXiv preprini arXiv: 2301v:2304.04055, 2023.
[29] Bolin Gao and Lacra Pavel. On the properties of the softmax function with application in game theory and reinforcement learning. arXiv preprint arXiv:/704.00803, 2017.
[30] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv: 101.00027, 2020.
[31] Zhihan Gao, Xingjian Shi, Hao Wang, Zhu, Yuyang Wang, Mu Li, and Yan Yeung. Earthformer: Exploring space-time time transformers for earth system forecasting, 2023.
[32] Sebastian Gehrmann, Tosin Adewumi, Karmanya Aggarwal, Pawan Sasanka Ammanamanchi, Aremu Anuoluwapo, Antoine Bosselut, Khyathi Raghavi Chandu, Miruna Clinciu, Dipanjan Das, Kaustubh Dhole, et al. The GEM benchmark: Natural language generation, its evaluation and metrics. arXiv preprint arXiv:2102.01672. 2021.
[33] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective on transformers. arXiv preprint
2023.
[34] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters in self-attention dynamics. Advances in Neural Information Processing Systems, 36, 2024.
[35] Borjan Geshkovski, Philippe Rigollet, and Domènec Ruiz-Balet. Measure-to-measure interpo- lation using transformers. arXiv preprint arXiv:2411.04551, 2024.
[36] Aaron Gokaslan and Vanya Cohen. OpenWebText Corpus. http::/8kkyLion77..tthaa
1o To/DpenneeTTextCrppsss 2019.
[37] Rafael C Gonzalez. Digital image processing. Pearson education india, 2009.
[38] IBM Granite Team. Granite 3.0 language models, 2024.
[39] Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse problems, 34(1):40 (5(())014004, 2017.
[40] Han Huang. Jiajia Yu, Jie Chen, and Rongjie Lai. Bridging mean-field games and normalizing flows with trajectory regularization. Journal of Computational Physics, 487:1 12155, 2023.
[41] Sonain Jamil, Md Jalil Piran, and Oh-Jin Kwon. A comprehensive survey of transformers for computer vision. Drones, 7(5):287, 2023.
[42] Di Jin, Zhijing Jin, Joey Tianyi Zhou, and Peter Szolovits. Is BERT really robust? strong baseline for natural language attack on text classification and entailment. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 8018-8025, 2020.
[43] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-58,, 2021.
[44] Kelvin Kan, François-Xavier Aubet, Tim Januschowski, Youngsuk Park, Konstantinos Benidis, Lars Ruthotto, and Jan Gasthaus. Multivariate quantile function forecaster. In International Conference on Artificial Intelligence and Statistics, pages 10603-10621. PMLR, 2022.
[45] Kelvin Kan, Xingjian Li, Benjamin J Zhang, Tuhin Sahai, Stanley Osher, Krishna Kumar, and Markos A Katsoulakis. Stability of transformers under layer normalization. arXiv preprint 1rxiv:2510.09904, 2025.
[46] Kelvin Kan, James G. Nagy, and Lars Ruthotto. LSEMINK: a modified Nowton-Kryloy method for log-sum-exp minimization. Electron. Trans. Numer. Anal., 60:618 08:618-635, 2024.
13

---

[47] Nikita Karagodin, Yury Polyanskiy, and Philippe Rigollet. Clustering in causal attention masking. arXiv preprint arXiv:2411.04990, 2024.
[48] Andrej Karpathy, Nanogpt. httppe///ggabbb.ocoo///rrpaaaaaa
2022.
[49] Jongseon Kim, Hyungjoon Kim, HyunGi Kim, Dongjun Lee, and Sungroh Yoon. A compre- hensive survey of time series forecasting: Architectural diversity and open challenges. arXiv preprint arXiv:241 1.05793, 2024.
[50] Martin Knott and Cyril S Smith. On the optimal mapping of distributions. Journal of Optimization Theory and Applications, 43:39-49, 1984.
[51] Jean-Michel Lasry and Pierre-Louis Lions. Mean field games. Japanese journal of mathemat- ics, 2(1))229-200)) 2007.
[52] Yann LeCun. The MNIST database of handwritten digits. http://yanm. lecun. com/exdb/mnistl/ 1998.
[53] Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invariant neural networks. In International Conference on Machine Leurning, pages 3744-3753. PMLR, 2019.
[54] Bei Li, Quan Du, Tao Zhou, Yi Jing, Shuhan Zhou, Xin Zeng, Tong Xiao, JingBo Zhu, Xuebo Liu, and Min Zhang. ODE transformer: An ordinary differential squation-inspiree model for sequence generation. Proceedings of the 60th Annual Meeling of the Association for Computational Linguistics (Volume 1: Long Pupers), pages 9335-8351, Dublin, Ireland, May 2022. Association for Computational Linguistics.
Smaranda Muresan, Preslav Nakov, and Aline Villavicencio, editors,
[55] Bei Li, Quan Du, Tao Zhou, Shuhan Zhou, Xin Zeng, Tong Xiao, and Jingho Zhu. ODE transformer: An ordinary differential cquation-inspired model for neural machine translation. arXiv preprint arXiv:2 104.0000 2021.
[56] Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. Transactions on Machine Learning Research, 2023.
[57] Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining, 2024.
[58] Yiping Lu, Zhuohan Li, Di He, Zhiqing Sun, Bin Dong, Tao Qin, Liwei Wang, and Tie-yan Liu. Understanding and improving transformer from multi-particle dynamic system point of view. In ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations. 2020.
[59] Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, pages 142-150, 2011.
[60] Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K. Gupta, and Aditya Grover. Climax: foundation model for weather and climate, 2023.
[61] Tung Nguyen, Rohan Shah. Hritik Bansal, Troy Arcomano, Romit Maulik, Veerabhadra Kotamarthi, Ian Foster, Sandeep Madireddy, and Aditya Grover. Scaling transformer neural networks for skillful and reliable medium-range weather forecasting, 2024.
[62] Hao Niu, Fengming Luo, Bo Yuan, Yi Zhang, and Jianyong Wang. transferring from neural ODE perspective. Electronics Letters, 60(17))e70015, 2024.
Efficient visual transformer
[63] Ikumi Okubo, Keisuke Sugiura, and Hiroki Matsutani. A cost-efficient FPGA-based CNN- transformer using neural ODE. IEEE Access, 2024.
[64] Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. OT-flow: Fast and accurate continuous normalizing flows via optimal transport. Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pages 9223-9232 2021.
14

---

[65] Derek Onken and Lars Ruthotto. Discretize-opimmize vs. optimize-diseretize for time-series regression and continuous normalizing flows. arXiv preprint arXiv:2005./3420, 2020.
[66] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in PyTorch. 2017.
[67] Lev Semenovich Pontryagin. Mathematical theory of optimal processes. Routledge, 2018.
[68] Charles Ruizhongtai Qi, Li Yi, Hao Su, and Leonidas J Guibas. Pointnet++: Deep hicrarchical feature learning on point sets in a metric space. Advances neural information processing systems, 30, 2017.
[69] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188:9, 2019.
[70] Jesse Roberts. How powerful are decoder-only transformer neural models? In 2024 Interna- tional Joint Conference Neural Networks (IJCNN), pages 1-8. IEEE, 2024.
[71] Bo-Kai Ruan, Hong-Han Shuai, and Wen-Huang Cheng. Vision transformers: state of the art and research challenges. arXiv preprint arXiv:2207.03041, 2022.
[72] Walter Rudin et al. Principles of mathematical analysis, volume 3. McGraw-hill New York, 1964.
[73] Lars Ruthotto and Eldad Haber. Deep neural networks motivated by partial differential equations. Journal of Mathematical Imaging and Vision, 62(3):352-36 2020.
[74] Michael E Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Sinkformers: Trans- formers with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics, pages 3515-3530. PMLR, 2022.
[75] Aman Sinha, Hongseok Namkoong, and John Duchi. Certifiable distributional robustness with principled adversarial training. In International Conference on Learning Representations, 2018.
[76] Matthew Staib and Stefanie Jegelka. Distributionally robust optimization and generalization in kernel methods. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[77] Elias M Stein and Rami Shakarchi. Real analysis: measure theory, integration, and Hilbert spaces. Princeton University Press, 2009.
[78] John Thickstun. The transformer model in equations. University of Washington: Seattle, WA, USA, 2021.
[79] Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. LaMDA: Language models for dialog applications. arXiv preprint arXiv:2201.08233 088339, 2022.
[80] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anee Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023.
[81] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.000 092089988. 2023.
[82] Mikaela Angelina Uy, Quang- Revisiting point cloud classification: A new benchmark dataset and classification model on real-world data. In Proceedings of the IEEF/CVF International Conference on Computer 2019.
Hieu Hieu Pham, Binh-Son Hua, Thanh Nguyen, and Sai-Kit Yeung.
Vision, pages
15

---

[83] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[84] Alexander Vidal, Samy Wu Fung, Luis Tenorio, Stanley Osher, and Levon Nurbekyan. Taming hyperparameter tuning in continuous normalizing flows using the JKO scheme. Scientific reports, 13(1):4501, 2023.
[85] Cédric Villani. Topics in optimal transportation, volume 58, American Mathematical Soc., 2021.
[86] Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2008.
[87] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.00 04768, 2020.
[88] Thomas Wang, Adam Roberts, Daniel Hesslow, Teven Le Scao, Hyung Won Chung, Iz Beltagy, Julien Launay, and Colin Raffel. What language model architecture and pretraining objective works best for zero-shot generalization? In International Conference on Muchine Learning, PMLR, 2022.
pages
[89] Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E Sarma, Michael M Bronstein, and Justin M Solomon. Dynamic graph enn for learning on point clouds. ACM Transactions on Graphics (tog), 38(5):1-12, 2019.
[90] Zhirong Wu, Shuran Song, Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao. 3D ShapeNets: A deep representation for volumetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1920, 2015.
[91] Chunlong Xia, Xinliang Wang, Feng Lv, Xin Hao, and Yifeng Shi. ViT-CoMer: Vision trans- former with convolutional multi-scale feature interaction for dense predictions. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 5493-5502, 2024.
[92] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pages 10524-10533. PMLR, 2020.
[93] Xovee Xu, Zhiyuan Wang. Fan Zhou, Ying Huang, Ting Zhong, and Goce Trajcevski. Dynamic transformer ODEs for large-scale reservoir inflow forecasting. Knowledge- Systems, 276:1 2023.
[94] Jingfeng Yang, Hongye Jin, Ruixiang Tang, Xiaotian Han, Qizhang Feng, Haoming Jiang, Shaochen Zhong, Bing Yin, and Xia Hu. Harnessing the power of LLMs in practice: A survey on chatgpt and beyond. ACM Transactions on Knowledge Discovery from Data, 18(6):1-32. 2024.
[95] Liu Yang and George Em Karniadakis. Potential flow generator with 2 optimal transport regularity for generative models. IEEE Transactions on Neural Networks and Learning Systems, 33(2):528-533 538, 2020.
[96] Liu Yang, Siting Liu, Tingwei Meng, and Stanley J Osher. In-context operator learning with data prompts for differential equation problems. Proceedings of the National Academy of Sciences, (2033))e2310142111
2023.
[97] Jiexia Ye, Weiqi Zhang, Ke Yi, Yongzi Yu, Ziyue Li, Jia Li, and Fugee Tsung. A survey of time series foundation models: Generalizing time series representation with large language model. arXiv preprint arXiv:2405.02355 2024.
[98] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-seqquene functions? In International Conference on Learning Representations, 2020.
16

---

[99] Benjamin J Zhang and Markos A Katsoulakis. A mean-fieed field games laboratory for generative modeling. arXiv preprint arXiv:2504.13534, 2023.
[100] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.
[101] Yaofeng Desmond Zhong, Tongtao Zhang, Amit Chakraborty, and Biswadip Dey. A neural ODE interpretation of transformer layers. In The Symbiosis of Deep Learning and Differential Equations II, 2022.
17

---

A Related Work
This section provides review of relevant work.
Continuous-time Architecture. There has been some applications of continuous- time formulations of Transformers. However, we note that they do not provide theoretical analysis to justify their modeling choice or to demonstrate its advantages. In contrast, our work proposes a foundational framework for understanding and designing effective Transformer architectures grounded in optimal control theory. Moreover, our work offers extensive theoretical maalysss-uuppon empirical results across diverse experiments- that motivates and substantiates the advantages of our model. Moreover, there is a key distinction between our and existing model formulations. In Tr-Transformer, we use the composition of all Transformer blocks to parametrize single dynamical system 7 governing the hidden states. To the best of our knowledge, the existing works use each Transformer block parametrize a dynamical system. For a Transformer with D Transformer blocks, the continuous- time model is represented as the output of D different dynamical systems. In particular, it is formulated as
by strong
(o( Xo.
(,() = -T),
for 1 SisD D 1,
(13)
(t) = (t); for E [0, 0 < i D - 1,
dt
where f. is the ith Transformer block parametrized by weights 0, and defined in and except that the fully-connecter layer (4) has no skip-connection.
This formulation is introduced in [5]. Here, we highlight several key differences between their work and ours. First, they only conduct the simple task of determining the parity of binary sequence in their work, rather than investigating its performance in general applications. More importantly, when their approach is applied, it fails improve performance over the vanilla Transformer and instead degrades it. In their experiments, the optimal transport regularization cannot improve the performance of their model when the sequence length exceeds eight. We observe similar issues when testing their model on other applications; see Appendix ! This is potentially due to their choice of formulation. Specifically, in (I3), as the model transitions from one Transformer block to the next, it effectively switches different dynamical system, introducing non mmoothosssssss the overall dynamics. This undermines the purpose of the optimal transport regularization and violates the regularity properties proved in Appendix F In contrast, our model is formulated using only one dynamical system. The resulting dynamics is smoother and thus inherently better suited to incorporate the optimal transport regularization. In particular,
modeling choice is consistent with the regularity properties proved in Appendices F.3 and [.6] This is also evident in our experimental results, while the regularization can always improve the generalization of OT-Transformer to significant extent, this is not the case with their model; in certain scenarios, the regularization may even degrade their model's performance. We also mention that, while (5) proposes alternative formulations for further investigation, it does not consider ours, highlighting the novelty and non-triviality of our approach.
Since then, there have been number of follow-up works that build on the formulation Equation 33 to
perform different tasks, including sequence generation [54] 55 58 [01), time series forecasting 16
93]. and image classification [62 63). Most of these methods only use the formulation [3 as
motivation, and none of them consider optimal transport regularization in their approach. Moreover, these models focused specific type of application and not
In order to access the performance of our OT-Transformer more comprehensively, we also include the existing Transformer formulation (13) as benchmark in our experiments; see Appendix for detailed experimental results. is referred to as "N-ODE Transformer" in our experimental results, following the terminology in (5).
prominent application of continuous-time neural networks is continuous normalizing flows (CNFs) [15]. CNFs use (6) to paramtrize invertible mappings between a standard Gaussian distribution and an unknown target distribution. The ill-posed nature of the CNF formulation can often add to the complexity and computational cost for solving problem. Optimal transport (OT) based regularization has prominent applications in CNFs and is a powerful tool improving accuracy and at times reducing cost. Among the infinitely many mappings between the two distributions.
OT-based CNFs.
18

---

target to find the optimal transport mapping. This is done by incorporating into the training objective regularization term(s) enforcing straight trajectories in (6). This renders the training problem well-posed 40 [99. The straight trajectories also offer numerical advantages, as they make the numerical integration of (6) more tractable.
OT-based CNFs
22
64 84 95
Mathematical Analysis. There have been works that theoretically analyze a continuous-time formulation of Transformers. [33 34), they show that continuous-time formulation can be interpreted as an interacting particle system, where each token can be perceived as a particle. They demonstrate that there is a clustering behavior among the tokens. Since then, there has been number of works that further investigate the dynamics of tokens through this interpretation, including [2] 12 35 47 to name few. However, we note that the aforementioned work is primarily theoretical and lacks evaluations beyond toy experiments. In (74), they show that, under some restriction on the weights, a continuous-time formulation of self-attention layers can be interpreted as a gradient flow. In 45, they use continuous-time training formulation to analyze the stability of Transformers under layer normalization. However, no experiments have been conducted following this analysis. To the best of our knowledge, existing theoretical analyses have not been conducted on continuous-time Transformers with OT regularization.
Relevance of Assumptions. We highlight the relevance of our assumptions and their alignment with recent developments the literature and industry. In recent survey [44 on LLM models, since the inception of GPT-style models, there has been a clear trend of decoder-only architectures taking popularity over encoder-only and encoder-decoder architectures due their superior performance despite having a simpler architecture, with many of the industry standard such as GPT models I
Llama 80 81 falling under the category. Similar observation also discussed in [28 88). Going
beyond LLMs, survey work such as [41] suggest encoder-only models are most popular for image classification tasks. While for time series problems, encoder-only architecture has been the norm with exploration of decoder-only architectures emerging, as pointed out in 49 197. While our proposed approach is applicable to any architecture, these trends underscore the importance of our theoretical guarantees on both encoder-only and decoder-only models, as they continue to define the state of the art across diverse domains and tasks.
B Comparison Between a Standard Transformer and OT-Transformer
To compare between our OT-Transformer and a standard Transformer, we provide side-by-side pscudocode of the forward propagation. This highlights the shared components, the OT- specific regularization and updates, and how Transforrrrrr can directly use an existing Transformer plug-and-play manner.
Implementation.
Algorithm Forward Propagation: Standard Transformer vs TT-Transformer Standard Transformer:
OT-Transformer::
Require: Input Xu, Transformer model fD f
Require: Input Xu, Transformer model f fpo...o o f1, terminal time T, regularization
parameter X, number of numerical integra-
tion steps M
Initialize: X(0) Xo, reg 0, At
1: for i 11,,..., do
1: for f -r, 2Ar do
Compute f(X(t))
% Adapt from standard Transformer
2: X; f((
X(t At) < X(t) + At ((((())))
4: reg + reg + At ((X(E)
3: end for
end for
4: return: output X
6: return output X(t), regularization À reg
Computational Overhead. Since OT-Transformer performs numerical integration (corresponding to the for loop in line I of Algorithm []. it incurs computational overhead compared to a standard Transformer. However, the computational cost can be mitigated by using a smaller model thanks to OT OT-Transformer's parameter efficiency; see Tables D and
19

---

Moreover, the OT regularization promotes highly regular hidden states (see Appendics F.2 F.7) D). This
allows TT-arroffrrmrr to further reduce computational cost: by using fewer integration steps (bigger A!) without sacrificing integration accuracy, see Table Consequently, OT-Transformer incurs less computational overhead than regular continuous-time Transformers. concrete runtime comparison for the NanoGPT experiment is reported in Table
Table 4: Average training time per iteration for the NanoGPT experiment.
Model
Time (ms)
Baseline
64.6
OT-Transformer
175.0
Standard Continuous-time Transformer
344.4
C Empirical Sensitivity Study of Hyperparameters
We perform additional NanoGPT experiments test the sensitivity of OT-Transformer with respect to its hyperparameters. The experiments are done over three random trials. The results are reported in Table We see from the results that OT-Transformer consistently outperforms the baseline across over two orders of magnitude of X. The performance of OT-Transformer is also very stable with respect to the number time steps. Moreover, tuning of hypeparrametess generally not required; even if the optimal À and number of integration steps (T/Al) are not used, OT-Transformer still outperforms the baseline.
Table 5: Ablation study on the strength of OT regularization (A) and integration time step (T/At) for the NanoGPT experiment. The experiments done over three random trials.
0.01
0.05
0.5
Baseline
Test loss
2.68 + 0.019
2.29 + 0.017
2.05 + 0.009
1.52 + 0.002
1.44 0.004
0.003
2.68 0.006
Time step
15
20
Baseline
loss
1.46 + 0.003
1.43 + 0.002
1.44 + 0.004
0.002
2 2.68
D Derivation of the Transformer Equation
We present the derivation of (3). the equation for the multihead self-attention layer. Recall that X; Rd denotes the input the (i 1)th Transformer block for i O.... D 1, where is
the number of tokens, d is their dimension, and D is the total number of Transformer blocks. And
X,, E Idd denotes the jth column of X,.
Using the notation defined Section the first four equations in 78 are given by
Q (b) (xx.)) = Qf x1.s, K (h( (xx.s) K/xis, (((h) (xx.) = Vixi,, where Q K v R
(14)
(x k.! (xxr:))
(h)
softmax,
(15)
V
u/,3
(1) (h) j,j'
(x
where
w Rdxk
h=1
7=1
(16)
u LayerNorm( x + u
(17)
Here, H is the number of attention heads, and softmax, denotes the softmax function applied on d-dimensional vector indexed by j'.
20

---

Substituting 14 into [5)
and (16 and plugging (6 into 17
we obtain
(Q) j. K
x
softmaxy' (Q/xis) ) (K!x
33 softmaxy'
(h)
V
V/
(18)
(K( (Q/x
softmaxy'
Vk
u LayerNorm
Xi.j h=1
(1) "3
Vixi
j'=1
(19)
LayerNorm
X
w
V Xi
(h)
h=1
Here, in the last step, we used the fact that
(h) jjs'
are scalars. We then plug 18 into
19)
and obtain
H
(K
(
Uuj = LayerNorm
Xi,j
w
V
X softmaxy'
V/
h=1
='=1
(K! Vi
= LayerNorr
X +
w softmax
(20)
h=1
where in the last step used the fact that X is the j'th column of X, Recall that in the main text, we omitted layer normalization for simplicity of exposition. The formulation (20) becomes the self-attention layer formulation (3) when omit the layer normalization function. Thus, we have derived the multihead self-attention layer formulation
E
Optimal Transport Background and Derivation of Training Problem
We review the relevant optimal transport background and then derive the training objective 8
Optimal Transport Background. We consider the space Raxd, to which the hidden states of OT-Transformer belong. Let (2 be the space Borel probability measure on Ra with finite second-order moments and Pu. P1 E (IRxxd) which specify the initial and terminal distributions. problem with quadratic cost [85] Chapter 1] is given by
The
inf /
W2( P1)
||X
Y| (X,,
(21)
TEI((pe.P))
xBuxd
where T(Po. P1) is the set of joint probability measures with Raxd Raxd with X- and Y- marginal distributions Pn and Pt. respectively. Here, the quadratic || - Y quantifies the cost of transporting unit of mass from location X location Y. We note that the Wasserstein space (P2(R) W2) equipped with the Wasserstein metric W2 a complete separable metric space [6 Theorem 6.18).
The Benamou-Brenier formulation
(1I is given by 6
inf f.P JOX [ |1 f) 2 p(, dXdt
(22)
subject to O,p(X, t) + V XX, (p(X (XX, t)) 0,
(23)
and p(X,0) (X): (XX 7) (X)).
(24)
This is dynamic formulation of 21.. which describes optimal transport dynamical system. In
particular, the probability density p kd x [0, - R20 evolves continuously over time under
the velocity field f RRdxn x [0. T] - dR as governed by the continuity equation (23). The
initial and terminal conditions 24 require that the probability density evolve from p(X, 0) Po (X)
21

---

to = (x().. The optimal velocity field has several important and favorable properties. Mass induced by the optimal velocity travels in straight lines constant speed [6]][Crrollary 7.22]. Moreover, under standard conditions on Po and pt, the optimal velocity field is unique (9 50
Since the velocity field f governs the movement of mass, given initial position Xn Po. the evolution of (() is governed by the ODE
d(((t) di
(X),)), for LE [0,T) with X(0) Xo.
(25)
Denote the solution operator of (25 by S Raxd x [0. + Raxd such that S(Xo, = X (().
Under suitable regularity conditions [emma 8.1.6), the solution of the continuity equation (23)
is given by p(.,t) = S(,, t) the pushforward of the probability measure by t(. Hence, 24) can be rewritten
inf (""
| Juxn 2 211/((())0) (Xo) dXodt
1
(26)
dx(() = f(X(t, for t e [0,T) X(O) Xo. S(-TT) P1.
subject to dt
Derivation of Training Problem. Under the setup of OT Transforfrrrr f is the Transformer blocks, and X,, and Po are the embedded input and its distribution, respectively. The probability measure p specifies the target distribution of the terminal state X(T), 'The terminal condition of (26) requires that the distributions of the target output y and corresponding terminal state X(T) to match. That is, it requires G(X(T), = for each pair (X()).y), where G' is the loss function defined in (8)
Thus, (26) becomes
inf Exa.y /
2 (1))
(27)
dx(t) f()(t), for t [0, X(() Xo. G(X(T),y) =.
subject to
dt
Here, we used Fubini's theorem Theorem 3.1] swap the order of the integrations, the ex- pectation is taken over the joint distribution of the cuput-output pairs (Xo, y) with an Xo-marginal distribution Po. Further, we make the following remarks on (27):
The optimization over f which is non non-parametric. In OT-Transformer, we parametrize f using Transformer blocks with weights 0, and the weights of the embedding and output layers are 7.
The optimization problem is intractable in general when the terminal condition G(X(T)). = 0 is imposed as a hard constraint. We relax this constraint by incorpo- rating as a penalty term in the objective function.
Thus, we obtain the training problem
(G(((T), v:n) '- (t) dt
min EX.Y 0
(28)
subject to dX(t) (()(),(;0 for LE [o, X(0) Xo-
dl.
Here, when we set μ == 3. under discretization (8 is equivalent to the OT-Transformer training problem (8) in the continuous-time setting subject to the dynamics 7 It is noteworthy that the hyperparameter j can be interpreted as the Lagrange multiplier for the terminal condition of (27). This reveals that the regularization hyperparameter À is inversely proportional to the Lagrange n² where μe" is the optimal Lagrange multiplier, and assuming that the Transformer blocks parametrizing f are sufficiently expressive, the solution of the OT-Transformer training problem (8) corresponds to the solution of the hard-constrained problem (27).
multiplier. When we have À = 1.
For clarity of presentation, slightly abuse notation by using Xo both random variable and dummy variable of integration.
For the commonly used cross- entsopy [Sctionn IJ and mean- error, the loss is zero when the model output equals the target output
22

---

Proofs of Theorems
In this section, we report in detail the assumptions and derivations of the theorems in Section
F1 Assumptions and Justification
In the following, we list the assumptions of the training problem ( on which our theoretical analysis is based. We then justify our assumptions by showing that they are satisfied in common applications.
1. the function G is proper, convex, and twice continuously differentiable in its first argument;
2. the function VG(:,) is Lipschitz continuous in both arguments, where the gradient is taken with respect to its first argument.
Here, the function G is the first term in the training objective g which is the composition of the output layer (5) and the loss function.
The assumptions are valid for encoder-only and decoder-only continuous-time Transformer training problems for classification, next-token prediction, and regression tasks. These settings apply to
many common Transformer architectures in the continuous time settings, including Vision Trans- formers (22] and language models including the GPT series [IO 69) 70 of OpenAI, PaLM [T8).
GLaM 23 and LaMDA [79] of Google, OPT 100) of Meta AI, and Granite of IBM 38
Next, we will demonstrate in detail why the assumptions hold under these settings. In this subsection, for clarity of presentation, we vectorize the terminal hidden states and denote their vectorizations by x(T) vec(X((T). Accordingly, (T) y)) by functions their vectorized forms (e.g., G(x(T),,)) This a slight abuse of notation, which will significantly simplify expressions and derivations.
denote functions originally defined on matrices (c.g..
On one hand, for regression tasks, the output layer (5) consists of linear transformation
(T),
(29)
where vho E Rexd are the weights of the output layer. The loss function is the mean-squared loss (MSE)
MSE y) - - y122.
(30)
Given a target output y E IRS, the loss function G is given by
G(x(T),y) = MSE (5)
(31)
2 x y
is easy to see that is proper, convex, and smooth (hence twice continuously differentiable). This satisfies the first assumption. Its gradient is given by
VG(((T), = v. (1)xx(T)-)) y));
(32)
which Lipschitz continuous in each of two arguments, with Lipschitz constants /22 and L, respectively, where
L | |1 12,
(33)
with || denoting the 2-operator norm of vb.. And we used the fact that || || ||2. This satisfics the second assumption.
On the other hand, for classification and next-token prediction tasks, the output layer is
eVox(T)
y = Softmax((s.x(77) 1 e
(34)
and the loss function is the cross-entropy loss
CrossEntropy(ý..) -y" log(ý).
(35)
23

---

Here, is the number of classes, 1. IR is vector of all ones, and y. E A which lie in the (c- di-deneniolall simplex, represent the model output and target output, respectively. The loss function is the log-sum-exp function plus a linear term [46]
eVx((T)
((x( ChesEBntropy/Ssttmxxxxxxxx))) = -y" log
1 (TVx
-y x (y T 1c) log (I (eeoi(t))
3ox(T) + log [ eTeex((T)
lincar term
log-sum-exp
Here, the last step, we used the fact that Tc) because y function is convex and smooth (hence twice continuously differentiable) 46 so is G. Moreover, G proper because value is always nonnegative and has nonempty effective domain. This shows that the first assumption is satisfied. Next, the gradient of G is
6-1 Since the lug-sum-exp
VG(((T),)) = + y + ab (iag(eV>x(2)) 1, e.ex)((T
= y + 2b
1 (T)
- wb (oftmax(tbxx(7).
Here diag(z) denotes diagonal matrix with diagonal entries equal z. straightforward to verify that VG is Lipschitz continuous with respect to its second arguement with Lipschitz constant
=
11 2l. Note that the softmax function Softmax(z)
I Lipschitz continuous with
Lipschitz constant
29
For any x (2),x2(T),
VV)( VG(x; ylll
-. + v softmax(M)xx - (-v y + v softmax((,xxxT)))
l (softmax (vix (7)) wwt. (T))) l|2
1 ll2 ||softmax (,,x (T)) softmax(t)xxx
11 lla || (T) ( (since softmax has Lipschitz constant
|| ll2 11 |x| (T) x2 (2)
L x1 (T) - x2 (2)
Therefore, VG is Lipschitz continuous both arguments, and the second assumption satisfied.
Our derivation also reveals that for regression, classification and next-token prediction tasks, VG is Lipschitz continuous in each of its two arguments, with the same Lipschitz constants in all cases. We summarize this finding in the following lemma.
Lemma 1. VG(..)) -) is Lipschitz continuous in each of its two arguments, with Lipschitz constants 2 and respectively, where L is the 2-operator norm of the output layer weights sbo
Mean Field Control Formulation
We first recall the pon-maraettric formulation of the training problem (10):
G(((T)))) A [' [lf(X(s), s) ds
min ( m
dX(tt f(X(t),t), for te [0,T].
(36)
subject to
dt
and X(0) Xo-
Here, for clarity of exposition, denote as the dummy variable for time integration.

---

The formulation 36
can be rewritten into mean field control problem given by
A / / 2
min /.P
y V Rdx 1) y, dXds,
g(p(:,;T)
(37)
subject to and
=0 =0, for tE [0,TT)
(38)
pX
(X),)), for e Rdxn
(39)
Here, mean field perspective is adopted by modeling the evolution of the density function p _dxn K [0. 2] + R2o. which characterizes the distribution of the hidden state for E [0, T]
given the target output In particular, the probability density of the hidden state evolves continuously
over time under the velocity field f, as governed by the continuity equation (38). The initial density
Pu Rdx= x R + Rzu defines the joint distribution of the target output and the hidden states at
t (that is, the embedded input The first term of the objective given by
(p( T)) pxx
(40)
F,3 Optimal Control Theory
We review optimal control theory that will be used our analysis.
We first recall the definition of the variational derivative of a functional. For test function w E 1,2 (Rdn+t), the variational derivative 4g of & with respect to p is defined as op
C(p + hw) - g(p)
-
8G(p)
lim h-+0
(XXy
(41)
h
JRC JReKn 8p
Evaluating (1 using the definition of (40). we have
89 8G(p(,,T) T))
= =
(42)
6p
Consider the mean field control problem 37 39 Under Lagrangian formulation, for hidden state
at location X and time with target output y, its potential function by Rddx [0. 7) -> IR is given by [5T)
ag(p(:,,TT) (X(T),)) 2/ f. (I(((()) öp
(X, inf f.P
(43)
X [ [I(((()) 8) ds
=inf inf ((X(T),)) +
Here, we used 42 in the last step.
Under standard regularity assumptions 225 |Section 10.3.1], the potential function is bounded and Moreover, when À > O, the potential function is the unique solution the Hamilton-Jacobi-Bellm (HJB) partial differentiation equation (PDE)
Lipschitz continuous in both arguments 25
x))) 0, 5,, H(V +
(44)
80(p(, (T) xX,
Φ X, T)
(45)
6p
where the gradient V, (X,,)t t) is taken with respect to the spatial variable. And the Hamiltonian И _dxnn - is given by
H(P) =us)= sup { (P.f) -
1
(46)
where P is the costate variable and (: denotes the Frobenius inner product. When À > 0, the objective strietly concave, and thus the supremum uniquely attained when Hamiltonian becomes
1P. The
1 H(P)
PPl.
(47)
2
25

---

Substituting
42
into
45
and
47
into
44
the HJB equation becomes
0, By
X. t)
X
t)
(48)
2x
(49)
Moreover, the derivation of the HJB equation reveals that the optimal velocity field f of the mean field control problem 37 to (39) attains the supremum the Hamiltonian when P = v t),
thus have
((X,
(X,1).
(50)
F.4 Well-posedness of Training Problem
We prove Theorem We first restate the theorem.
Theorem 1. There exists unique solution to the optimization problem (T0) if and only if
Proof. When À = 0, the Hamtiltonian 46 is degenerate and unbounded above, and the optimal velocity can have an arbitrarily large magnitude. This implies that the HJB PDE is not well-defined. degenerate mean field control problem that has infinitely many
As such, 37 39
reduces
solutions.
We note that À > 0 if and only if the Hamiltonian 46 is well-defined and admits a unique maximizer. Furthermore, the HJB PDE (44)-(45) is well-defined and only if the Hamiltonian is well-defined. In this case, the potential function is the unique solution to the HJB PDE [25]. Thus, the optimal velocity, which is given by the gradient of the potential function 500.
also unique.
F.5 Highly Regular Trajectory
Next, we show that the unique velocity guaranteed by Theorem I when À > 0 is highly regular. In particular, the Transfoffrmer produces exactly straight trajectories. Proposition 1. When À > O, under the optimal velocity f", each hidden state X(t) travels along straight trajectory at a constant speed.
Proof. We prove that the optimal velocity f fx(((tt,t) remains constant for t E [0, which induces straight trajectories and constant speed.
From the Pontryagin Maximum Principle
[67. we have that
dX(()
(PH(P(t)).
(51)
(t) (PH(((())) =
(52)
dt
P(T) = (T)
(53)
Here X(I) denote the trajectory induced by f", and P(() is the corresponding costate variable. The
gradient (G)) y) is taken with respect to the first argument. In the second step of 52 we used the fact that (P))) is independent of X.
From 52 we see that P(t) remains constant for all t [0, Thus, we have
P(t) P(T) VG(X(T).y) for all t [0, where applied (53) in the final step. Differentiating the formula of the Hamiltonian 47
(54)
we have
-p. VPH(P)= = 7.P
(55)
Substituting 54h and (55)
into (5T) we obtain
dX(t)
1
-
(56)
for E [0, 7]. This shows that the velocity for each hidden state is constant for all E [0, (,. This
implies each hidden state travels along straight line constant speed.
26

---

The regularized trajectories present stark contrast to those of the unregularized case. Thanks to such highly regular trajectories, the numerical integration of the dynamical system 7 particularly easy to perform. Practically, it requires fewer time steps to achieve accurate integration. Furthermore, since there are no fluctuations or abnormal magnitude changes in the velocity (as for the unregularized case), numerical stability during integration, and consequently during training, is significantly improved. This effect evident in our experiments, where the unregularized training exhibits instability, while the regularized training does not encounter such issues.
We note that the straight trajectories and constant speeds are not immediately apparent and are not guaranteed to hold in general. On one hand, the training problem is a relaxed optimal transport problem (see Appendix|E] details), which does not guarantee such regularity in general. On the other hand, this regularity is guaranteed only under our assumptions; see Appendix F.] The proof applies the Pontryagin Maximum Principle (67 and leverages the regularity assumptions on the loss function G: see Appendix
F.5 for details.
F.6 Regularity of Potential Function
Next, we show that for each given target output y, the hidden states trajectories never intersect.
Proposition 2. For two hidden states with different initial conditions Xo and Xo and the same target output y, their trajectories never intersect.
This proposition implies that the solution to the HJB PDE 48 (44) does not develop shocks. In other words, the gradient of the potential function V4y continuous space.
Proof. Integrating (56) from T, the trajectory of hidden state X(t) given by
((1) = X(T) (T' VG(X(T),))
(57)
for all t E [0, We first show that the trajectories corresponding to two different terminal states Xy and Xy at t T and the same target output cannot intersect. That is,
Xx + (T VG(X y) X + TT 1) VG(rr.y)
(58)
A
X
cannot hold for any / E [0, 7] X * Xy-
To this end, consider the function
-- G(X,y),
( = 2
X
+
(59)
A
which is strictly convex with respect to X for all E =0,,"), as G is convex with respect to X by assumption. Note that its gradient with respect X is
T V)(
(X =
(60)
À
which coincides the trajectory formulation in 57 By the strict convexity of F,, its gradient V is injective. Thus, (58) does not hold for any X Xr.
Next, we argue that for two hidden states with different initial conditions, the corresponding terminal states t T are different. Combining (57) and (60), at t = 0 we have
X(0)
(61)
Since F. is the gradient of strictly convex function, its inverse (VFo) defined by
X(T)
(62)
exists and is also the gradient of strictly convex function (I (Lemma of Appendix G Using the same injectivity arguments on (62) we have that for two hidden states with different initial conditions Xo and X,, their terminal states X and X are different. And we have established the previous paragraph that trajectories with different terminal states X,r and X, cannot intersect. Thus, we have proven that trajectories with different initial states cannot intersect.
D
27

---

F.7 Stable Forward Propagation
Theorem 2. Assume that the OT regularization hyperparaneter is such that X > TL2, where T is the terminal time and L is the 2-operator norm of the output layer weights (5) For embedded input-output pairs (X1(O).11) v1) and (X2(0). y2). the corresponding model outputs Yı and y2 satisfy
TL2 11 82112 (1 (1-T1²
- x + TL ly -Valla)
(63)
A
Proof. By Proposition [ and Proposition 2 the optimal velocity f remains constant along the trajectory (t). Thus, by (56). the trajectory is given by
dx(t) f (X()),
Ā V))
(64)
dt
Integrating (64) from / 0 to T, we obtain
( x()) T VG(X(T),
(65)
We now bound the difference between the terminal states X, (T) and (T) as follows
(i( - X2(T))lr = X, T (T).( X2(0) V
(66)
(X(0) (2(O)| T À (V) +v) (67)
-1) X2(O)lFF
+ (IV)
(68)
+ VG((x(T), - vo(
$ [|] - X2(O)|lF
T
VV((X vG(
(69)
T
[V)(Xx(T),)) yT)
A
(X X2(O)llF
T
+ A L' - X2(T))lF
(70)
T
+ X I.ly: - -2112, by Lemma
By the assumption that À TI,2, we have I4² < 1. Thus, becomes
(1- T |X (T) (2(T)ll |X (0) - x2 + TL |ly y2112
(71)
X
| -xa)) (1-T() [X) -2,(0]x +
TL
x ly y2l1z
(72)
Recall that for the case of regression, the model output (5) is given by ý (X(T); vb) = wh vec(()))
(73)
which is obviously Lipschitz continuous with respect to vec ((()) with Lipschitz constant L = |ll p. For the case of classification and next-token generation. the model output (5) is given by
(o(())
y = go (X():: bb.) = vec = 1 (X())
(74)
28

---

Since the softmax function is Lipschitz continuous with Lipschitz constant 1 22), the model output also Lipschitz continuous with respect
with the same Lipschitz constant ||4/oll
Thus, the corresponding model outputs y1 and Y2 satisfy
1 - V2l12 - L vec |eec(x (x())) veol vec((X(()))l
(IT (2))l
TL² A (L]X;(0) (0xx)(0))xx(0]1le) xx)]ll lyy velle) 7( TL
(01-21le))
by (72).
G
Auxiliary Lemma
We state and prove an auxiliary lemma, which used the proof of Proposition 2 The proof modifies the argument from [44 Proposition 1].
Appendix
Let D C R be open, H : D R be a strictly convex and continuously differentiable function and be grudient. Then h-1 exists and the gradient of a strictly convex functiin.
Lemma 2.
Proof. The strict convexity and continuous differentiability of H implies the existence h- and that Vh(x) is symmetric positive definite (SPD) for all x & 'D. Since h is continuously differentiable and Vh(x) is SPD for all x E D, by the Inverse Function Theorem 72 Theorem 9.24], h -1 is also continuously differentialle and
Vh (M(x)) = I for all x E DD
Here Is is the & x identity matrix. This implies Vh Vyy'(y) is also SPD for all y E h(D). Thus, by the symmetry of VA (y), we have
Oh 11: dy,
0[h-'1g for all i.j 1,2, ...,
(75)
Oy:
Here, y₁ denotes the jth entry of y. Let (vy) Ej Us Ioll (b-=) dtt Consider its partial derivative
On (y) = / [1 - ( Σ ay:
01h
Utt Oy:
(ty)dt
5-1
using 75
we get
f [h-")((ty)de
/ Σ
Ohh
yy
(ty)
Oyj
1=1
applying the chain rule (h~ (1~1) -E U/
alk a1 (ty). we obtain 8y1
=f( -h (- at (1n1a dt
performing integration by parts, we have
1b (ty)dt t[h (1 1=1 1=0 - 1n-itty
[/ (;(y), for i 1.,..... k.
Therefore, = Moreover & is strictly convex because its Hessian Vh -h-1(y) SPD for all y h(D). Therefore, -1 is the gradient of the strictly convex function d.
29

---

H Derivations for Robustness and Generalization
In this section, we provide detailed derivations for the robustness and generalization properties discussed in Section Section Recall the stability in forward propagation inequality
1 V2ll2 - C))x1 X2(0))l+ - y2lle,
(76)
where C- (1 "x
-1 L and C" =
(1 111--42 X x
TL? -1 r2²
H.1 Distributional robustness
We first provide proof of Theorem 3 the distributional robustness inequality
WT #1, T +w)== ÖWp(u,)) v), for any p 2 1, where Tg denotes the pushforward operator under the trained Transformer.
(77)
Proof. Define
vee(X1(0)) and X2 v0)( and note that [X, X2
||x X2112- Suppose x, ~ and X2 ~ v. Holding the output data to be the same, the stable forward propagation inequality gives | )- Txx))ll2 S C'|[x| X2ll2. By equivalence of p-norms we
have that there exists a constant (p)) such that
( - T(x2) Il < (( - TL2X - x - X211p-
(78)
For fixed p. let [(, v) to be the optimal coupling between μ and Then the pushforward of æ under the map (T @ ()x,, x2) = (T(x))
is a suboptimal coupling between measures
Iu and Tyy. Therefore, we have
WH HH, TPU) S / [ly: -2llddl ():
/
S
Taking the p-th each side, we obtain the desired result.
H.2 Derivation of Robustness Under Sampling Error Data Perturbation
Note that the data used for training always come from the underlying true distribution in theory; however, sampling error and perturbation error may occur when generating the training dataset. We show here that the forward model remains accurate when such errors are small.
To derive this property. we first specify the variables. Let μ and f denote the training distributions for the input and output data, and and the true distributions for the input and output data, respectively. Denote T # the pushforward operator under the trained Transformer.
Our goal is show that W, uu, 8) remains bounded. That is, the pushforward (induced by the trained Transformer) of the true distribution for input data closely approximates the true distribution of output data.
Applying triangle inequality to we obtain
WW ( ( B) + WT
TT)) S ČW,(p, v).
The use of Wasserstein metrics is erucial in this context. The derivation above relies on the fact that the Wasserstein distance Wp satisfies the triangle inequality property not shared by divergences such the KL divergence, for which this argument would hold. Thus, here we have
WT HU, D) < δ
Wpln, v)
W (T, +
(T,
(79)
input distribution
retraining loss
output distribution gap
30

---

Here, on the right hand side, the first term is the input distribution gap: it measures the difference between the training distribution and the true distribution of the input data. The third term measure the same for output data. We can bound both these terms using sample complexity results for Wasserstein metrics (p 2 1) [27], see also (82 below. Lastly, the second term in (79) quantifies the difference between the training distribution for output data and the pushforward (induced by the trained Transformer) of the training distribution for input data. This term is an informative surrogate for the training loss. When the training loss decreases, this term will tend to decrease. And when the training loss is O, this term will be 0 too.
Assuming the Transformer model is trained with sufficient accuracy, when the sampling error or perturbation error in the data is small, both the first and third terms will remain relatively small, showing that W (T gV, P) is not only bounded but also small. Thus, Equation (79) indicates robustness with respect to the training dataset. The model's prediction will remain close to the underlying true distribution even in the presence of mild sampling error or data corruption.
H.3 Out-of-Distributin Generalization
great benefit of having Lipschitz functions is that we may obtain performance guarantees on the quality of the model learned from finite training samples is at making predictions on samples. Suppose map T is trained on input samples X1,... XN I that form empirical measure åx. We use ideas from distributionally robust optimization (DRO) to produce performance bounds on -ut-of-sample performance and non-asymptotic generalization bounds 24 76
Denote (fx to be the Wasserstein ball of radius r around f.N, i.e., (aN) == (p E P(22) Wi (p. n I r)- For simplicity, consider the setting where f2 is bounded domain. Suppose we have out-of-distribution samples from distribution v E (,x)), we wish to obtain bounds on the test error of the Transformer T on the -iitt--distrbutioon samples, i.e., Eg (G(T(X): y)] -ET, [G(Z,v) v)], where Z = T() in terms of the training error E LN
[G (X), y)1 EygRo [G(Z, y)] where G is the terminal condition of the training objective X As Ω is a bounded domain, G is a Lipschitz function.
Observe that by Kantorovich duality, the 1- W-Wasserstein distance is given by the variational formula sup {Exx (x(()) E. [p())]]
Wi (i
DELip, (9)
where Lip, (52) is the set of 1 1-Lipschitz functions on Ω. Observe that with Theorem 3 can obtain a bound on the expected test error with respect to the distribution v E W,(An)
ET, v)] ET, [G( < Wi TYÂN)
(W,,)
(80)
=
which implies that
E. S Ep (G(IG v)] +1 --
(81)
This bound is deterministic and holds for any E W-( If statistical argument ensures this inclusion with high probability, then this becomes generalization guarantee. In particular, if v = μ. the distribution that generated the empirical distribution 14, then is computable as a function of samples and we can obtain aon-asympto generalization bound.
H.4 Generalization Bound via Concentration Inequalities
To turn the deterministic bound into a statistical generalization guarantee, first recall Ωoy to be the empirical measure constructed from the training data. Then μ W, with high probability- for a computable radius r-which is ensured by the concentration of the empirical measure in the Wasserstein distance. Indeed, under mild assumptions (e.g.. bounded domains or sub-Gaussian tails), we have sample complexity result [27]: there exists C O such that with probability at least 1 - 5,
log (1/() N
Wi(MAN) & 1// 1/2
= r(N,6).
(82)
31

---

Choosing the radius as function of å and the number of samples N accordingly. we obtain, by combining (0) and (82) that with probability at least 6,
By (((G((( < pp (G(T(X). + ((, 7
(83)
This is a non-asympiotic generalization bound for Lipschitz losses in terms of the Wasserstein distance between μ and JoN, linking distributional robustness of generalization performance. We emphasize this is only a crude application of DRO tools for analyzing Transformers that arise from optimal control tools. Finally, we note that Wasserstein distances are particularly well-suited for sample complexity analysis such (2). In contrast to KL divergence, which requires absolute continuity between distributions, Wasserstein distances remain well-defined even when comparing an empirical measure Äyーa discrete distribution-to possibly continuous distribution μ..
I Experimental Details and Results
We report the detailed experimental setups here. We adapted the code provided by 774 95,, main- taining the same default data processing setup, hyperparameters, and other experimental settings as used their implementations. Our implementation based on PyTorch [66] and experiments are conducted using NVIDIA A100 GPUs with 40GB of memory. Runtime varies by experiment. While Proposition D shows that the optimally trained model yiclds a straight hidden state trajectory with constant speed, in practice we use 8-20 integration steps and do not yet fully exploit this property to reduce runtime. In practice, fewer steps may suffice, which we leave for future work.
We also compared against N-ODE Transformer, an existing continuous-time Transformer formulation which is introduced in [S] and has been considered in other works. For details about the formulation and specific applications, see the discussion in Appendix [A] In the reported results, we refer to N-ODE Transformer with and without transport cost as unregularized N-ODE Transformer and regularized N-ODE Transformer, respectively. The goal here is to demonstrate that our proposed approach is not only theoretically sound, but also excels in performance when compared against the baseline method and similar approaches.
Point Cloud Classification. Given the ModelNet40 dataset, we experiment with the Set Trans- former model 53 It has encoder-decoder architecture and is specifically designed process unordered data, such as point clouds, ensuring that the output remains permutation invariant to its input. Following the setup of (74), we use Set Transformer (33 with two Induced Self Attention Blocks (ISABs) in the encoder, where each ISAB contains two Transformer blocks, and with a Pooling by Multihead Attention (PMA) Module in the decoder. For each instance, we uniformly sample 5000 points from each element in the dataset. For the continuous-time models, we use the same architecture except that we put fully-com fully-connected layer before the Transformer blocks so that the dimension is consistent for continuous-tin time dynamics. Also the hidden dimensions d and k of the ISABs are reduced from 256 200. This reduces the number of parameters for the ISABs by 24%. We use an Adam optimizer, with batch size 64, 200 training epochs, and learning rate of 1 x 10 For the regularized N-ODE Transformer and TT-Trassfommer, the regularization hyperparameters are X = 0.1 and X= = 1, respectively, as they provide the optimal performance in our tests. We use T = 1 and a total of 8 time steps for the numerical integration.
We perform the experiment over five random trials and report the best test accuracies in Table The unregularized continuous-time models encountered gradient explosion, resulting in NaN outputs, and the issue persists even with slight regularization. We found that the models never suffered from gradient explosion with sufficient regularization, indicating that transport cost effectively stabilizes the training process. Hence, we only report the performance of the regularized models. The bascline Set Transformer obtains an average test accuracy of 87.4%. The regularized N-ODE Transformer achieves an accuracy of 87.5%, indicating negligible improvement over the vanilla model. Our -r-Transfrrmrr shows a sizable improvement and reports an average 89, 89.9% test accuracy even with a smaller model. From the learning curves in Figure 2 we see that our model reports lower data-fitting loss for training data compared to the vanilla model, despite the inclusion of a regularization term. This example also highlights the generalizability of our model: despite using an encoder-decoder architecture, demonstrates clear advantages and noticeable performance improvements over the baseline.
32

---

Table 6: Number of parameters for the Transformer blocks, best and final test accuracies (with standard deviation) across five trials for the point cloud experiment.
Method/Exp.
Para. Count
Best Test Accuracy
Final Test Accuracy
Baseline
0.86M
87.4% + 0.45%
86.6% = 0.67% 86.7% 0.43%
Reg. N-ODE Trans.
0.65M
87.5% 0.51%
OT-Trans, (Ours)
0.65M
89.9% + 0.42%
89.3% + 0.69% %
0.9
2.5
0.8
Training Accurayy
0.7
Training Lose
2
0.6
1.5
0.5
0.4
Orig Reg N-ODE Trans.
0.3
0.5
OT Trans. (Ours)
0.2
25
50
75
100
125
150
175
200
25
50
75
100
125
150
175
200
0.9
2.5
0.8
Acmurayy
0.7
Test Loss
0.6
1.5
0.5
0.4
0.3
0.5
0.2
25
50
75
100
125
150
175
200
25
50
75
100
125
150
175
200
Epoch
Epoch
Figure Accuracy and data-fitting loss for the point cloud experiment (averaged over five trials)
MNIST Classification. Following 74,, the baseline model ViT has one Transformer block with single-head self-attention layer and no fully-connected layer. Since it has only one Transformer block, N-ODE Transformer and our OT-Transformer share the same formulation, and we report the results as OT Transssfforr.r
The -T-Trassfrrmer uses the same model architecture as the baseline model, except that the hidden dimensions d and & of the self-attention layer are reduced to 64 from 128. This reduces the number of parameters by over 80%. Adapting from the baseline setting, the patch size is 7 We use an Adam optimizer. The number of epochs is 45 and the batch size is 100. The learning rate is set to x 10-4 for the first 35 epochs, then decreased to 10-5 until the 41st epoch, which point is reduced to 5 x 10-6, For OT-Transformer, the regularization hyperparameter is X 0.01 as it provides the optimal performance our tests. We use 1 and a total of 20 time steps for numerical integration.
The experiments are conducted over five random trials. The best accuracies are reported in Table Tr-Transformer demonstrates significant improvements over the baseline in both accuracy and model efficiency. The baseline model achieved test accuracy of 93.0%. The unregularized OT-Transformer improves the test accuracy to 96.8%, although it uses much smaller model architecture. The transport cost regularization further improves the test accuracy to 97.1% while maintaining the same reduced parameter count. Notably, TT-Transfrrerr also exhibits significantly lower standard deviation across five trials when compared to the baseline and unregularized model, indicating enhanced stability and reliability in its performance. Interestingly, when we compare the learning curves of the unregularized regularized OT- Transssfforess in Figure we observe that including the transport cost regularization also reduces the training loss for data-fitte tata-fitting and accuracy.
33

---

Table 7: Number of parameters for the Transformer blocks, best and final test accuracies (with standard deviation) across five trials for the MNIST image classification experiment.
Method/Exp.
Para. Count
Best Test Accuracy
Final Test Accuracy
Baseline
93k
93.0% + 0.69%
93.0% = 0.67% 96.8% 0.25% 97.1 % + 0.15%
Unreg. OT-Trans.
18k
96.8% + 0.23%
OT-Trans, (Ours)
18k
97.1% + 0.16%
0.8
95%
0.7
Training Aucuray
0.6
90%
Twiining Loss
0.5
85%
0.3
80%
Orig
Unreg OT Trans.
0.2
75%
OT-TTans, (Ours)
10
30
30
0.8
95%
0.7
Aceuraty
90%
Test Lice
0.6
0.5
85%
0.4
80%
0.3
0.2
73%
0.1
10
20
30
40
10
30
40
Epoch
Epoch
Figure 3: Accuracy and data-fitting loss for the MNIST image classification experiment (averaged over five trials)
Cats and Dogs Classification. We again use VIT. The patch size is 16 x 16. We use an Adam optimizer. The learning rate 10-5 The number of epochs is 250, and the batch size 64. The hidden dimensions d and k are 128. For the baseline model. it has Transformer blocks. For the continuous-tim time models, the number Transformer blocks reduced 5; this reduces the number of parameters for the Transformer blocks by around 20%. For the regularized N-ODE Transformer and OT-Transformer, the regularization hyperparameters as they provide the optimal performance in our tests. We use T = 1 and a total of 20 time steps for the numerical integration.
0.005 and X 0.01, respectively,
We report in Table g the test accuracies over different seeded trials. We observe again that our OT-Transformer has the best performance and obtains test accuracy of 79.0%, improving from the baseline's 77.6%. The standard deviation of the test accuracy, at 0.31% is significantly lower than the baseline value of 0.86%, showing our proposed approach is more robust and reliable. We also observe that incorporating the transport cost regularization improves generalization and stability of OT-Transformer: without it, the average and standard deviation of test accuracy worsen to 78. 78.2% and 0.39%, respectively. Both the unregularized and regularized N-ODE Transformers report test accuracy of 75.6%, which worse than the baseline model, making them undesirable methods for the problem. Unlike our model, incorporating the regularization also has little effect on the performance N-ODE Transformer. This is likely due to the incompatibility of N-ODE Transformer and the regularizatio. We report the learning curves in Figure 4 When we compare the learning curves of the unregularized and regularized OT-Trranffommer, we see that including the transport cost regularization also improves the training loss for data- data-fitting and accuracy.
34

---

Table 8: Number of parameters for the Transformer blocks, best and final test accuracies (with standard deviation) across three trials for the cats and dogs image classification experiment.
Mxthod/Exp.
Para. Count
Best Test Accuracy
Final Test Accuracy
Baseline
1.77M
79.3% + 0.52%
77.6% + 0.86%
Unreg. N-ODE Trans.
1.48M
76.4% + 0.37%
75.6% + 0.48%
Reg. N-ODE Trans.
1,48M
76.4% + 0.30%
75.6% + 0.03%
Unreg. OT-Trans.
1.48M
78.8% = 0.63%
78.2% + 0.39%
OT-Trans, (Ours)
1.48M
79.5% + 0.46%
79.0% + 0.31%
85%
0.7
0.65
Triining Aneuney
75%
0.6
Training Leog:
1..55
65%
Orig
0.5
Unreg. N.ODE. Trans.
Reg. N-ODE Trans.
0.45
55%
Uureg. OT-Trans.
0.4
OT-'Trams. (Ours)
150
200
250
100
150
200
250
80%
0.7
75%
0.65
Test Lras
0.6
65%
0.55
0.5
55%
150
200
250
100
Epoch
Epoch
Figure 4: Accuracy and data tata-fitting loss for the cats and dogs image classification experiment
Sentiment Analysis. We use an identical baseline Transformer architecture as in 74] which has six layers of Transformer blocks. The Traanssfore counterpart has only 3 layers. reducing the number of parameters the Transformer blocks by half. We use an Adam optimizer with 15 epochs. The learning rate is 1 x 10t for the first 12 epochs and x 10 afterward. The batch size is 64. The hidden dimensions and & are 256. The batch size is 64. For both the regularized N-ODE Transformer and TT-Trassfommer, the regularization hyperparameter is X = 0.5, it provides the optimal performance in our tests. We use T = and a total of 8 time steps for the numerical integration.
We repeat the experiment for five random trials. In all trials, the unregularized N-ODE Transformer and OT-Transformer experienced issues with exploding gradients, resulting in NaN outputs. In order to estimate how the unregularized model would perform under more stable conditions, we impose slight transport cost with À 0.001. We note that the continuous-time models with slight and standard regularization completed all trials without issues. This shows the effectiveness of the transport cost regularization in stabilizing the training process and avoiding exploding gradients.
The best test accuracies are reported Table 9 The baseline architecture achieved a test accuracy of 83.9% The N-ODE Transformers with slight and standard regularization report test accuracy of 83.6% and 83.9%, respectively, which are not better than the baseline model. The N-ODE Transformer with slight regularization reports test accuracy of 83.6%. With standard regularization, the test accuracy slightly increases 83.9%. However, both results are not better than that of the baseline model. The OT-Transformer with slight regularization reported test accuracy of 82.7%. which is
35

---

subpar compared to the baseline model. On the other hand, the standard OT OT-nransformer achieves the best accuracy of 84.6%, which is 0.7% higher than the baseline model, in spite of using smaller model. The test accuracy is also 0.7% higher than that of the N-ODE Transformer's s. We note that with the incorporation of transport cost, the accuracy of N-ODE Transformer improved by only 0.3%. In contrast, the accuracy of OT-Transformer is boosted by 1.9%. Again, this is likely due to that our continuous-time formulation is inherently more suited for transport cost regularization than that of N-ODE Transformer.
The learning curves are reported Figure When we compare the results of the unregularized and regularized OT-Transformers, we see that the regularization effectively reduces overfitting by increasing training loss while simultaneously lowering test loss. Overall, we that the combination of our continuous-in inttimme formulation and transport cost regularization enhances parameter efficiency and generalization of Transformers.
Table 9: Number of parameters for the Transformer blocks, best and final test accuracies (with standard deviation) across five trials for for the sentiment analysis experiment. *: The unregularized continuous-time models experienced gradient explosion. And we estimate their performance by using a slight regularization À = 0.001.
Method/Exp.
Para. Count
Best Test Accuracy
Final Test Accuracy
Baseline
4.74M
83.9% + 0.26%
83.7% + 0.21 %
Unreg. N-ODE Trans.
2.37M
83.6% 0.40%*
83.4% + 0.40%*
Reg. N-ODE Trans.
2.37M
83.9% + 0.48%
83.5% + 0.84%
Unreg. OT-Trans.
2.37M
82.7% 0.38%*
82.1% + 0.89%*
OT-Trans (Ours)
2.37M
84.6% + 0.55 %
83.7% + 0.86%
90%
0.6
Training Accunicy
Trainitg Looe
80%
Orig
0.5
Unreg, N-ODE Trans.
70%
Reg. N-ODE Trans.
Unreg, OT Trans.
0.4
OT-Trans. (Ours)
60%
10
10
83%
90%
Accaricy
0.6
Teet Lowd
70%
0.5
22
60%
10
Epoch
Epoch
Figure 5: Accuracy and data-fitting loss for the sentiment analysis experiment
Text Generation: nanoGPT on Shakespeare (Character-level) Dataset. We use the nanoGPT implementation provided by 48 as the baseline model. The baseline architecture uses Transformer model with Transformer blocks, self-attention heads and hidden dimension of 384, resulting in total parameter count of 10.65 million. In contrast, for the OT-Transformer model, we reduce both the number of Transformer blocks and number of attention heads to and hidden size to 320, total this reduces the total parameter count to 6.16 million. We use 10 time steps for the numerical integration
36

---

and À = 1. We run total of 5000 training iterations. The learning rate decays from x 10 to 1 10 over the course of training. Model performances are measured using the loss in all tests. The reported results are reported in Table [0 and averaged over 3 random trials. We highlight two main findings from the results. First, our proposed model displays significant improvement in parameter efficiency compared to the baseline solution, this is consistent with earlier results from different classification examples. Second, our model outperforms the baseline model in both the best recorded test loss and final step test loss, more notably model ensures generalization while the baseline model overfits and shows a decline in performance, this coincides with our theoretical results and further highlights the importance of regularization.
Table 10: Number of parameters for the models, best and final test losses (with standard deviation) and additional LLM metries across three trials for for the nanoGPT experiment. Rouge and Bleu scores are omitted as less meaningful for character-level experiments; see 20 32
Method
Para. Count
Best Test Loss
Final Test Loss
Perplexity
BERT-P
BERT-R
BERT-FI
Baseline
10.65M
1.47 + 0.005
2.68 0.006
17.26
74.4 %
81.1 %
77.6 %
Ours
6.16M
1.44 0.004
1.44 + 0.005
4.40
74.8%
81.2%
77.9%
Orig
OT-Traas. (Ours)
3.5
2.5
3
2.5
2
1.5
1,000
2,000
3,000
4,000
5,000
0
1,000
2,000
3,000
4,000
5,000
Iteration
Iteraiinn
Figure 6: Training and test loss for the nanoGPT on Shakespeare (character-level) experiment.
Interpretation of Model Uncertainty through the Perplexity Metric We provide more detailed explanation of model robustness in terms of the prediction uncertainty of trained model. Note that the perplexity metric, derived from the test loss, measures the discrepancy between the predicted distribution and the underlying true distribution. as if it is equally uncertain among 10 plausible choices, even if the actual distribution is not uniform. Thus, a lower perplexity is usually more desirable practice. We derive Table TT from Table B including the additional perplexity metric. A visualization of the perplexity metric is also shown in Figure] 7 is evident that our model yields much smaller perplexity values across different noise levels, indicating much lower level of uncertainty in the model's predictions.
simple terms, a model with perplexity 10 behaves
Table 11:
Test loss and perplexity metrics of the nanoGPT model under noise, where noise is introduced as random text replacement.
Experiment
Metrie/ replace rate
0.0
0.005
0.01
0.05
Loss Base.
2.68
2.78
2.88
3.65
4.60
NanoGPT
Loss Ours
1.44
1.49
1.55
1.95
2.42
(text replace)
Perplexity Base.
14.59
16.12
17.81
38.47
99.48
Perplexity Ours
4.22
4.44
4.71
7.03
11.25
Text Generation: GPT-2 on Shakespeare (Word-level) Dataset. We next perform experiments using the more popular GPT-2 architecture on the Shakespeare (Word-level) dataset. We use the same implementation as in the first experiment, which is based [48). The goal for the experiment is to
37

---

Baseline (PPL=15) 15 Plausible Choices
OT-Transformer (PPL=4) -4 Plausible Choices
Baseline (PPL=99) -99 Plausible Choices
OT-Transformer (PPL#1 1) -11 Plausible Choices
Figure 7: Visualization of the perplexity metric for baseline and our OT- Transffsffr
test the generalizability of our proposed approach to large models of over 100M total parameters. The baseline architecture has total 12 layers. 12 attention heads and embedding size of 768, the models has 123. 7M parameters in total. For above mentioned reasons we the same exact architecture for OT-Transformer testing. Its also worth mentioning that the GPT-2 model is in fact overparametriieed for the specific task, which also means reducing the model size is less significant for the task and should not be the focus. We use 10 time steps for the numerical integration and set À = 1 for the example. We run total of 500 iterations for training. The learning rate decays from 6 x 10"4 to x 10-5 over the course of training. Model performances are measured using the test loss in all tests. The reported results are reported in Table 12 and averaged over 3 random trials.
The results indicate that our proposed approach comes out ahead in both the best documented test loss as well as the test loss in the last iteration. This shows the efficacy of our method extends to large-scale Transformer models. We also point out that while the baseline model sees fluctuations in performance as training progresses, with proper regularization our method is more stable as training progresses.
Table 12: Number of parameters for the models, best and final test losses (with standard deviation) across three trials for the GPT-2 and Shakespeare dataset experiment. We use test loss as the measure of model performance.
Method/Exp.
Para. Count
Best Test Loss
Final Test Loss
Baseline
123.7M
4.91 + 0.001
5.18 + 0.032
OT-Trans, (Ours)
123.7M
4.87 + 0.035
4.96 + 0.012
Text Generation: GPT-2 on Open WehText Dataset. Lastly, We conduct a large-scale exper- iment on text generation using GPT-2, trained on the Teen dataset [30]. The dataset is approximately 17GB in size, with the training data containing around 9 billion tokens and the test data containing about 4 million tokens. We use the GPT-2 model as the baseline, the Transformer architecture has in total 12 layers, 12 attention heads and embedding size of 768, as result, the model has 123.7M parameters in total. We use the same architecture to parametrize Tr-nsaofformrr. For -ranssosfrr,, we use 10 time steps for the numerical integration and set À = 1. We run total of 60, 00 iterations for training. The learning rate decays from 6 10-4 to 6 x 10-5 over the course of training. Given the large scale of the experiment, we perform only one seeded trial for each model. We use the example to test the performance of our method on not only large models but also large datasets.
38

---

11
Orig
11
10
OT-Trams (Ours)
10
Training
4
200
300
400
200
300
400
Iteration
Iteration
Figure Training and test loss for the GPT-2 on Shakespeare (word-level) experiment.
we see that our OT-Transformer outperforms the baseline GPT-2 by a sizable margin in both best test loss and final test loss. These results highlight the ability of our model to outperform established baselines in large-scale and realistic experimental settings.
We present the results in Table 13
Table 13: Number of parameters for the models, best and final test accuracies and LLM metrics for the GPT-2 and Wee dataset (9 billion tokens) experiment. Both the baseline and our OT--rranffrrmer use a model with 123.7 million parameters.
Loss
Perplexity
BLEU
ROUGE-1
ROUGE-2
ROUGE-L
BERT-P
Baseline
3.20
3.21
26.09
12.22%
57.30%
17.07%
35.17%
80.77% 81.91%
83.31 84.27%
82.02%
Ours
2.90
2.91
19.56
14.07%
59.4 59.46%
18 99%
38.05%
83.07%
Orig
OT Trans. (Ours)
Iteration (x10*)
Iteration (x10*)
Figure Training and test loss for the GPT-2 Open WebTett experiment. The Test loss is computed every 1,000 iterations.
39

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificcton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
The abstract and introduction clearly state the main claims, which are supported by the presented theoretical and numerical results.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
40

---

Justification: We have discussed in Section 4 that our assumptions for theoretical analysis may not fully align with the structure of general encoder-decoder Transformer architec- tures. However, we remark that our model performs well empirically on encoder-decoder architecture, see Section 5.1
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: All theorems and necessary assumptions are clearly documented in the main text, with detailed proofs to the theorems listed the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
41

---

Justification:
Details on experimental settings and results are thoroughly documented
in Appendix ensuring our numerical results can be reproduced.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Our code is available oltttppp//gttt..ccoo////l//r/rta
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.ce.
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https //nips cc/
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
42

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
The paper provides a comprehensive description of the experimental setup, including relevant parameters, implementation details, and evaluation metrics necessary to reproduce all experimental results presented. These details are included in Appendix
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yss]
Justification: We repeated our experiments over multiple trials, and report the mean and std in the main text and Appendix the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
43

---

Answer: [Yes]
Justification: We report the GPU specifications in the main text and Appendix | We also discussed runtime in Appendices B andd]
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: Yes, the research conducted in this paper conforms, in every respect, with the NeurlPS Code of Ethics as outlined in the provided guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Discussion on broader impacts are in Section 6
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (c.g., disinformation. generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out, For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
44

---

11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not pose any such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All creators or original owners of assets used in this paper have been properly credited. Our usage of these assets is in compliance with the license and terms of use.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Our work has been discussed and documented throughout the paper and Appendix in detail.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
45

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve research with human subjects nor crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does involve research with human subjects nor crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMSs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: No LLMs are used in any important, original, or non-standard component of the core methods in this research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer LLM policy for what should or should not be described.
46


Reviews:
[{"review_id": "QRFKPmu255", "reviewer": "Reviewer_fc21", "summary": "The paper introduces a novel framework that analyzes and augments Transformer models using optimal control theory and optimal transport (OT) regularization. The authors propose a novel regularization technique for training standard transformers as a continuous-time optimal control problem, where model depth corresponds to time and the loss is treated as a terminal cost. They train several versions of the so-called OT-Transformer architecture that adds regularization term to the training objective to design of the final layer. The method is plug-and-play, and does not require any complex modifications to the standard transformer architectures. The approach provides theoretical guarantees: in particular, stable forward propagation, robustness to input perturbations, which lead to improved model performance. Empirical evaluation obtained for various transformer architectures (encoders and decoders) on the small classification or language generation tasks shows superiority of the introduced approach in comparison to original transformer.", "questions": "My questions to the authors are as follows:\n\n1. I understand that in this work you focused mainly on training the models on toy examples to demonstrate the applicability of the proposed approach. However, even for such “toy” tasks it would be useful to provide a more qualitative evaluation of the results. Would you consider evaluating the texts generated by the original GPT models (trained on different datasets) with established LLM benchmarks or an “LLM‑as‑a‑judge” framework? For instance, for the Shakespeare‑based nanoGPT, you could adopt the evaluation protocol introduced in TinyStories: How Small Can Language Models Be and Still Speak Coherent English?, where ChatGPT was used to assess the grammar, logic, and consistency of the generated stories. For the larger GPT‑2 models trained on general‑domain data, could you also report results on simple benchmarks such as GLUE?\n\n2. Also I have a question regarding Figure 2, there, on the graph, we can observe an abrupt drop in training loss, and jump in training accuracy, while test loss starts coming upward on the 150 epoch. Is it connected to scheduler used during training. Can you, please, elaborate on this more?\n\n3. Also the computational overhead is not discussed, do the OT-regularization introduce some memory or speed overhead for the final model?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "The main strengths of the paper are as follows:\n\n1. A novel, clearly explained theoretical framework that allows to use optimal transport for training more robust Transformers-based models. The authors provide theoretical explanation for their method that seems clear and reasonable proving the stable forward propagation for the transformer-based models.\n2. The simplicity of integration. While the method is quite novel, its plug-and-play design allows to use it without significant modifications to the original Transformers that makes it a good candidate for further development in the research groups.\n3. Extensive experimental evaluation on different multimodal and language tasks and comparison with original transformers makes the theoretical research more solid. We can observe that the OT-Transformers trained in the same manner as the original ones outperforms the latter based on the accuracy, or the test loss.\n\nThe main weaknesses of the paper are as follows:\n1. While the empirical evaluation is extensive, its plug-and-play design is still limited to the very simple tasks, like MNIST or binary classification, language generation is evaluated only based on the test loss that is usually not a completely fair point for benchmarking the language models. While the loss can be low, the generations themselves can be of different quality, for example, the model can loop into repetitions, and we are not able to notice it based on the pure loss.\n2. In appendix G, the author notice the sensitivity of their model to the regularization parameter, however, it is noted only shallow, how this regularization term was adjusted (empirical study only)? Why this term turned out to be so important for the training stability is not properly discussed.", "was_revised": true, "final_justification": "I would keep my score. I believe that this work presents a novel more theoretical research, designed in a way, so that it can be applied to any Transformer-based model. The empirical evaluation is a little bit limited to small models, however, I don't see it as a disadvantage of the work, so, I believe it can be accepted for NeurIPS, and it will be a valuable and interesting paper for the community.", "rebuttal": "We sincerely thank the reviewer for their time, effort, and constructive feedback. They have helped us improve and strengthen our paper. We are encouraged that the reviewer acknowledged the novelty, theoretical significance, and simplicity of implementation of our work. We address the reviewer's concerns in bullet points as follows. We also conducted additional experiments to address the comments of the reviewers.\n___\n**Concern 1: Plug-and-play design is limited to simple tasks**\n\n**Response 1:** By plug-and-play, we mean **our model can be applied to any Transformer models** with only slight code modification. In particular, our model reuses existing Transformer architectures to parametrize a dynamical system (see Figure 1 and Equation 7). \n\n**Our approach is not limited to simple tasks and can be applied for any applications.** Our experiments cover problems with different scales, architectures, and applications. We specifically refer the reviewer to the GPT2 experiments in Appendix G, which shows that our model can handle realistic and hard tasks. \n\nImportantly, the plug-and-play nature, versatility, ease of adoption, and theoretically grounded performance improvement represent key advantages of our method.\n___\n**Concern 2: Experiment GPT with LLM benchmarks**\n\n**Response 2:** We thank the reviewer for the helpful suggestion, we have added additional evaluation results for both the NanoGPT (shakespeare character-level) and GPT2 (OpenWebText) experiments using common LLM benchmarks including Perplexity, Bleu, Rouge1, Rouge2, RougeL, BertP, BertR, BertF1.  For all metrics, our model consistently outperforms the baseline, demonstrating the strength of our proposed approach. We will also add these results to the camera-ready version of the paper. \n\n**GPT2 (OpenWebText) experiment**\n| Model    | Test Loss   | Perplexity | BLEU (%)  | ROUGE-1 (%) | ROUGE-2 (%) | ROUGE-L (%) | BERT-P (%) | BERT-R (%) | BERT-F1 (%) |\n| -------- | ---------- | ---------- | --------- | ----------- | ----------- | ----------- | ---------- | ---------- | ----------- |\n| Baseline | 3.19     | 26.09      | 12.22     | 57.30       | 17.07       | 35.17       | 80.77      | 83.31      | 82.02       |\n| Ours     | **2.91** | **19.56**  | **14.07** | **59.46**   | **18.99**   | **38.05**   | **81.91**  | **84.27**  | **83.07**   |\n\n**NanoGPT (Shakespeare character-level) experiment** (We omit the Rouge and Bleu scores here, as they are less meaningful for character-level experiments; see [2-3]) \n| Model    | Test loss | Perplexity   | BERT-P (%) | BERT-R (%) | BERT-F1 (%) |\n| -------- | --------- | ------------ | ---------- | ---------- | ----------- |\n| Baseline | 2.67      | 17.26    | 74.4      | 81.1      | 77.6       |\n| Ours     | **1.44**  | **4.40** | **74.8**  | **81.2**  | **77.9**   |\n___\n**Concern 3: Explain Figure 2's loss plunge**\n\n**Response 3:** We adopt this experiment from [1], where it is used as a benchmark task. We use the same set of hyperparameters as the original experiment. These include the number of epochs and learning rate scheduler. At the 150th epoch, the learning rate decreases from 1e-3 to 1e-4. The drop in training loss and corresponding jump in test accuracy coincide with this scheduled change. \n\nWe emphasize that the hyperparameters and training settings were originally chosen for the baseline and were not specifically tailored by us. Nevertheless, in all our seven extensive experiments, our model consistently achieves superior results even under suboptimal conditions.\n___\n**Concern 4: Does the regularizer introduce more memory or speed overhead**\n\n**Response 4:** We emphasize that the regularization adds negligible computational and memory cost. This is because the regularization term is computed alongside Equation 7; see the discussion in Section 3.\n___\n**Concern 5: Empirical sensitivity study of regularization term**\n\n**Response 5:** We observe that the model performance is stable with respect to different choices of $\\lambda$ and other hyperparameters. Tuning of hyperparameters is generally not required; even if the optimal $\\lambda$ and number of steps ($=\\frac{T}{\\Delta t}$) are not used, our model still outperforms the baseline.\n\nWe conduct additional NanoGPT experiments with different $\\lambda$'s and numbers of time steps ($=\\frac{T}{\\Delta t}$) for sensitivity study. The experiments are done over 3 random trials, and we report the average test loss and its standard deviation. We see that our model consistently outperforms the baseline across over two orders of magnitude of $\\lambda$. Our model performance is also very stable with respect to the number of time steps. These additional results will be added to the camera-ready version.\n\n| \\$\\lambda\\$ | 0.01         | 0.05         | 0.1          | 0.5          | 1                | 5            |Baseline            |\n| ----------- | ------------ | ------------ | ------------ | ------------ | ---------------- | ------------ | ------------ |\n| Test loss   | 2.68 ± 0.019 | 2.29 ± 0.017 | 2.05 ± 0.009 | 1.52 ± 0.002 | **1.44 ± 0.004** | 1.48 ± 0.003 | 2.68 ± 0.006 |\n\n| Time step (\\$T/\\Delta t\\$) | 1            | 5                | 10           | 15           | 20           |Baseline            |\n| -------------------------- | ------------ | ---------------- | ------------ | ------------ | ------------ | ------------ |\n| Test loss                  | 1.46 ± 0.003 | **1.43 ± 0.002** | 1.44 ± 0.004 | 1.46 ± 0.002 | 1.49 ± 0.008 | 2.68 ± 0.006 |\n___\n**Concern 6: Why this regularization term turned out to be so important for the training stability**\n\n**Response 6:** Indeed, this is one of the key points in our paper, and we have discussed this after Theorems 1 and 2 in Section 4, and full mathematical derivations are given in Appendices D2-D6. We here present a short version of our theory. We hope the additional discussion below helps clarify any confusion. \n\nWe prove that for standard Transformer and when $\\lambda=0$, the training problem is ill-posed in general. Specifically, there are infinitely many optimal solutions (trained Transformers) to the training problem. These include trained models with highly irregular (e.g., exploding) hidden states. This can explain why the model weights and gradients can explode during Transformer training especially when the model is deep; see the point cloud, sentiment analysis, NanoGPT, and GPT2 experiments in Appendix G.\n\nOn the other hand, we prove that when the OT-Transformer with $\\lambda>0$ is used, the optimal solution (trained Transformer) is unique. More importantly, we prove that the unique trained Transformer has highly regular hidden states; see Appendices D4-D6. Since **the OT regularization always leads to this unique and highly regular trained model, this can effectively prevent model weights and gradients from exploding and thus stabilize training.** These can explain the GPT-2 learning curves in Figure 8, where the Transformer architecture is deep and more prone to training instability, and our model's training and test curves are much less oscillatory than the baseline's.\n___\n[1] Sander, M. E., Ablin, P., Blondel, M., & Peyré, G. (2022, May). Sinkformers: Transformers with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics (pp. 3515-3530). PMLR.\n\n[2] Gehrmann, S., Adewumi, T., Aggarwal, K., Ammanamanchi, P. S., Aremu, A., Bosselut, A., ... & Zhou, J. (2021, August). The GEM Benchmark: Natural Language Generation, its Evaluation and Metrics. In Proceedings of the 1st Workshop on Natural Language Generation, Evaluation, and Metrics (GEM 2021) (pp. 96-120).\n\n[3] Denoual, E., & Lepage, Y. (2005). BLEU in characters: towards automatic MT evaluation in languages without word delimiters. In Companion Volume to the Proceedings of Conference including Posters/Demos and tutorial abstracts."}, {"review_id": "w3WpkK7Ogw", "reviewer": "Reviewer_ZUmk", "summary": "This paper formulates Transformer training as a continuous-time optimal control problem, leading to the proposal of OT-Transformer—a plug-and-play module that introduces an optimal transport (OT) based velocity regularization term. The authors provide theoretical results (well-posedness, stability bounds, distributional robustness) and validate the method across diverse domains including language, vision, and 3D point clouds.", "questions": "Are there any concrete empirical evidence supporting their theoretical claims, such as quantitative tests of stability or robustness under input perturbations? Additionally, how exactly is the continuous-time variable mapped to discrete Transformer layers—does each layer represent an equal time increment, or is the temporal analogy merely symbolic? It would also be helpful to clarify how the regularization strength λ is selected and whether the model's performance is sensitive to its tuning across different tasks or scales.", "limitations": "Despite its theoretical appeal, the paper lacks direct empirical validation of its core theoretical claims—such as forward stability or distributional robustness—which remain largely untested in controlled experiments. Moreover, the continuous-time formulation and its interpretation in relation to discrete Transformer architectures are insufficiently clarified, weakening the credibility of the dynamical systems analogy.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths：\nThis paper introduces a novel and theoretically grounded approach by casting Transformer training as a continuous-time optimal control problem, leading to a plug-and-play regularization scheme with provable stability and robustness properties. The proposed OT-Transformer is elegantly formulated, broadly applicable across modalities, and achieves consistent empirical improvements while reducing parameter counts, demonstrating both practical and theoretical value.\n\nWeaknesses:\nDespite its theoretical appeal, the paper lacks direct empirical validation of its core theoretical claims—such as forward stability or distributional robustness—which remain largely untested in controlled experiments. Moreover, the continuous-time formulation and its interpretation in relation to discrete Transformer architectures are insufficiently clarified, weakening the credibility of the dynamical systems analogy.", "was_revised": true, "final_justification": "I still have some concerns, 1) It is similar to the point raised by Reviewer fc21. 'While the empirical evaluation is extensive, its plug-and-play design is still limited to the very simple tasks, like MNIST or binary classification, language generation is evaluated only based on the test loss that is usually not a completely fair point for benchmarking the language models. While the loss can be low, the generations themselves can be of different quality, for example, the model can loop into repetitions, and we are not able to notice it based on the pure loss.' 2) The robustness statement in the title and conclusion of the article is not rigorous. The experiment lacks verification on the corresponding dataset. 3）The paper needs some major revisions or additions after submission, which is also a factor in my final scoring.", "rebuttal": "We sincerely thank the reviewer for their time, effort, and constructive feedback. They have helped us improve and strengthen our paper. We are encouraged that the reviewer acknowledged the novelty, elegance, theoretical significance, and practical value of our work. We address the reviewer's concerns in bullet points as follows. We also conducted additional experiments to address the comments of the reviewers.\n___\n**Concern 1: Lacks empirical validation for forward stability and distributional robustness**\n\n**Response 1:** We agree with the reviewer that additional experimentation can further support our theory, specifically Theorems 2 and 3 and Appendix F. We thank the reviewer for the helpful suggestion. We here conduct 4 additional experiments in which varying levels of random noise are injected into the test inputs. The noise was absent from the training data. This setup tests the model's **forward stability**, **distributional robustness**, and **out-of-distribution generalization**. **Our model consistently outperforms the baselines over all these tests across all noise levels.** Importantly, our model performs significantly better when the noise level is high. These additional results further support our theory and existing results.\n\nThe following results with full implementation details will be added to the camera-ready version of the paper.\n\n**Experiment 1: Point could classification with point dropout** (commonly used noise; see [5-6]) (average test accuracy and std)\n| Drop rate | 0.0               | 0.01              | 0.05              | 0.1               | 0.2               | 0.5               |\n| --------- | ----------------- | ----------------- | ----------------- | ----------------- | ----------------- | ----------------- |\n| Baseline  | 86.6% ± 0.45%     | 86.6% ± 0.48%     | 85.8% ± 0.60%     | 84.3% ± 0.69%     | 76.9% ± 0.88%     | 34.5% ± 1.94%     |\n| Ours      | **89.3% ± 0.55%** | **89.3% ± 0.55%** | **88.8% ± 0.34%** | **87.6% ± 0.62%** | **83.9% ± 0.80%** | **55.4% ± 4.87%** |\n\n**Experiment 1: Point could classification with point dropout** (performance drop, lower is better)\n| Drop rate | 0.01     | 0.05     | 0.1      | 0.2      | 0.5       |\n| --------- | -------- | -------- | -------- | -------- | --------- |\n| Baseline  | **0.0%**     | 0.8%     | 2.3%     | 9.7%     | 52.1%     |\n| Ours      | **0.0%** | **0.5%** | **1.7%** | **5.4%** | **33.9%** |\n\n**Experiment 2: NanoGPT with random text replacement** (follows [7]) (average test loss and std)\n| Replace rate | 0.0              | 0.005            | 0.01             | 0.05             | 0.1              |\n| ------------ | ---------------- | ---------------- | ---------------- | ---------------- | ---------------- |\n| Baseline     | 2.68 ± 0.004     | 2.78 ± 0.004     | 2.88 ± 0.004     | 3.65 ± 0.004     | 4.60 ± 0.005     |\n| Ours         | **1.44 ± 0.005** | **1.49 ± 0.004** | **1.55 ± 0.003** | **1.95 ± 0.010** | **2.42 ± 0.022** |\n\n**Experiment 2: NanoGPT with random text replacement** (performance drop, lower is better)\n| Replace rate | 0.005    | 0.01     | 0.05     | 0.1      |\n| ------------ | -------- | -------- | -------- | -------- |\n| Baseline     | 0.10     | 0.20     | 0.97     | 1.92     |\n| Ours         | **0.05** | **0.11** | **0.51** | **0.98** |\n\n**Experiment 3: MNIST with Gaussian noise** (commonly used noise; see [8]) (average test accuracy and std)\n| Noise level | 0.0                | 0.01               | 0.05               | 0.1                | 0.2                | 0.5                |\n| ----------- | ------------------ | ------------------ | ------------------ | ------------------ | ------------------ | ------------------ |\n| Baseline    | 92.97% ± 0.67%     | 92.99% ± 0.66%     | 92.96% ± 0.69%     | 92.76% ± 0.72%     | 91.70% ± 0.68%     | 80.64% ± 1.58%     |\n| Ours        | **97.05% ± 0.15%** | **97.05% ± 0.16%** | **96.99% ± 0.18%** | **96.89% ± 0.15%** | **96.39% ± 0.11%** | **90.10% ± 1.10%** |\n\n**Experiment 3: MNIST with Gaussian noise** (performance drop, lower is better)\n| Noise level | 0.01      | 0.05      | 0.1       | 0.2       | 0.5       |\n| ----------- | --------- | --------- | --------- | --------- | --------- |\n| Baseline    | **-0.02%**    | **0.01%**     | 0.21%     | 1.27%     | 12.33%    |\n| Ours        | 0.00% | 0.06% | **0.16%** | **0.66%** | **6.95%** |\n\n**Experiment 4: MNIST with uniform noise** (commonly used noise; see [8]) (average test accuracy and std) \n| Noise level | 0.0                | 0.01               | 0.05               | 0.1                | 0.2                | 0.5                |\n| ----------- | ------------------ | ------------------ | ------------------ | ------------------ | ------------------ | ------------------ |\n| Baseline    | 92.97% ± 0.67%     | 92.99% ± 0.65%     | 92.98% ± 0.63%     | 92.90% ± 0.58%     | 92.64% ± 0.57%     | 90.02% ± 0.45%     |\n| Ours        | **97.05% ± 0.15%** | **97.03% ± 0.14%** | **97.00% ± 0.15%** | **96.97% ± 0.14%** | **96.79% ± 0.16%** | **95.57% ± 0.11%** |\n\n**Experiment 4: MNIST with uniform noise** (performance drop, lower is better)\n| Noise level | 0.01  | 0.05      | 0.1       | 0.2       | 0.5       |\n| ----------- | ----- | --------- | --------- | --------- | --------- |\n| Baseline    | **0.02%** | **-0.02%**    | **0.07%**     | 0.33%     | 2.95%     |\n| Ours        | **0.02%** | 0.05% | 0.08% | **0.26%** | **1.48%** |\n___\n**Concern 2: The connection of continuous-time and discrete models are insufficiently clarified**\n\n**Response 2:** Transformer architectures include a skip-connection in each layer of their Transformer blocks. In particular, each layer can be written as \n$$ x_{t+1} = x_t + f(x_t; \\theta_t). $$\nThis can be seen as a forward Euler discretization of continuous dynamics. Specifically, we formulate the continuous dynamics as\n$$ \\frac{d x(t)}{dt} = f(x(t), t; \\theta), $$\nwhere $f$ is a Transformer model. This perspective is commonly adopted in continuous-time models and connects neural networks to dynamical systems. It enables the use of techniques from dynamical systems and has been central to the development of continuous-time models; see [1-4] for more details. **Due to space constraints, please refer to our responses 1 and 2 to Reviewer PviE for a pseudocode comparing the standard Transformer and our model.**\n\nAs a part of our contribution, we show that directly applying a continuous-time formulation to Transformers results in an ill-posed training problem, which can be a major issue for large-scale training problems. The trained Transformers from this ill-posed training problem are sub-par in performance, which we demonstrated by both 1. empirical experiments (in Sections 1 & 5, and Appendix G) and 2. inspecting the regularity and optimality conditions of the training problem, which are a Hamilton-Jacobi-Bellman partial differential equation coupled with a continuity equation (the proofs are given in Appendices D1-D4).\n\nOn the other hand, our novel theory demonstrates that the OT regularization leads to a highly regular trained model with generalization and robustness guarantees; see Section 4, and Appendices D5-D7 and F.\n___\n**Concern 3: How does the continuous hidden states (of our model) relate to the discrete hidden states (of standard Transformer)**\n\n**Response 3:** The initial continuous hidden states ${\\bf X}(0)$ is the same as the input of a standard Transformer. The terminal continuous hidden states ${\\bf X}(T)$ is used as the output of our model. When a forward Euler discretization scheme with step size 1 is used, our model reduces to a standard Transformer. In this case, the continuous and discrete hidden states are the same. Thus, our model can be seen as a generalization of the standard Transformer architecture, with a continuous-depth architecture. \n___\n**Concern 4: Clarify how the parameter $\\lambda$ is tuned, whether it is sensitive**\n\n**Response 4:** Due to space constraints, please refer to our response 5 to Reviewer fc21.\n___\n\n[1] Haber, E., & Ruthotto, L. (2017). Stable architectures for deep neural networks. Inverse problems, 34(1), 014004.\n\n[2] Chen, R. T., Rubanova, Y., Bettencourt, J., & Duvenaud, D. K. (2018). Neural ordinary differential equations. Advances in neural information processing systems, 31.\n\n[3] Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., & Duvenaud, D. (2019) FFJORD: Free-Form Continuous Dynamics for Scalable Reversible Generative Models. In International Conference on Learning Representations.\n\n[4] Ruthotto, L. (2024). Differential equations for continuous-time deep learning. arXiv preprint arXiv:2401.03965.\n\n[5] Wang, Y., Sun, Y., Liu, Z., Sarma, S. E., Bronstein, M. M., & Solomon, J. M. (2019). Dynamic graph cnn for learning on point clouds. ACM Transactions on Graphics (tog), 38(5), 1-12.\n\n[6] Qi, C. R., Yi, L., Su, H., & Guibas, L. J. (2017). Pointnet++: Deep hierarchical feature learning on point sets in a metric space. Advances in neural information processing systems, 30.\n\n[7] Jin, D., Jin, Z., Zhou, J. T., & Szolovits, P. (2020, April). Is bert really robust? a strong baseline for natural language attack on text classification and entailment. In Proceedings of the AAAI conference on artificial intelligence (Vol. 34, No. 05, pp. 8018-8025).\n\n[8] Rafael C. Gonzalez and Richard E. Woods. 2006. Digital Image Processing (3rd Edition). Prentice-Hall, Inc., USA."}, {"review_id": "ExvYdImpOe", "reviewer": "Reviewer_SwmS", "summary": "In this paper, the authors bring insights from optimal control theory to Transformer model training. With a plug-and-play modification, Transformers are formulated in a continuous-time setting and incorporate regularization during training. Seven experiments covering point-cloud, image, sentiment, and text generation tasks show that OT-Transformer either matches or exceeds the baseline while often using fewer parameters.", "questions": "1. Adding evaluations across both smaller and larger model sizes, with and without OT-Transformer, may help assess the effectiveness of the proposed method.\n\n2. Fixing the hyperparameters and using fixed random seeds would improve reproducibility and reliability.\n\n3. Providing a guideline on how to reduce the model size when applying OT-Transformer would be helpful.", "limitations": "See Weaknesses.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "1. The authors claim that the proposed method is better at avoiding overfitting. However, it is difficult to determine whether the evaluation results support this claim, as the model size is reduced. The improved performance or reduced overfitting may be due to the smaller number of parameters rather than the new training algorithm.\n\n2. The authors also claim that stable forward propagation results in robustness to input perturbations and out-of-distribution generalization. However, there is no evaluation provided to support this claim.\n\n3. There is no pseudocode, which makes it difficult to understand how the formulation is implemented and whether it is correct.\n\n4. There is no discussion or ablation study on the selection of hyperparameters such as lambda, T, and the total number of time steps. The choice of hyperparameters appears cherry-picked, as the number of training epochs and the learning rate vary significantly across datasets. Additionally, there is no general rule provided for how or why the model size should be reduced for OT-Transformer: in the MNIST experiment, the model size is reduced by 5x, while in the Shakespeare dataset, there is no reduction.", "was_revised": true, "final_justification": "During the rebuttal period, the authors have answered most of my questions. These include:\n\n1. The authors have provided pseudo code that helps clarify the implementation of the proposed algorithm.\n2. The authors have pointed out that the selection of hyperparameters is based on previous works.\n3. A series of experiments have been conducted to support the claims on robustness and out-of-distribution generalization.\n4. The authors have reported the training time compared to both the original model and the Standard Continuous-Time Transformer, providing sufficient information to evaluate the pros and cons.\n5. Although I am still not convinced by the discussion on model size, I believe it does not significantly affect the contribution of this work.\n\nFor the above reasons, I am willing to raise my score to borderline accept.", "rebuttal": "We sincerely thank the reviewer for their time, effort, and constructive feedback. They have helped us improve and strengthen our paper. We address the reviewer's concerns in bullet points as follows. We also conducted additional experiments to address the comments of the reviewers.\n___\n**Concern 1: Difficult to determine whether the evaluation results support that our model is better at avoiding overfitting, since the model size is reduced. The improved performance/reduced overfitting may be due to reduced number of parameters**\n\n**Response 1: Our seven diverse experimental results (in Appendix G) suggest the opposite. Our model does not suffer from overfitting issues. Our model is more parameter-efficient. We explain through the following details.**\n\n1. In our experiments, our model **achieves both lower training and test losses (on unseen data) consistently across different examples** (see Appendix G), despite having fewer parameters. **These results show that the performance improvement is due to parameter efficiency, not overfitting.**\n\n2. We point the reviewer to the two GPT2 experiments presented, in which our model and the baseline have the same number of parameters, yet our model consistently performs better. This directly supports that our model is better at avoiding overfitting.\n\n3. We remark that reducing model size to demonstrate parameter efficiency is a common practice in the continuous-time modeling literature—for example, continuous normalizing flows [3-5].\n\n4. To further show that we do not reduce model size to avoid overfitting, we conducted a new NanoGPT experiment, where OT-Transformer uses the same model size as the baseline, and we report the overall results below. Under repeated trials, **our models outperforms the baseline, whether the model size is matched or reduced.**\n\n**NanoGPT experiment**\n| Model    | Model size | Test loss        |\n| -------- | ---------- | ---------------- |\n| Baseline | 10.65M     | 2.68 ± 0.006     |\n| Ours     | 10.65M     | **1.44 ± 0.007** |\n| Ours     | 6.16M      | **1.44 ± 0.005** |\n\n___\n**Concern 2: The authors also claim that stable forward propagation results in robustness to input perturbations and out-of-distribution generalization**\n\n**Response 2: These are not merely claims; they are rigorously and mathematically proven results.** See Theorems 1-3 and Appendix F. We prove that our training formulation leads to stable forward propagation, which in turn leads to improved robustness and generalization. Specifically, we derive deterministic generalization guarantees when the training or test distributions are perturbed. Moreover, we derive a non-asymptotic and statistical generalization guarantee for out-of-distribution test datasets. Importantly, these guarantees can be controlled through the strength of the OT regularization.\n\nOn the other hand, these bounds do not hold in general for a standard Transformer setup without any type of optimal transport regularization; see Theorems 1 and 2 in Section 4. Because the training is ill-posed and can lead to models with highly irregular (e.g. exploding) hidden states. This can also explain why we have a more stable training loss curve in the GPT2 WebText experiment in Figure 8, where the model is deep and more prone to training instability. Because **our training formulation always leads to a unique model with highly regular hidden states**, this can prevent the hidden states and network weights from exploding during training and resulting in NaNs; see the derivations in Appendices D3-D6.\n\nAlso, our extensive results across seven experiments show that our model has improved robustness and generalization; see Tables 1 and 2 and Appendix G where our approach consistently leads to improved performance, lower variance and better training stability.  \n___\n**Concern 3: No evaluation provided to support that stable forward propagation results in robustness and out-of-distribution generalization**\n\n**Response 3:** We appreciate this useful feedback. **This concern is addressed by our four additional experiments, where our model consistently demonstrates superior robustness and out-of-distribution generalization over the baseline in all settings.** Due to space constraints, please refer to our response 1 to Reviewer ZUmk.\n___\n**Concern 4: No psuedocode to explain**\n\n**Response 4:** We thank the reviewer for the helpful suggestion, we have added a pseudocode that helps clarify and explain our proposed approach. Due to space constraints, please refer to our responses 1 and 2 to Reviewer PviE.\n___\n**Concern 5: Discussion/ablation study on selecting hyperparameters**\n\n**Response 5:** Additional ablation study is performed to address the concern. Due to space constraints, please refer to our response 5 of Reviewer fc21.\n___\n**Concern 6: Hyperparameters appear to be cherry-picked**\n\n**Response 6: We did not cherry-pick any hyperparameters.** The hyperparameters used in the Transformer models and experiments—such as the number of epochs, step size, batch size—follow the original settings from [1] and [2]. While these hyperparameters may appear different for different examples, they are determined by the original authors and not us. \n\nWe also want to highlight that even though the hyperparameters used are not tailored for our approach, we consistently obtained superior results across seven diverse experiments, demonstrating the strength of our framework.\n___\n**Concern 7: Use fixed random seeds**\n\n**Response 7:** As discussed in our paper, we use fixed random seed = 1,2,3,4,5 for our experiments to ensure reproducibility. Moreover, the reported results are averaged over multiple trials with different seeds to reduce variance and ensure our claim is accurate. \n\nThe full training details for all examples are reported in Appendix G. The reader should be able to reproduce our results using the information shared. We will also release our source code on Github upon publication of this work to further ensure reproducibility.\n___\n**Concern 8: Provide guidelines on how to reduce the model size**\n\n**Response 8:** We reduce the number of parameters to demonstrate that our model is more parameter-efficient. This approach follows common practice in the continuous-time modeling literature; see [3-5] for example.\n\nIt is not necessary to reduce the model size. Specifically, using the same model size as the baseline can also achieve improved performance; see our GPT2 experiments in Appendix G. We also perform additional experiments on NanoGPT, where our model consistently outperforms the baseline, whether it is using the same or reduced model size. For the results, please refer to our response 1 to you.\n___\n[1] Sander, M. E., Ablin, P., Blondel, M., & Peyré, G. (2022, May). Sinkformers: Transformers with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics (pp. 3515-3530). PMLR.\n\n[2] Karpathy, A. (2023). NanoGPT [Source code]. GitHub. \n\n[3] Chen, R. T., Rubanova, Y., Bettencourt, J., & Duvenaud, D. K. (2018). Neural ordinary differential equations. Advances in neural information processing systems, 31.\n\n[4] Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., & Duvenaud, D. (2019). FFJORD: Free-Form Continuous Dynamics for Scalable Reversible Generative Models. In International Conference on Learning Representations.\n\n[5]  Massaroli, S., Poli, M., Park, J., Yamashita, A., & Asama, H. (2020). Dissecting neural odes. Advances in neural information processing systems, 33, 3952-3963."}, {"review_id": "SrJuior35Y", "reviewer": "Reviewer_PviE", "summary": "This paper investigates the optimal control perspective of Transformer architectures. Building on this perspective, the authors introduces an OT-regularizer which seems to help the performance of the Transformer training in practice.", "questions": "See above.", "limitations": "See above.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "The perspective presented in this paper is nice and fresh. Also, I appreciate the OT-based regularizer working nicely in practice.\n\nA few suggestions/questions:\n\n1. First, I would appreciate as a reader if the authors really highlight the main difference resulted by the OC-based perspective. I get it from reading the paper, but it would really clear up the main contributions if the authors really walk through an example (e.g. GPT training) how the OC-based perspective together with OT-regularize leads to a different training framework.\n\n2. Also, I would like to see the validation loss curve for OT-transformer and compare it with standard Transformer to see how the training dynamics differ.\n\n3. Along the similar vein as 1), please show how the OT-based regularizer is implemented in practical Transformer training. You should clearly present the equation of the regularizer, or pseudocode or some sort.\n\n4. Please share the full training details in the appendix. (Optimizer hyperparameters etc)\n\n5. I noticed that for some of the experiments you chose smaller models for OT-Transformer, but for GPT2 training, you used the same size transformers. Is that any particular reason for this? I wonder what happens if you use smaller models for GPT2 training as well?", "was_revised": true, "final_justification": "The authors clarified the relevance of the paper. I'd like to maintain my positive assessment.", "rebuttal": "We sincerely thank the reviewer for their time, effort, and constructive feedback. They have helped us improve and strengthen our paper. We also thank them for acknowledging the novelty of our work. We address the reviewer's concerns in bullet points as follows. We also conducted additional experiments to address the comments of the reviewers.\n___\n**Concerns 1 and 2: How OT regularizer lead to a different training framework? Show how the model is implemented, by equation or pseudocode**\n\n**Responses 1 and 2:** Figure 1 of our paper gives an overview of our approach, and how it can be adapted from existing transformer architecture. To further clarify this, we here provide a pseudocode in algorithm format that details how our framework can be easily implemented. \n\nSince uploading images and links are prohibited, we have made our best effort to present the pseudocode here. We will include the pseudocode (in proper LaTeX algorithm formatting, with the standard Transformer and our model compared side-by-side) to the camera-ready version of the paper.\n\n**Pseudocode: Forward Propagation of Standard Transformer vs OT-Transformer**\n\nInputs: Input $\\mathbf{X}_0$, Transformer model $f = f_D \\circ \\dots \\circ f_1$, terminal time $T$, regularization parameter $\\lambda$, number of numerical integration steps $M$\n\nDefine: Step size $\\Delta t = \\frac{T}{M}$\n\n**Standard Transformer**\n\nFor $i = 1, 2, \\dots, D$:\n\n  $\\mathbf{X}{i} \\leftarrow \\mathbf{X}{i-1} + f_i(\\mathbf{X}_{i-1})$\n\nReturn: output $\\mathbf{X}_D$\n\n**OT-Transformer**\n\nInitialize: $\\mathbf{X}(0) = \\mathbf{X}_0$, $\\text{reg} = 0$\n\nFor $t = \\Delta t, 2 \\Delta t, \\dots, T$:\n\n  Compute $f(\\mathbf{X}(t))$ *Adapt from standard Transformer*\n\n  $\\mathbf{X}(t + \\Delta t) \\leftarrow \\mathbf{X}(t) + \\Delta t \\cdot f(\\mathbf{X}(t))$\n\n  $\\text{reg} \\leftarrow \\text{reg} + \\Delta t \\cdot \\|\\|f(\\mathbf{X}(t)) \\|\\|_F^2$\n\nReturn: output $\\mathbf{X}(t)$, regularization $\\lambda \\cdot \\text{reg}$\n___\n**Concern 3: Show validation loss curve**\n\n**Response 3:** Validation curves for all tested examples are reported in Appendix G, specifically, see Figures 2-8.\n___\n**Concern 4: Share full training details**\n\n**Response 4:** The full training details for all examples are reported in Appendix G. The reader should be able to reproduce our results using the information shared. We will also release our source code on Github upon publication of this work to further ensure reproducibility.\n___\n**Concern 5: Why GPT-2 uses a model of the same size**\n\n**Response 5:** Overall, our model consistently outperforms the baseline models, whether it uses the same or a reduced model size.\n\nSince the GPT-2 model is of larger size (>100M parameters), we use it to demonstrate that our approach can also generalize to larger models and datasets.\n\nFor the other models, we reduce the number of parameters to demonstrate that our model is more parameter-efficient. This approach follows common practice in the continuous-time modeling literature; see [1-3] for example.\n\nAdditionally, we conducted a new NanoGPT experiment in which OT-Transformer uses the same model size. The overall results, reported below, clearly show that our approach achieves better performance with both full-sized and reduced-sized models.\n\n**NanoGPT experiment**\n| Model    | Model size | Test loss        |\n| -------- | ---------- | ---------------- |\n| Baseline | 10.65M     | 2.68 ± 0.006     |\n| Ours     | 10.65M     | **1.44 ± 0.007** |\n| Ours     | 6.16M      | **1.44 ± 0.005** |\n___\n[1] Chen, R. T., Rubanova, Y., Bettencourt, J., & Duvenaud, D. K. (2018). Neural ordinary differential equations. Advances in neural information processing systems, 31.\n\n[2] Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., & Duvenaud, D. (2019). FFJORD: Free-Form Continuous Dynamics for Scalable Reversible Generative Models. In International Conference on Learning Representations.\n\n[3]  Massaroli, S., Poli, M., Park, J., Yamashita, A., & Asama, H. (2020). Dissecting neural odes. Advances in neural information processing systems, 33, 3952-3963."}]

Author rebuttal:

Meta-review: This paper proposes a novel plug-and-play model called OT-Transformer, which formulates Transformers in a continuous-time setting and adds an optimal transport (OT) regularization into the training loss. With a large enough regularization parameter $\lambda > TL^2$, the model comes with a property dubbed stable forward propagation (Theorem 2), which states that if two input–target output pairs are similar to each other, then the output of the optimally trained model will also be similar; the property is shown to enhance generalization and robustness (Theorem 3). The proposed model is tested with seven experiments across different modalities and tasks, exhibiting consistent improvements over baselines.

Most reviewers appreciated the novelty of the proposed framework. Reviewers ZUmk and SwmS raised a concern that the robustness and generalization implications of the stable forward propagation property were not properly tested in the paper. With the authors’ additional experiments during the rebuttal phase, I believe that the issue is addressed. After a careful review of the entire discussion, I found that most other questions and concerns were well-addressed. Some reviewer commented that the evaluations were done only on relatively small models, but they also added that this is not viewed as a disadvantage; I also agree with this view, given the diversity of experiments presented.

Overall, the paper presents a novel plug-and-play framework that can be potentially impactful. I recommend acceptance. As noted in the Author Final Remarks, please include the additional LLM benchmark results and robustness results, which will further strengthen the paper. Please also include the runtime analysis (in the response to Reviewer xcGF), as this will better reflect the time–performance trade-off.
Decision: Accept (poster)