Title: Global Convergence in Training Large-Scale Transformers
Venue: NeurIPS 2024
Accepted: yes
Keywords: Transformer

gradient flow

mean-field analysis

global convergence
Abstract: Despite the widespread success of Transformers across various domains, their optimization guarantees in large-scale model settings are not well-understood. This paper rigorously analyzes the convergence properties of gradient flow in training Transformers with weight decay regularization. First, we construct the mean-field limit of large-scale Transformers, showing that as the model width and depth go to infinity, gradient flow converges to the Wasserstein gradient flow, which is represented by a partial differential equation. Then, we demonstrate that the gradient flow reaches a global minimum consistent with the PDE solution when the weight decay regularization parameter is sufficiently small. Our analysis is based on a series of novel mean-field techniques that adapt to Transformers. Compared with existing tools for deep networks (Lu et al., 2020) that demand homogeneity and global Lipschitz smoothness, we utilize a refined analysis assuming only $\textit{partial homogeneity}$ and $\textit{local Lipschitz smoothness}$. These new techniques may be of independent interest.

Global Convergence in Training Large-Scale Transformers

Abstract

Despite the widespread success of Transformers across various domains, their optimization guarantees in large-scale model settings are not well-understood. This paper rigorously analyzes the convergence properties of gradient flow in training Transformers with weight decay regularization. First, we construct the mean-field limit of large-scale Transformers, showing that as the model width and depth go to infinity, gradient flow converges to the Wasserstein gradient flow, which is represented by a partial differential equation. Then, we demonstrate that the gradient flow reaches a global minimum consistent with the PDE solution when the weight decay regularization parameter is sufficiently small. Our analysis is based on a series of novel mean-field techniques that adapt to Transformers. Compared with existing tools for deep networks [47] that demand homogeneity and global Lipschitz smoothness, we utilize a refined analysis assuming only partial homogeneity and local Lipschitz smoothness . These new techniques are of independent interest.

1 Introduction

Transformers have revolutionized the field of deep learning since their introduction in [66]. These models are distinguished by their immense scales, often comprising billions of parameters to achieve state-of-the-art performance. Notably, this massive parameterization enables them to excel in a variety of domains, notably in natural language processing [21, 55, 65] and vision tasks [20, 36], where they have significantly advanced the frontiers of machine learning.

Despite the widespread adoption of Transformer models, our understanding of their optimization guarantees is still in its early stages. One particularly intriguing phenomenon is that as the size of model increases, training algorithms typically converges globally despite the highly nonconvex landscape of the training objective function. Remarkably, it remains somewhat enigmatic how gradient-based approaches can consistently succeed when training large-scale Transformers.

Notably, there have been several recent works showing the global convergence of training overparameterized neural networks [51, 16, 28, 14, 47, 22, 23, 35, 3, 25, 76]. In particular, several works [47, 22, 23] studied the setting with deep neural networks with skip connections. By studying the connections between the network with discretization in the parameter space and a corresponding ordinary differential equation system [71, 12, 43], these works demonstrated global convergence guarantees of wide and deep neural networks based on a mean-field analysis. However, these results are established based on certain homogeneity and/or global Lipschitz smoothness properties of the

neural network, which are not applicable to Transformer models. Therefore, it remains an open question how gradient-based methods can effectively train large-scale Transformers.

1.1 Our contribution

In this work, we bridge the gap between Transformer theory and practice by demonstrating the global convergence of Transformer training optimization via gradient flow in a large-scale model regime. We analyze the mean-field limit of the Transformer model, which is characterized by the distribution of model parameters, shifting the focus from parameter space to distributional dynamics in the Wasserstein metric [16]. This approach yields two key theorems:

We show the closeness between practical discrete Transformers trained by gradient flow and continuous Transformers whose parameter distribution follows a partial differential equation of the Wasserstein gradient flow (Theorem 3.1). Our result demonstrates that large-scale discrete Transformers can be approximated by its mean-field limit and the approximation error can be expressed in terms of the width and depth of the Transformer models.

This approximation facilitates our analysis of the global convergence (Theorem 4.1) of discrete Transformer models. By leveraging the universal approximation capabilities of either the self-attention or feed-forward layers, we demonstrate that a basic gradient flow method can reliably find a global optimum, despite the highly non-convex landscape of the training objective.

We also highlight our novel contributions to Transformer theory through the development of these two core results:

The assumption on activation regularity conditions (Assumption 2) is less stringent compared to those usually found in studies of two-layer neural networks [51, 16, 28, 14] or deep ResNet networks [22, 23, 47]. In particular, many existing approximation guarantees reply on a Lipschitz continuity property of the network gradients, which limits the mean-filed study to neural networks with smooth activation functions. In comparison, our analysis relaxes this assumption and only requires local Lipschitz continuity of the gradient in expectation. This relaxation broadens the applicability of our approach and ensures that our result can cover more practical Transformer architectures.

Our model differs from the ResNet models in [47, 22, 23, 13], as those models incorporate only a single identical encoder within each evolutionary block. Unlike the typical theoretical configurations, our model employs two distinct encoders f and h that alternate throughout the network's depth. More importantly, despite the distinct encoders used, the continuous limit of our model uniformly interprets the encoder as an average of f and h , providing a rigorous validation of concepts proposed in [47] and [67] from a new perspective.

Our global convergence guarantee for training Transformer models is also broadly applicable: our assumption (Assumption 4) ensures global convergence by relying on the universal approximation capabilities of either the self-attention or the feed-forward encoder. Additionally, we incorporate a more flexible framework by adopting partial 1 -homogeneity for only a subset of the parameters, in contrast to the full parameter homogeneity required in studies such as [47]. This modification enables the use of softmax and sigmoid activation, expanding beyond the hardmax and ReLU restricted by full homogeneity.

Additional related works. See Appendix B for a detailed discussion.

Notations. For any α ∈ R d , dim( α ) refers to its dimension d . For any B ∈ R d × d , its trace is denoted by Tr( B ) . For any positive integer n , Let [ n ] = { 1 , 2 , . . . , n } . Let 0 d denote the d -dimension vector of all zeros. Let W p ( µ, ν ) denote the Wassersteinp distance between two probability measures µ, ν ∈ P ( R d ) for p ≥ 1 . For a matrix A = ( a 1 , a 2 , . . . , a n ) , define its vectorization version as vec[ A ] := ( a ⊤ 1 , a ⊤ 2 , . . . , a ⊤ n ) ⊤ . Let δ ( · ) denote the Dirac mass and 1 {·} be the indicator function. Let supp( · ) denote the support of any distribution. Let ∥·∥ = ∥·∥ 2 denote the l 2 norm and ∥·∥ max denote the maximum norm. For any subsets D 1 , D 2 in Euclidean space, define C ( D 1 , D 2 ) as the collection of functions that map D 1 to D 2 and are continuous over D 1 . Define the Bounded Lipschitz norm for any measure µ ∈ M ( R d ) as ∥ µ ∥ BL := sup { ∫ fdµ : f : R d → R , sup | f |≤ 1 , f is 1 -Lipschitz } .

2 Transformer model

In this section, we describe our deep Transformer model with each data input as a sequence, and the gradient flow algorithm used for training.

2.1 Data setting

In our paper, the data input is both general and straightforward: an input sequence H ∈ R D × ( N +1) consisting of N + 1 tokens, each with dimension D . We consider the setting where each input sequence H is associated with a label y ( H ) ∈ R , where y ( H ) is the target function we aim to learn. Furthermore, we assume that each instance H is i.i.d. drawn from a population distribution µ .

Relation to in-context learning (ICL) Our data setting is versatile and applicable to any task involving sequential input. It particularly suits the in-context learning (ICL) scenario [6, 10, 75], where models are capable of making accurate predictions on new data when prompted with training examples from the same pool. For clarity, consider the input sequence H ∈ R D × ( N +1) formatted as follows:

Here, { x i } i ∈ [ N ] are the input vectors, each associated with a corresponding label { y i } i ∈ [ N ] . The last token, x N +1 is the test input for which a prediction is made. The third row contains the customized and fixed positional encoding vectors { p i } i ∈ [ N ] , which typically include ones, zeros, and indicators denoting the token for prediction. The label for the query point x N +1 is then given by y N +1 = y ( H ) in our terminology. ICL operates in a zero-shot fashion, without any updates to the model's parameters, highlighting a unique and powerful capability of these systems to adapt and generalize based on the provided context alone. In [6], the authors demonstrate that fixed Transformers can approximate in-context penalized generalized linear regression to any desired degree.

2.2 Model

We follow a common configuration of Transformer architectures [6, 38, 40, 48, 73] where each Transformer block consists of two distinct layers: a self-attention mechanism layer and a token-wise feed-forward neural network layer, both equipped with skip connections. We assume that both layers consist of the average of M heads, treated uniformly as the width across all blocks for simplicity. The formulation for a matrix input Z ∈ R D × ( N +1) and a given residual step size η > 0 is as follows: Each residual self-attention layer is represented by

and each residual feed-forward neural network layer is defined by

for parameter vectors θ and w in the Euclidean space. The encoders for the self-attention and feed-forward layers are denoted as f : R D × ( N +1) → R D × ( N +1) and h : R D × ( N +1) → R D × ( N +1) , respectively. The self-attention encoder f formulation, commonly adopting a multiplicative or dot-product approach as detailed in [8, 38, 48, 64, 66, 73], can be exemplified by

where W V , W K , W Q ∈ R s × D , and W O ∈ R D × s . This formulation can be reparametrized to

where V, W ∈ R D × D , θ = vec[ V, W ] . The activation σ A typically uses column-wise softmax, but component-wise ReLU is also viable, as in [6]. For the feed-forward layer, an example of the encoder is h ( Z, w ) = W 2 σ M ( W 1 Z ) , as detailed in [6, 38, 73], where w = vec[ W 1 , W 2 ] and the activation σ M is component-wise ReLU. Alternatively, setting h ≡ 0 results in a Transformer block that comprises only the self-attention layer, referred to as 'attention-only' Transformers, as discussed in [6, 46, 49, 66].

Next, we analyze a Transformer network composed of L Transformer blocks, referring to L as the depth of the model. In this paper, we introduce an additional term, η , in (2.1) and (2.2) to simulate the model's evolution in a residual manner. We set the step size η as ∆ t/ 2 , where ∆ t = 1 /L. As L increases, ∆ t approaches zero, allowing Transformer blocks to incrementally contribute to the model's overall progression. The structure of the network is then defined as follows:

for each t = 0 , ∆ t, . . . , ( L -1)∆ t with ̂ T Θ ( H, 0) = H . We abbreviate the subscript t = 0 , ∆ t, . . . , ( L -1)∆ t by t and j = 1 , 2 , . . . , M by j for simplicity. Here, Θ = { θ t,j , w t,j } t,j denotes all parameters in the Transformer model.

Throughout this paper, we treat D and N as bounded finite values, while M and L are treated as diverging, aligning with the setting of large-scale Transformers.

2.3 Gradient flow

For the l 2 regularization with λ > 0 , we consider training the constructed Transformer model using the following λ -regularized risk objective:

with the population squared risk function defined as

In Section 3.2, we will show that l 2 -regularization on the parameter norms is essential for the wellposedness of the (Wasserstein) gradient flow to control parameter growth under our mild assumptions, even with a very small λ > 0 . Similar strategies that consider necessary l 2 regularization are employed in [23] and [70]. Then, drawing on the methodologies in [6, 32, 46], our model processes the final output through a simple read-out function, Read[ · ] , extracting the ( d +1 , N +1) -th entry of its input. We propose that this read-out layer can be expanded to any linear mapping with bounded parameter norm without affecting the validity of our theoretical results.

To minimize the objective function (2.5), we implement the standard gradient flow method as follows:

Step 1. Initially, for each t = 0 , ∆ t, . . . , ( L -1)∆ t , we sample M particles θ (0) t,j , w (0) t,j with j ∈ [ M ] independently from ρ 0 ( θ, w | t ) , where ρ 0 is a pre-defined distribution with bounded support.

Step 2. Then, we update all parameters θ ( τ ) t,j , w ( τ ) t,j in the set Θ ( τ ) = { θ ( τ ) t,j , w ( τ ) t,j } t,j using gradient flow (scaled by ML ), which is defined as follows:

Define the function ̂ R ( H ; Θ) = 1 2 (Read[ ̂ T Θ ( H, 1)] -y ( H )) 2 , and the partial derivative ̂ p Θ ( H,t ) = ∂ ̂ R ( H ; Θ) /∂ ̂ T Θ ( H,t ) ⊤ for each t = 0 , ∆ t/ 2 , ∆ t, . . . , ( L -1)∆ t, ( L -1 / 2)∆ t . Refer to Appendix C.4 for the explicit formula of ̂ p Θ ( H,t ) . Using the chain rule, we derive the explicit form of the gradient flow as follows:

where

3 Approximation by the mean-field limit

In this section, we present a rigorous approximation result that bridges Transformer models in (2.4) with their mean-field limit as continuous Transformers. Thus, the width M and depth L in our proposed model are treated as discretization of this continuous limit in the parameter space.

3.1 Assumptions

In addition, we introduce the norm ∥·∥ 2 -col as the maximum l 2 norm across all columns of a matrix. We proceed under several mild assumptions related to the data distribution and the encoders f and h .

Assumption 1 (Data regularity) . There exists some universal constant B > 0 such that, for any H ∈ supp( µ ) , we have max {∥ H ∥ 2 -col , y ( H ) } ≤ B . In addition, a universal constant K y > 0 ensures that y ( H ) is K y -Lipschitz continuous for ∥·∥ F over H ∈ supp( µ ) .

Remark Assumption 1 is irrelevant to the Transformer model, and is only a fairly mild assumption on the data.

Assumption 2 (Transformer particle growth bound) . We assume that the gradient of f ( T, θ ) and h ( T, w ) exists. Furthermore, we have

for some continuous, monotonically increasing functions ϕ P , ϕ T for every coordinate, and a universal constant K > 0 . Similarly, if we replace f with h and θ with w , the same conditions apply.

Remark There are three key observations for Assumption 2. Firstly, it incorporates the ∥·∥ 2 -col norm, which is particularly useful for handling sequential inputs where each column represents a token. Secondly, as we consider higher-order multiplications between data and parameters, this assumption accommodates a broader range of self-attention encoders, such as the one in (2.3) with softmax or ReLU activation (where the derivative is defined as ReLU ′ ( x ) = 1 { x > 0 } ). Lastly, a particularly interesting and frontier question is identifying the function ϕ T , and we have listed related literature in Appendix B.

Assumption 3 (Locally Lipschitz continuous gradient in expectation) . Besides Assumption 2, for any L T > 0 and any L T -Lipschitz continuous functions T 1 = T 1 ( H ) and T 2 = T 2 ( H ) , for every i ∈ [ N +1] , we have

for K T = max { sup H ∥ T 1 ∥ 2 -col , sup H ∥ T 2 ∥ 2 -col } , K P = max {∥ θ ∥ , ∥ θ ′ ∥} , and some continuous functions ϕ PT , ϕ TP , ϕ PP , ϕ TT that are monotonically increasing for every coordinate. Similarly, if we replace f with h and θ with w , the same conditions apply.

Remark Assumption 3 states that functions are locally Lipschitz continuous in expectation, suitable for encoders that utilize ReLU functions and have second-order derivatives almost everywhere. This assumption is naturally satisfied if the activation has a locally Lipschitz continuous gradient.

Define P 2 as the set of probability measures endowed with the Wasserstein2 distance, where the Lipschitz continuity with respect to the depth holds, i.e. there exists some universal constant C ρ > 0 such that ∥ ρ ( · , t ) -ρ ( · , t ′ ) ∥ BL ≤ C ρ | t -t ′ | for any t, t ′ ∈ [0 , 1] .

Choice of ρ 0 Suppose ρ 0 ∈ P 2 satisfies that for any t ∈ [0 , 1] , the support of ρ 0 ( · , · , t ) is contained within the set { ( θ, w ) : ∥ θ ∥ 2 + ∥ w ∥ 2 ≤ R 2 } for a universal constant R . Additionally, for each t ∈ [0 , 1] , it holds that ∫ θ,w ρ 0 ( θ, w, t ) d ( θ, w ) = 1 . This condition suits common bounded support distributions, and a natural choice is a uniform distribution across a disk with radius R for each t ∈ [0 , 1] .

We would like to clarify that verifying Assumptions 2 and 3 for concrete examples of Transformer architectures with smooth activation functions is fairly intuitive, and the proof is mainly based on a series of tedious calculations. We give a concrete proposition with its brief proof in Appendix G.

3.2 Continuous Transformer and Wasserstein gradient flow

Drawing inspiration from [47] and [67], which suggest that deep residual networks behave like ensembles of residual networks locally, we apply a similar manipulation to formulate the continuous version of (2.4). Consider the following continuous version T ρ ( H,t ) ∈ R D × ( N +1) , governed by the following continuous ODE that averages the two encoders :

In (3.1), each encoder f or h is conceptualized as a particle, and we consider the distribution of these particles denoted as ρ ( θ, w, t ) . For any ρ ∈ P 2 that have a bounded support, the well-posedness of T ρ ( H,t ) that satisfies the Transformer ODE (3.1) is shown in Proposition C.1. Transitioning to the framework with continuous Transformers, our objective shifts to minimizing the l 2 risk function with regularization on the second moment of ρ as follows:

with

Define p ρ ( H,t ) ∈ R D × ( N +1) , the partial derivative of R ( ρ ) relative to T ρ ( H,t ) at a local query point H , as the solution derived in Appendix C.4 using the classical adjoint sensitivity method [58]:

Using this, we can compute the functional derivative to ρ as follows:

The following Proposition claims that δQ δρ is indeed the derivative with respect to ρ (specifically, the Fréchet derivative [30]) for the functional Q ( ρ ) .

Proposition 3.1 (Functional derivative to ρ ) . Under Assumptions 1 and 2, for any pair ρ, ν ∈ P 2 that have bounded supports, we have

where δQ δρ is defined in (3.4) , and ⟨ δQ δρ , ν -ρ ⟩ = ∫ 1 0 ∫ ( θ,w ) δQ δρ · ( ν -ρ ) d ( θ, w ) dt ∈ R .

Now, we are in a position to display the gradient flow of ρ in the Wasserstein metric [16], given by a McKean-Vlasov type equation [4, 37, 54, 56]. Specifically, we study the following partial differential equation of the distribution ρ ( τ ) ( θ, w, t ) :

where ρ (0) = ρ 0 , div is the divergence operator, and the gradient functions are defined as

Propositions D.1 and 3.2 provide the well-posedness of both gradient flow and Wasserstein gradient flow respectively. In both propositions, a λ > 0 is essential to stabilize the optimization process by controlling both the maximum and average norms across all parameters. If λ is set to 0 , it is only possible to establish the well-posedness of (3.5) over a finite maximal interval [47]. Similar adjustments to regularize the risk function are also noted in [23].

Proposition 3.2 (Existence and uniqueness of Wasserstein gradient flow) . Under Assumptions 1 and 2, there exists a unique solution ( ρ ( τ ) ) τ ≥ 0 ∈ P 2 × R with ρ (0) = ρ 0 for (3.5) . Additionally, for any τ ≥ 0 , we have i. ρ ( τ ) has a bounded support { θ, w : ∥ θ ∥ 2 + ∥ w ∥ 2 ≤ R τ } × [0 , 1] , where R τ = ( R + 1) exp( C R τ ) -1 for some constant C R that only depends on N,D,λ and the parameters of the assumptions.

ii. ∫ 1 0 ( ∥ θ ∥ 2 + ∥ w ∥ 2 ) ρ ( τ ) ( θ, w, t ) d ( θ, w ) dt ≤ A 2 0 , where A 0 := R 2 + λ -1 (2 B 2 +2 B 2 exp( K (1+ R + R 2 )) 2 ) .

3.3 Approximation of large-scale Transformer

In this section, we discuss the general results associated with approximating our discrete Transformer model to its mean-field limit. First, we highlight that the minimization of the risk function with discretization, whether or not regularization is included, closely approximates the minimal risk achievable by continuous models.

Proposition 3.3 (Global minimum approximation of discretization) . Under Assumptions 1 and 2, we define P 2 ,r as the set of distributions in P 2 concentrated on { ( θ, w ) : ∥ θ ∥ 2 + ∥ w ∥ 2 ≤ r 2 } × [0 , 1] . for any r > 0 . Then there exists a constant C dependent on N,D,r and the parameters of the assumptions such that

Proposition 3.3 specifies that the distributions under consideration must have bounded support. While it is typically challenging to confirm whether the minimal risk is indeed achieved on a distribution with bounded support, this assumption is justified as λ regulates parameter norms, implicitly encourages solutions residing in a compact region of the parameter space.

We now present the main theorem concerning the convergence of the gradient flow process to the Wasserstein gradient flow as outlined in (3.5). The proof with detailed explanation of the techniques used in Theorem 3.1 is provided in Appendix D.

Theorem 3.1 (Gradient flow approximation of discretization) . Define the empirical distribution as ˆ ρ ( τ ) := 1 ML ∑ t ∑ M j =1 δ ( θ ( τ ) t,j , w ( τ ) t,j , t ) for any τ ≥ 0 . Under Assumptions 1-3, we

have that (ˆ ρ ( τ ) ) τ ≥ 0 weakly converges to ( ρ ( τ ) ) τ ≥ 0 almost surely along any sequence such that L → ∞ , M/ log L → ∞ . Moreover, for any fixed τ > 0 and any δ > 0 , with probability at least 1 -3 exp( -δ ) with respect to the parameter initialization Θ (0) , we have

for some constant C that depends on on N,D,τ,λ and the parameters of the assumptions.

Theorem 3.1 significantly advances our understanding by controlling the difference regarding both the Transformer output, the risk function, and the regularized risk function. It's noted that the difference bound in the model's approximation may increase, possibly exponentially [22, 23, 51], as the time horizon extends. As argued in [51], such behavior may be inherent to the systems being modeled.

Additionally, the technical uniqueness and innovation of this theorem contrast sharply with previous results from overparametrized ResNet models. Our analysis distinguishes itself in two ways. First, our discrete Transformer model (2.4) uniquely splits the averaged encoder ( f + h ) / 2 into two distinct blocks with encoders f and h . Second, we demonstrate uniform error control over any finite time interval [0 , τ ] , enabling continuous monitoring of maximum error across the gradient flow's trajectory. In contrast, models in prior studies such as [22, 23] restricts the error analysis to a specific s ∈ [0 , τ ] .

4 Global convergence of gradient flow

In this section, we explore the optimization problem for gradient flow in the context of the discrete Transformer model, focusing on our general global convergence results.

4.1 An additional assumption

To ensure the global convergence of gradient flow for our discrete Transformer model, we introduce the following assumption. While influenced by the work in [16, 22, 23, 47], our assumption is uniquely tailored to the context of Transformers:

Assumption 4. There exists a pair ( g, α ) ∈ { ( f, θ ) , ( h, w ) } with a partition α = ( α 1 , α 2 ) such that

We emphasize that the universal kernel property, as discussed in [52], closely relates to the universal approximation abilities. Under our assumption, we require the universal approximation capabilities of either the self-attention encoder or the feed-forward encoder. In Appendix G, we provide a concrete example of Transformer architectures and verify the validity of Assumption 4.

The universal kernel property of the feed-forward layer encoder h is well-established, particularly in two-layer neural network contexts [74]. Conversely, the universal approximation abilities of self-attention layers is a frontier research area, which, while not extensively covered in this paper, holds significant potential. Often labeled as 'memorization capacity", this area is recently explored across multiple studies [27, 31, 38, 39, 49, 63, 73]. The interconnection between approximation abilities and memorization capacities is established in [38]. Notably, [49] investigated the expressive capabilities of one single multi-head softmax self-attention layer, thereby potentially validating our assumptions.

Finally, we posit that the universal kernel applies to α 2 within a compact set, as the function's scale can be moderated by the homogeneous part α 1 . In scenarios where α 2 and K are absent, our

assumption simplifies to that in [47], characterized by complete homogeneity. Conversely, in the absence of the α 1 component, our framework aligns with [23] which necessitates a more stringent support condition for K , as detailed later in Theorem 4.1.

4.2 Global convergence result

In this section, we establish the convergence properties of the optimization task for discrete Transformers through gradient flow dynamics.

Theorem 4.1 (Global convergence up to λ ) . Suppose that Assumptions 1-4 hold, and the Wasserstein gradient flow ( ρ ( τ ) ) τ ≥ 0 weakly converges to some ρ ∞ ∈ P 2 . If for some universal constant R ∞ > 1 , the following two conditions hold:

( ρ ( τ ) ) τ ≥ 0 is concentrated on { θ, w : ∥ θ ∥ 2 + ∥ w ∥ 2 ≤ R 2 ∞ } × [0 , 1] when τ is sufficiently large.

If Assumption 4 holds with ( g, α ) = ( f, θ ) , we assume there exists a t ∗ ∈ [0 , 1] such that the connected set supp( ρ ∞ ( · , t ∗ )) ⊃ D×K×{ w 0 } , for some w 0 ∈ R dim( w ) and D ⊂ R dim( θ 1 ) that separates { θ 1 : ∥ θ 1 ∥ = 1 /R ∞ } and { θ 1 : ∥ θ 1 ∥ = R ∞ } .

ii ′ . If Assumption 4 holds with ( g, α ) = ( h, w ) , we assume there exists a t ∗ ∈ [0 , 1] such that the connected set supp( ρ ∞ ( · , t ∗ )) ⊃ ×{ θ 0 } × K × D , for some θ 0 ∈ R dim( θ ) and D ⊂ R dim( w 1 ) that separates { w 1 : ∥ w 1 ∥ = 1 /R ∞ } and { w 1 : ∥ w 1 ∥ = R ∞ } .

Then, for any ϵ > 0 , there exists some τ 0 > 0 such that

with probability at least 1 -3 exp( -δ ) with respect to the parameter initialization Θ (0) for any δ > 0 . Here, C 1 is some constant dependent only on N,D,τ 0 , λ and the parameters of the assumptions, while C 2 depends only on N,D,R ∞ and the parameters of the assumptions.

Theorem 4.1 depicts the behavior of the risk function ̂ R (Θ ( τ ) ) as the training duration τ is sufficiently large. Specifically, ̂ R (Θ ( τ ) ) asymptotically approaches zero as both L →∞ and M/ log L →∞ , with an additional term that scales with λ . This additional term attributes to the incorporation of a λ -weighted penalty on the norm of the parameters in our training objective ̂ Q . Consequently, by selecting an appropriately small λ > 0 , the risk approximates zero, demonstrating global convergence to the minimum of ̂ R .

In addition, Theorem 4.1 posits some additional assumptions: the weak convergence of ρ ( τ ) , the longtime uniform boundedness, and the separation property for α 1 with the support expansion of α 2 to K . Similar assumptions are made in the literature of deep model optimization theory [22, 23, 47]. While these types of assumptions are typically challenging to justify, we provide high-level justifications for them in Appendix C.5, deferring detailed verification to future research.

We then present a corollary that directly follows from Theorem 4.1:

Corollary 4.1. Continuing with the notations and assumptions from Theorem 4.1, suppose λ ≤ C λ ϵ for some universal constant C λ > 0 . Then, for any δ > 0 , constants τ 0 , L 0 , K 0 > 0 can be found such that:

with probability at least 1 -3 exp( -δ ) with respect to the parameter initialization Θ (0) . Notably, if C λ ≤ (2 C 2 ) -1 , the upper bound in (4.1) is less than or equal to ϵ .

Corollary 4.1 claims that with a fixed δ > 0 , for any ϵ > 0 , we can achieve an order of ϵ -close approximation with sufficiently large L and M . Though our result is asymptotic and does not involve an explicit rate, it is the first of its kind and lays the groundwork for future theoretical optimization guarantees for Transformers.

5 Proof ideas of main theorems

Given the technical nature of this paper, this section presents the key ideas behind the proof of our main novel results, along with an outline of the proof preparation.

Idea for Theorem 3.1 This convergence is described in two parts. First, the finite-time result (points (i)-(iii)) uses propagation of chaos [62] to analyze how differences evolve over time, comparing the evolution of parameter particles in discrete and continuous dynamics. The approximation bound is derived using a third auxiliary dynamic ("nonlinear dynamics"), involving the triangle inequality and Grönwall's inequality, which allows us to bound output differences over time.

Second, weak convergence of the empirical distribution process relies on optimal transport theory and stability results for Wasserstein gradient flows [4], focusing on the convergence of momentum fields [4, 60]. This also requires bounding the gradient differences between discrete and continuous Transformers as they approach the mean-field limit. See Appendix D for a detailed illustration, including a description of each main step.

Idea for Theorem 4.1 We first establish the continuity of the functional gradient δQ δρ | ρ ∞ , ensuring it remains constant if the derivative with respect to β is constant over a region. Next, we derive the key bound for Q ( ρ ∞ ) , which is proportional to λ , by analyzing the functional energy Q 's landscape through its derivatives.

Finally, we show that the finite-time risk can approach this bound. Achieving ϵ -level loss requires Q ( ρ ( τ 0 ) ) ≤ ϵ for some large τ 0 . Applying Theorem 3.1, we show ̂ Q ( τ 0 ) becomes sufficiently small, and since ̂ Q ( ρ ( τ ) ) is non-increasing, it remains small for τ ≥ τ 0 . See Appendix E for a detailed illustration, including a description of each main step.

Proof preparation for the main theorems Appendix C.3 lists several useful lemmas essential to the main results. Specifically, Lemmas C.1-C.6 ensure the boundedness of key components and bound the output differences between discrete and continuous Transformers under different parameter settings. This boundedness is non-trivial due to the mild Assumptions 2 and 3 that fit the Transformer architecture. The technical lemmas in Appendix C.3 form the foundation for all subsequent proofs. Before introducing the nonlinear dynamics used to bound parameter differences under non-i.i.d. settings, these lemmas first establish an important oracle approximation bound result (Lemma D.6) with i.i.d. parameter settings. Additionally, they serve as key tools for bounding the (functional) gradient differences between Transformer dynamics, as shown in Lemmas D.1-D.3, which are essential for proving the approximation bound in Theorem 3.1.

In Theorem 4.1, Lemma E.1 plays a key role by demonstrating that for any ρ , there is always a nearly descent direction around ρ for Q ( ρ ) , implying that all local minima are nearly global. This motivates further landscape analysis for bounding δQ δρ | ρ ∞ in the main theorem.

6 Conclusion

We conclude by summarizing our key contributions and suggesting future research directions. This paper establishes the global convergence of large-scale Transformer models through gradient flow dynamics, providing a thorough theoretical foundation. Our analysis, focused on the mean-field limit with infinite width and depth, shifts optimization from parameter space to distributional probability measures. We present two main theorems: one confirming the close approximation between discrete and continuous gradient flows, and another demonstrating global convergence, highlighting that basic optimization methods can successfully navigate complex landscapes to find optimal solutions. The techniques and results from this study lay the groundwork for further exploration into Transformer optimization. Future work could explore direct gradient descent with specific focus on step sizes, and expand on the in-context learning approximation capabilities of Transformers, as initiated by [6]. Additionally, it's crucial to rigorously assess under what conditions can self-attention layers serve as universal kernels to enhance our theoretical understanding, and to determine the generalization error bounds of Transformers trained on finite samples. These directions promise to deepen the theoretical and practical insights into Transformer models.

Acknowledgments and Disclosure of Funding

We thank the anonymous reviewers for their helpful comments. Yuan Cao is partially supported by NSFC 12301657 and Hong Kong RGC-ECS 27308624. Mengdi Wang acknowledges the support by NSF IIS-2107304, NSF CPS-2312093, ONR 1006977 and Genmab. Han Liu's research is partially supported by the NIH R01LM01372201. Jason M. Klusowski was supported in part by the National Science Foundation through CAREER DMS-2239448, DMS-2054808 and HDR TRIPODS CCF-1934924. Jianqing Fan's research was partially supported by NSF grants DMS-2210833, DMS-2053832, and ONR grant N00014-22-1-2340.

References

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations , 2022.

Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in Neural Information Processing Systems , 2019.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning , pages 242-252, 2019.

Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. In metric spaces and in the space of probability measures. In Gradient Flows , 2005.

Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning , pages 322-332, 2019.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Workshop on Efficient Systems for Foundation Models @ ICML2023 , 2023.

Raphaël Barboni, Gabriel Peyré, and François-Xavier Vialard. On global convergence of resnets: From finite to infinite width using linear parameterization. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

James Bernhard. Alternatives to the scaled dot product for attention in the transformer neural network architecture, 2023.

Blake Bordelon, Hamza Tahir Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head transformer dynamics, 2024.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 1877-1901. Curran Associates, Inc., 2020.

Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems , 2019.

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems , 31, 2018.

Yihang Chen, Fanghui Liu, Yiping Lu, Grigorios Chrysos, and Volkan Cevher. Generalization of scaled deep resnets in the mean-field regime. In The Twelfth International Conference on Learning Representations , 2024.

Zixiang Chen, Yuan Cao, Quanquan Gu, and Tong Zhang. A generalized neural tangent kernel analysis for two-layer neural networks. arXiv: Learning , 2020.

Jingpu Cheng, Qianxiao Li, Ting Lin, and Zuowei Shen. Interpolation, approximation and controllability of deep neural networks. arXiv preprint arXiv:2309.06015 , 2023.

Lénaïc Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. In Proceedings of the 32nd International Conference on Neural Information Processing Systems , NIPS'18, page 3040-3050, Red Hook, NY, USA, 2018. Curran Associates Inc.

Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems , 2019.

George V. Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of Control, Signals and Systems , 2:303-314, 1989.

George Dasoulas, Kevin Scaman, and Aladin Virmaux. Lipschitz normalization for selfattention layers with application to graph neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 2456-2466. PMLR, 18-24 Jul 2021.

Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xiao Wang, Carlos Riquelme Ruiz, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd Van Steenkiste, Gamaleldin Fathy Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Collier, Alexey A. Gritsenko, Vighnesh Birodkar, Cristina Nader Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetic, Dustin Tran, Thomas Kipf, Mario Lucic, Xiaohua Zhai, Daniel Keysers, Jeremiah J. Harmsen, and Neil Houlsby. Scaling vision transformers to 22 billion parameters. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 7480-7512. PMLR, 23-29 Jul 2023.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In North American Chapter of the Association for Computational Linguistics , 2019.

Zhiyan Ding, Shi Chen, Qin Li, and Stephen Wright. On the global convergence of gradient descent for multi-layer resnets in the mean-field regime, 2021.

Zhiyan Ding, Shi Chen, Qin Li, and Stephen Wright. Overparameterization of deep resnet: Zero loss and mean-field analysis. Journal of Machine Learning Research , 23:48-1, 2022.

Alexey Dosovitskiy. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning , pages 1675-1685, 2019.

Weinan E, Jiequn Han, and Qianxiao Li. A mean-field optimal control formulation of deep learning. Research in the Mathematical Sciences , 6(1):1-41, 2019.

Benjamin L. Edelman, Surbhi Goel, Sham M. Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms, 2022.

Cong Fang, Hanze Dong, and Tong Zhang. Over parameterized two-level neural networks can learn near optimal feature representations. ArXiv , abs/1910.11508, 2019.

Cong Fang, Yihong Gu, Weizhong Zhang, and Tong Zhang. Convex formulation of overparameterized deep neural networks. IEEE Transactions on Information Theory , 68(8):5340-5352, 2022.

Bela A. Frigyik, Santosh Srivastava, and Maya R. Gupta. Introduction to functional derivatives. UWEETechReport 2008-0001, University of Washington Department of Electrical Engineering, 2008.

Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? a study through the random features lens. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? a case study on learning with representations. In The Twelfth International Conference on Learning Representations , 2024.

Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249 , 2023.

Ken ichi Funahashi. On the approximate realization of continuous mappings by neural networks. Neural Networks , 2:183-192, 1989.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems , pages 8571-8580, 2018.

Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 37822-37836. Curran Associates, Inc., 2022.

Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokkerplanck equation. SIAM Journal on Mathematical Analysis , 29, 04 2000.

Tokio Kajitsuka and Issei Sato. Are transformers with one layer self-attention using low-rank weight matrices universal approximators?, 2024.

Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention, 2021.

Junghwan Kim, Michelle Kim, and Barzan Mozafari. Provable memorization capacity of transformers. In The Eleventh International Conference on Learning Representations , 2023.

Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. In Forty-first International Conference on Machine Learning , 2024.

Qianxiao Li, Long Chen, Cheng Tai, and Weinan E. Maximum principle based algorithms for deep learning. Journal of Machine Learning Research , 18(165):1-29, 2018.

Qianxiao Li and Shuji Hao. An optimal control approach to deep learning and applications to discrete-weight neural networks. In International Conference on Machine Learning , pages 2985-2994. PMLR, 2018.

Qianxiao Li, Ting Lin, and Zuowei Shen. Deep learning via dynamical systems: An approximation perspective. Journal of the European Mathematical Society , 25(5):1671-1709, 2022.

Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In International Conference on Machine Learning , pages 19689-19729. PMLR, 2023.

Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. In The Twelfth International Conference on Learning Representations , 2024.

Yiping Lu, Chao Ma, Yulong Lu, Jianfeng Lu, and Lexing Ying. A mean field analysis of deep ResNet and beyond: Towards provably optimization via overparameterization from depth. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 6426-6436. PMLR, 13-18 Jul 2020.

Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention-based neural machine translation. In Lluís Màrquez, Chris Callison-Burch, and Jian Su, editors, Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing , pages 1412-1421, Lisbon, Portugal, September 2015. Association for Computational Linguistics.

Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multihead attention in transformers. In The Twelfth International Conference on Learning Representations , 2024.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory , 2019.

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665E7671, 2018.

Charles A. Micchelli, Yuesheng Xu, and Haizhang Zhang. Universal kernels. J. Mach. Learn. Res. , 7:2651-2667, 2006.

Atsushi Nitanda and Taiji Suzuki. Stochastic particle gradient descent for infinite ensembles, 2017.

Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Particle dual averaging: optimization of mean field neural network with global convergence rate analysis*. Journal of Statistical Mechanics: Theory and Experiment , 2022(11):114010, nov 2022.

OpenAI. Gpt-4 technical report. ArXiv , abs/2303.08774, 2023.

Felix Otto. The geometry of dissipative evolution equations: The porous medium equation. Communications in Partial Differential Equations , 26(1-2):101-174, 2001.

Allan Pinkus. Approximation theory of the mlp model in neural networks. Acta Numerica , 8:143 - 195, 1999.

L. S. Pontryagin, V. G. Boltyanskii, R. V . Gamkrelidze, and E. F. Mishechenko. The mathematical theory of optimal processes. Zamm-zeitschrift Fur Angewandte Mathematik Und Mechanik , 43:514-515, 1963.

H. Risken. The Fokker-Planck Equation: Methods of Solution and Applications . Springer, 1996.

Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY , 55(5863):94, 2015.

Filippo Santambrogio. {Euclidean, metric, and Wasserstein} gradient flows: an overview. Bulletin of Mathematical Sciences , 7:87-154, 2016.

Alain-Sol Sznitman. Topics in propagation of chaos. In Ecole d'Eté de Probabilités de Saint-Flour XIX-1989 , pages 165-251. Springer, 1991.

Shokichi Takakura and Taiji Suzuki. Approximation and estimation ability of transformers for sequence-to-sequence functions with infinite dimensional input. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 33416-33447. PMLR, 23-29 Jul 2023.

Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention for transformer models. In International Conference on Machine Learning , 2020.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothé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. ArXiv , abs/2302.13971, 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł 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.

Andreas Veit, Michael J Wilber, and Serge Belongie. Residual networks behave like ensembles of relatively shallow networks. Advances in neural information processing systems , 29, 2016.

C. Villani. Optimal Transport, Old and New , volume 338 of Grundlehren der Mathematischen Wissenschaften . Springer-Verlag, 2008.

James Vuckovic, Aristide Baratin, and Rémi Tachet des Combes. A mathematical theory of attention. ArXiv , abs/2007.02876, 2020.

Colin Wei, J. Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets v.s. their induced kernel. In Neural Information Processing Systems , 2018.

E Weinan. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics , 1(5):1-11, 2017.

E Weinan, Chao Ma, and Lei Wu. Machine learning from a continuous viewpoint, i. Science China Mathematics , 63:2233 - 2266, 2019.

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? ArXiv , abs/1912.10077, 2019.

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM , 64(3):107115, 2021.

Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models incontext. In R0-FoMo:Robustness of Few-shot and Zero-shot Learning in Large Foundation Models , 2023.

Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes overparameterized deep ReLU networks. Machine Learning , Oct 2019.

A Overview of Appendix

The appendix is organized as follows:

Appendix B: Additional related works are discussed.

Appendix C:

-In Appendix C.1, additional notations and preliminary details are introduced.

-In Appendix C.2, we show the proof of Proposition C.1, concerning the existence and uniqueness of the continuous Transformer ODE (3.1).

-In Appendix C.3, useful lemmas for the main proofs are detailed.

-In Appendix C.4, the explicit formulas for p ρ ( H,t ) and ̂ p Θ ( H,t ) are explored via the adjoint sensitivity method.

-In Appendix C.5, high-level explanations are provided to substantiate the assumptions made in Theorem 4.1.

Appendix D: Includes proofs of main results from Section 3.

Appendix E: Includes proofs of main results from Section 4.

Appendix F: Includes proofs of all auxiliary technical results mentioned in Appendix C-E.

Appendix G: The verification of Assumptions 2-4 for a concrete example of Transformer architectures is provided.

Appendix H: We provide simple experiment results and discuss the widths and depths of Vision Transformers that guarantee convergence, achieving near-zero training loss and 100% training accuracy on the CIFAR-10 dataset.

B Additional related work

Theory of Transformers. Some very recent works have studied theoretical properties of Transformer models from different aspects. [75, 33] studied the in-context learning guarantees for single-layer Transformers to perform linear regression predictions after being trained with linear regression example tasks. [1, 6, 32] studied the in-context learning capability of Transformers through the function approximation point of view, and demonstrated that there exists Transformers with specific parameter configurations that can perform particular in-context learning tasks. [36, 45] investigated how single-layer Transformers can be trained to learn simple image models and topic models respectively.

The most closely related work to ours is [41], which is the only study we are aware of that addresses the general and universal in-context learning capability of large-scale Transformers through optimization dynamics. It shows that a two-layer MLP followed by a linear attention layer can approximate functions in a general Barron space sufficiently well as the Transformer width increases. Additionally, its corresponding mean-field dynamics, via Wasserstein gradient flow, converges to global minima for in-context feature learning. Another work exploring the mean-field limit of Transformers is [9], which examines the limit as the depth, key-query length, and number of heads increase to infinity.

In addition, we have noticed a lot of theoretical interest in identifying the optimal choice of ϕ T in Assumption 2, i.e., the Lipschitz constant of the Jacobian matrix of the self-attention term. For instance, [19] suggest that ϕ T can be bounded by √ N/D +poly( ∥ T ∥ F ) , where poly( · ) denotes a polynomial function. In the context of l 2 self-attention, [39] find ϕ T to be √ N log N/D , which notably does not depend on ∥ T ∥ F , and [69] demonstrate that for l 1 distance metrics in attention layers, ϕ T could be √ D log N .

Global convergence of fully connected neural networks. A line of recent works have studied the global convergence of (stochastic) gradient descent in training overparameterized neural networks in the mean-field regime [16, 51, 50, 70, 28, 29]. They consider the limit of the neural network as the width of the network at each layer goes to infinity, and models the limit of the network as a functional of the distribution of network parameters. A separate line of works also established the global convergence guarantees for training overparameterized neural networks in the 'neural tangent kernel' regime [35, 3, 25, 76, 17, 2, 5, 11], where the gradient descent training iterates are

asymptotically equivalent to the training iterates of kernel regression based on the neural tangent kernel.

Connection between ordinary differential equation models and infinite-depth ResNets. Our work is also closely related to the recent literature aiming to understand ResNets by analyzing their connections to ordinary differential equations [71, 12, 43, 42, 72, 26, 22, 47, 23, 44, 7, 15, 13]. Specifically, [71, 12, 44, 15] studied the approximation of flow-based networks via discrete networks. [43, 42, 72, 26, 22, 47, 23, 44, 7] studied the optimization of the infinite-depth and infinite-width ResNets. [13] studied the generalization properties of the ResNet trained in the mean-field regime.

C Proof setup

C.1 Additional technical notations

Define

Thus, δQ δρ could be expressed as

Additionally, we can combine G f with G h , and ̂ G f with ̂ G h to reformulate as

and

Remark 1. To facilitate the proof, we restate Assumptions 2 and 3 for g ( T, β ) . Under Assumption 2, the gradient of g ( T, β ) respect to T and β exists. Additionally, we have Under Assumption 2:

Under Assumption 3: For any L T > 0 and any L T -Lipschitz continuous functions T 1 = T 1 ( H ) and T 2 = T 2 ( H ) , for every i ∈ [ N +1] , we have

Verifying all these results above only needs the basic triangle inequality of general norms, so we omit the trivial proof. We will apply them directly throughout the proofs of results. Additionally, we omit writing L T for simplicity in the proof, as all functions applied to Assumption 3 will be Lipschitz continuous with some universally bounded Lipschitz constant.

Next, we introduce some additional technical notations. Denote the identity matrix with d -dimension as I d . Define the sample space Ω := R dim β × [0 , 1] , and P (Ω) as the probability measure space

defined on Ω . For any Θ = { β t,j } t/ ∆ t +1 ∈ [ L ] ,j ∈ [ M ] and ˜ Θ = { ˜ β t,j } t/ ∆ t +1 ∈ [ L ] ,j ∈ [ M ] , define d (Θ , ˜ Θ) = ∑ t ∑ M j =1 ∥ β t,j -˜ β t,j ∥ . Define the local risk function as

Define the nested family of compact subsets ( P r ) r> 0 as

For any ρ, ν ∈ P (Ω) and p ≥ 1 , define the l p distance ∥ ρ -ν ∥ p as

Specifically, when p = 1 , we have

for any ρ, ν ∈ P 2 concentrated on P r . For simplicity, any H discussed throughout this paper is assumed to lie within supp( µ ) .

C.2 Transformer ODE existence and uniqueness

In this section, we establish the existence and uniqueness of the solution T ρ ( H,t ) to the ODE presented in (3.1) for any H , given that ρ ∈ P 2 is concentrated on a bounded support, specifically, P r for some r > 0 . This following proposition forms the cornerstone of the subsequent technical analyses:

Proposition C.1 (Existence and uniqueness of Transformer ODE) . Under Assumptions 1 and 2, for any ρ ∈ P 2 that has a bounded support, there exists a unique solution of (3.1) on t ∈ [0 , 1] that is Lipschitz continuous with respect to ( H,t ) .

Initially, we demonstrate that the integral ∫ β ρ ( β, t ) dβ is bounded. According to the definition P 2 , it follows that

for any t, t ′ ∈ [0 , 1] . Integrating (C.4) over t 2 ∈ [0 , 1] obtains

For the remainder of the technical proof, we will employ (C.5) without additional elaboration.

Proof of Proposition C.1. Step I: Create a small neighboring area with local Lipschitz continuity

Consider any vector β such that ∥ β ∥ ≤ r . Define F ( T, t ) := ∫ β g ( T, β ) ρ ( β, t ) dβ . For T within the rectangle { T : ∥ T -H ∥ max ≤ δ } , where δ > 0 is bounded, both ∥ T ∥ 2 -col and ∥ T ∥ F are also bounded. Given Assumption 2 (i), g ( T, β ) is universally bounded by some constant K δ,r . Moreover,

under Assumption 2 (ii) and (iii), g ( T, β ) is Lipschitz continuity with some constant L δ,r . Hence, within the rectangle { T : ∥ T -H ∥ max ≤ δ } × [0 , 1] the following properties hold:

which indicates that F ( T, t ) is continuous with respect to t within the rectangle { T : ∥ T -H ∥ max ≤ δ } × [0 , 1] .

Moreover, within the bounded region { T : ∥ T -H ∥ max ≤ δ } , Assumption 2 (iii) ensures that ∥∇ vec[ T ] vec[ g ( T, β )] ∥ 2 bounded. Consequently, g ( T, β ) is Lipschitz-continuous with respect to T for ∥·∥ F . Denote this Lipschitz constant by L ′ δ,r . Therefore, for any T, T ′ ∈ { T : ∥ T -H ∥ max ≤ δ } and t ∈ [0 , 1] , it follows that

which deduces the Lipschitz continuity of F ( T, t ) with respect to T for ∥·∥ F .

Step II: Show that the maximal existence interval is infinite by repeatedly using the PicardLindelöf Theorem

Invoking the Picard-Lindelöf Theorem, there exists some ϵ > 0 such that the initial value problem

has a unique solution on t ∈ [0 , ϵ ] . Given that this claim holds for any H , the standard ODE Extensibility Theorem guarantees a continuation of T ( t ) to a maximal interval of existence, denoted as [0 , t max ] .

Assume by contradiction that t max < 1 . From (3.1) and Assumption 2(i), for any t ∈ [0 , t max ] we see that

Therefore, by the Grönwall's inequality, we have

This presents a contradiction to the notion that t max < 1 . This is because, By reapplying the local Picard-Lindelöf Theorem using the state T ( H,t max ) as the new initial condition, we can extend the interval of existence beyond t max . Consequently, we must conclude that t max = 1 , and the existence and uniqueness follows.

Step III: Show that the Lipschitz continuity with respect to ( H,t )

In the final part of our proof, we demonstrate that T ρ ( H,t ) is Lipschitz continuous with respect to ( H,t ) for H ∈ supp( µ ) and any t ∈ [0 , 1] . Given that T ρ ( H,t ) is universally bounded within H ∈ supp( µ ) and any t ∈ [0 , 1] , we only need to focus on establishing its Lipschitz continuity with respect to H and t separately. The Lipschitz continuity with respect to t is derived from

for any t 1 , t 2 ∈ [0 , 1] . Given Assumption 2 (iii), we have

for any H,H ′ ∈ supp( µ ) . Define L H := ϕ T ( N,D,B exp( K (1 + C ρ / 2 + r + r 2 )))(1 + r + r 2 ) . Utilizing Grönwall's inequality, we establish:

given that T ρ ( · , 0) serves as the identity mapping. Consequently, T ρ ( H,t ) demonstrates Lipschitz continuity with respect to H ∈ supp( µ ) .

C.3 Useful technical lemmas

Lemma C.1 (Continuous Transformer output bound) . Under Assumption 2, for any distribution ρ ∈ P (Ω) where ∫ 1 0 ∫ β || β || 2 2 ρ ( β, t ) dβdt ≤ A 2 for some constant A > 0 and for any t ∈ [0 , 1] , we have

Lemma C.2 (Continuous Transformer difference bound) . Under Assumption 2 and given H , for any ρ, ν ∈ P 2 that satisfy ∫ 1 0 ∫ β || β || 2 2 ρ ( β, t ) dβdt ≤ A 2 and have bounded supports P r for some constants A,r > 0 , we have that

Here, the universal constant C r only depends on N,D,r,A , and the parameters of the assumptions. Lemma C.3 (Continuous Transformer gradient component bound) . Under Assumption 1 and 2, for any ρ ∈ P (Ω) where supp( ρ ) ⊂ P r and ∫ 1 0 ∫ β ∥ β ∥ 2 ρ ( β, t ) dβdt ≤ A 2 , we have

Lemma C.4 (Discrete Transformer bound) . Under Assumptions 2, for any Θ where 1 ML ∑ t ∑ M j =1 ∥ β ∥ 2 ≤ A 2 for some universal constant A > 0 and at any t = 0 , ∆ t/ 2 , ∆ t, . . . , ( L -1 / 2)∆ t, 1 , we have

Lemma C.5 (Discrete Transformer difference bound) . Under Assumption 1 and 2, for any H , let Θ = { β t,j } t/ ∆ t +1 ∈ [ L ] ,j ∈ [ M ] , ˜ Θ = { ˜ β t,j } t/ ∆ t +1 ∈ [ L ] ,j ∈ [ M ] such that max {∥ β t,j ∥ , ∥ ˜ β t,j ∥} ≤ r for any t = 0 , . . . , ( L -1)∆ t, j = 1 , . . . , M . Then, we have that

Here the universal constant C r only depends on N,D,r , and the parameters of the assumptions.

LemmaC.6 (Discrete Transformer gradient component bound) . Under Assumption 1 and 2, for any Θ such that sup t,j ∥ β t,j ∥ 2 ≤ r 2 and 1 ML ∑ t ∑ M j =1 ∥ β ∥ 2 ≤ A 2 , we have, for any t = 0 , ∆ t/ 2 , . . . , ( L -1 / 2)∆ t, 1

where B T = B exp( K (1 + A + A 2 )) .

Lemma C.7 (Norm average concentration) . Under Assumption 2, consider a parameter setting Θ = { β t,j } t/ ∆ t +1 ∈ [ L ] ,j ∈ [ M ] i.i.d. drawn from { ρ ( β | t ) } t/ ∆ t +1 ∈ [ L ] ,j ∈ [ M ] where ρ ∈ P 2 is concentrated on P r and satisfies ∫ β ρ ( β, t ) dβ = 1 for every t ∈ [0 , 1] . Then, with probability at least 1 -exp( -δ ) with respect to the parameter initialization Θ (0) , we have

for any δ > 0 . Here, the ≲ notation hides the dependencies on r and the parameters specified in the assumption.

Lemma C.8 (Matrix product difference bound) . Suppose that for some d > 0 , the matrices A 1 , A 2 , . . . , A L and B 1 , B 2 , . . . , B L satisfy the following conditions:

For each l = 1 , . . . , L , the norms of the matrices are bounded as ∥ A l ∥ ≤ 1 + a l , ∥ B l ∥ ≤ 1 + b l , where a l , b l > 0 .

The product of the increments for each matrix is bounded by ∏ L l =1 1 + max { a l , b l } ≤ C for some constant C > 0 .

Under these conditions, it holds that

C.4 Solution of adjoint ODE

In this section, we define the partial derivative

without specifying its explicit formula. Denote the derivative of T ρ ( H, 1) to T ρ ( H,t ) (after vectorization) by the Jacobian J ρ ( H,t ) ∈ R ( N +1) D × ( N +1) D , and assume that ˙ Jρ ( H,t ) exists for any t ∈ [0 , 1] . Then [58] shows that J ρ ( H,t ) satisfies the adjoint equation of the ODE.

for any t ∈ [0 , 1] . By applying the chain rule and exchanging the order of the derivative and integral, we have, for any t ∈ [0 , 1] , that

Hence, by taking the derivative with respect to t , we obtain that

with the solution

On the other hand, we have

where E read is a D × ( N +1) zero matrix except 1 at the ( d +1 , N +1) -th entry. Moreover, from (C.12) and (C.13), we see that

Additionally, we could explicitly derive the formula for ̂ p Θ ( H,t ) = ∂ ̂ R ( H ;Θ) ∂ ̂ T Θ ( H,t ) for the discrete Transformer. By applying the chain rule multiple times across each layer with the encoder either f or h , for any t = 0 , ∆ t, . . . , ( L -1)∆ t, 1 , we have

and

C.5 Explanation for assumptions made in Theorem 4.1

The justification of the two assumptions outlined in Theorem 4 warrants careful consideration. While we provide only high-level justifications, they underpin significant aspects of our theoretical framework.

For the first assumption, we argue that the regularization parameter λ , which penalizes the magnitude of the parameter norms, implicitly promotes solutions that are confined to a compact subset of the parameter space. This rationale is conceptual and requires that regularization effectively constrains the growth of the parameter norms, thereby localizing the solutions.

The second assumption concerns the separation property. It is naturally satisfied as long as the origin 0 dim θ +dim w remains an interior point of supp( ρ ∞ ( · , t )) . This condition is relatively mild and

is generally satisfied. The challenge arises in verifying that supp( ρ ∞ ( · , t )) for the α 2 component extends to encompass the entire space K . While direct confirmation is elusive, it is suggested by [23] initially spans K , this expansive support property is maintained at any finite time. Thus, we conjecture that the condition holds under these circumstances, providing a basis for this assumption.

D Proofs of main results in Section 3

D.1 Proof of Theorem 3.1

This convergence is detailed in two parts. First, the finite time result, as stated in points (i)-(iii), utilizes a concept in probability theory known as propagation of chaos [62] to examine how differences evolve uniformly across a given time interval. In the context of our model, this involves comparing how parameter particles evolve under discrete versus continuous dynamics.

Specifically, the approximation bound is derived using a third auxiliary dynamic, termed the "nonlinear dynamics," by bounding the dynamic difference over the entire finite time interval. This process involves applying the triangle inequality to each component and concluding with a Grönwall's inequality. Since the Transformer output can be bounded by the dynamic difference, we can then bound the output difference at any specific time, along with the difference regarding different time for the same dynamic. By applying a probability union bound on a dense set of L 2 points, we can extend this to bound the maximal difference over any time interval.

Secondly, the weak convergence of the empirical distribution process leverages optimal transport theory alongside abstract stability results for Wasserstein gradient flows [4]. This argument involves detailed analysis of the discretization of particle distributions in space, particularly focusing on obtaining the convergence of the sequence of momentum fields [4, 60] that could directly leads to the result. To obtain the convergence of the momentum field sequence, we also need to bound the parameter gradient difference between discrete and continuous Transformers as the mean-field limit.

Preparatory Step: Nonlinear dynamics

We first define some auxiliary quantities and differential equations that are useful for the proof. For any gradient flow parameter setting Θ ( τ ) = { β ( τ ) t,j } t,j , from its definition (2.7), we could rewrite the dynamics as

for any gradient flow time τ > 0 , depth index t = 0 , ∆ t, . . . , ( L -1)∆ t and width index j = 1 , . . . , M . For simplicity, any mentioned constant only depends on N,D,τ,λ and the parameters of the assumptions, and we abbreviate the subscript t = 0 , ∆ t, . . . , ( L -1)∆ t, j = 1 , . . . , M by t, j throughout the proof.

Inspired by the 'propagation of chaos" idea [62], we could define the 'nonlinear dynamics" with the same initialization setting ˜ Θ ( τ ) = { ˜ β ( τ ) t,j } t,j , i.e.

for any t, j . Here, ( ρ ( s ) ) s ≥ 0 is the solution to the Wasserstein gradient flow (3.5), of which the uniqueness is implied by Proposition 3.2. Since (D.2) is just the particle flow of (3.5), its existence and uniqueness are guaranteed by Proposition 3.2.

Observing that { β t,j } t,j are independent due to the dynamics only involving ( ρ ( s ) ) s ≥ 0 with i.i.d initialization over ρ 0 , we can consider { ˜ β ( s ) t,j } t,j as i.i.d. samples drawn from { ρ ( s ) } t,j . In addition, from Propositions 3.2 and D.1, for any t, j we have max {∥ β t,j ∥ , ∥ ˜ β t,j ∥} ≤ R τ , where R τ is defined as in these propositions and does not depend on M and L .

Preparatory Step: Bound the gradient difference regarding parameter settings

As the second preparatory step for the proof of Theorem 3.1, we present the following three lemmas that will be helpful:

Lemma D.1 (Continuous gradient difference bound) . Suppose Assumptions 1-3 hold. If we have ρ, ν ∈ P 2 concentrated on P r for some r > 0 , and β, ˜ β such that max {∥ β ∥ , ∥ ˜ β ∥} ≤ r , then

for any t ∈ [0 , 1] . Here, the constant C G only depends on N,D,r , and the parameters of the assumptions.

Lemma D.2 (Discrete gradient difference bound) . Under Assumption 1-3, for any Θ = { β t,j } t/ ∆ t +1 ∈ [ L ] ,j ∈ [ M ] , ˜ Θ = { ˜ β t,j } t/ ∆ t +1 ∈ [ L ] ,j ∈ [ M ] such that max {∥ β t,j ∥ , ∥ ˜ β t,j ∥} ≤ r for any t = 0 , . . . , ( L -1)∆ t, j = 1 , . . . , M . Then we have that

for any β, ˜ β ∈ { β : ∥ β ∥ ≤ r } and t = 0 , ∆ t, . . . , ( L -1)∆ t . Here, the constant C G only depends on N,D,r and the parameters of the assumptions and d (Θ , ˜ Θ) is defined as ∑ t ∑ M j =1 ∥ β t,j -˜ β t,j ∥ .

Lemma D.3 (Oracle gradient approximation with discretization) . Under Assumptions 1-3, suppose that the parameter setting Θ is i.i.d. drawn from { ρ ( β | t ) } t/ ∆ t +1 ∈ [ L ] ,j ∈ [ M ] for some ρ ∈ P 2 concentrated on P r and satisfies that ∫ β ρ ( β, t ) dβ = 1 for any t ∈ [0 , 1] . Then with probability at least 1 -4 exp( -δ ) with respect to the parameter initialization Θ (0) , we have

for any β ∈ { β : ∥ β ∥ ≤ r } , t = 0 , ∆ t, . . . , ( L -1)∆ t, 1 and any δ > 0 . Here, ≲ hides the dependencies on N,D,r and the parameters of the assumptions.

Proof of Theorem 3.1. Our proof consists of several steps outlined below:

Step I: Show the W 2 continuity of parameter (sample) distributions

Our analysis commences with the bound for 0 < s 1 < s 2 < τ , we have

where each particle at time s 1 is paired with its position at time s 2 , leveraging the Jensen's inequality. Recalling the identity

shown in Proposition D.1, it follows that

where the last inequality uses (D.30). Since A 2 0 ≲ 1 + λ -1 , we see that W 2 (ˆ ρ ( s 1 ) , ˆ ρ ( s 2 ) ) ≤ C (1 + λ )( s 2 -s 1 ) 1 / 2 for some constant C dependent on the parameters listed in the result. Similarly, we have

where ˜ β ( s 1 ) ∼ ρ ( s 1 ) ( β, t ) is embedded with its future position at ˜ β ( s 2 ) . The last step is feasible by setting C large enough, noticing that Q ( ρ 0 ) ≤ λA 2 0 / 2 . To summarize, we have

for some constant C > 0 dependent on the parameters listed in the result.

Step II: Bound the difference between gradient flow dynamics and non-linear dynamics

Next, we aim to bound ∆( s ) := sup s ′ ∈ [0 ,s ] sup t,j ∥ β ( s ′ ) t,j -˜ β ( s ′ ) t,j ∥ for any s ∈ [0 , τ ] . Taking the difference of (D.1) and (D.2), we obtain that

where the final inequality stems from Lemma D.2, employing a constant C G dependent on the parameters listed in the result. By taking the supremacy over t, j in (D.4), and considering 1 ML d (Θ ( s ) , ˜ Θ ( s ) ) ≤ sup t,j ∥ β ( τ ) t,j -˜ β ( τ ) t,j ∥ for any s ≥ 0 , we derive

Further supremacy taken over s ∈ [0 , τ ] yields:

where we define ˜ ∆( τ ) := sup t,j sup s ∈ [0 ,τ ] ∥ ∥ ∥ ̂ G ( ˜ β ( s ) t,j , ˜ Θ ( s ) , t ) -G ( ˜ β ( s ) t,j , ρ ( s ) , t ) ∥ ∥ ∥ for simplicity of notation. Apply the Grönwall's inequality to (D.5) yields

It remains to bound ˜ ∆( τ ) to bound ∆( τ ) . It's worth noting that by Lemmas D.1 and D.2, for any s 1 , s 2 ∈ [0 , τ ] and t, j , we have

(D.7)

for some constant C ∆ dependent on the parameters listed in the result. Moreover, for any t, j , we have

where the universal boundedness of ∥ G ( ˜ β ( s ) t,j , ˜ Θ ( s ) , t ) -λ ˜ β ( s ) t,j ∥ can be readily derived from the third result of Lemma C.3 alongside Assumption 2 (ii). Therefore, we obtain

for any s 1 , s 2 ∈ [0 , τ ] .

Define τ n = nτ/L 2 for n = 1 , 2 , . . . , L 2 . Leveraging Lemma D.3 and applying the union bound over n ∈ [ L 2 ] (updating the δ in the lemma to δ +2log L ), we obtain, with a probability of at least 1 -exp( -δ ) with respect to the parameter initialization Θ (0) :

Furthermore, leveraging (D.9), we deduce

Returning to (D.6), we establish that with a probability of at least 1 -exp( -δ ) with respect to the parameter initialization Θ (0) ,

We denote the event where the above inequality holds as E 1 , thus we have P ( E 1 ) ≥ 1 -exp( -δ ) .

Step III: Prove the finite time results

Now, we are poised to demonstrate the results in Theorem 3.1 that concern supremacy over s ∈ [0 , τ ] . The verification of Lemmas C.4 reveals the existence of a universal constant B τ := B exp( K (1 + R τ + R 2 τ )) such that

for any H and t ∈ [0 , 1] .

Utilizing Lemma D.6 and applying the union bound over n ∈ [ L 2 ] , we observe

Additionally, note that for any s 1 , s 2 ∈ [0 , τ ] , we derive from Lemmas C.2 and C.5 that

where the last inequality is derived utilizing (D.3) and (D.8). Consequently, we have

with probability at by 1 -exp( -δ ) . We denote the event where the above inequality holds as E 2 , thus we have P ( E 2 ) ≥ 1 -exp( -δ ) . Considering that { ˜ β ( s ) t,j } t,j could be regarded as i.i.d. samples

drawn from { ρ ( s ) } t,j , employing a similar method with the concentration guarantee from Lemma C.7, we can readily deduce the existence of an event E 3 with P ( E 3 ) ≥ 1 -exp( -δ ) such that, under E 3 , we have

Now, let's analyze the scenario under the probability event E 1 ∩ E 2 ∩ E 3 with P ( E 1 ∩ E 2 ∩ E 3 ) ≥ 1 -3 exp( -δ ) . Lemma C.5 demonstrates that

This further implies, by (D.11), that

Since max {∥ T ρ ( τ ) ( H,t ) ∥ 2 -col , ∥ ̂ T Θ ( τ ) ( H,t ) ∥ 2 -col , ∥ ̂ T ˜ Θ ( τ ) ( H,t ) ∥ 2 -col } is universally bounded, (D.13) immediately indicates that

and

Step IV: Prove the weakly convergence

For the remainder of the proof, we adopt a similar approach as in the proof of Theorem 2.6 in [16]. We denote ˆ ρ as ˆ ρ M,L for any given M and L . It's essential to note that we treat ˆ ρ M,L as probability measures in this step of the proof.

Recalling (D.3), for any s 1 , s 2 ∈ [0 , τ ] , we have

for some constant C dependent on the parameters listed in the result. We observe that the family of curves ( s ↦→ ˆ ρ ( s ) M,L ) M,L is equicontinuous in W 2 on [0 , τ ] , uniformly in M,L . Additionally, the family (ˆ ρ M,L ) M,L lies within a W 2 ball, thus weakly precompact. As the weak topology is weaker than the topology induced by W 2 , according to the Arzelà-Ascoli theorem, along any sequence where L →∞ and log L/M →∞ , we can identify a subsequence that converges weakly to a certain process ( ν ( s ) ) s ≥ 0 ∈ P 2 × R , concentrated on P R τ at all times. In the subsequent analysis, we solely focus on this subsequence, still denoted as (ˆ ρ M,L ) M,L .

For any t ∈ [0 , 1] , let's define the sequence ( E ⊤ M,L ) M,L of momentum fields, which is a vectorvalued measure on [0 , τ ] × Ω , denoted by E M,L := ̂ G ( β, Θ ( s ) M,L , t )ˆ ρ ( s ) ( β, t ) ds . We also define E := G ( β, ν ( s ) , t ) ν ( s ) ( β, t ) ds .

Considering that both ˆ ρ M,L and ν are concentrated on P R τ , we also have uniform convergence in the Bounded Lipschitz metric. Hence, for any bounded and Lipschitz function φ : [0 , τ ] × R dim β → R dim β , it holds

uniformly among s ∈ [0 , τ ] along the sequence.

Note that

for some constant C E dependent on the parameters listed in the result and with probability at least 1 -exp( -δ ) with respect to the parameter initialization Θ (0) for any δ > 0 . Here, the third inequality of (D.16) utilizes the third result of Lemma D.3, and the fourth inequality uses the second result of Lemma D.3, following a similar process in Step II to achieve supremacy over s ∈ [0 , τ ] .

From (D.16), we infer that ∣ ∣ ∣ ∫ τ 0 ∫ 1 0 ∫ β φ · d ( E M,L -E ) ∣ ∣ ∣ → 0 almost surely along the sequence. Hence, E M,L converges weakly to E almost surely along the sequence, and the particle gradient flow for ( ν ( τ ) ) τ ≥ 0 almost surely satisfies (2.7) on [0 , τ ] for any arbitrarily given τ > 0 . According to the Fokker-Planck equation without noise involved [59], we conclude that ( ν ( τ ) ) τ ≥ 0 almost surely satisfies (3.5). Consequently, the uniqueness stated in Proposition 3.2 ensures that ( ν ( τ ) ) τ ≥ 0 = ( ρ ( τ ) ) τ ≥ 0 almost surely.

D.2 Proof of Proposition 3.1

Proof. Suppose that the Fréchet derivative δR δρ indeed exists, we establish

Therefore, it suffices to show that the Fréchet derivative of L with respect to ρ is

Denote ρ η = ρ + η ( ν -ρ ) . We provide the following lemma to bound T ρ η ( H, 1) -T ρ ( H, 1) by expanding the first-order derivative as follows

Lemma D.4 (First-order derivative of Transformer output) . Under Assumption 2, for any H and ρ, ν ∈ P 2 that have bounded supports, we have

where ρ η := ρ + η ( ν -ρ ) .

Given (D.18), we observe from the solution of p ρ (C.13) and (C.14) that

Hence, by applying (D.19) to the risk function, we obtain

which indicates (D.17) and concludes the proof.

D.3 Proof of well-posedness of Wasserstein gradient flow

Proof. Following a similar idea as Proposition 2.5 of [16], we leverage the general theory of Wasserstein gradient flow developed in [4]. Define the functional family Q r ( ρ ) as

For any r > 0 , let's consider any admissible transport γ ∈ P Ω × Ω concentrated on P r . By definition, both of its marginals, denoted by ρ 1 and ρ 2 , are concentrated on P r . We define the transport cost for γ as

for p ≥ 1 . Additionally, we denote the transport interpolation as ρ γ α := ((1 -α ) ρ 1 + αρ 2 ) # γ . Our proof consists of several steps outlined below.

Step I: Show that Q r is proper and continuous for W 2 on its closed domain:

Note that the parameters r, D, N, λ remain fixed throughout this proof step, so we hide the constant dependencies on them. For any ( β, t ) ∈ P r , we have Q r ( δ ( β,t ) ) = 1 2 E µ [(Read[ H +∆ tg ( H,β )] -y ( H )) 2 ] + λ 2 ∥ β ∥ 2 < ∞ . This indicates that Q r is proper. Moreover, for any ρ, ν ∈ P 2 whose bounded support belong to P r , Lemma C.1 ensures that

and Lemma C.2 guarantees that

for any H . Therefore, we have

Furthermore, since both ρ and ν have bounded support, ∥ β ∥ 2 is Lipschitz continuous with respect to ( β, t ) . Therefore, by the Kantorovich-Rubinstein Theorem (see Theorem 5.10 of [68], for example), we have

Combining (D.20) and (D.21), we obtain that Q ( ρ ) -Q ( ν ) = O ( W 2 ( ρ, ν )) . Therefore, Q r is continuous for W 2 on its closed domain.

Step II: Show that α ↦→ Q ( ρ γ α ) /C 2 2 ( γ ) is differentiable and has a Lipschitz continuous derivative Let's denote h ( α ) := Q r ( ρ γ α ) . Lemma C.3 ensures that for any ρ ∈ P 2 with bounded support belonging to P r , we have bounded ∥ δQ δρ ( · , · ) ∥ F on P r . Therefore, Q r ( ρ γ α ) is differentiable with respect to t , and the derivative reads

Then, it suffices to show that h ′ ( α ) is Lipschitz continuous. To accomplish this, we first propose the following lemma for later use:

Lemma D.5 (Locally Lipschitz of ρ for the gradient) . Under Assumptions 1 and 2, for any ρ, ν ∈ P 2 concentrated on P r , there exists some constant L r depending on r, N, D and parameters of the assumptions such that

Returning to the lemma proof, for α 1 , α 2 ∈ [0 , 1] , by the triangle inequality we have | h ′ ( α 1 ) -h ′ ( α 2 ) |≤ J 1 + J 2 where

(D.23)

by Lemma D.5. The final inequality of (D.23) applies Hölder's inequality to obtain that C 2 1 ( γ ) ≤ C 2 2 ( γ ) . Furthermore,

where L ′ r := sup ( β,t ) ∈ P r ∥ δQ δρ | ρ = ρ γ α 2 ( β, t ) ∥ < ∞ from Lemma C.3. Combining (D.23) and (D.24) leads us to the result that h ′ ( α ) /C 2 2 ( γ ) is Lipschitz continuous.

Step III: Show the well-posedness of Wasserstein gradient flow at some finite time

We follow a similar approach to the proof of Proposition 2.5 in [16]. Since h ′ ( α ) is λ h × C 2 2 ( γ ) -Lipschitz continuous with respect to α for some λ h , the well-posedness of the Wasserstein gradient flow for Q r with the velocity field constrained on P r is a corollary of Theorem 11.2.2 of [4]. Specifically, there exists a unique curve ( ρ ( τ ) r ) τ ≥ 0 continuous in P 2 such that:

for ρ ( τ ) r -a.e. Given the initialization ρ 0 concentrated on P R , for any r > R , the unique ρ ( τ ) r exhibits a first exit time denoted as

By defining this exit time, for any ¯ r > r and τ ∈ [0 , τ r ] , we observe v r ( τ )( β, t ) = G ( β, ρ ( τ ) r , t ) and v ¯ r ( τ )( β, t ) = G ( β, ρ ( τ ) ¯ r , t ) . Due to uniqueness, we infer ρ ( τ ) r = ρ ( τ ) ¯ r on τ ∈ [0 , τ r ] . Considering ρ ( τ ) r as the solution to (3.5), we establish the existence and uniqueness of the Wasserstein gradient flow for Q over [0 , τ r ] .

Step IV: Show the well-posedness of Wasserstein gradient flow at all time

To establish the Wasserstein gradient flow's definition for τ ≥ 0 , it's necessary to demonstrate that lim r →∞ τ r = ∞ . For any r > R , according to the energy identity in Theorem 11.2.1 of [4], on [0 , t r ] , we observe that τ ↦→ Q ( ρ ( τ ) ) is non-increasing. Specifically, this represents

Therefore, for any τ ∈ [0 , τ r ] , utilizing Lemma C.1, we have

where

Thus, we have ∫ 1 0 ∫ β ∥ β ∥ 2 ρ ( τ ) ( β, t ) ≤ 2 λ Q ( ρ ( τ ) ) ≤ R 2 + λ -1 ( 2 B 2 + 2 B 2 exp ( K (1 + R + R 2 ) ) 2 ) = A 2 0 . According to Assumption (ii) and Lemma C.3, for any ( β, t ) ∈ P r , we have

(D.27) where C R := √ N +1 ϕ P ( B exp( K (1 + A 0 + A 2 0 )))( B + B exp( K (1 + A 0 + A 2 0 ))) exp ( ϕ T ( N,D, √ N +1 KB exp( K (1 + A 0 + A 2 0 ))(1 + A 0 + A 2 0 ) ) . Applying (D.27) to the gradient flow equation

for τ ≥ 0 , we obtain d ∥ β ( τ ) ∥ dτ = ⟨-v ( τ ) r ( β,t ) , ∥ β ( τ ) ∥⟩ ∥ β ( τ ) ∥ ≤ ∥ v ( τ ) r ( β, t ) -λβ ∥ ≤ C R (1 + ∥ β ( τ ) ∥ ) . This indicates

by the Grönwall's inequality. Therefore, for any T > 0 , ρ ( T ) is concentrated on P ( R +1)exp( C R T ) -1 , implying that for r > ( R +1)exp( C R T ) , we have τ r > T . Hence, we conclude lim r →∞ τ r = ∞ , establishing the existence of a unique Wasserstein gradient flow from (3.5) over τ > 0 .

Eventually, we establish the three properties listed in Proposition 3.2 for ρ ( τ ) . By (3.5), for any t ∈ [0 , 1] , we have

indicated by the Divergence Theorem as ρ ( s ) has bounded support. Next, for any τ ≥ 0 , (D.28) shows that ρ ( τ ) is concentrated on P R τ . Moreover, (D.25) now holds for any τ > 0 , implying ∫ 1 0 ∫ β ∥ β ∥ 2 ρ ( τ ) ( β, t ) ≤ A 2 0 for any τ ≥ 0 .

D.4 Proof of well-posedness of gradient flow

Proposition D.1 (Existence and uniqueness of gradient flow) . Under Assumptions 1-3, for any initialization of Θ (0) i.i.d. drawn from { ρ 0 ( θ, w | t ) } t,j , there exists a unique solution (Θ ( τ ) ) τ ≥ 0 for (2.7) . Additionally, for any τ ≥ 0 , we have i. Θ ( τ ) has a bounded support, meaning sup t,j ( ∥ θ ( τ ) t,j ∥ 2 2 + ∥ w ( τ ) t,j ∥ 2 2 ) ≤ R τ .

Here, R τ and A 0 are defined as in Proposition 3.2.

Proof. The local Lipschitz continuity established in Lemma D.2 directly implies the continuity of ̂ G β ( β t,j , Θ , t ) with respect to Θ . Since { ML · ̂ G β ( β t,j , Θ , t ) } t/ ∆ t +1 ∈ [ L ] ,j ∈ [ M ] serves as the gradient of ̂ Q (Θ) , it follows that ̂ Q (Θ) is continuously differentiable, indicating the local semiconvexity of ̂ Q (Θ) . Specifically, for any Θ , there exists some κ > 0 such that ̂ Q (Θ) + κ ∑ t ∑ M j =1 ( ∥ θ t,j ∥ 2 2 +

∥ w t,j ∥ 2 2 ) is convex within a small neighborhood of Θ . The existence and uniqueness of a gradient flow over the maximal interval [0 , τ max ] is a standard result (see Section 2.1 of [61]).

For any τ ∈ [0 , τ max ] , it holds that

where the last inequality follows from Jensen's inequality. (D.29) establishes that ̂ Q (Θ ( τ ) ) is both upper and lower bounded, and Θ ( τ ) exhibits a bounded curve length over any the time interval [0 , τ max ] . By compactness, if τ max is finite, then Θ ( τ max ) exists and thus must exist beyond τ max , which leads to contradiction. Therefore, τ max = ∞ , and the well-posedness of the gradient flow for τ ≥ 0 consequently follows. Additionally, (D.29) shows that for any τ ≥ 0 ,

The last inequality of (D.30) follows from Lemma C.4, thereby showing that 1 ML ∑ t ∑ j =1 ∥ β t,j ∥ 2 2 ≤ A 2 0 for any τ ≥ 0 .

As the final part of our proof, we demonstrate that the norm of any entry of Θ is bounded at any given time τ ≥ 0 . Note that

Applying (D.31) to the gradient flow

for τ ≥ 0 , we have d ∥ β ( τ ) t,j ∥ dτ = ⟨-̂ G ( β ( τ ) t,j , Θ ( τ ) ,t ) , ∥ β ( τ ) ∥⟩ ∥ β ( τ ) ∥ ≤ ∥ ̂ G ( β ( τ ) t,j , Θ ( τ ) , t ) -λβ ∥ ≤ C R (1+ ∥ β ( τ ) ∥ ) . This indicates

D.5 Proof of Proposition 3.3

Before commencing the proof, we introduce two proxy Transformer procedures in addition to ¯ T Θ and ˜ T ρ . The first proxy, denoted as ¯ T Θ , involves moving the layers with the encoder h slightly forward by ∆ t/ 2 in the depth index. This adjustment results in a discrete Transformer with only L layers, where each layer has a step size of ∆ t and an encoder of ( f + h ) / 2 , represented by g . Specifically, ¯ T Θ can be written as

The second proxy, denoted as ˜ T ρ , extends the width to infinity by letting M → ∞ , effectively replacing the average with an integral:

We let all four Transformers share the same initial state ̂ T Θ ( H, 0) = ¯ T Θ ( H, 0) = ˜ T ρ ( H, 0) = T ρ ( H, 0) = H .

We first present the following lemma, considering parameters i.i.d. drawn from some distribution ρ ∈ P 2 with bounded support:

Lemma D.6 (Oracle approximation of discretization) . Under Assumptions 1 and 2, suppose that the parameter setting Θ is i.i.d. drawn from { ρ ( θ, w | t ) } t,j for some ρ ∈ P 2 concentrated on { ( θ, w ) : ∥ θ ∥ 2 + ∥ w ∥ 2 ≤ r 2 } × [0 , 1] and satisfies that ∫ θ,w ρ ( θ, w, t ) d ( θ, w ) = 1 for any t ∈ [0 , 1] . Then with probability at least 1 -exp( -δ ) with respect to the parameter initialization Θ (0) , we have

for any H , t = 0 , ∆ t, . . . , ( L -1)∆ t, 1 and any δ > 0 . Here, ≲ hides the dependencies on N,D,r and the parameters of the assumptions.

Proof of Proposition 3.3. Since ρ ∈ P 2 ,r has a bounded support for any ρ , there exists some ρ ∗ ∈ P 2 ,r such that R ( ρ ∗ ) = inf ρ ∈ P 2 ,r R ( ρ ) . According to Lemma D.6, we can find a specific Θ such that

where C depends on N,D,r , and the parameters of the assumptions. Moreover, from Lemma D.6, we ensure that each entry β t,j of Θ satisfies ∥ β t,j ∥ ≤ r . Verification of Lemmas C.1 and C.4 on T ρ and ̂ T Θ respectively leads to their uniform boundedness, i.e., sup t T ρ ( H,t ) ≲ 1 and sup t ̂ T Θ ( H,t ) ≲ 1 . Therefore, we have

Here, ≲ hides the dependencies on N,D,r and the parameters of the assumptions. The result then follows.

The proof for the energy functional Q (and ̂ Q ) follows a similar approach. There exists some ρ ∗ ∈ P 2 ,r such that Q ( ρ ∗ ) = inf ρ ∈ P 2 ,r Q ( ρ ) . From Lemmas D.6 and C.7, we can find a specific Θ such that

by setting C large enough. Verification of Lemmas C.1 and C.4 on T ρ and ̂ T Θ respectively leads to their uniform boundedness. Hence, we have

E Proofs of main results in Section 4

For simplicity, we assume that Assumption 4 holds with ( g, α ) = ( f, θ ) . The proof for the case of ( g, α ) = ( h, w ) is symmetric, involving a simple substitution of f with h and θ with w .

E.1 Proofs of Theorem 4.1 and Corollary 4.1

Our proof of Theorem 4.1 consists of three parts, each focusing on bounding differences related to the energy functional Q or ̂ Q .

The first step establishes the continuity of the functional gradient δQ δρ | ρ ∞ . This ensures that if the derivative with respect to β for the functional gradient is constant over a region, then the functional gradient remains constant within that region.

The second step provides the key bound for Q ( ρ ∞ ) , which is proportional to λ . This involves a detailed analysis of Q 's landscape by bounding its derivatives.

After obtaining the bound for Q ( ρ ∞ ) , the final steps are to show that the finite-time risk can approach this bound. Achieving a loss as small as ϵ requires Q ( ρ ( τ 0 ) ) ≤ ϵ for some sufficiently large τ 0 . We then apply Theorem 3.1, with constant dependency on τ 0 , to show that ̂ Q ( τ 0 ) becomes sufficiently small. Since ̂ Q ( ρ ( τ ) ) is non-increasing, ̂ Q ( τ ) remains small for all τ ≥ τ 0 .

Preparatory Step: Landscape analysis

First, the following lemma suggests that as long as the risk R ( ρ ) remains positive, a descent direction for Q ( ρ ) can be constructed at any depth index, provided that λ is sufficiently small. This implies that by adjusting λ , one can influence the gradient flow to effectively reduce Q ( ρ ) .

Lemma E.1 (Landscape of Q ( ρ ) ) . Suppose that Assumptions 1-4 hold. For any ρ ∈ P 2 concentrated on P r with some r > 0 , any w 0 ∈ R dim w , and any t ∗ ∈ [0 , 1] such that ∫ β ρ ( β, t ∗ ) ≥ 1 / 2 , there exists a ν ∈ P ( R dim β ) such that i. for any β ∈ supp( ν ) , we have 1 /B r ≤ ∥ β ∥ ≤ B r

ii. for any ( θ 1 , θ 2 , w ) , we have θ 2 ∈ K and w = w 0 .

Here, B r , C 1 , C 2 are constants that depends on N,d,r and the parameters of the assumptions.

Given the t ∗ and w 0 specified in the theorem, Lemma E.1 indicates that there exists some

such that ∫ β δQ δρ | ρ ∞ ( β, t ∗ ) ( ν ( β ) -ρ ∞ ( β | t ∗ ) ) dβ ≤ C 1 λ -C 2 R ( ρ ) , where B R ∞ , C 1 , C 2 are constants dependent on N,d,R ∞ and the parameters of the assumptions.

In addition, for any ρ ∈ P 2 , we define the following two functional derivatives:

It is obvious that δQ δρ ≡ ( δQ f δρ + δQ h δρ ) / 2 .

Proof of Theorem 4.1. Step I: Show that δQ δρ | ρ ∞ ( β, t ) is continuous with respect to ( β, t )

In Step I of the proof of Lemma E.1, we establish that ρ ∞ ∈ P 2 , with a bounded support P R ∞ , implies p ρ ∞ ( H,t ) is C p -Lipschitz continuous, where C p is a constant dependent solely on N,D,R ∞ , and the parameters in our assumptions.

Next, we would like to show that δQ δρ ρ ∞ ( β, t ) is continuous with respect to ( β, t ) ∈ R div β × [0 , 1] . Let's focus on the region ( β, t ) ∈ P r for any r > R ∞ , so that ρ ∞ is also concentrated on P r . It's noteworthy that for any bounded support ( β, t ) ∈ P r , Lemma C.1 and Assumption 2 (i) ensure the universal boundedness of g ( T ρ ( H,t ) , β ) , and Lemma C.3 ensures the universal boundedness of p ρ ( H,t ) for any H ∈ supp( µ ) , with the constants depending solely on N,D,r , and the parameters of the assumptions.

Combining the Lipschitz continuity of p ρ ∞ ( H,t ) and T ρ ∞ ( H,t ) with respect to ( H,t ) , as shown in Proposition C.2, along with the Lipschitz continuity of g ( T, β ) with respect to ( T, β ) when ∥ T ∥ F is universally bounded (as guaranteed by Assumption 2 (ii) and (iii)), and their universal boundedness, we derive that Tr ([ g ( T ρ ∞ ( H,t ) , β ) ] ⊤ p ρ ∞ ( H,t ) ) is C G -Lipschitz continuous for ∥·∥ 2 with respect to ( β, t ) ∈ P r for some constant C G that depends only on N,D,r and the parameters of the assumptions. Since the Lipschitz constant C G is independent of the choice of H , we see that Tr ([ g ( T ρ ∞ ( H,t ) , β ) ] ⊤ p ρ ∞ ( H,t ) ) is uniformly continuous across all H ∈ supp( µ ) with respect to ( β, t ) ∈ P r . Consequently, we have that

is continuous with respect to ( β, t ) ∈ P r . Since the choice of r is arbitrary, we conclude that δQ δρ | ρ ∞ ( β, t ) is continuous with respect to ( β, t ) ∈ R div β × [0 , 1] .

Step II: Show that Q ( ρ ∞ ) ≲ λ with further landscape analysis

In the first part of the proof, we will adopt a similar approach to Theorem 3.9 of [47] to demonstrate that the stationary point of the Wasserstein gradient flow, denoted ρ ∞ , satisfies Q ( ρ ∞ ) ≲ λ . It's worth noting that in [47], the authors assume λ = 0 and conclude R ( ρ ∞ ) = Q ( ρ ∞ ) = 0 , but this claim relies on assuming the global existence of the Wasserstein gradient flow rather than proving it directly.

Based on the pivotal findings from [53] regarding the stationary points in the Wasserstein space, we infer that the stationary point ρ ∞ of the Wasserstein gradient flow (3.5), i.e.

must satisfy ∇ β δQ δρ | ρ ∞ = 0 almost everywhere over supp( ρ ∞ ) . This further indicates that ∇ β δQ δρ | ρ ∞ ( θ 1 , θ 2 , w, t ∗ ) = 0 almost everywhere over supp( ρ ∞ ( · , t ∗ )) . The fact ρ ( · , t ∗ ) is a connected set, coupled with the continuity of the Frechét differential δQ δρ | ρ ∞ with respect to β , implies that, δQ δρ | ρ ∞ ( θ 1 , θ 2 , w, t ∗ ) = C for some constant C over ( β, t ∗ ) ∈ supp( ρ ( · , t ∗ )) .

Given the separation assumption on the support of ρ ∞ ( · , t ∗ ) , we ensure that for any ( θ 1 , θ 2 ) ∈ ( B dim θ 1 (0 , B R ∞ ) /B dim θ 1 (0 , 1 /B R ∞ ) ) × K , there exists c ∈ R , 1 /R ∞ B R ∞ ≤ | c |≤ R ∞ B R ∞

such that ( cθ 1 , θ 2 , w 0 , t ∗ ) ∈ supp( ρ ∞ ) . Combined with Assumption 4 (i), which implies the 1 -homogeneity of f ( T, θ 1 , θ 2 ) with respect to θ 1 , we have

Hence, given that ∇ β δQ δρ | ρ ∞ ( · , t ∗ ) = 0 almost everywhere over supp( ρ ∞ ( · , t ∗ )) , it also holds that

which implies

Given the condition that ∥ ( cθ 1 , θ 2 , w 0 , t ∗ ) -( θ 1 , θ 2 , w 0 , t ∗ ) ∥ ≤ R 2 ∞ B R ∞ , and recalling that δQ δρ | ρ ∞ ( θ 1 , θ 2 , w, t ∗ ) ≡ C across ( β, t ∗ ) ∈ supp( ρ ( · , t ∗ )) , (E.2) further indicates that

for any ( θ 1 , θ 2 ) ∈ ( B dim θ 1 (0 , B R ∞ ) /B dim θ 1 (0 , 1 /B R ∞ ) ) ×K . Hence, by Lemma E.1 we have

Therefore, we have R ( ρ ∞ ) ≤ C 1 +2( R 4 ∞ B 2 R ∞ + R 3 ∞ B R ∞ ) C 2 λ , and

which completes the first part of our proof.

Step III: Bound the difference between Q ( ρ ( τ ) ) and Q ( ρ ∞ ) when τ is large

Proposition 3.2 establishes that the second moment for ρ ( τ ) is uniformly bounded across all τ ≥ 0 :

where A 0 is defined as in Proposition 3.2. Therefore, the weak convergence of probability measures ( ρ ( τ ) ) τ ≥ 0 is equivalent to the convergence in the Wasserstein2 distance, i.e.

When τ is sufficiently large, ρ ( τ ) concentrates on P R ∞ . Therefore, according to Lemma C.2, there exists a constant C 0 depending solely on N , D , R ∞ , and the parameters of the assumptions that

for any H and t ∈ [0 , 1] when τ is sufficiently large. Note that Lemma C.1 shows that

for any H and t ∈ [0 , 1] . Thus, from (E.7), we have

where the second inequality incorporates the Kantorovich-Rubinstein Theorem (see Theorem 5.10 of [68], for example) and the 2 R ∞ -Lipschitz continuity of ∥ β ∥ 2 over the region ( β, t ) ∈ P R ∞ . Combining equations (E.6) and (E.8), we deduce that for any ϵ > 0 , there exists some τ 0 > 0 such that | Q ( ρ ∞ ) -Q ( ρ ( τ 0 ) ) |≤ ϵ .

Step IV: Complete the proof by bounding the difference between ̂ Q (Θ ( τ ) ) and Q ( ρ ( τ ) ) when τ is large

The final step can be seen as a direct corollary of the approximation result in Theorem 3.1. According to Theorem 3.1, there exists a constant C 1 dependent on N , D , τ 0 , λ , and the parameters specified in the assumptions, such that

with probability at least 1 -3 exp( -δ ) with respect to the parameter initialization Θ (0) for any δ > 0 . Combining the outcomes from the preceding steps, we obtain

Note that

so the sequence ( ̂ Q (Θ ( τ ) )) τ ≥ 0 is non-decreasing. Hence, for any τ ≥ τ 0 ,

which completes the proof, recalling that C 1 depends only on N,D,τ 0 , λ and the parameters of the assumptions, and C 2 depends only on N,D,R ∞ , and the parameters of the assumptions.

Proof of Corollary 4.1. Given the choice of λ By Theorem 4.1, there exists some τ 0 > 0 such that

The result holds by setting L and M/ log L sufficiently large to ensure that

F Proofs of auxiliary results

F.1 Proof of Lemma D.1

Proof. Lemma C.1 confirms that ∥ T ρ ( H,t ) ∥ 2 -col and ∥ T ν ( H,t ) ∥ 2 -col are bounded uniformly by B T = B exp( K (1 + r + r 2 )) . Considering the definition (C.2), it suffices to demonstrate that

where

Here, the symbol ≲ hides dependencies on N , D , r , and the parameters of the assumptions. To bound J 1 , consider that

The third inequality in Equation (F.1) is derived from Lemma C.3, while the fourth inequality relies on Assumption 3 (i) and (iii). Lastly, bounding ∥ T ρ ( H,t ) -T ρ ( H,t ) ∥ 2 -col by W 1 ( ρ, ν ) is achieved with Lemma C.2.

On the other hand, to bound J 2 , we have

In (F.2), the third inequality relies on Assumption 2 (i). Consequently, to establish J 2 ≲ exp( C G W 1 ( ρ, ν )) -1 , it is adequate to demonstrate that ∥ p ν ( H,t ) -p ρ ( H,t ) ∥ F ≤ I 1 + I 2 ≲ W 1 ( ρ, ν ) , where

From Lemma C.2, it is trivial that | Read[ T ρ ( H, 1) -T ν ( H, 1)] |≤ ∥ T ρ ( H, 1) -T ν ( H, 1) ∥ F ≲ W 1 ( ρ, ν ) . Thus, I 1 ≲ W 1 ( ρ, ν ) given the boundedness of ∥∇ vec[ T ] vec[ g ( T ν ( H,t ) , β )] ∥ as provided

in Assumption 2 (iii). To bound I 2 , we have

for some constant C r dependent on the parameters listed in the result setting. Here, the first inequality in (F.3) stems from | Read[ T ρ ( H, 1)] -y ( H ) |≤ B T + B , the second inequality is ensured by the boundedness of ∥∇ vec[ T ] vec[ g ( T ν ( H,t ) , β )] ∥ as stated in Assumption 2 (iii), and the fourth inequality is provided by Assumption 3 (iv). The last inequality in (F.3) arises from Assumption 2 (iii) and Lemma C.2. By combining Equation (F.3) with the bounds of J 1 and I 1 , we deduce that J 1 + J 2 ≲ exp( C G W 1 ( ρ, ν )) -1 + ∥ β -˜ β ∥ for some constant C G dependent on the parameters listed in the result, thereby completing the proof.

F.2 Proof of Lemma D.2

Proof. Lemma C.4 demonstrates that ∥ ̂ T Θ ( H,t ) ∥ 2 -col and ∥ ̂ T ˜ Θ ( H,t ) ∥ 2 -col are bounded by B T = B exp( K (1 + A + A 2 )) for any H and t ∈ [0 , 1] . We begin by bounding

To demonstrate that ∥ ̂ G ( β, Θ , t ) -̂ G ( β, ˜ Θ , t ) ∥ ≤ C G 1 ML d (Θ , ˜ Θ) , it suffices to show

and

as the other part for h and w follows a similar proof approach.

To bound J 1 , by Assumption 3 (i) we have

where the fourth inequality utilizes Lemma C.5 and the last inequality uses Lemma C.6. To bound J 2 , by Assumption 2 (ii), we have

Hence, it suffices to show that ∥ ̂ p Θ ( H,t +∆ t/ 2) -̂ p ˜ Θ ( H,t +∆ t/ 2) ∥ F ≲ 1 ML d (Θ , ˜ Θ) to establish J 2 ≤ 1 ML d (Θ , ˜ Θ) . Recalling the formula ̂ p Θ in (C.16), we have ∥ ̂ p Θ ( H,t + ∆ t/ 2) -̂ p ˜ Θ ( H,t + ∆ t/ 2) ∥ F ≤ I 1 + I 2 , where

and

where the first inequality applies Lemma C.8, and the second inequality relies on Assumption 2 (iii) and Assumption 3 (ii). Therefore, we conclude that I 2 ≲ 1 ML d (Θ , ˜ Θ) . By combining the bounds of I 1 and I 2 , we observe that Equation (F.5) holds, thereby establishing ∥ ̂ G ( β, Θ , t ) -̂ G ( β, ˜ Θ , t ) ∥ ≤ C G 1 ML d (Θ , ˜ Θ) for some C G dependent on N , d , r , and the parameters of the assumptions.

Therefore, we only need to show that and

to obtain ∥ ̂ G ( β, ˜ Θ , t ) -̂ G ( ˜ β, ˜ Θ , t ) ∥ ≤ C G (1 + λ ) ∥ β -˜ β ∥ . Here, we only establish the inequality above for f and θ , as the proof of the other inequality follows a similar pattern. Note that by Assumption 3 (iii), we have

where the last inequality applies Lemma C.6. Therefore, we conclude that ∥ ̂ G ( β, ˜ Θ , t ) -̂ G ( ˜ β, ˜ Θ , t ) ∥ ≤ C G (1 + λ ) ∥ β -˜ β ∥ , completing the proof.

F.3 Proof of Lemma D.3

Proof. Lemmas C.1 and C.4 establish that max ∥ T ρ ( H,t ) ∥ 2 -col , ∥ ̂ T Θ ( H,t ) ∥ 2 -col ≤ B T := B exp( K (1 + r + r 2 )) . According to Lemma D.6, there exists an event E with P ( E ) ≥ 1 -exp( -δ ) such that under E , we have

for any H and t = 0 , ∆ t, . . . , ( L -1)∆ t, 1 . Following the same proof procedure as in Lemma D.6, with ρ replaced by ̂ ρ , and bounding only J 1 and J 3 in the proof (as there is no need to utilize Hoeffding's inequality to bridge the difference due to a finite width M ), we could obtain the bound

Wepresent the proof only for the case involving ρ . The bounding of ∥ ̂ G ( β, Θ , t ) -G ( β, ̂ ρ, t ) ∥ F can be derived analogously by substituting ρ with ̂ ρ using Equation (F.9), and skipping the process of bounding ∥ D 1 -D 2 ∥ where D 1 and D 2 will be defined later. The bounding of ∥ G ( β, ̂ ρ, t ) -G ( β, ρ, t ) ∥ can be straightforwardly achieved by combining the results obtained from the other two cases.

By the definitions of the gradients in Equations (C.2) and (C.3), we observe that ∥ ∥ ∥ ̂ G ( β, Θ , t ) -G ( β, ρ, t ) ∥ ∥ ∥ ≲ ∥ ∥ ∥ ̂ G f ( θ, Θ , t ) -G f ( θ, ρ, t ) ∥ ∥ ∥ + ∥ ∥ ∥ ̂ G h ( w, Θ , t ) -G h ( w,ρ, t ) ∥ ∥ ∥ . We will focus on showing that ∥ ∥ ∥ 2( ̂ G f ( θ, Θ , t ) -G f ( θ, ρ, t )) ∥ ∥ ∥ ≲ L -1 + √ δ +log( L +1) M , as the other part of the proof follows a similar approach.

Let's define the following quantities A 1 , A 2 , B 1 , B 2 , C 1 , C 2 :

The universal boundedness of ∥ A 1 ∥ and ∥ A 2 ∥ is implied by Assumption 2 (ii), while the universal boundedness of | B 1 | and | B 2 | is implied by Assumption 1. Additionally, ∥ C 1 ∥ and ∥ C 2 ∥ can be

bounded via Assumption 2 (iii), and one can refer to the proofs of Lemmas C.3 and C.6 for detailed explanations.

We claim that to obtain the result, it suffices to show that

There exists some event E 2 with P ( E 2 ) ≥ 1 -exp( -δ ) such that ∥ C 1 -C 2 ∥ ≲ L -1 + √ δ +log( L +1) M under E ∩ E 2 .

This is because if we can establish the above statements, then under the event E ∩ E 2 with P ( E ∩ E 2 ) ≥ 1 -2 exp( -δ ) , we obtain J = ∥ ∥ ∥ 2( ̂ G f ( θ, Θ , t ) -G f ( θ, ρ, t )) ∥ ∥ ∥ ≲ L -1 + √ δ +log( L +1) M . Given the similarity in proof for ∥ ∥ ∥ 2( ̂ G h ( w, Θ , t ) -G h ( w,ρ, t )) ∥ ∥ ∥ , we deduce that with probability at least 1 -4 exp( -δ ) with respect to the parameter initialization Θ (0) , we have ∥ ∥ ∥ ̂ G ( β, Θ , t ) -G ( β, ρ, t ) ∥ ∥ ∥ ≲ L -1 + √ δ +log( L +1) M . The remainder of the proof focuses on bounding the quantities in statements (i)-(iii).

Proof for statement (i): By Assumption 3 (iv), we have ∥ A 1 -A 2 ∥ ≤ ϕ TT ( N,D, √ N +1 B T , r ) ∥ ̂ T Θ ( H,t ) -T ρ ( H,t ) ∥ F ≲ L -1 + √ δ +log( L +1) M under event E .

Proof for statement (ii): Under event E , it is obvious to see ∥ B 1 -B 2 ∥ ≤ ∥ ̂ T Θ ( H,t ) -T ρ ( H,t ) ∥ F ≲ L -1 + √ δ +log( L +1) M .

Proof for statement (iii): We further define the following quantities

max {∥∇ vec[ T ] vec[ f ( ̂ T Θ ( H,s ) , θ s,j )] ∥ , ∥∇ vec[ T ] vec[ h ( ̂ T Θ ( H,s +∆ t/ 2) , w s,j )] ∥} ≤ B J for some constant B J dependent on the parameters listed in the result. This implies that each column of ∇ vec[ T ] vec[ f ( ̂ T Θ ( H,s ) , θ s,j )] or ∇ vec[ T ] vec[ h ( ̂ T Θ ( H,s + ∆ t/ 2) , w s,j )] has l 2 norm upper bounded by B J as well. Applying Hoeffding's inequality to each column of ∇ vec[ T ] vec[ f ( ̂ T Θ ( H,s ) , θ s,j )] and ∇ vec[ T ] vec[ h ( ̂ T Θ ( H,s + ∆ t/ 2) , w s,j )] , and subsequently calculating the union bound across all columns yields:

and

for any z > 0 . For (F.10) and (F.11), we further the consider the union bound across all s = 0 , ∆ t, . . . , ( L -1)∆ , and let z = B J √ 2 M ( δ +log(2( N +1) DL )) , which implies that with probability at least 1 -exp( -δ ) with respect to the parameter initialization Θ (0) , we have

and

bounded by B J √ 2 M ( N +1) D ( δ +log(2( N +1) DL )) ≤ C J √ δ +log( L +1) M for any s = 0 , ∆ t, . . . , ( L -1)∆ . Here, C J is some constant that only depends on N,D,r and the parameters of the assumptions. Denote this probability event by E 2 , and we have P ( E 2 ) ≥ 1 -exp( -δ ) . Under E 2 , by Lemma C.8, we have

∥

D

1

-

D

2

∥

For any s = 0 , ∆ t, . . . , ( L -1)∆ , we define

and

Since Assumption 2 (iii) indicates that max {∥ A s,j ∥ , ∥ B s,j ∥} ≲ ∆ t = L -1 , we have

Applying Lemma C.8 once more, we have

Since Assumption 2 (iii) ensures the boundedness of ∥ D 4 ∥ , we have

Therefore, to show that ∥ D 3 -D 4 ∥ ≲ L -1 + √ δ +log( L +1) M , it suffices to show that

By Assumption 3 (iv) and Lemma D.6, we have

and

where the last inequality employs Assumption 2 (i). Therefore, we conclude that

(F.15)

where the third inequality uses Assumption 3 (iv), and the last inequality relies on the Lipschitz continuity as demonstrated in Proposition C.1. Combining (F.12), (F.13), and (F.15) yields ∥ C 1 -

F.4 Proof of Lemma D.4

Proof. Define F ρ ( T ¯ ρ , t ) = ∫ β vec[ g ( T ¯ ρ ( H,t ) , β )] ρ ( β, t ) dβ for any ρ, ¯ ρ ∈ P 2 . From Taylor's expansion, we have

of which the last equality holds as Lemma C.2 shows that ∥ T ρ η ( H,t ) -T ρ ( H,t ) ∥ F = O ( W 2 ( ρ η , ρ )) = O ( η ) , where we hide the constant dependence on B,K,N,r . Therefore, we have

which leads to (D.18).

F.5 Proof of Lemma D.5

Proof. Fix ( β, t ) ∈ P r . Lemma C.2 implies that

for some constant C r dependent on the parameters listed in the result. Our goal is to regulate the difference between ˙ p ρ ( H,t ) and ˙ p ν ( H,t ) to control the the propagation of ∥ p ρ ( H, · ) -p ν ( H, · ) ∥ .

Note that by (C.14) and Assumption 2,

where and

The third inequality of (F.17) uses Lemma C.1 to obtain

and

with Assumption 2(iii) to bound the norm of the Jacobian matrix with L r, 1 . The last inequality of (F.17) employs Lemma C.3 to bound ∥ p ν ( H,t ) ∥ F with L r, 2 . Applying the Grönwall's inequality to (F.17), we obtain

Since

Thus, we complete the proof of the first result.

By Lemma C.3, under Assumption 1 we have

In addition, by Lemma C.1, under Assumption 2 (i) we have

Therefore, for the gradient function δQ δρ , by Lemma C.3 we have

Hence, it suffices to show that for any H such that ∥ H ∥ 2 -col ≤ B , we have ∥ g ( T ρ ( H,t ) , β ) -g ( T ν ( H,t ) , β ) ∥ 2 -col ≤ L r, 5 ∥ ρ -ν ∥ 1 for some L r, 5 > 0 in order to obtain the second result of this lemma. By Assumption 2 (iii), we see that

where the second inequality again uses Lemma (C.2). Combining (F.19) and F.20 completes the proof of the second result.

F.6 Proof of Lemma D.6

Proof. Denote the empirical distribution of ¯ T Θ and ˜ T ρ by

respectively. It's straightforward to verify that ¯ ρ and ˜ ρ meet the conditions outlined in Lemma C.1, and T ¯ ρ = ¯ T Θ and T ˜ ρ = ˜ T ρ . Hence, Lemmas C.1 and C.4 indicate that max {∥ ̂ T Θ ( H,t ) ∥ 2 -col , ∥ ¯ T Θ ( H,t ) ∥ 2 -col , ∥ ˜ T ρ ( H,t ) ∥ 2 -col } ≤ B exp( K (1 + r + r 2 )) for any H and t ∈ [0 , 1] . We then define B T := B exp( K (1 + r + r 2 )) .

The following decomposition equation holds:

Our proof will bound J 1 , J 2 and J 3 , possibly in a probabilistic manner, to obtain the desired result. Bounding J 1 : Note that according to Assumption 2 (i),

Bounding J 2 : For any t = 0 , ∆ t, . . . , ( L -1)∆ , i ∈ [ N + 1] , j ∈ [ M ] , we have ∥ g ( ¯ T Θ ( H,t ) , β t,j ) : ,i ∥ ≤ B T . Hence, by the Hoeffding's inequality, for any z > 0 we have

By the union bound over i ∈ [ N +1] and t = 0 , ∆ t, . . . , ( L -1)∆ , the above inequality implies

Denote the event such that

by E δ . (F.26) directly indicates P ( E δ ) ≥ 1 -exp( -δ ) .

Suppose that the high probability event E δ occurs. Let's denote B T √ 2 M -1 ( δ +log(( N +1) L )) by B δ for brevity. From Assumption 2 (iii), it follows that for any t = 0 , ∆ t, . . . , ( L -1)∆ ,

Repeatedly applying Equation (F.27) yields

for some constant C 2 > 0 dependent on the parameters listed in the result.

Bounding J 3 : It's worth noting that the convergence proof with a convergence rate of O (∆ t ) = O ( 1 L ) for ˜ T ρ ( H,t ) , the first-order Euler method for T ρ ( H,t ) , is non-standard. This departure from convention arises because we do not assume the boundedness of the second-order derivative d 2 T ρ ( H,t ) dt 2 , instead relying on the continuity of ρ ( · , t ) with respect to the depth index t . In this proof, O ( · ) hides dependencies on N , D , r , and the parameters of the assumptions.

From the definition of ˜ T ρ and T ρ , we have

To bound I 1 , we use (3.1) to get

Given that Proposition C.1 establishes the Lipschitz continuity of T ρ ( H,t ) with respect to t under the condition that ρ ∈ P 2 has a bounded support, we can conclude:

for some constant C 3 , 1 > 0 dependent on the parameters listed in the result. Furthermore, Lemma C.1 and Proposition C.1 demonstrate that g ( T ρ ( H,t ) , β ) is both bounded and Lipschitz continuous with respect to β . Thus, we have

for some constant C 3 , 2 > 0 dependent on the parameters listed in the result. Substituting (F.31) and (F.32) into (F.30), we find that there exists a constant C 3 , 5 dependent on the parameters listed in the result such that I 1 ≤ C 3 , 5 ∆ t 2 .

Additionally, Assumption 2 (iii) implies that

Therefore, by bounding I 1 + I 2 , we obtain the following inequality

which implies, after being used multiple times, that

for any t = 0 , ∆ t, . . . , ( L -1)∆ t, 1 and some constant C 3 dependent on the parameters listed in the result. Combining (F.24), (F.28), and (F.35) yields the desired result.

F.7 Proof of Lemma E.1

Proof. Let ˜ µ ( t ) be the measure induced by T ρ ( H,t ) with H ∼ µ , and ˜ µ be ˜ µ ( t ∗ ) . By verifying Lemma C.1, we establish that ∥ T ρ ( H,t ) ∥ 2 -col ≤ B T := B exp( K (1 + r + r 2 )) for any H and

t ∈ [0 , 1] . Consequently, supp( ˜ µ ( t )) ⊂ { T : ∥ T ∥ 2 -col ≤ B T } for any t ∈ [0 , 1] . The remainder of the proof involves four steps:

Step I: Show that p ρ ( H,t ) is Lipschitz continuous with respect to ( H,t )

In the proof of this proposition, when referring to the Lipschitz continuity of a function, we imply its Lipschitz continuity for H within the support of µ . Recall that we have shown in Lemma C.3 that p ρ ( H,t ) is universally bounded and Lipschitz continuous for ∥·∥ F and any t ∈ [0 , 1] .

From the formula of p ρ in (C.14), for any H,H ′ , we have

| Read[ T ρ ( H ′ , 1)] -y ( H ′ ) | ︸ ︷︷ ︸ I 3

From Proposition C.1, we observe that p ρ ( H,t ) is Lipschitz continuous. Since y ( · ) is K y -Lipschitz continuous, as given in Assumption 4 (iii), we obtain that I 1 ≲ ∥ H -H ′ ∥ F . Moreover, from Assumption 1 and Lemma C.1, we see that I 3 is universally bounded. Therefore, to show that p ρ ( H,t ) is Lipschitz continuous for ∥·∥ F , it suffices to demonstrate that I 2 ≲ 1 and I 4 ≲ ∥ H -H ′ ∥ F .

From Assumption 2 (iii), we have

Thus, dividing I 4 by the uniformly bounded part I 2 , we obtain

where the final inequality holds by Assumption 3 (iv).

By the Lipschitz continuity of T ρ ( H,t ) for ∥·∥ F as stated in Proposition C.1, we have sup t ∈ [0 , 1] ∥ T ρ ( H,t ) -T ρ ( H ′ , t ) ∥ F ≲ ∥ H -H ′ ∥ F . Hence, utilizing the universal boundedness of the last equation in (F.37), we derive I 4 ≲ ∥ H -H ′ ∥ F .

Considering all assertions regarding I 1 , I 2 , I 3 and I 4 , we conclude that p ρ ( H,t ) is C p -Lipschitz continuous with respect to H for some constant C p dependent on the parameters listed in the result that is sufficiently large. The Lipschitz continuity of p ρ ( H,t ) with respect to t could be easily derived from the boundedness of the Jacobian matrix, as asserted in Assumption 2 (iii). Moreover, since p ρ ( H,t ) is universally bounded shown in Lemma C.3, we conclude that p ρ ( H,t ) is Lipschitz continuous with respect to ( H,t ) for some universal Lipschitz constant C p dependent on the parameters listed in the result that is sufficiently large.

Step II: Prepare bounds related to p ρ for later use

(C.14) implies that p ρ solves the adjoint equation

with

implied by Assumption 2 (iii). Therefore, the Grönwall's inequality directly indicates that

for any t ∈ [0 , 1] .

Step III: Construct the descent direction ν

By the well-posedness of the ODE solution to (3.1) as shown in Proposition C.1, the solution map T ρ ( H, · ) is invertible. Hence, for any ρ ∈ P 2 , there exists a continuous inverse map T -1 t such that T -1 t ( T ρ ( H,t )) = H for any H and t ∈ [0 , 1] . Let's define the following function ¯ F ( T ) to approximate

The function ¯ F ( T ) , arising from the composition of p ρ ( · , t ) and T -1 t ∗ ( · ) , exhibits continuity over T ∈ supp( ˜ µ ) owing to the continuous nature of p ρ ( · , t ) and T -1 t ∗ ( · ) . Therefore, Assumption 4 (ii) could be applied to ¯ F . Since f ( · , θ ) is a universal kernel constrained on θ ∈ R dim θ 1 ×K [52], and ¯ F ( T ) is continuous with respect to T , there exists a sequence { ( c k , θ k ) } k ≥ 0 ⊂ ( R × R dim θ 1 ×K ) N such that

given some ϵ > 0 and any T such that ∥ T ∥ 2 -col ≤ B T . Notably, since f ( T, θ k 1 , θ k 2 ) = -f ( T, -θ k 1 , θ k 2 ) , we could assume without loss of generality that c k ≥ 0 for any k . Furthermore, there exists a constant k ϵ such that

We define C ( ϵ ) := ∑ k ϵ k =1 c k , and ¯ ν ( β ) ∈ P ( R dim β ) as the probability distribution such that, given ¯ β ∼ ¯ ν and for any k ≥ 0 , ¯ β has probability c k /C ( ϵ ) of being (2 C ( ϵ ) θ k 1 , θ 2 , 0) . Then, (F.40) transforms into where

From (F.41), we claim that there exists some R ( ϵ ) such that

¯ ν ( { β : 1 /R ( ϵ ) ≤ ∥ β ∥ ≤ R ( ϵ ) } ) ≥ 1 / 2 ,

∥ ∥ ∥ ∫ ∥ β ∥ >R ( ϵ ) g ( T, β )¯ ν ( β ) dβ ∥ ∥ ∥ max ≤ ϵ/ 3 .

We are now in the position to define the descent direction ν . By defining ν ∈ P ( R dim β ) as the measure obtained by truncating any part outside 1 /R ( ϵ ) ≤ ∥ β ∥ ≤ R ( ϵ ) from ¯ ν , and scaling the measure function by 1 / ¯ ν ( { β : 1 /R ( ϵ ) ≤ ∥ β ∥ ≤ R ( ϵ ) } ) ≤ 2 , we can establish that

for any T such that ∥ T ∥ 2 -col ≤ B T . A straightforward deduction from (F.42) is that for any T such that ∥ T ∥ 2 -col ≤ B T , we have ∥ F ( T ) -∫ β g ( T, β ) ν ( β ) dβ ∥ F ≤ ϵ √ ( N +1) D . It is clear that ν has

a bounded support as { β : 1 /R ( ϵ ) ≤ ∥ β ∥ ≤ R ( ϵ ) } . We will determine the value of ϵ later, ensuring it based only on N , D , r , and the parameters of the assumptions.

Utilizing the gradient definition in (3.4) and ∫ β g ( T, β ) ρ ( β | t ∗ ) dβ = F ( T ) + p ρ ( T -1 t ∗ ( T ) , t ∗ ) , we obtain

For ρ ∈ P 2 concentrated on P r , we observe

Hence, to bound J 1 , we have

where C 3 := ( N +1)( B T + B )( N +1) D exp( C 0 ) .

To bound J 2 , we have

Combining (F.43) and (F.44), by choosing ϵ = 1 4 exp( -2 C 0 ) /C 3 , which only depends on N , D , r , and the parameters of the assumptions, we have

F.8 Proof of Lemma C.1

Proof. By applying the Cauchy-Schwarz inequality, we trivially obtain ∫ 1 0 ∫ β || β || 2 ρ ( β, t ) dβdt ≤ A . Thus, leveraging (3.1) and Assumption 2 (i), we can infer

Therefore, by the Grönwall's inequality, we have

F.9 Proof of Lemma C.2

Proof. As per Lemma C.1, the boundedness of ∥ T ν ( H,t ) ∥ 2 -col and ∥ T ρ ( H,t ) ∥ 2 -col is established by a constant C := ∥ H ∥ 2 -col exp( K (1 + A + A 2 )) > 0 for all t ∈ [0 , 1] . Consequently, from (3.1), this implies

for any t 1 , t 2 ∈ [0 , 1] . Therefore, T ν ( H,t ) is (1 + C ρ / 2) KC (1 + r + r 2 ) -Lipschitz with respect to t for ∥·∥ 2 -col , and thus √ N +1(1 + C ρ / 2) KC (1 + r + r 2 ) -Lipschitz with respect to t for ∥·∥ F . Note that by (3.1),

We then bound J 1 and J 2 using the following two lemmas separately. Firstly, since ∥ T ρ ∥ F ≤ √ N +1 C and ∥ T ν ∥ F ≤ √ N +1 C , we have by Assumption 2 that

Therefore, by (F.48) we have

Secondly, we aim to bound the integral J 2 given Assumption 2 and ∥ T ν ( H,t ) ∥ 2 -col ≤ C on t ∈ [0 , 1] . Again by Assumption 2 we have

Since T ν ( H,s ) is N +1(1+ C ρ / 2) KC (1+ r + r 2 ) -Lipschitz with respect to s for ∥·∥ F (as shown in (F.46)), by (F.51), we obtain that g ( T ν ( H,s ) , β ) is √ N +1(1+ C ρ / 2) KCϕ T ( N,D, √ N +1 C )(1+ r + r 2 ) 2 -Lipschitz with respect to s . Thus, g ( T ν ( H,s ) , β ) is C ′ -Lipschitz with respect to ( s, β ) , where C ′ = √ N +1(1 + C ρ / 2) KCϕ T ( N,D, √ N +1 C )(1 + r + r 2 ) 2 + √ N +1 ϕ P ( C )(1 + r ) . This indicates, by the Kantorovich-Rubinstein Theorem (see Theorem 5.10 of [68], for example), that

Define C ∗ := max { C ′ , ϕ T ( N,D, N +1 C )(1 + r + r 2 ) } . By combining (F.49) and (F.52), we have

Applying the Grönwall's inequality then shows

F.10 Proof of Lemma C.3

Proof. Let's consider a fixed ( β, t ) pair within P r . Given Lemma C.1, we establish ∥ T ρ ( H,t ) ∥ 2 -col ≤ B exp( K (1 + A + A 2 )) . Consequently, under Assumption 1, we deduce

On the other hand, from (C.14), we have

Combining (F.55) and (F.56), we could obtain

F.11 Proof of Lemma C.4

Proof. By employing the Cauchy-Schwarz inequality, we readily observe that 1 ML ∑ t ∑ M j =1 ∥ β ∥ ≤ A . Consequently, leveraging (2.4) and Assumption 2, we ascertain that for any t = 0 , ∆ t, . . . , ( L -1)∆ t , we obtain:

F.12 Proof of Lemma C.5

Proof. Lemma C.4 shows that ∥ ˜ T Θ ( H,t ) ∥ 2 -col and ∥ ˜ T ˜ Θ ( H,t ) ∥ 2 -col are bounded by B T := B exp( K (1 + A + A 2 )) for any H and t ∈ [0 , 1] . From (2.4), for any t = 0 , . . . , ( L -1)∆ t , from Assumption 2 (ii) and (iii) we have

for some constant C and C depending only N,D,r and assumptions. Similarly, we have

1 2 ∥ ̂ T Θ ( H,t +∆ t ) -̂ T ˜ Θ ( H,t +∆ t ) ∥ F ≤∥ ̂ T Θ ( H,t +∆ t/ 2) -̂ T ˜ Θ ( H,t +∆ t/ 2) ∥ F (1 + C 1 (∆ t/ 2))

M

∑

j

=1

+(∆

Combining (F.58) and (F.59), we derive

where C 3 is a constant depending solely on N , D , r , and the parameters of the assumptions. Iterating (F.60) multiple times yields

t/

2)

C

2

M

-

1

∥

w

t,j

-

w

t,j

˜

∥

.

(F.59)

F.13 Proof of Lemma C.6

Proof. By verifying that Θ satisfies the conditions outlined in Lemma C.4, we establish ∥ ̂ T Θ ( H,t ) ∥ 2 -col ≤ B exp( K (1 + A + A 2 )) . The first two results stem from Assumption 2 (i) and (ii), with recognition that ∥ T ∥ F ≤ √ N +1 ∥ T ∥ 2 -col for any T ∈ R D × ( N +1) . As for the third result, consider t = 0 , ∆ t, . . . , ( L -1)∆ t, 1 , where

where the first inequality arises from the fact that the matrix 2 -norm is greater equal than the norm of any of its columns, and the fourth inequality follows from Assumption 2 (iii).

F.14 Proof of Lemma C.7

Proof. Note that ∥ β t,j ∥ ≤ r for any t = 0 , ∆ t, . . . , ( L -1)∆ t and j ∈ [ M ] with its expectation denoted as ∫ β ∥ β t,j ∥ ρ ( β, t ) dβ . Applying Hoeffding's inequality yields, for any z > 0 and t = 0 , ∆ t, . . . , ( L -1)∆ t

By applying the union bound over t = 0 , ∆ t, . . . , ( L -1)∆ t , the inequality above implies

In addition, we have

Combining (F.62) and (F.63), and setting z = r √ 2 M -1 ( δ +log(2 L )) completes the proof.

F.15 Proof of Lemma C.8

Proof. The proof will be trivial by noting the equality

Hence, we have

G Assumption verification for a concrete example

In this section, we consider

with the collection of parameters θ = vec[ V, W ] , where softmax denotes the column-wise softmax function. Moreover, consider

with the collection of parameters w = vec[ W 1 , W 2 ] . Here, HuberizedReLU denotes the entry-wise HuberizedReLU activation function defined as

Then, we can consider a Transformer model defined by equations (2.1), (2.2), and (2.4) in the paper, where the functions f and h are specified above. We suppose that this Transformer model is applied to a learning task with data that satisfies Assumption 1. We have the following proposition.

Proposition G.1. Consider the Transformer model defined by equations (2.1) , (2.2) , and (2.4) , with f ( Z, θ ) and h ( Z, w ) defined in (G.1) and (G.2) respectively. Then Assumptions 2-4 all hold.

Proof. We omit the detailed derivations for the function h ( Z, w ) , which corresponds to the MLP part, in our verification of Assumptions 2 and 3, as h ( Z, w ) satisfies Assumptions 2 and 3 is relatively more intuitive, especially given the proofs for f ( Z, θ ) .

Denote Z = ( z 1 , . . . , z N +1 ) ∈ R D × N +1 . Then the function f can be rewritten as

where f ( Z, θ ) : ,i = ∑ N +1 j =1 P ij V z j and P i, : = softmax( Z ⊤ Wz i ) . Next, we calculate the derivatives of f ( Z, θ ) : ,i with respect to Z and θ as follows:

For Z : the Jacobian J ∈ R ( N +1) D × ( N +1) D is J = ( J ij ) 1 ≤ i,j ≤ N +1 , where J ij = ∂f : ,i ∂z j ( Z, θ ) ∈ R D × D . After calculation, we obtain

where Q i := diag( P i : ) -P ⊤ i : P i : , E ij is the matrix with zeros everywhere except one the ( i, j ) -th entry, and δ ij is the Kronecker delta ( 1 if i = j , 0 otherwise).

For θ : Define A i = Z ⊤ Wz i . After calculation, we have

Thus, by the chain rule, we have

where ( z ⊤ j , . . . , z ⊤ j ) ⊤ contains D copies of z j . We then verify the assumptions one by one. For Assumption 2 (i), we have

where the second-to-the-last inequality follows by the fact that ℓ 2 -norm can be upper bounded by the ℓ 1 -norm, and the last inequality follows by the fact that each column of the softmax output has an ℓ 1 -norm equaling one. Therefore, the first condition in Assumption 2 with K = 1 is verified for the function f in (G.1).

For h in (G.2), we have

where the second inequality follows by the property of HuberizedReLU that | HuberizedReLU( x ) |≤ | x | . This demonstrates that Assumption 2 (i) with K = 1 holds for h in (G.2) as well.

For Assumption 2 (ii), (G.5) leads to

Moreover, (G.4) leads to

Moreover, we have

Combining the two equations above gives Assumption 2 (ii) with ϕ P ( ∥ T ∥ 2 -col ) = ∥ T ∥ 2 -col + 2 ∥ T ∥ 2 2 -col .

For Assumption 2 (iii), we have

For any 1 ≤ i, j ≤ N +1 , we have

Hence, we have

The above equation demonstrates that for f , Assumption 2 (iii) holds with ϕ T ( N,D, ∥ T ∥ F ) = ( N + 1)(1 + 2 ∥ T ∥ 2 F ) . We have verified Assumption 2 for the attention layer encoder f . The verification for h is similar and easier.

Next, we verify Assumption 3. Given that we are currently considering the example where the encoder employs a smooth univariate activation function, we can prove stronger results by removing the expectation E µ .

(i) and (iii): Given the calculation of derivatives in (G.4) and (G.5) we have presented above, we first show that

is locally Lipschitz continuous with respect to Z and θ . By (G.3) and the chain rule, we can derive that

and the local Lipschitz continuity is then obvious given the boundedness of ∇ W kl P ij and ∂P ij ∂z j with respect to necessary parameters.

As we prove that P ij is locally Lipschitz continuous, given that then each component in (G.4) and (G.5) is locally Lipschitz continuity with respect to both Z and θ , and is obviously bounded by an increasing function of N,D, ∥ θ ∥ , K T , L T . Then the local Lipschitz continuity is straightforward as they are all sufficiently smooth.

(ii) and (iv): Because the norm of the difference of two Jacobian matrices ∥ J 1 -J 2 ∥ 2 is bounded by √ ∑ 1 ≤ i,j ≤ N +1 ∥ J 1 ij -J 2 ij ∥ 2 2 , it suffices to show that J ij is locally Lipschitz continuous with respect to both θ and Z. Again each component of J ij that depends on Z or θ, i.e. Z, Q i , W, P ij , is bounded by an increasing function of N,D,K P , L T , K T , ∥ θ ∥ , and is locally Lipschitz continuous given sufficient smoothness. Hence, (ii) and (iv) also hold.

For Assumption 4, we consider the pair ( g, α ) = ( h, w ) , and the partition α = ( α 1 , α 2 ) with α 1 = W 2 , α 2 = W 1 . We also let a compact set K = { W 1 : ∥ W 1 ∥ ≤ 1 } . Then Assumption 4 (i) on the partial 1 -homogeneity property straightforwardly holds:

Regarding Assumption 4 (ii) on the universal kernel property, we first note that according to the choice ( g, α ) = ( h, w ) , this assumption is purely an assumption on the MLP part of the Transformer. Here we give the detailed proof as follows.

First of all, according to the classic universal approximation theory (see the wiki page of 'universal approximation theorem' and [34, 18, 57] for more details), we know that two-layer fully-connected

networks with non-polynomial activation functions and without any constraints on its parameters are universal approximates.

Therefore, we know that the function class span { W 2 ReLU 2 ( W 1 T ) : W 1 ∈ R dim( W 1 ) , W 2 ∈ R dim( W 2 ) } is dense in C ( ∥ T ∥ 2 -col ≤ B, R D × ( N +1) ) . Moreover, by the definition of HuberizedReLU, for any B > 0 and any ̂ W 1 , ̂ W 2 , there exist small constant c such that c · ̂ W 1 ∈ K , c · ∥ ̂ W ∥ 1 ≤ B -1 , and

where the second equation follows by the positive 2 -homogeneity of ReLU 2 activation. This implies that

Therefore, we conclude that span { W 2 HyberizedReLU( W 1 T ) : W 2 ∈ R dim( W 2 ) ×K} is dense in C ( ∥ T ∥ 2 -col ≤ B, R D × ( N +1) ) . This finishes the validation of Assumption 4.

H Experiments

As discussed in Sections 3 and 4, our mean-field approximation results and global convergence results are asymptotic guarantees requiring exponentially large number of heads M and number of layers L . Such results are due to the nature of mean-field type analysis. In practice, we frequently observe that global convergence can be achieved by Transformer models of reasonable sizes. In this section, we run simple experiments on training Vision Transformers (ViT) [24] on the CIFAR-10 datasets to demonstrate global convergence in practical applications.

We train Vision Transformers with different numbers of heads and layers. In all our experiments, we split each CIFAR-10 image into four patches and then pass the patches into Vision Transformer models. We keep the dimension of each attention head to be 128 . The output of each self-attention layer is passed through a single-hidden-layer feedforward component with 128 hidden neurons and GeLU activation. Both the self-attention and feedforward components include skip connections. We implement dropout in the self-attention layers as well as the feedforward layers with a dropout probability of 0.1. The model is attached to a linear classifier.

In all experiments, we train the ViT models using Adam for 200 epochs with a mini-batch size 512 . We set the initial learning rate to be 1 e -4 , and implement a cosine annealing learning rate schedule. We do not use any data augmentation or explicit regularization techniques, so that global convergence for large enough models implies close-to-zero training loss and close to 100% training accuracy.

In the first set of experiments, we fix the depth of the ViT to 6 layers (i.e., there are six self-attention layers, each followed by a single-hidden-layer feedforward component). We train such Vision Transformers with the numbers of heads per layer ranging from 4 to 40 , and record the training loss and training accuracy throughout training. The results are given in Figure 1. Based on the results, it is clear that for ViT models with more than 20 heads can achieve close-to-zero training loss and close to 100% training accuracy, demonstrating global convergence on the CIFAR-10 training data.

In the second set of experiments, we fix number of heads in the ViT model per layer to 8. We train such Vision Transformers with depths ranging from 2 to 20 , and record the training loss and training accuracy throughout training. The results are given in Figure 2. Again, the results indicate that ViT models with more than 16 layers can achieve close-to-zero training loss and close to 100% training accuracy on the CIFAR-10 dataset, implying global convergence.

We note that all these experiments are conducted on a standard GPU card. We can observe clear global convergence when the Vision Transformer is sufficiently wide or deep, but still within reasonable scales. This indicates that, although our theoretical guarantees require extremely large numbers of heads and layers due to the limitations of the mean-field technical tools, global convergence can be achieved by Transformers of reasonable sizes in practice.

(a) Training loss

(b) Training accuracy

2.0

4 heads, 6 layers

100

8heads, 6 layers

12 heads, 6 layers

06

16 heads, 6 layers

1.5

20 heads,6 layers

80

24 heads, 6 layers

2

training loss

28 heads, 6 layers

accurac

70

32 heads,6 layers

4 heads, 6 layers

1.0

36heads,6 layers

40 heads, 6 layers

training

60

8heads, 6 layers

12 heads, 6 layers

16 heads, 6 layers

50

20 heads, 6 layers

24 heads, 6 layers

0.5

40

28 heads, 6 layers

32 heads,6 layers

30

36 heads, 6 layers

0.0

40 heads, 6 layers

20

0

T

25

50

75

100

125

150

175

200

0

T

25

50

75

T

100

125

150

175

200

(a) Training loss

(b) Training accuracy

8heads,2 layers

100

2.0

8heads, 4 layers

8heads, 6 layers

06

8heads, 8 layers

8 heads, 10 layers

80

1.5

8 heads, 12 layers

2

SSO

8 heads, 14 layers

accurac

70

trainingl

1.0

8heads,18 layers

8 heads, 16 layers

8 heads, 20 layers

training

60

8heads,2 layers

8heads, 4 layers

8heads,6 layers

8heads,8layers

50

8 heads,10 layers

8 heads, 12 layers

0.5

40

8 heads, 14 layers

8 heads, 16 layers

30

8 heads, 18 layers

0.0

20

8 heads,20 layers

0

T

25

50

75

100

125

150

175

200

0

T

25

50

75

100

125

150

175

200

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction accurately summarize the main contributions and scope of the paper (Section 1.1), clearly outlining the theoretical advancements while aligning with the stated assumptions.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in 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]

Justification: We operate under noiseless label settings, as noisy conditions typically preclude achieving zero risk. The limitations of Assumption 4 and the assumptions in Theorem 4.1 are discussed in the respective sections. Although some assumptions are currently untestable, we provide high-level justifications for their validity.

Guidelines:

The answer NA means that the paper has no 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 a 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 on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a 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 important 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 necessary assumptions (Assumptions 1-4) are stated in the main content prior to presenting the theoretical results, with complete and correct proofs included in the supplementary material.

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We clearly disclose all the information for reproducing experimental results.

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 a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to 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 to 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 case authors are welcome to describe 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 to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: This paper includes only simple experiments in the supplemental material to verify the main theoretical results. Therefore, the access to the data and code is not central to the paper's contribution.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we 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 a 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/public/guides/CodeSubmissionPolicy ) 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.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a 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) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: This paper specifies all necessary details to understand the results.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: This paper reports appropriate information about the statistical significance.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence 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).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It 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 a 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 in 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 computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: This paper provides sufficient details on the computer resources used for the simulation studies. The task size is small and can be easily reproduced using a standard GPU.

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 as 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 it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research presented in this paper fully complies with the NeurIPS Code of Ethics. All ethical guidelines pertinent to the theoretical nature of the work have been adhered to rigorously.

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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or 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: [NA]

Justification: The paper does not discuss societal impacts as it strictly addresses theoretical optimization guarantees for Transformer models, focusing solely on providing rigorous proofs without exploring practical applications.

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 (e.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 not tied to particular applications, let alone deployments. However, if there is a 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, it 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 as 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 or unintentional) misuse of the technology.

If 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper is theoretical and does not involve the release of data or models. Therefore, there are no risks for misuse or dual-use associated with this research.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The public CIFAR-10 datasets are explicitly mentioned and properly respected.

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 a 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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper is theoretical and does not involve the use of existing assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an 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: The paper is theoretical and does not involve crowdsourcing nor research with human subjects.

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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper is theoretical and does not involve crowdsourcing nor research with human subjects.

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 to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper considers the theoretical mean-field limit of Transformers where width and depth go to infinity and studies the approximation error and convergence properties of gradient flow with weight decay. Both a residual self-attention layer and a residual feedforward layer are approximated by an ODE which averages the two encoders, whose solution models distribution of the parameters of both blocks throughout the depth of the Transformer. It is shown that the dynamics of the discretized model approximates the Wasserstein gradient flow under some  regularity assumptions. Moreover, it is shown under partial homogeneity and universal kernel assumptions that if the Wasserstein flow weakly converges to a stationary distribution, then the discrete model must also converge with arbitrarily small risk.\n\nSTRENGTHS:\n* The paper is a nontrivial extension of existing mean-field analyses of two-layer neural networks and ResNets to Transformers, which is challenging due to the existence of both feedforward and attention layers, the latter of which is new in the literature.\n* The assumptions for the networks and input data are quite general and can encompass e.g. the softmax mechanism and in-context learning settings.\n* The analyses and techniques are sufficiently novel and detailed and serve as an initial characterization of the dynamics of large scale Transformers, in particular the global convergence analysis of Section 4.\n\nWEAKNESSES:\n* The obtained error bounds (Theorem 3.1) are so large that they likely have little practical relevance beyond the initial moments of training. In particular, the hidden constants are in the worst case super-super-exponential with respect to the time horizon. Specifically, the bound is exponential in $\\phi_T(N,D,C\\times B_\\tau)$ where $B_\\tau$ is exponential in $R_\\tau$, which in turn is exponential in $\\tau$. While bounds exponentially diverging in time are common in the literature (which is a priori expected from an ODE discretization argument) as mentioned in  the paper, the dependency in Theorem 3.1 is much worse and some effort needs to be made to at least justify this. For example, one exponentation seems to be removable if e.g. $\\phi_T$ is bounded. Currently the dependency in the hidden constants are not made obvious without going through the proofs in the Appendix, which I feel should be addressed more transparently.\n* The convergence analysis (Theorem 4.1) also suffers from the same problem. The analysis requires a time horizon $\\tau_0$ large enough so that $W_2(\\rho^{(\\tau)}, \\rho_\\infty)$ is exponentially small (w.r.t. $R_\\infty$) compared to the desired error $\\epsilon$. This horizon is then fed into the approximation bound of Theorem 3.1, again yielding very large constants unless convergence is achieved in the very early stages of training. (This issue seems less critical since it only affects the $C_1$ term, although the $C_2\\lambda$ term is still super-exponential in $R_\\infty$.)\n* The recurring rate $L^{-1}+\\sqrt{\\log L/M}$ is not shown to be tight, leading me to further question the utility of the provided bounds. Why is the scaling for $L,M$ different? How does the rate compare with the analysis for ResNets or other ODE systems with 'width' & 'depth' dimensions?\n* A nitpick: the terms 'universal constant' or 'universal bound' are used quite frequently, however they should be reserved for constants that do not depend on *any* problem parameters.\n\nQUESTIONS:\n* Can the dependency on time horizon be alleviated with stronger model assumptions?\n* Is the $L^{-1}+\\sqrt{\\log L/M}$ rate tight/expected? (see Weakness)\n* The widths of both the feedforward layer and the attention layer are both set equal to $M$, however the number of heads cannot typically be expected to be very large compared to the width of feedforward layers. Does the analysis generalize to when they are different?\n* Some recent papers [1,2] have also studied mean-field limits of Transformers from different perspectives, it would be nice to add a comparison in the related works section (although the latter is a contemporary work).\n\n[1] https://openreview.net/forum?id=xm2lU7tteQ\n\n[2] https://arxiv.org/abs/2405.15712\n\nLIMITATIONS:\nSee Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors theoretically investigate the mean field limit of residual transformer models. As depth and width go to infinity, thanks to the residual structure, the forward pass can be modeled as an ordinary differential equation, and the training gradient flow converges to a Wasserstein gradient flow. The author show well-posedness of this PDE, provide non-asymptotic time uniform bounds on quantities of interest and show global convergence to the global minimum of the empirical risk as $L,m/log(L) \\to+\\infty$ and $\\lambda \\to 0$.\n\nSTRENGTHS:\nThe tackled problem is extremely relevant and the derived results are sound. Given the size of the work, the authors did a good job overall in organizing the contents.\n\nWEAKNESSES:\n1. The content is intense and the presentation doesn't help a novel reader. I believe polishing of the manuscript would improve the quality. \nMore concretely, I would add a small paragraph in the main manuscript summarizing the main ideas behind the novelties. In the appendix \t\t\tinstead it is really easy to get lost in a long proof, loosing track of the main goal. I suggest to include a small paragraph at the beginning of \t\teach section with the main steps required for the final proof, to guide the reader and make even the technical part more accessible \t\t\t(something on the line of what the authors did in D.3 would be perfect). \n\n2. While this did not and will not impact my score, for manuscripts of this size I believe at least a really small experimental section (even in the appendix) to test the claims would make the results more sound.\n\nQUESTIONS:\nOn the technical side: \n1. While I understand the need of assumption i) in theorem 4.1, it is not obvious to me how stringent it is. Already for an extremely simple problem like $Q(\\rho) = \\int \\rho \\log(\\rho) + \\lambda \\int x^2 d\\rho$ the minimizer is a Gaussian, that has unbounded support. \tGiven the last example, the heuristic justification given by the authors seems to not be sufficient. Do the authors see a way to transfer this condition on the growth rate of $R(\\Theta)$ at infinity?\n2. The immediate extension of the deterministic case would be to consider the full Fokker-Plank, as it may emerge as a model of stochastic optimization. Already in the simplest case, the PDE under consideration would be $\\partial_\\tau \\rho^{(\\tau)} = div(\\rho^{(\\tau)} \\nabla U) + \\sigma \\Delta \\rho^{(\\tau)}$. The stationary distribution in this case would be $\\rho \\propto e^{-U/\\sigma}$, that has unbounded support for $U$ defined on all parameter space. This suggest that already in the simplest case (just pure isotropic diffusion added), the extension of the results may not be straightforward. I would appreciate the authors' comment on this, and I would like to know if they see an easy way to weaken that hypothesis as it's something one has no control on.\n3. While mentioned in the proofs, it is not clear in the main manuscript what $\\delta$ is. For example, in theorem 4.1 the statement is \"with high probability\", but it is not mentioned on what. I would add two words in the theorems/lemmas that are in high probability to specify with respect to what. \n4. It sounds strange to me that theorem 3.1 gives a bound on a compact set, while theorem 4.1 on an unbonded time horizon, with what appears to be exactly the same constant. Is the constant in theorem 4.1 really the same of the one in 3.1 or it includes other terms?\n\nLIMITATIONS:\nThe authors correctly assessed the limitations in the manuscript, I have no further suggestions.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-06"}

comments: SUMMARY:
This paper analyzes gradient flow on Transformer networks. It is shown that for wide and deep Transformers, gradient flow converges to the Wasserstein gradient flow and reaches a global minimum.

STRENGTHS:
This is a well-written paper and it seems the results are strong and clean. However, I have very little background in this area and cannot verify the correctness of the statements. I hope the AC can find another qualified reviewer and delete this review if possible.

WEAKNESSES:
N/A

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies transformers in their mean field limit. They use a ResNet architecture with infinite depth. The residual blocks are made of two steps: one transformer step and one standard MLP. They also use infinite width for both steps.\nThen, they show that the gradient flow in this limit is well posed, by using the notion of Wasserstein gradient flows since the mean field model is defined on space of measures. They also show consistency with the original discrete transformer model.\nThen, they study the global convergence properties of the gradient flow: they prove a sort of global convergence result, which requires some hypotheses.\n\nSTRENGTHS:\nThe writing of the paper is quite clear.\n\nProving mean field convergence results has been studied a lot in the literature on theoretical deep learning. The result of the authors is interesting because they deal with the transformer architecture which ends up being a sum of two maps in the mean field representation. The authors can prove similar results in this context, which are new to the best of my knowledge of the literature.\n\nWEAKNESSES:\nOne of the main weakness for a submission to NeurIPS is the format of the paper. The 9 pages of the main paper are devoted to a (well done) explanation of the main results. There is no concrete ideas of proof in the main paper. More than 40 pages of technicalities are necessary to validate their result. \n\nThe novelty of the results is the fact that transformer architecture is taken into account in the mean-field limit. However, the stated results do not show to what extent this architecture helps in obtaining the results. In fact, the transformer architecture appears more as something that hinders the standard proof on global convergence. \nThe paper does not shed light on the interest/peculiarity of the transformer architecture.\nI would have appreciated if some simple experiments could be done to assess the range of validity of the results and in particular testing the importance of the transformer architecture.\n\nThe hypothesis in section 4 to obtain a global convergence result seems not checkable in practice. In my opinion, it is not a global convergence result but rather an \"if\" theorem. Although I know that this assumption has been put forward in other previous papers, it is still an very demanding hypothesis.\n\nQUESTIONS:\nWhat are specific features of the transformer architecture that are necessary to make the results valid? In other words, can you extend your results to more general architecture than transformers?\n\nThe justification of the hypothesis on the separability property in C5 is difficult to understand. Can you elaborate more on this assumption and try to report precisely on the progress of the literature on this assumption?\n\nI do think it is not fair to claim global convergence for such kind of \"if\" theorem with a hypothesis that is uncheckable. I suggest the authors to rephrase their contributions.\n\nLIMITATIONS:\nNot applicable.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a rigorous theoretical analysis of the global convergence properties of gradient flow in training Transformers in the setting of in-context learning. By constructing the mean-field limit, the authors show that as the model's width and depth increase to infinity, the gradient flow converges to a Wasserstein gradient flow. This convergence ensures that the training process reaches a global minimum when the weight decay regularization parameter is sufficiently small. The paper introduces techniques adapted to Transformers, leveraging partial homogeneity and local Lipschitz smoothness to demonstrate the close approximation between discrete and continuous models and to establish the global convergence of gradient flow.\n\nSTRENGTHS:\nStrength:\n\n1. The paper studies the Transformer model in the context of in-context learning. Although the mean-field approach has been applied to infinite width and depth ResNet models before, the Transformer model is currently the most widely used and thus worthy of detailed study. This paper extends the analysis by incorporating two distinct encoders, filling the gap in understanding the Transformer model in infinite limits.\n\n2. The paper introduces new techniques by extending previous approaches used for ResNet models. The authors refine the analysis by assuming only partial homogeneity and local Lipschitz smoothness, which is a key extension. The idea of considering the continuous dynamics as a function of the average behavior of f and h is new and interesting.\n\n3. The final result has a mild requirement for the dependence of depth and width, where one would only need depth $L = \\Theta(\\epsilon^{-1})$ and the number of Transformer block $M = \\Theta(\\epsilon^{-2} \\log (\\epsilon^{-1}))$ to achieve a loss as small as $\\epsilon$.  (However, the reviewer has some questions about the result, so this strength can still be questionable.)\n\n4. The paper is well-organized and clearly written. The authors effectively communicate the ideas, contributions, approaches, and limitations of their work, making it accessible and comprehensible to the reader.\n\nWEAKNESSES:\n## Potential fallacy in the claim of Corollary 4.1:\n\n1. The reviewer is convinced by most results presented in the paper, but has concerns about the claim in Corollary 4.1. The main issue arises from the fact that the authors treat many objects as constants, even though they clearly depend on some key quantities. In short, the reviewer thinks that $C_1$ in Theorem 4.1 depends significantly on $\\epsilon$, which makes the claim in Corollary 4.1 questionable.  \nThe concern stems from the observation that $C$ in Theorem 3.1 may depend exponentially on $\\tau$. Following the proof of Theorem 4.1, to achieve a loss as small as $\\epsilon$, one would need $Q(\\rho^{(\\tau_0)}) \\le \\epsilon$. This introduces a non-trivial dependence of $\\tau_0$ on $\\epsilon$. Given that $C$ may depend exponentially on $\\tau_0$, $C_1$ in Theorem 4.1 would thus depend exponentially on $\\tau_0$, leading to a complex dependence on $\\epsilon$. This dependence cannot be ignored, and treating $C_1$ as a simple constant to achieve Corollary 4.1 appears incorrect. (The reviewer believes this is how Corollary 4.1 is currently derived.)  \nBased on the above, the reviewer thinks that \n    1. The current approach in proving Corollary 4.1 is flawed.\n    2. A naive application of Theorem 4.1 might only lead to asymptotic bounds for $L$ and $M$, instead of the linear and quadratic terms presented nicely in the current version of Corollary 4.1.\n\n## Limitations of settings and results\n\nThough the reviewer agrees that studying and achieving a good understanding of the case where both width and depth go to infinity is valuable, there are some limitations in the results presented in the paper.\n\n1. The main result, Theorem 4.1, depends on Assumptions 1-4. Although the reviewer considers Assumptions 1-4 to be reasonable and likely to hold under certain settings, the authors do not make it clear when these assumptions would hold. Consequently, the setting of the Transformer model might differ from classical Transformer models (e.g., it might only hold under the reparameterized case (2.3) instead of keeping all value, key, query, and output matrices). The reviewer will address this point more in the Questions section.\n\n2. The main result, Theorem 4.1, relies heavily on a few additional assumptions that are not justified in the paper:\n    1. the Wasserstein gradient flow weakly converges to some distribution $\\rho_{\\infty}$\n    2. the uniform boundedness of $\\rho^{(\\tau)}$ for large time $\\tau$\n    3. the separation property for $\\alpha_1$ with the support expansion of $\\alpha_2$ to $K$.\n\nThe reviewer acknowledges that some of these assumptions are also adopted in other mean-field approaches and appreciates the authors' efforts to justify them intuitively in Section C.5. However, they are not examined carefully and may harm the validity of the statement. For example, for the second assumption, the reviewer agrees that a large regularization $\\lambda$ can implicitly bound the norm of the solution. However, this might require a large $\\lambda$, which could make the results less meaningful. Notice that the current bound shown in Proposition 3.2 implies that $R_{\\tau}$ may depend exponentially on $\\tau$, thus it is not clear that a small $\\lambda$ can result in a bounded solution. And based on Theorem 4.1, a small $\\lambda$ would be required to achieve a small loss.\n\n3. The result does not provide a non-asymptotic bound for the required training time. In other words, there is no guarantee that $\\hat{Q}(\\Theta^{(\\tau)})$ will be small for any specific $\\tau$. This is because there is no control over the convergence rate of $\\rho^{(\\tau)}$ to $\\rho^{\\infty}$, which would require a stronger assumption beyond the current weak convergence assumption of $\\rho^{(\\tau)}$.\n\n4. The loss and training dynamics are only considered for the population risk under $\\mu$. There is no consideration of the effect of an empirical dataset. This is a clear limitation, but the reviewer can accept this as an initial step toward studying the Transformer model. \n\n## Minor:\n\n1.  Some parts of the statements are not well organized. For example, Assumption 2 is an assumption about norm bounds of function $f$ and its gradient. The line \"define ReLU'(x) = ...\" has no direct connection with the assumption. The reviewer thinks that the authors are trying to argue that by using ReLU'(x) as the gradient of the ReLU function, $f$ defined in (2.3) would satisfy Assumption 2. If this is the case, the reviewer suggests the authors separate the statement into 1) what the assumption of $f$ is. and 2) when the assumption would hold. Another minor point is that the Fréchet derivative still acts differently from the gradient, so some careful treatment might be needed if the authors claim that the result in Theorem 4.1 applies to Transformers with ReLU activation. \n\nTypos:\n1. Line 111, it should be \"h(Z, w) = W_2 \\sigma_M(W_1 Z)\" instead of $W_1 H$. \n2. Line 156, in Assumption 2, $\\Psi_j$ in the description should be $\\Phi_P$. \n3. Line 86: \"~~the~~ our deep Transformer model\"\n4. Line 901, For the line that bounds $R(\\Theta)$, there is no $C_{\\lambda}$ in the first line, and the transition from the first line to the second line is currently shown as $1/2 \\epsilon + ... \\le 1/4 \\epsilon + ...$\n\nQUESTIONS:\n1. Please address the above concern for Corollary 4.1 carefully. The answer may greatly affect the final rating from the reviewer.\n\n2. Questions related to assumption 1-4:\n    1. For the purpose of confirmation, have you shown in which cases Assumption 1-4 are met for the Transformer model? My current understanding is that they would be met when you reparameterize $f$ as in (2.3) but not when you keep the W_V, W_K, W_Q and W_O matrices (at least Assumption 2 won't be satisfied in this case). Is that correct? The review is not against the reparameterization and considers this is still a Transformer model, but want to understand the difference and the claim of the paper more carefully. \n    2. As the focus of the paper is on the Transformer model, the main results are based on the structure defined in (2.1) (2.2), (2.4) and Assumptions 1-4 over f and h. The review wonders if it is possible to provide a separate corollary for a commonly used form that satisfy Assumption 1-4, such as (2.3), making it clearer on an example where minimum assumptions are used. Or in other words, are there technical difficulties in showing Assumption 1-4 for (2.3)? Is the universal approximation capabilities the key obstacle?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitation of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-29"}