Global Convergence in Training Large-Scale Transformers

Authors: cheng  gao
yuan  cao
zihao  li
yihan  he
mengdi  wang
han  liu
jason m klusowski
jianqing  fan
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 06:08:09

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.

Main paper: 
 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 ii. 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.iii. 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  
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.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:H = [h 1 , h 2 , . . . , h N +1 ] = x 1 x 2 . . . x N x N +1 y 1 y 2 . . . y N 0 p 1 p 2 . . . p N p N +1 i.i.d.∼ µ, y N +1 = y(H).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.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 byEQUATIONand each residual feed-forward neural network layer is defined byMLP w1,w2,...,w M (Z, η) = Z + ηM -1 M j=1 h(Z, w j ) (2.2)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 byf (Z, θ) = W O W V Zσ A (W K Z) ⊤ W Q Z ,where W V , W K , W Q ∈ R s×D , and W O ∈ R D×s . This formulation can be reparametrized toEQUATIONwhere 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:T Θ (H, t + ∆t/2) = Attn θt,1,...,θ t,M ( T Θ (H, t), ∆t/2) T Θ (H, t + ∆t) = MLP wt,1,...,w t,M ( T Θ (H, t + ∆t/2), ∆t/2) (2.4)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.For the l 2 regularization with λ > 0, we consider training the constructed Transformer model using the following λ-regularized risk objective:EQUATIONwith the population squared risk function defined asR(Θ) = E µ 1 2Read[ T Θ (H, 1)] -y(H).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 θEQUATIONt,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 M L), which is defined as follows:EQUATIONDefine 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:EQUATIONt,j , Θ (τ ) , t).(2.7)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 Wasserstein-2 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. 
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:EQUATIONIn (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:EQUATIONwithEQUATIONDefine 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] :vec[p ρ (H, t)] ⊤ = Read[T ρ (H, 1)]-y(H) exp 1 t β ∇ vec[T ] vec[g(T ρ (H, t), β)ρ(β, t)dβdt] DN +d+1,:. Using this, we can compute the functional derivative to ρ as follows:δQ δρ (θ, w, t) = E µ Tr f (T ρ (H, t), θ) + h(T ρ (H, t), w) 2 ⊤ p ρ (H, t) + λ 2 (∥θ∥ 2 2 + ∥w∥ 2 2 ). (3.4)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 haveQ(ρ + η(ν -ρ)) = Q(ρ) + η δQ δρ , ν -ρ + o(η),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):EQUATIONwhere ρ (0) = ρ 0 , div is the divergence operator, and the gradient functions are defined asG f (θ, ρ, t) = 1 2 E µ ∇ θ Tr f (T ρ (H, t), θ) ⊤ p ρ (H, t) + λθ, G h (w, ρ, t) = 1 2 E µ ∇ w Tr h(T ρ (H, t), w) ⊤ p ρ (H, t) + λw.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 (2B 2 +2B 2 exp(K(1+ R + R 2 )) 2 ).iii. (θ,w) ρ (τ ) (θ, w, t)d(θ, w) = 1 for any t ∈ [0, 1].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 thatinf Θ R(Θ) ≤ inf ρ∈P 2,r R(ρ) + C L -1 + log(L + 1) M , inf Θ Q(Θ) ≤ inf ρ∈P 2,r Q(ρ) + C(1 + λ) L -1 + log(L + 1) M .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 ρ(τ) := 1EQUATIONt,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 havei. sup s∈[0,τ ] |Read[ T Θ (s) (H, t)] -Read[T ρ (s) (H, t)]|≤ C L -1 + δ+log(L+1) M ii. sup s∈[0,τ ] | R(Θ (s) ) -R(ρ (s) )|≤ C L -1 + δ+log(L+1) M iii. sup s∈[0,τ ] | Q(Θ (s) ) -Q(ρ (s) )|≤ C L -1 + δ+log(L+1) Mfor 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, τ ]. 
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.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 i. (Partial 1-homogeneity) for any T ∈ R D×(N +1) and c ∈ R, we have g(T, cα1 , α 2 ) = cf (T, α 1 , α 2 ). ii. (Universal kernel) a compact set K ⊂ R dim(α2) ensures that the span of g(•, α) : α ∈ R dim(α1) × K is dense in C(∥T ∥ 2-col ≤ B, R D×(N +1) ) for any B > 0.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.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:i. (ρ (τ ) ) τ ≥0 is concentrated on {θ, w : ∥θ∥ 2 + ∥w∥ 2 ≤ R 2 ∞ } × [0, 1] when τ is sufficiently large.ii. 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 at * ∈ [0, 1] such that the connected set supp(ρ ∞ (•, t * )) ⊃ ×{θ 0 } × K × D, for some θ 0 ∈ R dim(θ) and D ⊂ R dim(w1) that separates {w 1 : ∥w 1 ∥ = 1/R ∞ } and {w 1 : ∥w 1 ∥ = R ∞ }.Then, for any ϵ > 0, there exists some τ 0 > 0 such thatsup τ ≥τ0 R(Θ (τ ) ) ≤ ϵ + C 1 L -1 + δ + log(L + 1) M + C 2 λ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:sup τ ≥τ0,L>L0,M/log L>K0 R(Θ (τ ) ) ≤ (1/2 + C 2 C λ )ϵ (4.1)with probability at least 1 -3 exp(-δ) with respect to the parameter initialization Θ (0) . Notably, if C λ ≤ (2C 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. 
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. 
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. 
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). 
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? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA]• 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.  
 
 Appendix: The appendix is organized as follows:• Appendix B: Additional related works are discussed.• Appendix C:- • Appendix D: Includes proofs of main results from Section 3.In Appendix C.• 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.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.Defineβ := (θ, w) ⊤ , g(T, β) := f (T, θ) + h(T, w) 2 .Thus, δQ δρ could be expressed asδQ δρ (β, t) = E µ Tr g(T ρ (H, t), β) ⊤ p ρ (H, t) + λ 2 ∥β∥ 2 2 . (C.1)Additionally, we can combine G f with G h , and G f with G h to reformulate asEQUATIONandG(β, Θ, t) = E µ ∇ β Tr f ( T Θ (H, t), θ)/2 h( T Θ (H, t + ∆t/2), w)/2 ⊤ p Θ (H, t + ∆t/2) p Θ (H, t + ∆t) + λβ. (C.3)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:i. ∥g(T, β)∥ 2-col ≤ K∥T ∥ 2-col (1 + ∥β∥ + ∥β∥ 2 )ii. For every i ∈ [N + 1], we have ∥∇ β g(T, β):,i ∥ 2 ≤ ϕ P (∥T ∥ 2-col )(1 + ∥β∥) iii. ∥∇ vec[T ] vec[g(T, β)]∥ 2 ≤ ϕ T (N, D, ∥T ∥ F )(1 + ∥β∥ + ∥β∥ 2 )Under Assumption 3: For any L T > 0 and any L T -Lipschitz continuous functions T 1 = T 1 (H) andT 2 = T 2 (H), for every i ∈ [N + 1], we have i. E µ ∥∇ θ g(T 1 , β) :,i -∇ θ g(T 2 , β) :,i ∥ 2 ≤ ϕ P T (∥θ∥, K T , L T ) sup H ∥T 1 -T 2 ∥ 2-col , ii. E µ ∥∇ vec[T ] vec[g(T 1 , β)] -∇ vec[T ] vec[g(T 1 , β ′ )]∥ 2 ≤ ϕ T P (N, D, sup H ∥T 1 ∥ F , K P , L T )∥θ -θ ′ ∥ iii. E µ ∥∇ θ g(T 1 , β) :,i -∇ θ g(T 1 , β ′ ) :,i ∥ 2 ≤ ϕ P P (K P , sup H ∥T 1 ∥ 2-col , L T )∥θ -θ ′ ∥, iv. E µ ∥∇ vec[T ] vec[g(T 1 , β)] -∇ vec[T ] vec[g(T 2 , β)]∥ 2 ≤ ϕ T T (N, D, K T , ∥θ∥, L T ) sup H ∥T 1 -T 2 ∥ FVerifying 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. .Define the nested family of compact subsets (P r ) r>0 asP r := {β : ∥β∥ ≤ r} × [0, 1], ∀r > 0.For any ρ, ν ∈ P(Ω) and p ≥ 1, define the l p distance ∥ρ -ν∥ p as∥ρ -ν∥ p = 1 0 β |ρ(x) -ν(x)| p dβdt 1/p .Specifically, when p = 1, we haveW 1 (ρ, ν) = sup 1 0 β f (ρ -ν)dβdt : f is 1-Lipschitz, f(0, 0) = 0 ≤ sup 1 0 β |f ||ρ -ν|dβdt : f is 1-Lipschitz, f(0, 0) = 0 ≤ (r + 1)∥ρ -ν∥ 1for any ρ, ν ∈ P 2 concentrated on P r . For simplicity, any H discussed throughout this paper is assumed to lie within supp(µ).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 (C.4)| β ρ(β, t)dβ - β ρ(β, t ′ )dβ|≤ C ρ |t -t ′ | for any t, t ′ ∈ [0, 1]. Integrating (C.4) over t 2 ∈ [0, 1] obtains β ρ(β, t)dβ =1 + β ρ(β, t)dβ - 1 0 β ρ(β, t ′ )dβdt ′ ≤1 + C ρ 1 0 |t -t ′ |dt ′ ≤1 + C ρ /2. (C.5)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:|F (T, t 1 ) -F (T, t 2 )|≤ max{K δ,r , L δ,r }∥ρ(•, t 1 ) -ρ(•, t 2 )∥ BL ≤ C ρ max{K δ,r , L δ,r }|t 1 -t 2 |. (C.6) 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|F (T, t) -F (T ′ , t)|≤ L ′ δ,r ∥T -T ′ ∥ F β ρ(β, t)dβ ≤ L ′ δ,r (1 + C ρ /2)∥T -T ′ ∥ F , (C.7)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 Picard-Lindelöf TheoremInvoking the Picard-Lindelöf Theorem, there exists some ϵ > 0 such that the initial value problemṪ (H, t) = F (T, t), T (H, 0) = H 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 thatd dt ∥T ρ (H, t)∥ 2-col ≤ ∥ Ṫρ (H, t)∥ 2-col = ∥ β g(T ρ (H, t), β)ρ(β, t)dβ∥ 2-col ≤ β ∥g(T ρ (H, t), β)∥ 2-col ρ(β, t)dβ ≤ β K(1 + ||β|| 2 +||β|| 2 2 )ρ(β, t)∥T ρ (H, t)∥ 2-col dβ.(C.8) Therefore, by the Grönwall's inequality, we have∥T ρ (H, t max )∥ 2-col ≤ ∥T ρ (H, 0)∥ 2-col exp( ⊤ 0 β K(1 + ||β|| 2 +||β|| 2 2 )ρ(β, s)dβds) ≤ ∥H∥ 2-col exp(K(1 + C ρ /2 + r + r 2 )t max ) < ∞.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∥T ρ (H, t 1 )-T ρ (H, t 2 )∥ 2-col ≤ t2 t1 β ∥g(T ρ (H, t), β∥ 2-col ρ(β, t)dβdt ≤ (1+C ρ /2)KC(1+r+r 2 )(t 2 -t 1 ), (C.9) for any t 1 , t 2 ∈ [0, 1]. Given Assumption 2 (iii), we have ∥T ρ (H, t) -T ρ (H ′ , t)∥ F ≤ ⊤ 0 β ∥g(T ρ (H, s), β) -g(T ρ (H ′ , s), β)∥ F ρ(β, s)dβds ≤ ⊤ 0 β ϕ T (N, D, B exp(K(1 + C ρ /2 + r + r 2 )))(1 + r + r 2 )∥T ρ (H, s) -T ρ (H ′ , s)∥ F ρ(β, s)dβds (C.10) 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:∥T ρ (H, t) -T ρ (H ′ , t)∥ F ≤ ∥H -H ′ ∥ F exp(L H 1 0 β ρ(β, s)dβds) = exp(L H )∥H -H ′ ∥ Fgiven that T ρ (•, 0) serves as the identity mapping. Consequently, T ρ (H, t) demonstrates Lipschitz continuity with respect to H ∈ supp(µ).Lemma C.1 (Continuous Transformer output bound). Under Assumption 2, for any distribution ρ ∈ P(Ω) where1 0 β ||β|| 2 2 ρ(β, t)dβdt ≤ A 2 for some constant A > 0 and for any t ∈ [0, 1], we have ∥T ρ (H, t)∥ 2-col ≤ ∥H∥ 2-col exp(K(1 + A + A 2 )).Lemma C.2 (Continuous Transformer difference bound). Under Assumption 2 and given H, for any ρ, ν ∈ P 2 that satisfy1 0 β ||β|| 2 2 ρ(β, t)dβdt ≤ A 2and have bounded supports P r for some constants A, r > 0, we have thatsup t∈[0,1] ∥T ρ (H, t) -T ν (H, t)∥ F ≤ C r W 1 (ρ, ν).Here, the universal constant C r only depends on N, D, r, A, and the parameters of the assumptions.Under Assumption 1 and 2, for any ρ ∈ P(Ω) where supp(ρ) ⊂ P r and1 0 β ∥β∥ 2 ρ(β, t)dβdt ≤ A 2 , we have sup (β,t)∈Pr ∥g(T ρ (H, t), β)∥ F ≤ √ N + 1KB exp(K(1 + A + A 2 ))(1 + r + r 2 ), sup (β,t)∈Pr ∥p ρ (H, t)∥ F ≤ (B + B exp(K(1 + A + A 2 ))) exp ϕ T (N, D, √ N + 1KB exp(K(1 + A + A 2 ))(1 + A + A 2 ) , sup (β,t)∈Pr δQ δρ (β, t) ≤ √ N + 1KB exp(K(1 + A + A 2 ))(1 + r + r 2 )(B + B exp(K(1 + A + A 2 ))) exp ϕ T (N, D, √ N + 1KB exp(K(1 + A + A 2 ))(1 + A + A 2 ) + λ 2 r 2 .Lemma C.4 (Discrete Transformer bound). Under Assumptions 2, for any Θ where1 M L t M j=1 ∥β∥ 2 ≤ A 2for some universal constant A > 0 and at any t = 0, ∆t/2, ∆t, . . . , (L -1/2)∆t, 1, we have∥ T Θ (H, t)∥ 2-col ≤ ∥H∥ 2-col exp(K(1 + A + A 2 )).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∥ T Θ (H, t) -T Θ (H, t)∥ F ≤ C r 1 M L d(Θ, Θ).Here the universal constant C r only depends on N, D, r, and the parameters of the assumptions.Under Assumption 1 and 2, for any Θ such that sup t,j ∥β t,j ∥ 2 ≤ r 2 and 1M L t M j=1 ∥β∥ 2 ≤ A 2 , we have, for any t = 0, ∆t/2, . . . , (L- 1/2)∆t, 1 sup (θ,t)∈Pr max{∥f ( T Θ (H, t), θ)∥ F , ∥h( T Θ (H, t), w)∥ F , ∥g( T Θ (H, t), β)∥ F } ≤ √ N + 1KB T (1 + r + r 2 ), sup (β,t)∈Pr sup i∈[N +1] max{ ∇ θ f ( T Θ (H, t), θ) :,i , ∇ w h( T Θ (H, t), w) :,i , ∇ β g( T Θ (H, t), β) :,i } ≤ ϕ P (B T )(1 + r), sup (β,t)∈Pr ∥ p Θ (H, t)∥ F ≤ (B + B T ) exp ϕ T (N, D, √ N + 1KB T )(1 + A + A 2 ) ,whereB 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| 1 M L t M j=1 ∥β t,j ∥ 2 - 1 0 β ∥β∥ 2 ρ(β, t)dβdt|≲ L -1 + δ + log(L + 1) M .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:1. 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.2. The product of the increments for each matrix is bounded byL l=1 1 + max{a l , b l } ≤ C for some constant C > 0.Under these conditions, it holds that∥ L l=1 A l - L l=1 B l ∥ ≤ C L l=1 ∥A l -B l ∥.In this section, we define the partial derivativep ρ (H, t) := ∂R(H; ρ) ∂T ρ (H, t) ⊤ ∈ R D×(N +1)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.Jρ (H, t) = -J ρ (H, t)∇ vec[T ] vec[ β g(T ρ (H, t), β)ρ(β, t)dβ] (C.11)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 vec[p ρ (H, t)] ⊤ = ∂R(H; ρ) ∂vec[T ρ (H, 1)] ∂vec[T ρ (H, 1)] ∂vec[T ρ (H, t)] = vec[ ∂R(H; ρ) ∂T ρ (H, 1) ⊤ ] ⊤ ∂vec[T ρ (H, 1)] ∂vec[T ρ (H, t)] = vec[p ρ (H, 1)] ⊤ J ρ (H, t)Hence, by taking the derivative with respect to t, we obtain thatvec[ ṗρ (H, t)] ⊤ = -vec[p ρ (H, t)] ⊤ β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβ with the solution vec[p ρ (H, t)] ⊤ = vec[p ρ (H, 1)] ⊤ exp 1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)]ρ(β, s)dβds . (C.12)On the other hand, we havep ρ (H, 1) = ∂R(H; ρ) ∂T ρ (H, 1) = (Read[T ρ (H, 1)] -y(H))E read , (C.13)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 thatvec[p ρ (H, t)] ⊤ = (Read[T ρ (H, 1)]-y(H))•exp 1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)]ρ(β, s)dβds DN +d+1,:.(C.14)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 havevec[ p Θ (H, t)] = (Read[ T Θ (H, 1)] -y(H)) (s-t)/∆t+1∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] (s-t)/∆t+1∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w s,j )] DN +d+1,: , (C.15) and vec[ p Θ (H, t + ∆t/2)] = (Read[ T Θ (H, 1)] -y(H)) (s-t)/∆t+2∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] (s-t)/∆t+1∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w s,j )]DN +d+1,:.(C.16)C.5 Explanation for assumptions made in Theorem 4.1The 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θ+dimw 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 3D.1 Proof of Theorem 3.1This 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.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β (τ ) t,j = β (0) - τ 0 G(β (s) t,j , Θ (s) , t)ds (D.1)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. β (τ ) t,j = β (0) - τ 0 G( β (s) t,j , ρ (s) , t)ds, β (0) t,j = β (0) t,j (D.2)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.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, thenG(β, ρ, t) -G( β, ν, t) ≤ C G exp(C G W 1 (ρ, ν)) -1 + (1 + λ)∥β -β∥ .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∥ G(β, Θ, t) -G( β, Θ, t)∥ ≤ C G 1 M L d(Θ, Θ) + (1 + λ)∥β -β∥ .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 haveG(β, Θ, t) -G(β, ρ, t) ≲ L -1 + δ + log(L + 1) M , ∥G(β, ρ, t) -G(β, ρ, t)∥ ≲ L -1 + δ + log(L + 1) M , G(β, Θ, t) -G(β, ρ, t) ≲ L -1for 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) distributionsOur analysis commences with the bound for 0 < s 1 < s 2 < τ , we haveW 2 (ρ (s1) , ρ(s2) ) 2 ≤ 1 M L t M j=1 |β (s1) t,j -β (s2) t,j | 2 ≤ (s 2 -s 1 ) M L t M j=1 s2 s1 ∥ G(β (s) t,j , Θ (s) , t)∥ 2 ds,where each particle at time s 1 is paired with its position at time s 2 , leveraging the Jensen's inequality.Recalling the identityd Q(Θ (s) ) ds = 1 M L t M j=1 ∥ G(β (s) t,j , Θ (s) , t)∥ 2shown in Proposition D.1, it follows thatW 2 (ρ (s1) , ρ(s2) ) ≤ (s 2 -s 1 ) 1/2 Q 1/2 (Θ (0) ) ≤ λ 2 A 2 0 ,where the last inequality uses (D.30). SinceA 2 0 ≲ 1 + λ -1 , we see that W 2 (ρ (s1) , ρ(s2) ) ≤ C(1 + λ)(s 2 -s 1 ) 1/2for some constant C dependent on the parameters listed in the result. Similarly, we haveW 2 (ρ (s1) , ρ (s2) ) 2 ≤ E β (s2) -β (s1) 2 ≤ (s 2 -s 1 ) s2 s1 1 0 β G(β, ρ (s) , t) 2 dβdtds ≤ (s 2 -s 1 )Q(ρ 0 ) ≲ (s 2 -s 1 )C(1+λ)where β (s1) ∼ ρ (s1) (β, t) is embedded with its future position at β (s2) . The last step is feasible by setting C large enough, noticing that Q(ρ 0 ) ≤ λA 2 0 /2. To summarize, we havemax W 2 (ρ (s1) , ρ(s2) ), W 2 (ρ (s1) , ρ (s2) ) ≤ C(1 + λ) √ s 2 -s 1 (D.3)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∥β (τ ) t,j -β (τ ) t,j ∥ ≤ τ 0 G(β (s) t,j , Θ (s) , t) -G( β (s) t,j , ρ (s) , t) ds ≤ τ 0 G(β (s) t,j , Θ (s) , t) -G( β (s) t,j , Θ (s) , t) ds + τ 0 G( β (s) t,j , Θ (s) , t) -G( β (s) t,j , ρ (s) , t) ds ≤C G τ 0 1 M L d(Θ (s) , Θ (s) ) + (1 + λ)∥β (s) -β (s) ∥ ds + τ 0 G( β (s) t,j , Θ (s) , t) -G( β (s) t,j , ρ (s) , t) ds ≤C G τ 0 1 M L d(Θ (s) , Θ (s) ) + (1 + λ)∥β (s) -β (s) ∥ ds + sup s∈[0,τ ] G( β (s) t,j , Θ (s) , t) -G( β (s) t,j , ρ (s) , t) (D.4)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 considering1 M L d(Θ (s) , Θ (s) ) ≤ sup t,j ∥β (τ ) t,j -β (τ ) t,j ∥ for any s ≥ 0, we derive sup t,j ∥β (τ ) t,j -β (τ ) t,j ∥ ≤ C G (2+λ) τ 0 sup t,j ∥β (s) -β (s) ∥ds+sup t,j sup s∈[0,τ ] G( β (s) t,j , Θ (s) , t) -G( β (s) t,j , ρ (s) , t)Further supremacy taken over s ∈ [0, τ ] yields:∆(τ ) ≤ ∆(τ ) + C G (2 + λ) τ 0 ∆(s)ds, (D.5)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∆(τ ) ≤ exp C G (2 + λ)τ ∆(τ ). (D.6)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 haveG( β (s2) t,j , Θ (s2) , t) -G( β (s2) t,j , ρ (s2) , t) -G( β (s1) t,j , Θ (s1) , t) -G( β (s1) t,j , ρ (s1) , t) ≤ G( β (s1) t,j , Θ (s1) , t) -G( β (s2) t,j , Θ (s2) , t) + G( β (s1) t,j , ρ (s1) , t) -G( β (s2) t,j , ρ (s2) , t) ≤C ∆ (exp(C ∆ W 1 (ρ (s1) , ρ (s2) )) -1) + C ∆ 1 M L d( Θ (s1) , Θ (s2) ) + C ∆ (1 + λ)∥ β (s1) t,j -β (s2) t,j ∥ ≲ exp(C ∆ C √ s 2 -s 1 ) -1 + (1 + λ) sup t,j ∥ β (s1) t,j -β (s2) t,j ∥ ≲ √ s 2 -s 1 + sup t,j ∥ β (s1) t,j -β (s2)t,j ∥ (D.7) for some constant C ∆ dependent on the parameters listed in the result. Moreover, for any t, j, we have∥ β (s1) t,j -β (s2) t,j ∥ ≤ s2 s1 ∥G( β (s) t,j , Θ (s) , t)∥ ≤ √ s 2 -s 1 s2 s1 ∥G( β (s) t,j , Θ (s) , t)∥ 2 ≲ (1+λ 2 ) √ s 2 -s 1 ,(D.8) 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 obtainG( β (s2) t,j , Θ (s2) , t) -G( β (s2) t,j , ρ (s2) , t) -G( β (s1) t,j , Θ (s1) , t) -G( β (s1) t,j , ρ (s1) , t) ≲ (1+λ+λ 2 +λ 3 ) √ s 2 -s 1 (D.9) 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 δ + 2 log L), we obtain, with a probability of at least 1 -exp(-δ) with respect to the parameter initialization Θ (0) :sup s∈{τn} n∈[L 2 ] G( β (s) t,j , Θ (s) , t) -G( β (s) t,j , ρ (s) , t) ≲(1 + λ + λ 2 + λ 3 )L -1 + δ + log L + log(L + 1) M ≲L -1 + δ + log(L + 1)M Furthermore, leveraging (D.9), we deducesup t,j sup s∈[0,τ ] G( β (s) t,j , Θ (s) , t) -G( β (s) t,j , ρ (s) , t) ≲ sup t,j sup s∈{τn} n∈[L 2 ] G( β (s) t,j , Θ (s) , t) -G( β (s) t,j , ρ (s) , t) + sup |s1-s2|≤τ /L 2 (1 + λ + λ 2 + λ 3 ) √ s 2 -s 1 ≲L -1 + δ + log(L + 1) M .Returning to (D.6), we establish that with a probability of at least 1 -exp(-δ) with respect to the parameter initializationΘ (0) , sup s∈[0,τ ] sup t,j ∥β (s) t,j -β (s) t,j ∥ = ∆(τ ) ≲ L -1 + δ + log(L + 1) M .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 resultsNow, 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 constantB τ := B exp(K(1 + R τ + R 2 τ )) such that max{∥T ρ (τ ) (H, t)∥ 2-col , ∥ T Θ (τ ) (H, t)∥ 2-col , ∥ T Θ (τ ) (H, t)∥ 2-col } ≤ B τfor any H and t ∈ [0, 1].Utilizing Lemma D.6 and applying the union bound over n ∈ [L 2 ], we observesup s∈{τn} n∈[L 2 ] ∥ T Θ (s) (H, t) -T ρ (s) (H, t)∥ F ≲ L -1 + δ + log(L + 1) M .Additionally, note that for any s 1 , s 2 ∈ [0, τ ], we derive from Lemmas C.2 and C.5 that∥ T Θ (s 1 ) (H, t) -T ρ (s 1 ) (H, t)∥ F -∥ T Θ (s 2 ) (H, t) -T ρ (s 2 ) (H, t)∥ F ≤∥ T Θ (s 1 ) (H, t) -T Θ (s 2 ) (H, t)∥ F + ∥T ρ (s 1 ) (H, t) -T ρ (s 2 ) (H, t)∥ F ≲ sup t,j ∥ β (s1) t,j -β (s2) t,j ∥ + W 1 (ρ (s1) , ρ (s2) ) ≲ √ s 2 -s 1 (D.10)where the last inequality is derived utilizing (D.3) and (D.8). Consequently, we havesup s∈[0,τ ] ∥ T Θ (s) (H, t) -T ρ (s) (H, t)∥ F ≲ sup s∈{τn} n∈[L 2 ] ∥ T Θ (s) (H, t) -T ρ (s) (H, t)∥ F + sup |s2-s1|≤τ /L 2 √ s 2 -s 1 ≲ L -1 + δ + log(L + 1) M (D.11) 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 havesup s∈[0,τ ] | 1 M L t M j=1 ∥ β (s) t,j ∥ 2 - 1 0 β ∥β∥ 2 ρ (s) (β, t)dβdt|≲ L -1 + δ + log(L + 1) M .Now, let's analyze the scenario under the probability eventE 1 ∩ E 2 ∩ E 3 with P(E 1 ∩ E 2 ∩ E 3 ) ≥ 1 -3 exp(-δ). Lemma C.5 demonstrates that sup s∈[0,τ ] |Read[ T Θ (s) (H, t)] -Read[ T Θ (s) (H, t)]|≤ sup s∈[0,τ ] ∥ T Θ (s) (H, t) -T Θ (s) (H, t)∥ F ≲ sup s∈[0,τ ] sup t,j ∥β (s) t,j -β (s) t,j ∥ ≲ L -1 + δ + log(L + 1) M .(D.12)This further implies, by (D.11), thatsup s∈[0,τ ] |Read[ T Θ (s) (H, t)] -Read[T ρ (s) (H, t)]| ≤ sup s∈[0,τ ] |Read[ T Θ (s) (H, t)] -Read[ T Θ (s) (H, t)]|+ sup s∈[0,τ ] ∥ T Θ (s) (H, t) -T ρ (s) (H, t)∥ F ≲L -1 + δ + log(L + 1) M , (D.13) 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 sup s∈[0,τ ] | R(Θ (s) )-R(ρ (s) )|≲ sup s∈[0,τ ] |Read[ T Θ (s) (H, t)]-Read[T ρ (s) (H, t)]|≲ L -1 + δ + log(L + 1) M , (D.14) and sup s∈[0,τ ] | Q(Θ (s) ) -Q(ρ (s) )|≤ sup s∈[0,τ ] | R(Θ (s) ) -R(ρ (s) )|+ sup s∈[0,τ ] | 1 M L t M j=1 ∥ β (s) t,j ∥ 2 - 1 0 β ∥β∥ 2 ρ (s) (β, t)dβdt| ≲L -1 + δ + log(L + 1) M .(D.15)Step IV: Prove the weakly convergenceFor 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 haveW 2 (ρ (s1) M,L , ρ(s2) M,L ) ≤ C(1 + λ) √ s 2 -s 1for 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 φEQUATIONM,L , t) -G(β, ν (s) , t) ρ(s) (β, t)dβdtds 16 ) 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. 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.EQUATIONM,L , t) -G(β, ρ(s) M,L , t) ρ(s) (β, t)dβdtds +∥φ∥ max τ 0 1 0 β G(β, ρ(s) , t) -G(β, ν (s) , t) ρ(s) (β, t)dβdtds + sup s∈[0,τ ] ∥ρ (s) -ν s ∥ BL ≲L -1 + sup s∈[0,τ ] ∥ρ (s) -ν s ∥ BL + sup s∈[0,τ ],∥β∥≤Rτ G(β, ρ(s) , t) -G(β, ν (s) , t) ≲L -1 + δ + log(L + 1) M + sup s∈[0,τ ] ∥ρ (s) -ν s ∥ BL (D.Proof. Suppose that the Fréchet derivative δR δρ indeed exists, we establish δQ δρ(θ, w, t) = δR δρ (θ, w, t) + λ 2 (∥θ∥ 2 2 + ∥w∥ 2 2 ).Therefore, it suffices to show that the Fréchet derivative of L with respect to ρ is δR δρ(β, t) = E µ Tr g(T ρ (H, t), β) ⊤ p ρ (H, t) . (D.17) 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 havevec[T ρη (H, 1) -T ρ (H, 1)] =η 1 0 β exp 1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)]ρ(β, s)dβ) • vec[g(T ρ (H, t), β)](ν -ρ)(β, t)dβdt + o(η), (D.18) where ρ η := ρ + η(ν -ρ).Given (D.18), we observe from the solution of p ρ (C.13) and (C. 14) thatRead[T ρ (H, 1)] -y(H) Read[T ρη (H, 1) -T ρ (H, 1)] =vec[p ρ (H, 1)] ⊤ vec[T ρη (H, 1) -T ρ (H, 1)] =η 1 0 β vec[p ρ (H, t)] ⊤ vec[g(T ρ (H, t), β)](ρ -ν)(β, t)dβdt + o(η) =η 1 0 β Tr g(T ρ (H, t), β) ⊤ p ρ (H, t) (ρ -ν)(β, t)dβdt + o(η), (D.19)Hence, by applying (D. 19 ) to the risk function, we obtainR(ρ η ) -R(ρ) = 1 2 E µ Read[T ρη (H, 1)] -y(H) 2 -Read[T ρ (H, 1)] -y(H) 2 = E µ Read[T ρ (H, 1)] -y(H) Read[T ρ (H, 1) -T ρη (H, 1)] + Read[T ρ (H, 1) -T ρη (H, 1)]O(Read[T ρη (H, 1) -T ρ (H, 1)]) = η δR δρ , ν -ρ + o(η),which indicates (D.17) and concludes the proof.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 (ρ) asQ r (ρ) = Q(ρ) ρ(P r ) = 1, ∞ otherwise.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 C p (γ) := ( |x -y| p dγ(x, y)) 1/p 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 haveQ 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∥T ρ (H, t) + T ν (H, t)∥ F = O(1),and Lemma C.2 guarantees that∥T ρ (H, t) -T ν (H, t)∥ F = O(W 1 (ρ, ν)) = O(W 2 (ρ, ν))for any H. Therefore, we haveR(ν) -R(ρ) = 1 2 E µ Read[T ν (H, 1)] -y(H) 2 -Read[T ρ (H, 1)] -y(H) 2 = E µ Read[T ρ (H, 1)] -y(H) Read[T ρ (H, 1) -T ν (H, 1)] + Read[T ρ (H, 1) -T ν (H, 1)]O(Read[T ν (H, 1) -T ρ (H, 1)]) = O(W 2 (ρ, ν)).(D.20)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λ 2 β ∥β∥ 2 (ρ -ν)dβdt = O(W 1 (ρ, ν)) = O(W 2 (ρ, ν)). (D.21)Combining (D.20) and (D.21), we obtain thatQ(ρ) -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 readsh ′ (α) =⟨ δQ δρ ρ=ρ γ α , d dα ρ γ α ⟩ = d δQ δρ ρ=ρ γ α (1 -α)(β 1 , t 1 ) + α(β 2 , t 2 ) (β 1 , t 1 ) -(β 2 , t 2 ) dγ((β 1 , t 1 ), (β 2 , t 2 )) .(D.22) 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 supt∈[0,1] ∥p ρ (H, t) -p ν (H, t)∥ F ≤ L r ∥ρ -ν∥ 1 , sup (β,t)∈Pr δQ δρ ρ (β, t) - δQ δρ ν (β, t) ≤ L r ∥ρ -ν∥ 1 .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 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, Step III: Show the well-posedness of Wasserstein gradient flow at some finite timeJ 1 := d δQ δρ ρ=ρ γ α 1 (1 -α 1 )(β 1 , t 1 ) + α 1 (β 2 , t 2 ) (β 1 , t 1 ) -(β 2 , t 2 ) dγ((β 1 , t 1 ), (β 2 , t 2 )) -d δQ δρ ρ=ρ γ α 2 (1 -α 1 )(β 1 , t 1 ) + α 1 (β 2 , t 2 ) (β 1 , t 1 ) -(β 2 , t 2 ) dγ((β 1 , t 1 ), (β 2 , t 2 )) ≤ sup (β,t)∈Pr δQ δρ ρ (β, t) - δQ δρ ν (β, t) ∥(β 1 , t 1 ) -(β 2 , t 2 )∥ 1 dγ((β 1 , t 1 ), (β 2 , t 2 )) ≤L r ∥ρ γ α1 -ρ γ α2 ∥ 1 • C 1 (γ) ≤L r C 2 1 (γ)|α 1 -α 2 | ≤L r C 2 2 (γ)|α 1 -α 2 | (D.J 2 := d δQ δρ ρ=ρ γ α 2 (1 -α 1 )(β 1 , t 1 ) + α 1 (β 2 , t 2 ) (β 1 , t 1 ) -(β 2 , t 2 ) dγ((β 1 , t 1 ), (β 2 , t 2 )) -d δQ δρ ρ=ρ γ α 2 (1 -α 2 )(β 1 , t 1 ) + α 2 (β 2 , t 2 ) (β 1 , t 1 ) -(β 2 , t 2 ) dγ((β 1 , t 1 ), (β 2 , t 2 )) ≤ sup (β,t)∈Pr δQ δρ ρ=ρ γ α 2 (β, t) α 1 -α 2 ∥(β 1 , t 1 ) -(β 2 , t 2 )∥ 2 dγ((β 1 , t 1 ), (β 2 , t 2 )) ≤L ′ r C 2 2 (γ)|α 1 -α 2 | (D.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:dρ (τ ) r dτ = div β (ρ (τ ) r v (τ ) r )whereEQUATIONr , t), (β, t) ∈ P r , 0, otherwise.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τ r := inf{τ > 0 : ρ (τ ) r (P r ) < 1}.By defining this exit time, for any r > r and τEQUATIONEQUATIONr , 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 representsdQ(ρ (τ ) ) dτ = 1 0 β dQ dρ ρ=ρ (τ ) , div β (ρ (τ ) G(β, ρ (τ ) , t)) = 1 0 β G(β, ρ (τ ) , t), div β (ρ (τ ) G(β, ρ (τ ) , t)) dβdt = 1 0 β ρ (τ ) ∥G(β, ρ (τ ) , t)∥ 2 2 dβdt ≤ 0. (D.25)Therefore, for any τ ∈ [0, τ r ], utilizing Lemma C.1, we haveQ(ρ (τ ) ) ≤ Q(ρ 0 ) =E µ 1 2 Read[T ρ0 (H, 1)] -y(H) 2 + λ 2 1 0 β ∥β∥ 2 ρ 0 (β, t)dβdt ≤ 1 2 E µ [(∥T ρ0 (H, 1)∥ 2-col + B) 2 ] + λR 2 2 ≤E µ [(∥T ρ0 (H, 1)∥ 2 2-col + B 2 ] + λR 2 2 ≤B 2 + B 2 exp K(1 + R + R 2 ) 2 + λR 2 2 (D.26)Thus, we have1 0 β ∥β∥ 2 ρ (τ ) (β, t) ≤ 2 λ Q(ρ (τ ) ) ≤ R 2 + λ -1 2B 2 + 2B 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∥v (τ ) r (β, t) -λβ∥ = ∥G(β, ρ (τ ) , t) -λβ∥ ≤ N +1 i=1 ∇ β g(T ρ (H, t), β) :,i p ρ (H, t) :,i ≤ sup i∈[N +1] ∥∇ β g(T ρ (H, t), β) :,i ∥ N +1 i=1 ∥p ρ (H, t) :,i ∥ ≤ √ N + 1 sup i∈[N +1] ∥∇ β g(T ρ (H, t), β) :,i ∥ ∥p ρ (H, t)∥ F ≤ √ N + 1ϕ P (∥T ρ (H, t)∥ 2-col )∥p ρ (H, t)∥ F (1 + ∥β∥) ≤C R (1 + ∥β∥), (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 + 1KB exp(K(1 + A 0 + A 2 0 ))(1 + A 0 + A 2 0) . Applying (D.27) to the gradient flow equationdβ (τ ) dτ = -v (τ ) r (β, t), β (0) = β for τ ≥ 0, we obtain d∥β (τ ) ∥ dτ = ⟨-v (τ ) r (β,t),∥β (τ ) ∥⟩ ∥β (τ ) ∥ ≤ ∥v (τ ) r (β, t) -λβ∥ ≤ C R (1 + ∥β (τ ) ∥). This indicates ∥β (τ ) ∥ ≤ (∥β∥ + 1) exp(C R τ ) -1 ≤ (R + 1) exp(C R τ ) -1 (D.28)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β ρ (τ ) (β, t)dβ = β ρ (0) (β, t)dβ + τ 0 β div β (ρ (s) G (s) (β, ρ (s) , t))dβ ds =1 + τ 0 0 • ds = 1indicated 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, implying1 0 β ∥β∥ 2 ρ (τ ) (β, t) ≤ A 2 0 for any τ ≥ 0.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 τ .ii.1 M L t M j=1 (∥θ (τ ) t,j ∥ 2 2 + ∥w (τ ) t,j ∥ 2 2 ) ≤ A 2 0 .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 {M L • 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 thatQ(Θ (0) ) ≥ Q(Θ (0) ) -Q(Θ (τ ) ) = τ 0 t M j=1 ⟨ d Q(Θ) dβ t,j Θ=Θ (τ ) , M L d Q(Θ) dβ t,j Θ=Θ (τ ) ⟩dτ = τ 0 1 M L t M j=1 ∥ G β (β (τ ) t,j , Θ (τ ) , t)∥ 2 dτ ≥ 1 M Lτ t M j=1 τ 0 ∥ G β (β (τ ) t,j , Θ (τ ) , t)∥dτ 2 ,(D.29) 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,1 M L t M j=1 ∥β t,j ∥ 2 2 ≤ 2λ -1 Q(Θ (τ ) ) ≤ 2λ -1 Q(Θ (0) ) =λ -1 E µ Read[ T Θ (0) (H, 1)] -y(H) 2 + 1 M L t M j=1 ∥β (0) t,j ∥ 2 2 ≤λ -1 E µ [(∥ T Θ (0) (H, 1)∥ 2-col + B) 2 ] + R 2 ≤2λ -1 E µ [(∥ T Θ (0) (H, 1)∥ 2 2-col + B 2 ] + R 2 ≤R 2 + λ -1 2B 2 + 2B 2 exp K(1 + R + R 2 ) 2 =A 2 0(D.30) The last inequality of (D.30) follows from Lemma C.4, thereby showing that 1M L t M j=1 ∥β t,j ∥ 2 2 ≤ A 20 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 ∥ G(β, Θ (τ ) , t) -λβ∥≤ N +1 i=1 ∇ θ f ( T Θ (H, t), θ) :,i p Θ (H, t) :,i /2 + ∇ w h( T Θ (H, t + ∆/2), w) :,i p Θ (H, t + ∆t/2) :,i /2 ≤ sup i∈[N +1] (∥∇ θ f (T ρ (H, t), θ) :,i ∥ N +1 i=1 ∥p ρ (H, t) :,i ∥/2 + ∥∇ w h(T ρ (H, t), w) :,i ∥ N +1 i=1 ∥p ρ (H, t + ∆t/2) :,i ∥/2) ≤ √ N + 1 sup i∈[N +1] ∥∇ θ f (T ρ (H, t), θ) :,i ∥ ∥p ρ (H, t)∥ F /2 + ∥∇ w h(T ρ (H, t + ∆t/2), w) :,i ∥ ∥p ρ (H, t + ∆t/2)∥ F /2 ≤ √ N + 1 max{ϕ P (∥T ρ (H, t)∥ 2-col ), ϕ P (∥T ρ (H, t + ∆t/2)∥ 2-col )} max{∥p ρ (H, t)∥ F , ∥p ρ (H, t + ∆t/2)∥ F }(1 + ∥β∥) ≤C R (1 + ∥β∥),(D.31) Applying (D.31) to the gradient flowdβ (τ ) t,j dτ = -G(β (τ ) t,j , Θ (τ ) , t)for τ ≥ 0, we haved∥β (τ ) t,j ∥ dτ = ⟨-G(β (τ ) t,j ,Θ (τ ) ,t),∥β (τ ) ∥⟩ ∥β (τ ) ∥ ≤ ∥ G(β (τ ) t,j , Θ (τ ) , t)-λβ∥ ≤ C R (1+∥β (τ ) ∥). This indicates ∥β (τ ) ∥ ≤ (∥β∥ + 1) exp(C R τ ) -1 ≤ (R + 1) exp(C R τ ) -1 = R τ .(D.32)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 asTΘ (H, t + ∆t) = TΘ (H, t) + ∆tM -1 M j=1 f ( TΘ (H, t), θ t,j ) + M j=1 h( TΘ (H, t), w t,j ) = TΘ (H, t) + ∆tM -1 M j=1g( TΘ (H, t), β t,j ).(D.33)The second proxy, denoted as T ρ , extends the width to infinity by letting M → ∞, effectively replacing the average with an integral:T ρ (H, t + ∆t) = T ρ (H, t) + ∆t β g( T ρ (H, t), β)ρ(β|t)dβ. (D.34)We let all four Transformers share the same initial stateT Θ (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∥ T Θ (H, t) -T ρ (H, t)∥ F ≲ L -1 + δ + log(L + 1) M .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∥ T Θ (H, t) -T ρ * (H, t)∥ F ≤ C L -1 + log(L + 1) M ,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| R(Θ) -R(ρ * )|≤E µ [|Read[ T Θ (H, 1) -T Θ (H, 1)]|•||Read[ T Θ (H, 1) + T Θ (H, 1)] + 2y(H)|] ≲L -1 + log(L + 1) M .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   ∥ T Θ (H, t) -T ρ * (H, t)∥ F ≤ C L -1 + log(L+1) M , | 1 M L t M j=1 ∥β t,j ∥ 2 - 1 0 β ∥β∥ 2 ρ(β, t)dβdt|≤ C L -1 + log(L+1) M .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| Q(Θ)-Q(ρ * )|≤ | R(Θ)-R(ρ * )|+λ| 1 M L t M j=1 ∥β t,j ∥ 2 - 1 0 β ∥β∥ 2 ρ(β, t)dβdt|≤ C(1+λ) L -1 + log(L + 1) M .The result thus follows.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.1Our 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 .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 dimw , 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 .iii.β δQ δρ (β, t * ) ν(β) -ρ(β|t * ) dβ ≤ C 1 λ -C 2 R(ρ)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ν ∈ P B dimθ1 (0, B R∞ )/B dimθ1 (0, 1/B R∞ ) × K × {w 0 } 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:δQ f δρ (θ, w, t) = E µ Tr f (T ρ (H, t), θ) ⊤ p ρ (H, t) + λ∥θ∥ 2 2 δQ h δρ (θ, w, t) = E µ Tr h(T ρ (H, t), w) ⊤ p ρ (H, t) + λ∥w∥ 2 2 .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 respectto (β, 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 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 analysisIn 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. ) ∈ 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 1homogeneity of f (T, θ 1 , θ 2 ) with respect to θ 1 , we haveδQ f δρ ρ∞ (cθ 1 , θ 2 , w 0 , t * ) = c δQ f δρ ρ∞ (θ 1 , θ 2 , w 0 , t * ) + (|c|-1)λ∥θ 1 ∥ 2 . δQ h δρ ρ∞ (cθ 1 , θ 2 , w 0 , t * ) = δQ h δρ ρ∞ (θ 1 , θ 2 , w 0 , t * ). (E.1)Hence, given that ∇ β δQ δρ | ρ∞ (•, t * ) = 0 almost everywhere over supp(ρ ∞ (•, t * )), it also holds that EQUATION∇ (θ1,θ2,w) δQ δρ ρ∞ (θ 1 , θ 2 , w 0 , t * ) ≤ (R 2 ∞ B R∞ + R ∞ )λ. (E.2) Given the condition that ∥(cθ 1 , θ 2 , w 0 , t * ) -(θ 1 , θ 2 , w 0 , t * )∥ ≤ R 2 ∞ B R∞ ,(θ 1 , θ 2 , w 0 , t * ) -C ≤ (R 4 ∞ B 2 R∞ + R 3 ∞ B R∞ )λ. (E.3) 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 C 1 λ -C 2 R(ρ ∞ ) ≥ β δQ δρ ρ∞ (β, t * ) ν(β) -ρ ∞ (β|t * ) dβ = β δQ δρ ρ∞ (β, t * ) -C ν(β) -ρ ∞ (β|t * ) dβ ≥ - β (R 4 ∞ B 2 R∞ + R 3 ∞ B R∞ )λ ν(β) + ρ ∞ (β|t * ) dβ ≥2(R 4 ∞ B 2 R∞ + R 3 ∞ B R∞ )λ. (E.4) Therefore, we have R(ρ ∞ ) ≤ C1+2(R 4 ∞ B 2 R∞ +R 3 ∞ B R∞ ) C2λ, andQ(ρ ∞ ) ≤ R(ρ ∞ ) + 1 0 β ∥β∥ 2 ρ(β, t)dβdt ≤ C 1 + 2(R 4 ∞ B 2 R∞ + R 3 ∞ B R∞ ) C 2 + R 2 ∞ λ, (E.5)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:1 0 β ∥β∥ 2 ρ (τ ) (β, t)dβdt ≤ A 2 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 Wasserstein-2 distance, i.e.lim τ →∞ W 2 (ρ (τ ) , ρ ∞ ) = 0. (E.6)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∥T ρ∞ (H, t) -T ρ (τ ) (H, t)∥ F ≤ C 0 W 2 (ρ (τ ) , ρ ∞ ) (E.7)for any H and t ∈ [0, 1] when τ is sufficiently large. Note that Lemma C.1 shows thatmax{∥T ρ (τ ) (H, t)∥ 2-col , ∥T ρ∞ (H, t)∥ 2-col } ≤ B exp(K(1 + R ∞ + R 2 ∞)) =: B T for any H and t ∈ [0, 1]. Thus, from (E.7), we have|Q(ρ ∞ ) -Q(ρ (τ ) )|≤ 1 2 E µ |Read[|T ρ (τ ) (H, 1)]|+|T ρ∞ (H, 1)]|+2|y(H)| ∥T ρ∞ (H, 1) -T ρ (τ ) (H, 1)∥ F + λ 2 1 0 β ∥β∥ 2 (ρ ∞ -ρ (τ ) )(β, t)dβdt ≤(B T + B)C 0 W 2 (ρ (τ ) , ρ ∞ ) + λ 2 R ∞ W 1 (ρ (τ ) , ρ ∞ ) ≤((B T + B)C 0 + λ 2 R ∞ )W 2 (ρ (τ ) , ρ ∞ ), (E.8) where the second inequality incorporates the Kantorovich-Rubinstein Theorem (see Theorem 5.10 of [68] , for example) and the 2R ∞ -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 largeThe 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| Q(Θ (τ0) ) -Q(ρ (τ0)) )|≤ C 1 L -1 + δ + log(L + 1) Mwith probability at least 1 -3 exp(-δ) with respect to the parameter initialization Θ (0) for any δ > 0.Combining the outcomes from the preceding steps, we obtainQ(Θ (τ0) ) ≤ ϵ + C 1 L -1 + δ + log(L + 1) M + C 2 λ. (E.9)Note thatd dτ Q(Θ (τ ) ) = t M j=1 ⟨ d Q(Θ) dβ t,j Θ=Θ (τ ) , -M L d Q(Θ) dβ t,j Θ=Θ (τ ) ⟩ = - 1 M L t M j=1 ∥ G β (β (τ ) t,j , Θ (τ ) , t)∥ 2 ≤ 0, so the sequence ( Q(Θ (τ ) )) τ ≥0 is non-decreasing. Hence, for any τ ≥ τ 0 , R(Θ (τ ) ) ≤ Q(Θ (τ ) ) ≤ Q(Θ (τ0) ) ≤ ϵ + C 1 L -1 + δ + log(L + 1) M + C 2 λ,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 thatsup τ ≥τ0 R(Θ (τ ) ) ≤ϵ/2 + C 1 L -1 + δ + log(L + 1) M + C 2 C λ λ ≤(1/4 + C 2 C λ )ϵ + C 1 L -1 + δ + log(L + 1) M .The result holds by setting L and M/log L sufficiently large to ensure thatC 1 L -1 + δ + log(L + 1) M ≤ C 1 L -1 + 2(1 + δ) log(L + 1) M ≤ ϵ/4.F. E µ ∇ β Tr g(T ρ (H, t), β) ⊤ p ρ (H, t) -∇ β Tr g(T ν (H, t), β) ⊤ p ν (H, t) ≤ J 1 + J 2 ,whereJ 1 := E µ ∇ β Tr (g(T ν (H, t), β) -g(T ρ (H, t), β)) ⊤ p ν (H, t) ≲ ∥β -β∥, J 2 := E µ ∇ β Tr g(T ρ (H, t), β) ⊤ (p ν (H, t) -p ρ (H, t)) ≲ exp(C G W 1 (ρ, ν)) -1.Here, the symbol ≲ hides dependencies on N , D, r, and the parameters of the assumptions. To bound J 1 , consider thatJ 1 ≤ sup i∈[N +1] E µ g(T ν (H, t), β) :,i -g(T ρ (H, t), β) :,i N +1 i=1 ∥p ν (H, t) :,i ∥ ≤ √ N + 1 sup i∈[N +1] E µ g(T ν (H, t), β) :,i -g(T ρ (H, t), β) :,i ∥p ν (H, t)∥ F ≲ sup i∈[N +1] E µ g(T ν (H, t), β) :,i -g(T ρ (H, t), β) :,i ≤ϕ P T (r, B T )∥T ρ (H, t) -T ρ (H, t)∥ 2-col + ϕ P P (r, B T )∥β -β∥ ≲W 1 (ρ, ν) + ∥β -β∥. (F.1)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 haveJ 2 ≤ √ N + 1 sup i∈[N +1] E µ ∥g(T ρ (H, t), β) :,i ∥ ∥p ν (H, t) -p ρ (H, t)∥ F ≤ √ N + 1∥(T ρ (H, t), β)∥ 2-col ∥p ν (H, t) -p ρ (H, t)∥ F ≲∥p ν (H, t) -p ρ (H, t)∥ F . (F.2)In (F.2), the third inequality relies on Assumption 2 (i). Consequently, to establishJ 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 (ρ, ν),whereI 1 =|Read[T ρ (H, 1) -T ν (H, 1)]| exp 1 t β ∇ vec[T ] vec[g(T ν (H, s), β)]ρ(β, s)dβds , I 2 =|Read[T ρ (H, 1)] -y(H)| exp 1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)]ρ(β, s)dβds -exp 1 t β ∇ vec[T ] vec[g(T ν (H, s), β)]ρ(β, s)dβds . 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 providedin Assumption 2 (iii). To bound I 2 , we haveI 2 ≲ exp 1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)]ρ(β, s)dβds -exp 1 t β ∇ vec[T ] vec[g(T ν (H, s), β)]ρ(β, s)dβds ≲ exp 1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)]ρ(β, s)dβds - 1 t β ∇ vec[T ] vec[g(T ν (H, s), β)]ρ(β, s)dβds -I dimvec[T ]≲ exp1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)] -∇ vec[T ] vec[g(T ν (H, s), β)] ρ(β, s)dβds -1 ≲ exp ϕ T T (N, D, B T , r)∥g(T ρ (H, t), β) -g(T ν (H, t), β)∥ F -1 ≲ exp C r ϕ T T (N, D, B T , r)ϕ T (N, D, √ N + 1B T )(1 + r + r 2 )W 1 (ρ, ν) -1 (F.3) 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 thatJ 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 ∥ G(β, Θ, t) -G(β, Θ, t)∥ = 1 2 E µ ∇ θ Tr f ( T Θ (H, t), θ) -f ( T Θ (H, t), θ) ⊤ p Θ (H, t + ∆t/2) ⊤ + 1 2 E µ ∇ θ Trf ( T Θ (H, t), θ) ⊤ p Θ (H, t + ∆t/2) -p Θ (H, t + ∆t/2) ⊤ , 1 2 E µ ∇ w Tr h( T Θ (H, t + ∆t/2), w) -h( T Θ (H, t + ∆t/2), w) ⊤ p Θ (H, t) ⊤ + 1 2 E µ ∇ w Trh( T Θ (H, t + ∆t/2), w) ⊤ p Θ (H, t) -p Θ (H, t) ⊤ ⊤ . To demonstrate that ∥ G(β, Θ, t) -G(β, Θ, t)∥ ≤ C G 1 M L d(Θ, Θ), it suffices to show J 1 := E µ ∇ θ Tr f ( T Θ (H, t), θ) -f ( T Θ (H, t), θ) ⊤ p Θ (H, t + ∆t/2) ≤ C G 1 M L d(Θ, Θ),(F.4) andJ 2 := E µ ∇ θ Trf ( T Θ (H, t), θ) ⊤ p Θ (H, t + ∆t/2) -p Θ (H, t + ∆t/2) ≤ C G 1 M L d(Θ, Θ), (F.5) as the other part for h and w follows a similar proof approach.To bound J 1 , by Assumption 3 (i) we haveJ 1 ≤ N +1 i=1 E µ ∇ θ f ( T Θ (H, t), θ) :,i -f ( T Θ (H, t), θ) :,i ∥ p Θ (H, t + ∆t/2) :,i ∥ ≤E µ sup i∈[N +1] ∇ θ f ( T Θ (H, t), θ) :,i -f ( T Θ (H, t), θ) :,i N +1 i=1 ∥ p Θ (H, t + ∆t/2) :,i ∥ ≤ϕ T (r, B T )∥ T Θ (H, t) -T Θ (H, t)∥ 2-col √ N + 1∥ p Θ (H, t + ∆t/2)∥ F ≲ϕ T (r, B T ) √ N + 1∥ p Θ (H, t + ∆t/2)∥ F 1 M L d(Θ, Θ) ≲ 1 M L d(Θ, Θ), (F.6)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 haveJ 2 ≤ √ N + 1E µ sup i∈[N +1] f ( T Θ (H, t), θ) :,i ∥ p Θ (H, t + ∆t/2) -p Θ (H, t + ∆t/2)∥ F = √ N + 1ϕ(B T )(1 + r)E µ ∥ p Θ (H, t + ∆t/2) -p Θ (H, t + ∆t/2)∥ F . (F.7)Hence, it suffices to show that∥ p Θ (H, t + ∆t/2) -p Θ (H, t + ∆t/2)∥ F ≲ 1 M L d(Θ, Θ) to establish J 2 ≤ 1 M L 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 ,whereI 1 =(Read[ T Θ (H, 1) -T Θ (H, 1)] (s-t)/∆t+2∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] (s-t)/∆t+1∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w s,j )]DN +d+1,:≤ 1 M L d(Θ, Θ) (s-t)/∆t+2∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] (s-t)/∆t+1∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w s,j )] ≤ 1 M L d(Θ, Θ) exp (∆t/2) (s-t)/∆t+1∈[(1-t)/∆t] M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] +(∆t/2) (s-t)/∆t+1∈[(1-t)/∆t] M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, t), w s,j )] ≤ 1 M L d(Θ, Θ) exp ϕ T (N, D, √ N + 1KB T )(1 + r + r 2 ) ≲ 1 M L d(Θ, Θ),andI 2 =|Read[ T Θ (H, 1)] + y(H)| (s-t)/∆t+2∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] (s-t)/∆t+1∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w s,j )] - (s-t)/∆t+2∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, t), θ s,j )] (s-t)/∆t+1∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, t + ∆/2), w s,j )]DN +d+1,:≤(B + B T ) • exp (∆t/2) (s-t)/∆t+1∈[(1-t)/∆t] M -1 M j=1 max{ ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] , ∇ vec[T ] vec[f ( T Θ (H, t), θ s,j )]} +(∆t/2) (s-t)/∆t+1∈[(1-t)/∆t] M -1 M j=1 max{ ∇ vec[T ] vec[h( T Θ (H, t), w s,j )] , ∇ vec[T ] vec[h( T Θ (H, t), w s,j )]} • (∆t/2) (s-t)/∆t+1∈[(1-t)/∆t] j∈[M ] M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] -M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, t), θ s,j )] + M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, t), w s,j )] -M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, t), w s,j )] ≤(B + B T ) exp ϕ T (N, D, √ N + 1KB T )(1 + r + r 2 ) ϕ T P (N, D, √ N + 1B T , r) 1 M L d(Θ, Θ) ≲ 1 M L d(Θ, Θ).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 M L d(Θ, Θ). By combining the bounds of I 1 and I 2 , we observe that Equation (F.5) holds, thereby establishing d, r , and the parameters of the assumptions.∥ G(β, Θ, t) -G(β, Θ, t)∥ ≤ C G 1 M L d(Θ, Θ) for some C G dependent on N ,It remains to prove that ∥ G(β, Θ, t) -G( β, Θ, t)∥ ≤ C G (1 + λ)∥β -β∥. Note that ∥ G(β, Θ, t) -G( β, Θ, t)∥ ≤λ∥β -β∥ + 1 2 E µ ∇ θ Tr f ( T Θ (H, t), θ) -f ( T Θ (H, t), θ) ⊤ p Θ (H, t + ∆t/2) ⊤ , 1 2 E µ ∇ w Tr h( T Θ (H, t + ∆t/2), w) -h( T Θ (H, t + ∆t/2), w) ⊤ p Θ (H, t) ⊤ ⊤.Therefore, we only need to show thatEQUATIONE µ ∇ w Tr h( T Θ (H, t), w) -h( T Θ (H, t), w) ⊤ p Θ (H, t) ≤ C G ∥w -w∥, 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 haveE µ ∇ θ Tr f ( T Θ (H, t), θ) -f ( T Θ (H, t), θ) ⊤ p Θ (H, t + ∆t/2) ≤ sup i∈[N +1] E µ ∇ θ f ( T Θ (H, t), θ) -f ( T Θ (H, t), θ) :,i √ N + 1∥ p Θ (H, t + ∆t/2)∥ F ≤ϕ P P (r, B T )∥θ -θ∥ √ N + 1∥ p Θ (H, t + ∆t/2)∥ F ≲∥θ -θ∥.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∥ T Θ (H, t) -T ρ (H, t)∥ F ≲ L -1 + δ + log(L + 1) M . (F.8)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∥ T Θ (H, t) -T ρ (H, t)∥ F ≲ L -1 . (F.9)We present 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 thatG(β, Θ, 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.A 1 , A 2 , B 1 , B 2 , C 1 , C 2 : A 1 :=∇ θ vec[f ( T Θ (H, t), θ)], A 2 := ∇ θ vec[f (T ρ (H, t), θ)] B 1 :=Read[ T Θ (H, 1) -y(H)], B 2 := Read[T ρ (H, 1) -y(H)] C 1 := (s-t)/∆t+2∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] (s-t)/∆t+1∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w s,j )]DN +d+1,:C 2 := exp 1 t β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβdt DN +d+1,:The universal boundedness of ∥A := 2( G f (θ, Θ, t) -G f (θ, ρ, t)) as J = ∥A 1 B 1 C 1 -A 2 B 2 C 2 ∥ ≤∥(A 1 -A 2 )B 2 C 2 ∥ + ∥A 1 (B 1 -B 2 )C 2 ∥ + ∥A 1 B 1 (C 1 -C 2 )∥ ≤∥A 1 -A 2 ∥∥B 2 ∥∥C 2 ∥ + ∥A 1 ∥∥B 1 -B 2 ∥∥C 2 ∥ + ∥A 1 ∥∥B 1 ∥∥C 1 -C 2 ∥ ≲∥A 1 -A 2 ∥ + ∥B 1 -B 2 ∥ + ∥C 1 -C 2 ∥.We claim that to obtain the result, it suffices to show that i.∥A 1 -A 2 ∥ ≲ L -1 + δ+log(L+1) M under event E. ii. ∥B 1 -B 2 ∥ ≲ L -1 + δ+log(L+1) M under event E.iii. There exists some eventE 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 eventE∩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).By Assumption 3 (iv), we have∥A 1 -A 2 ∥ ≤ ϕ T T (N, D, √ N + 1B T , r)∥ T Θ (H, t) -T ρ (H, t)∥ F ≲ L -1 + δ+log(L+1) Munder 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 quantitiesD 1 := (s-t)/∆t+2∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] (s-t)/∆t+1∈[(1-t)/∆t] j∈[M ] I dimvec[T ] + (∆t/2)M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w s,j )] D 2 := (s-t)/∆t+2∈[(1-t)/∆t] I dimvec[T ] + (∆t/2) β ∇ vec[T ] vec[f ( T Θ (H, s), θ)]ρ(β|s)dβ (s-t)/∆t+1∈[(1-t)/∆t] I dimvec[T ] + (∆t/2) β ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w)]ρ(β|s)dβ D 3 := exp (s-t)/∆t+2∈[(1-t)/∆t] (∆t/2) β ∇ vec[T ] vec[f ( T Θ (H, s), θ)]ρ(β, s)dβ + (s-t)/∆t+1∈[(1-t)/∆t] (∆t/2) β ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w)]ρ(β, s)dβ D 4 := exp 1 t β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβdt Note that ∥C 1 -C 2 ∥ ≤ ∥D 1 -D 2 ∥ + ∥D 2 -D 3 ∥ + ∥D 3 -D 4 ∥. Assumption 2 (iii)indicates that for any s = 0, ∆t, . . . , (L -1)∆ and j = 1, . . . , M , we havemax{∥∇ 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:P M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] - β ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )]ρ(β|s)dβ ≥ (N + 1)Dz ≤2(N + 1)D exp(- z 2 2B 2 J M ) (F.10) and P M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, t), w s,j )] - β ∇ vec[T ] vec[h( T Θ (H, t), w s,j )]ρ(β|s)dβ ≥ (N + 1)Dz ≤2(N + 1)D exp(- z 2 2B 2 J M )(F.11) 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 2M (δ + log(2(N + 1)DL)), which implies that with probability at least 1 -exp(-δ) with respect to the parameter initialization Θ (0) , we have∥M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] - β ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )]ρ(β|s)dβ∥ and ∥M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, t), w s,j )] - β ∇ vec[T ] vec[h( T Θ (H, t), w s,j )]ρ(β|s)dβ∥ bounded by B J 2M (N + 1)D(δ + log(2(N + 1)DL)) ≤ C J δ+log(L+1) Mfor 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 ∥ ≤ 1 2L (s-t)/∆t+2∈[(1-t)/∆t] M -1 M j=1 ∇ vec[T ] vec[f ( T Θ (H, s), θ s,j )] - β ∇ vec[T ] vec[f ( T Θ (H, s), θ)]ρ(β, s)dβ + (s-t)/∆t+1∈[(1-t)/∆t] M -1 M j=1 ∇ vec[T ] vec[h( T Θ (H, s), w s,j )] - β ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w)]ρ(β, s)dβ ≤C J δ + log(L + 1) M .(F.12) For any s = 0, ∆t, . . . , (L -1)∆, we defineA s,j = (∆t/2) β ∇ vec[T ] vec[f ( T Θ (H, s), θ)]ρ(β|s)dβ and B s,j = (∆t/2) β ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w)]ρ(β|s)dβ.Since Assumption 2 (iii) indicates that max{∥A s,j ∥, ∥B s,j ∥} ≲ ∆t = L -1 , we have∥exp(A s,j ) -I -A s,j ∥ ≲ ∥A s,j ∥ 2 , ∥exp(B s,j ) -I -B s,j ∥ ≲ ∥B s,j ∥ 2 .Applying Lemma C.8 once more, we have∥D 2 -D 3 ∥ ≤ (s-t)/∆t+2∈[(1-t)/∆t] ∥A s,j ∥ 2 + (s-t)/∆t+1∈[(1-t)/∆t] ∥B s,j ∥ 2 ≲ L -1 . (F.13)Since Assumption 2 (iii) ensures the boundedness of ∥D 4 ∥, we have∥D 3 -D 4 ∥ ≲ exp (s-t)/∆t+2∈[(1-t)/∆t] (∆t/2) β ∇ vec[T ] vec[f ( T Θ (H, s), θ)]ρ(β, s)dβ + (s-t)/∆t+1∈[(1-t)/∆t] (∆t/2) β ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w)]ρ(β, s)dβ - 1 t β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβdt -1 . (F.14) Therefore, to show that ∥D 3 -D 4 ∥ ≲ L -1 + δ+log(L+1) M, it suffices to show thatJ 34 := (s-t)/∆t+2∈[(1-t)/∆t] (∆t/2) β ∇ vec[T ] vec[f ( T Θ (H, s), θ)]ρ(β, s)dβ + (s-t)/∆t+1∈[(1-t)/∆t] (∆t/2) β ∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w)]ρ(β, s)dβ - 1 t β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβdt ≲ L -1 + δ + log(L + 1) M .By Assumption 3 (iv) and Lemma D.6, we have∇ vec[T ] vec[f ( T Θ (H, s), θ)] -∇ vec[T ] vec[f (T ρ (H, s), θ)] ≲ ∥ T Θ (H, s)-T ρ (H, s)∥ F ≲ L -1 + δ + log(L + 1) M ,and∇ vec[T ] vec[h( T Θ (H, s + ∆t/2), w)] -∇ vec[T ] vec[h(T ρ (H, s), w)] ≲ ∥ T Θ (H, s + ∆t/2) -T ρ (H, s)∥ F ≲∥ T Θ (H, s + ∆t/2) -T Θ (H, s)∥ F + ∥ T Θ (H, s) -T ρ (H, s)∥ F ≲L -1 + δ + log(L + 1) M + L -1 ∥M -1 M j=1 h( T Θ (H, s), w s,j )∥ F ≲L -1 + δ + log(L + 1) M ,where the last inequality employs Assumption 2 (i). Therefore, we conclude that 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 -∥D 3 -D 4 ∥ ≲ J 34 ≲L -1 + δ + log(L + 1) M + ∥(∆t/2) β ∇ vec[T ] vec[f (T ρ (H, t), θ)]ρ(β, t)dβ∥ (s-t)/∆t+1∈[(1-t)/∆t] (∆t/2) β ∇ vec[T ] vec[f (T ρ (H, s), θ)]ρ(β, s)dβ + (s-t)/∆t+1∈[(1-t)/∆t] (∆t/2) β ∇ vec[T ] vec[h(T ρ (H, s)(H, s), w)]ρ(β, s)dβ - 1 t β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβdt ≲L -1 + δ + log(L + 1) M + sup |s1-s2|≤∆t ∇ vec[T ] vec[g(T ρ (H, s 1 ), β)] -∇ vec[T ] vec[g(T ρ (H, s 2 ), β)] ≲L -1 + δ + log(L + 1) M + sup |s1-s2|≤∆t ∥T ρ (H, s 1 ) -T ρ (H, s 2 )∥ F ≲L -1 + δ + log(L + 1) M (F.C 2 ∥ ≤ ∥D 1 -D 2 ∥ + ∥D 2 -D 3 ∥ + ∥D 3 -D 4 ∥ ≲ L -1 + δ+log(L+1) M . F.4 Proof of Lemma D.4Proof. Define F ρ (T ρ, t) = β vec[g(T ρ(H, t), β)]ρ(β, t)dβ for any ρ, ρ ∈ P 2 . From Taylor's expansion, we havevec[ Ṫρη (H, t) -Ṫρ (H, t)] = F ρ (T ρη , t) -F ρ (T ρ , t) + F ρη (T ρη , t) -F ρ (T ρη , t) = β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβ vec[T ρη (H, t) -T ρ (H, t)] + η β vec[g(T ρη (H, t), β)](ν -ρ)(β, t)dβ + o(η) = β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβ vec[T ρη (H, t) -T ρ (H, t)] + η β vec[g(T ρ (H, t), β)](ν -ρ)(β, t)dβ + η β ∇ vec[T ] vec[g(T ρ (H, t), β)](ν -ρ)(β, t)dβ vec[T ρη (H, t) -T ρ (H, t)] + o(η) = β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβ vec[T ρη (H, t) -T ρ (H, t)] + η β vec[g(T ρ (H, t), β)](ν -ρ)(β, t)dβ + o(η), 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 16 ) which leads to (D.18).d dt exp - ⊤ 0 β ∇ vec[T ] vec[g(T ρ (H, s), β)]ρ(β, s)dβ) vec[T ρη (H, t) -T ρ (H, t)] = exp - ⊤ 0 β ∇ vec[T ] vec[g(T ρ (H, s), β)]ρ(β, s)dβ) vec[ Ṫρη (H, t) -Ṫρ (H, t)] - β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβ vec[T ρη (H, t) -T ρ (H, t)] = exp - ⊤ 0 β ∇ vec[T ] vec[g(T ρ (H, s), β)]ρ(β, s)dβ) η β vec[g(T ρ (H, t), β)](ν -ρ)(β, t)dβ + o(η) , (F.Proof. Fix (β, t) ∈ P r . Lemma C.2 implies that∥p ρ (H, 1)-p ν (H, 1)∥ F = |Read(T ρ (H, 1)-Read(T ν (H, 1)|≤ C r W 1 (ρ, ν) ≤ C r (r +1)∥ρ-ν∥ 1for some constant C r dependent on the parameters listed in the result. Our goal is to regulate the difference between ṗρ (H, t) and ṗν (H, t) to control the the propagation of ∥p ρ (H, •) -p ν (H, •)∥.d dt ∥p ρ (H, t) -p ν (H, t)∥ F ≤∥ ṗρ (H, t) -ṗν (H, t)∥ F = vec[p ρ (H, t)] ⊤ β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβ -vec[p ν (H, t)] ⊤ β ∇ vec[T ] vec[g(T ν (H, t), β)]ν(β, t)dβ ≤∥p ρ (H, t) -p ν (H, t)∥ F β ∇ vec[T ] vec[g(T ρ (H, t), β)] ρ(β, t)dβ +∥p ν (H, t)∥ F β ∇ vec[T ] vec[g(T ν (H, t), β)] (ρ -ν)(β, t)dβ ≤L r,1 ∥p ρ (H, t) -p ν (H, t)∥ F β ρ(β, t)dβ + ∥p ν (H, t)∥ F ∥ρ(•, t) -ν(•, t)∥ 1 ≤L r,1 ∥p ρ (H, t) -p ν (H, t)∥ F β ρ(β, t)dβ + L r,2 ∥ρ(•, t) -ν(•, t)∥ 1 (F.17) where L r,1 = ϕ T (N, D, √ N + 1B exp(K(1 + r + r 2 )))(1 + r + r 2 ) and L r,2 = (B +B exp(K(1+r +r 2 ))) exp ϕ T (N, D, √ N + 1KB exp(K(1+r +r 2 ))(1+r +r 2 ) .The third inequality of (F.17) uses Lemma C.1 to obtain∥T ρ (H, t)∥ F ≤ √ N + 1∥T ρ (H, t)∥ 2-col ≤ √ N + 1B exp(K(1 + r + r 2 ))and∥T ν (H, t)∥ F ≤ √ N + 1∥T ν (H, t)∥ 2-col ≤ √ N + 1B exp(K(1 + r + r 2 ))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∥p ρ (H, t) -p ν (H, t)∥ F ≤C r (r + 1) exp(L r,1 )∥ρ -ν∥ 1 + 1 0 L r,1 L r,2 exp(L r,1 1 t β ρ(β, t)dβds)∥ρ(•, t) -ν(•, t)∥ 1 dt ≤C r (r + 1) exp(L r,1 )∥ρ -ν∥ 1 + L r,1 L r,2 exp(L r,1 ) 1 0 ∥ρ(•, t) -ν(•, t)∥ 1 dt =(C r + L r,1 L r,2 ) exp(L r,1 ))∥ρ -ν∥ 1 (F.18) Since 1 0 ∥ρ(•, t) -ν(•, t)∥ 1 dt = 1 0 β |ρ(β, t) -ν(β, t)|dβdt = ∥ρ -ν∥ 1 .Thus, we complete the proof of the first result.By Lemma C. ≤ E µ [∥g(T ρ (H, t), β) -g(T ν (H, t), β)∥ F ∥p ρ (H, t)∥ F ] + E µ [∥g(T ν (H, t), β)∥ F ∥p ρ (H, t) -p ν (H, t)∥ F ] ≤ L r,3 E µ [∥g(T ρ (H, t), β) -g(T ν (H, t), β)∥ F ] + L r,4 (C r + L r,1 L r,2 )∥ρ -ν∥ 1 ≤ √ N + 1L r,3 E µ [∥g(T ρ (H, t), β) -g(T ν (H, t), β)∥ 2-col ] + L r,4 (C r + L r,1 L r,2 )∥ρ -ν∥ 1 .(F. 19 ) 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 ∥g(T ρ (H, t), β) -g(T ν (H, t), β)∥ 2-col ≤ϕ T (N, D, max{∥T ρ (H, t)∥ F , ∥T ν (H, t)∥ F })(1 + r + r 2 )∥T ρ (H, t) -T ρ (H, t)∥ 2 ≤ϕ T (N, D, √ N + 1KB exp(K(1 + r + r 2 )))(1 + r + r 2 )C r W 1 (ρ, ν) ≤ϕ T (N, D, √ N + 1KB exp(K(1 + r + r 2 )))(1 + r + r 2 )C r (1 + r)∥ρ -:= B exp(K(1 + r + r 2 )).The following decomposition equation holds:∥ T Θ (H, t)-T ρ (H, t)∥ F ≤ ∥ T Θ (H, t) -TΘ (H, t)∥ F J1 + ∥ TΘ (H, t) -T ρ (H, t)∥ F J2 + ∥ T ρ (H, t) -T ρ (H, t)∥ F J3 ,(F.21) 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), ∥ T Θ (H, t) + (∆t/2)M -1 M j=1 f ( T Θ (H, t), θ t,j )∥ 2-col ≤∥ T Θ (H, t)∥ 2-col + (∆t/2)M -1 M j=1 ∥f ( T Θ (H, t), θ t,j )∥ 2-col ≤∥ T Θ (H, t)∥ 2-col + (K∆t/2)M -1 M j=1 ∥ T Θ (H, t)∥ 2-col (1 + ∥θ t,j ∥ + ∥θ t,j ∥ 2 ) ≤B T (1 + (K∆t/2)(1 + r + r 2 )) ≤B T (1 + (K/2)(1 + r + r 2 )) Denote B T (1 + (K/2)(1 + r + r 2 ))= T Θ (H, t) + (∆t/2)M -1 M j=1 f ( T Θ (H, t), θ t,j ) + (∆t/2)M -1 M j=1 h T Θ (H, t) +(∆t/2)M -1 M j=1 f ( T Θ (H, t), θ t,j ), w t,j (F.22)Then, from the formula of (F.22) and Assumption 2 (iii), we see that for any t = 0, ∆t, . . . , (L -1)∆,∥ T Θ (H, t + ∆t) -TΘ (H, t + ∆t)∥ F ≤∥ T Θ (H, t) -TΘ (H, t)∥ F 1 + (∆t/2)M -1 M j=1 ϕ T (N, D, √ N + 1B T )(1 + ∥θ t,j ∥ + ∥θ t,j ∥ 2 ) +(∆t/2)M -1 M j=1 ϕ T (N, D, √ N + 1 BT )(1 + ∥w t,j ∥ + ∥w t,j ∥ 2 ) ∥ T Θ (H, t) + (∆t/2)M -1 M j=1 f ( T Θ (H, t), θ t,j ) -TΘ (H, t)∥ F ≤∥ T Θ (H, t) -TΘ (H, t)∥ F 1 + ∆tM -1 M j=1 ϕ T (N, D, √ N + 1 BT )(1 + ∥β t,j ∥ + ∥β t,j ∥ 2 ) + √ N + 1B T (K∆t/2)(1 + r + r 2 )(∆t/2)M -1 M j=1 ϕ T (N, D, √ N + 1 BT )(1 + ∥w t,j ∥ + ∥w t,j ∥ 2 ) ≤∥ T Θ (H, t) -TΘ (H, t)∥ F 1 + ∆tϕ T (N, D, √ N + 1 BT )(1 + r + r 2 ) + √ N + 1ϕ T (N, D, √ N + 1 BT )B T (K∆t 2 /4)(1 + r + r 2 ) 2(F.23) Therefore, applying (F.23), we deduce that for any t = 0, ∆t, . . . , (L -1)∆t, 1, we have:∥ T Θ (H, t)-TΘ (H, t)∥ 2-col ≤ √ N + 1B T (K∆t 2 /4)(1+r+r 2 ) exp(ϕ T (N, D, √ N + 1 BT )(1 + r + r 2 )) ∆t C 1 L -1 (F.24) where C 1 := √ N + 1B T K(1 + r + r 2 ) exp(ϕ T (N, D, √ N + 1 BT )(1 + r + r 2 ))/4.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 haveP(∥M -1 M j=1 g( TΘ (H, t), β t,j ) :,i - β g( TΘ (H, t), β t,j ) :,i ρ(β|t)dβ∥ ≥ z) ≤ 2 exp(- z 2 2B 2 T M ).By the union bound over i ∈ [N + 1] and t = 0, ∆t, . . . , (L -1)∆, the above inequality implies Suppose that the high probability event E δ occurs. Let's denote B T 2M -1 (δ + log((N + 1)L)) by B δ for brevity. From Assumption 2 (iii), it follows that for any t = 0, ∆t, . . . , (L -1)∆,P(sup t ∥M -1 M j=1 g( TΘ (H, t), β t,j )- β g( TΘ (H, t), β t,j )ρ(β|t)dβ∥ 2-col ≥ z) ≤ 2(N +1)L exp(- z 2 2B 2 T M ).∥ TΘ (H, t + ∆t) -T ρ (H, t + ∆t)∥ F ≤∥ TΘ (H, t) -T ρ (H, t)∥ F +∆t∥M -1 M j=1 g( TΘ (H, t), β t,j ) - β g( TΘ (H, t), β t,j )ρ(β|t)dβ∥ F +∆t∥ β (g( TΘ (H, t), β t,j ) -g( T ρ (H, t), β t,j ))ρ(β|t)dβ∥ F ≤∥ TΘ (H, t) -T ρ (H, t)∥ F + ∆t √ N + 1B δ + +∆t β ∥g( TΘ (H, t), β t,j ) -g( T ρ (H, t), β t,j )∥ F ρ(β|t)dβ ≤∥ TΘ (H, t) -T ρ (H, t)∥ F (1 + ∆tϕ T (N, D, √ N + 1B T )(1 + r + r 2 )) +∆t √ N + 1B δ .(F.27) Repeatedly applying Equation (F.27) yields∥ T Θ (H, t) -TΘ (H, t)∥ F ≤ √ N + 1B δ ϕ T (N, D, √ N + 1B T )(1 + r + r 2 ) exp(ϕ T (N, D, √ N + 1B T )(1 + r + r 2 ))=C 2 M -1 (δ + log((N + 1)L)).(F.28) 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( 1L ) 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 derivatived 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 ∥ T ρ (H, t + ∆t) -T ρ (H, t + ∆t)∥ F ≤∥ T ρ (H, t) -T ρ (H, t)∥ F + ∥T ρ (H, t + ∆t) -T ρ (H, t) -∆t β g(T ρ (H, t), β)ρ(β|t)dβ∥ F I1 + ∆t ∥ β (g(T ρ (H, t), β) -g( T ρ (H, t), β))ρ(β|t)dβ∥ F I2 . (F.∥ β g(T ρ (H, t), β)(ρ(β|s) -ρ(β|t))dβ∥ F = sup s∈[t,t+∆t] ∥ β g(T ρ (H, t), β)(ρ(β, s) -ρ(β, t))dβ∥ F ≤ sup s∈[t,t+∆t] C 3,2 ∥ρ(•, s) -ρ(•, t)∥ BL ≤C 3,2 C ρ ∆t (F.32) 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 thatI 2 ≤ β ∥g(T ρ (H, t), β) -g( T ρ (H, t), β)∥ F ρ(β|t)dβ ≤ϕ T (N, D, √ N + 1B T )(1 + r + r 2 )∥T ρ (H, t), β) -T ρ (H, t), β)∥ F (F.33)Therefore, by bounding I 1 + I 2 , we obtain the following inequalityEQUATIONF.34) which implies, after being used multiple times, that∥T ρ (H, t), β) -T ρ (H, t), β)∥ F ≤ C 3,5 ∆t 2 (1 + C 3,6 ∆t) L+1 -1 C 3,6 ∆t ≤ C 3 L -1 (F.35)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. I 4 ≲ ∥H -H ′ ∥ F .From Assumption 2 (iii), we haveI 2 ≤ exp 1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)] ρ(β, s)dβds ≤ exp ϕ T (N, D, √ N + 1B T )(1+r+r 2 ) .Thus, dividing I 4 by the uniformly bounded part I 2 , we obtainI 4 ≤ exp 1 t β ∇ vec[T ] vec[g(T ρ (H ′ , s), β)] -∇ vec[T ] vec[g(T ρ (H, s), β)] ρ(β, s)dβds -I divvec[T ]≤ exp1 t β ∇ vec[T ] vec[g(T ρ (H ′ , s), β)] -∇ vec[T ] vec[g(T ρ (H, s), β)] ρ(β, s)dβds -1 ≤ exp ϕ T T (N, D, √ N + 1B T , r) sup t∈[0,1] ∥T ρ (H, t) -T ρ (H ′ , t)∥ F -1,(F.37) 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 havesup 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 deriveI 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 equationvec[ ṗρ (H, t)] ⊤ = -vec[p ρ (H, t)] ⊤ β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβ (F.38) with β ∇ vec[T ] vec[g(T ρ (H, t), β)]ρ(β, t)dβ ≤ (1 + C ρ /2)ϕ T (N, D, √ N + 1B T )(1 + r + r 2 ) =: C 0implied by Assumption 2 (iii). Therefore, the Grönwall's inequality directly indicates thatE µ ∥p ρ (H, t)∥ 2 F ≥ exp(-2C 0 t)E µ ∥p ρ (H, 1)∥ 2 F ≥ exp(-2C 0 )E µ [|Read[T ρ (H, 1)] -y(H)| 2 ] ≥ 2 exp(-2C 0 )R(ρ), E µ ∥p ρ (H, t)∥ 2 F ≤ exp(2C 0 t)E µ ∥p ρ (H, 1)∥ 2 F ≤ exp(2C 0 )E µ [|Read[T ρ (H, 1)] -y(H)| 2 ] ≤ 2 exp(2C 0 )R(ρ), 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 F (T ) = -p ρ (T -1 t * (T ), t * ) + β g(T, β)ρ(β|t * )dβ - h(T, w 0 ) 2 ,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 F (T ) - ∞ k=1 c k f (T, θ k ) max ≤ ϵ/3 (F.39)given some ϵ > 0 and anyT 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 F (T ) - whereF (T ) == -p ρ (T -1 t * (T ), t * ) + β g(T, β)ρ(β|t * )dβ.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 thatF (T ) - β g(T, β)ν(β)dβ max ≤ ϵ. (F.42)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.Step IV: Upper bound β δQ δρ (β, t * ) ν(β) -ρ(β|t * ) dβ to complete the proof Utilizing the gradient definition in (3.4) and β g(T, β)ρ(β|t * )dβ = F (T ) + p ρ (T -1 t * (T ), t * ), we obtainβ δQ δρ (β, t * ) ν(β) -ρ(β|t * ) dβ =E µ β Tr g(T ρ (H, t * ), β) ⊤ p ρ (H, t * ) ν(β) -ρ(β|t * ) dβ + λ 2 β ∥β∥ 2 ν(β) -ρ(β|t * ) dβ ≤E T ∼ µ(t) Tr g(T, β) ⊤ v(β) -ρ(β|t * ) p ρ (T -1 t * (T ), t * ) + λ 2 (R(ϵ) + r) = E T ∼ µ(t) Tr F (T ) - β g(T, β)ν(β)dβ ⊤ p ρ (T -1 t * (T ), t * ) J1 -E T ∼ µ(t) Tr p ρ (T -1 t * (T ), t * ) ⊤ p ρ (T -1 t * (T ), t * ) J2 + λ 2 (R(ϵ) + r)For ρ ∈ P 2 concentrated on P r , we observeR(ρ) ≤ 1 2 E µ [ |Read[T ρ (H, 1)]|+|y(H)| 2 ] ≤ 1 2 (B T + B) 2 .Hence, to bound J 1 , we haveJ 1 ≤E T ∼ µ(t) ∥F (T ) - β g(T, β)ν(β)dβ∥ F • ∥p ρ (T -1 t * (T ), t * )∥ F dt ≤ (N + 1)DϵE T ∼ µ(t) ∥p ρ (T -1 t * (T ), t * )∥ F dt ≤(N + 1)Dϵ E T ∼ µ(t) ∥p ρ (T -1 t * (T ), t * )∥ 2 F 1/2 dt =(N + 1)Dϵ E µ ∥p ρ (H, t)∥ 2 F 1/2 dt ≤ √ 2(N + 1)D exp(C 0 )R(ρ) 1/2 ϵ ≤(N + 1)(B T + B)(N + 1)D exp(C 0 )R(ρ)ϵ =C 3 R(ρ)ϵ, (F.43) where C 3 := (N + 1)(B T + B)(N + 1)D exp(C 0 ). To bound J 2 , we have J 2 =E T ∼ µ(t) ∥p ρ (T -1 t (T ), t)∥ 2 F dt =E µ ∥p ρ (H, t)∥ 2 F dt ≥E µ exp(-2C 0 )∥p ρ (H, 1)∥ 2 F dt ≥ 1 2 exp(-2C 0 )R(ρ).(F.44)Combining (F.43) and (F.44), by choosing ϵ = 1 4 exp(-2C 0 )/C 3 , which only depends on N , D, r, and the parameters of the assumptions, we have β δQ δρ (β, t * ) ν(β) -ρ(β|t * ) dβ ≤ - 1 4 exp(-2C 0 )R(ρ) + λ 2 R( 1 4 exp(-2C 0 )/C 3 ) + r . Setting B r = R( 1 F.∥T ρ (H, t)∥ 2-col ≤ ∥T ρ (H, 0)∥ 2-col exp( 1 0 β K(1 + ||β|| 2 +||β|| 2 2 )ρ(β, t)dβdt) ≤ ∥H∥ 2-col exp(K(1 + A + A 2 )).F. (H, t) -T ν (H, t)∥ F =∥ ⊤ 0 Ṫρ (H, s) -Ṫν (H, s)ds∥ F =∥ ⊤ 0 β g(T ρ (H, s), β)ρ(β, t)dβds - ⊤ 0 β g(T ν (H, s), β)ν(β, s)dβds∥ F ≤ ⊤ 0 β ∥g(T ρ (H, s), β) -g(T ν (H, s), β)∥ F ρ(β, s)dβds J1 + ∥ ⊤ 0 β g(T ν (H, s), β) ρ -ν (β, s)dβds∥ F J2 (F.47)We then bound J 1 and J 2 using the following two lemmas separately. Firstly, since∥T ρ ∥ F ≤ √ N + 1C and ∥T ν ∥ F ≤ √ N + 1C, we have by Assumption 2 that ∥g(T ρ (H, s), β) -g(T ν (H, s), β)∥ F ≤ sup ∥T ∥ F ≤ √ N +1C ∥∇ vec[T ] vec[g(T, β)]∥ 2 ∥T ρ (H, s) -T ν (H, s)∥ F ≤ ϕ T (N, D, √ N + 1C)(1 + r + r 2 )∆(H, s) (F.48)Therefore, by (F.48) we have J 1 ≤ ϕ T (N, D, √ N + 1C)(1 + C ρ /2)(1 + r + r 2 ) ⊤ 0 ∆(H,EQUATION(T ν (H, s), β) is √ N + 1(1+C ρ /2)KCϕ T (N, D, √ N + 1C)(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 + 1C)(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), thatJ 2 = ∥ 1 0 D g(T ν (H, s), β) ρ -ν (β, s)dβds∥ F ≤ C ′ W 1 (ρ, ν). (F.52) Define C * := max{C ′ , ϕ T (N, D, √ N + 1C)(1 + r + r 2 )}.By combining (F.49) and (F.52), we have∆(H, t) ≤ C * ⊤ 0 ∆(H, s)ds + C * W 1 (ρ, ν). (F.53)Applying the Grönwall's inequality then shows∆(H, t) ≤ C * exp(C * t)W 1 (ρ, ν). (F.54) Specifically, we have ∥T ρ (H, 1) -T ν (H, 1)∥ F ≤ C * exp(C * )W 1 (ρ, ν). F.10 Proof of Lemma C.3Proof. 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∥g(T ρ (H, t), β)∥ F ≤ √ N + 1∥g(T ρ (H, t), β)∥ 2-col ≤ √ N + 1K∥T ρ (H, t)∥ 2-col (1 + r + r 2 ) ≤ √ N + 1KB exp(K(1 + A + A 2 ))(1 + r + r 2 ).(F.55)On the other hand, from (C.14), we have ∥p ρ (H, t)∥ F ≤ |Read(T ρ (H, 1)) -y(H)|• exp 1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)ρ(β, s)dβds] 2 ≤ (B + B exp(K(1 + A + A 2 ))) exp 1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)] ρ(β, s)dβds ≤ (B + B exp(K(1 + A + A 2 ))) exp 1 t β ∇ vec[T ] vec[g(T ρ (H, s), β)] ρ(β, s)dβds ≤ (B + B exp(K(1 + A + A 2 ))) exp 1 0 β ϕ T (N, D, ∥T ρ (H, s)∥ F )(1 + ∥β∥ + ∥β∥ 2 )ρ(β, s)dβds ≤ (B + B exp(K(1 + A + A 2 ))) exp ϕ T (N, D, √ N + 1KB exp(K(1 + A + A 2 ))(1 + A + A 2 ) (F.(H, t), β)∥ F ∥p ρ (H, t)∥ F + λ 2 r 2 ≤ √ N + 1KB exp(K(1 + A + A 2 ))(1 + r + r 2 )(B + B exp(K(1 + A + A 2 ))) exp ϕ T (N, D, √ N + 1KB exp(K(1 + A + A 2 ))(1 + A + A 2 ) + λ 2 r 2 .for any t = 0, ∆t, . . . , 1.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, whereEQUATIONwhere 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.7Proof. 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)∆tP(|M -1 M j=1 ∥β t,j ∥ 2 - β ∥β∥ 2 ρ(β, t)dβ|≥ z) ≤ 2 exp(- z 2 2r 2 M ).By applying the union bound over t = 0, ∆t, . . . , (L -1)∆t, the inequality above impliesP(| 1 M L t M j=1 ∥β t,j ∥ 2 - 1 L t β ∥β∥ 2 ρ(β, t)dβ|≥ z) ≤P(sup t |M -1 M j=1 ∥β t,j ∥ 2 - β ∥β∥ 2 ρ(β, t)dβ|≥ z) ≤2L exp(- z 2 2r 2 M ). (F.62)In addition, we have| 1 L t β ∥β∥ 2 ρ(β, t)dβ - 1 0 β ∥β∥ 2 ρ(β, t)dβdt|≤ sup |t-s|≤∆t | β ∥β∥ 2 ρ(β, t)dβ - β ∥β∥ 2 ρ(β, s)dβ| ≤r 2 L -1 sup |t-s|≤∆t ∥ρ(•, t) -ρ(•, s)∥ BL ≤r 2 C ρ L -1 .(F.63) Combining (F.62) and (F.63), and setting z = r 2M -1 (δ + log(2L)) completes the proof. Proof. The proof will be trivial by noting the equalityL l=1 A l - L l=1 B l = L l=1 l-1 s=1 B s (A l -B l ) L s=l+1 A s .Hence, we have∥ L l=1 A l - L l=1 B l ∥ ≤ L l=1 ∥ l-1 s=1 B s ∥ • ∥A l -B l ∥ • ∥ L s=l+1 A s ∥ ≤ C L l=1 ∥A l -B l ∥.In this section, we considerf (Z, θ) = V Zsoftmax(Z ⊤ W Z) (G.1)with the collection of parameters θ = vec[V, W ], where softmax denotes the column-wise softmax function. Moreover, considerh(Z, w) = W 2 HuberizedReLU(W 1 H) (G.2)with the collection of parameters w = vec[W 1 , W 2 ]. Here, HuberizedReLU denotes the entry-wise HuberizedReLU activation function defined asHuberizedReLU(x) =      0, if z ≤ 0; z 2 /2, if z ∈ [0, 1]; z -1/2, if z ≥ 1.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 f (Z, θ) = V Zsoftmax(Z ⊤ W Z) = (f (Z, θ) :,i ) 1≤i≤N +1 , where f (Z, θ) :,i = N +1 j=1 P ij V z j and P i,: = softmax(Z ⊤ W z 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 ∂zj (Z, θ) ∈ R D×D . After calculation, we obtainJ ij = V ZQ i [E ji Z ⊤ W + Z ⊤ W ⊤ δ ij ] + P ij V,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). 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∥f (T, θ)∥ 2-col = ∥V T softmax(T ⊤ W T )∥ 2-col ≤∥V ∥ 2 • ∥T ∥ 2 • ∥softmax(T ⊤ W T )∥ 2-col ≤∥θ∥ 2 • ∥T ∥ 2-col • ∥softmax(T ⊤ W T )∥ 1-col ≤∥θ∥ 2 • ∥T ∥ 2-col ,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. For Assumption 2 (ii), (G.5) leads to∥∇ vec[V ] f (T, θ) :,i ∥ 2 ≤ N +1 j=1 ∥P ij T ⊤ :,j ∥ 2 ≤ N +1 j=1 P ij ∥T :,j ∥ 2 ≤ ∥T ∥ 2-col .Moreover, (G.4) leads to ∥J ij ∥ 2 .∇ vec[W ] f (T,For any 1 ≤ i, j ≤ N + 1, we have∥J ij ∥ 2 ≤P ij ∥V ∥ 2 + ∥T ∥ 2 ∥Q i ∥ 2 ∥E ji T ⊤ W + T ⊤ W ⊤ δ ij ∥ 2 ∥V ∥ 2 ≤∥V ∥ 2 1 + 2∥T ∥ 2 2 ∥W ∥ 2 ≤∥θ∥ 2 + 2∥T ∥ 2 F ∥θ∥ 2 2 ≤(1 + 2∥T ∥ 2 F )(1 + ∥θ∥ 2 + ∥θ∥ 2 2). Hence, we haveEQUATION2 ). 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 thatP ij = softmax(Z ⊤ W z i )is locally Lipschitz continuous with respect to Z and θ. By (G.3) and the chain rule, we can derive that∇ W kl P ij = N +1 j=1 Z ki Z jl P ij (δ ik -P il ), ∂P ij ∂z j = Q i [E ji z ⊤ j W + z ⊤ j W ⊤ δ ij ].and the local Lipschitz continuity is then obvious given the boundedness of ∇ W kl P ij and ∂Pij ∂zj 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:h(T, W 1 , c • W 2 ) = c • W 2 HuberizedReLU(W 1 H) = c • h(T, W 1 , W 2 ).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 • 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.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).exp(-2C 0 )/C 3 ), C 1 = R( 1 4 exp(-2C 0 )/C 3 ) + r /2, and C 2 = 14 exp(-2C 0 ) completes the proof. 


Decision: The authors present a rigorous theoretical analysis of global convergence in an in-context learning setting for certain transformer-type neural network architectures. This is accomplished by a mean-field theory approach: one takes width and depth to infinity, then studies the approximation error and convergence properties of gradient flow with weight decay. The authors show that the dynamics of the discretized model approximates the Wasserstein gradient flow under some regularity assumptions, and moreover that under partial homogeneity and universal kernel assumptions, if the Wasserstein flow weakly converges to a stationary distribution, then the discrete model must also converge with arbitrarily small risk.

After robust discussion, reviewers uniformly acknowledge the nontrivial technical contributions of the authors to the mean-field study of neural networks, which so far have concentrated on ResNet architectures. Reviewers also emphasized several drawbacks of the work: the lack of practical insights into the transformer architecture afforded by the theory; certain imprecisions in presentation of the networks to which the authors’ analysis applies, which, after protracted discussion were eventually more-or-less clarified; the difficulty of checking the authors’ strong assumptions on model/data necessary to obtain global convergence results in practice (which, it must be noted, are inherited from prior work on mean-field theory of neural networks). However, on balance, reviewers argued that the paper's contributions to mean-field theory should be seen as outweighing these concerns. The AC concurs with this assessment and recommends acceptance. The authors are encouraged to incorporate all feedback from the reviews on presentation into the revision, especially the outputs of the discussion with `iBME`.

Reviews: review: paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: ## Potential fallacy in the claim of Corollary 4.1:

1. 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.  
The 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.)  
Based on the above, the reviewer thinks that 
    1. The current approach in proving Corollary 4.1 is flawed.
    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.

## Limitations of settings and results

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

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

2. The main result, Theorem 4.1, relies heavily on a few additional assumptions that are not justified in the paper:
    1. the Wasserstein gradient flow weakly converges to some distribution $\rho_{\infty}$
    2. the uniform boundedness of $\rho^{(\tau)}$ for large time $\tau$
    3. the separation property for $\alpha_1$ with the support expansion of $\alpha_2$ to $K$.

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

3. 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)}$.

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

## Minor:

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

Typos:
1. Line 111, it should be "h(Z, w) = W_2 \sigma_M(W_1 Z)" instead of $W_1 H$. 
2. Line 156, in Assumption 2, $\Psi_j$ in the description should be $\Phi_P$. 
3. Line 86: "~~the~~ our deep Transformer model"
4. 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 + ...$
questions: questions: 1. Please address the above concern for Corollary 4.1 carefully. The answer may greatly affect the final rating from the reviewer.

2. Questions related to assumption 1-4:
    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. 
    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?
limitations: limitations: The authors have adequately addressed the limitation of the paper.
paper_summary: paper_summary: This 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.
Then, 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.
Then, they study the global convergence properties of the gradient flow: they prove a sort of global convergence result, which requires some hypotheses.
strength_weakness: strength_weakness: One 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. 

The 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. 
The paper does not shed light on the interest/peculiarity of the transformer architecture.
I 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.

The 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.
questions: questions: What 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?

The 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?

I 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.
limitations: limitations: Not applicable.
paper_summary: paper_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.
strength_weakness: strength_weakness: N/A
questions: questions: N/A
limitations: limitations: N/A
paper_summary: paper_summary: The 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$.
strength_weakness: strength_weakness: 1. The content is intense and the presentation doesn't help a novel reader. I believe polishing of the manuscript would improve the quality. 
More concretely, I would add a small paragraph in the main manuscript summarizing the main ideas behind the novelties. In the appendix 			instead 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 		each section with the main steps required for the final proof, to guide the reader and make even the technical part more accessible 			(something on the line of what the authors did in D.3 would be perfect). 

2. 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.
questions: questions: On the technical side: 
1. 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. 	Given 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?
2. 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.
3. 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. 
4. 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?
limitations: limitations: The authors correctly assessed the limitations in the manuscript, I have no further suggestions.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: * 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.
* 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$.)
* 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?
* 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.
questions: questions: * Can the dependency on time horizon be alleviated with stronger model assumptions?
* Is the $L^{-1}+\sqrt{\log L/M}$ rate tight/expected? (see Weakness)
* 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?
* 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).

[1] https://openreview.net/forum?id=xm2lU7tteQ

[2] https://arxiv.org/abs/2405.15712
limitations: limitations: See Weaknesses.

Comments: comment: The reviewer thanks the authors for their further efforts. While the reviewer is not able to examine the full proof within the rebuttal period, the reviewer will take a closer look into it in the coming days. Based on the provided proof, the reviewer believes that it can be incorporated into the revised paper and is therefore willing to increase the score, provided no major concerns arise upon closer inspection. As a side note, the reviewer finds it interesting to explore how the parameters specific to the Transformer architecture ($\phi_P, \phi_T$ and similar quantities in Assumptions 1-4) influence the convergence properties in the final theorem. Once again, the reviewer appreciates the authors' efforts during the rebuttal period to address the concerns raised.
Thank you for your suggestions on the experiments and for increasing your score! We will make sure to add experiments to study the size of parameters that can guarantee to global convergence.
We hope that by presenting the details that you have requested, we can convince you that our theory can be applied to fairly practical Transformer models, and that the volume of work required should not be an issue. We are dedicated to add additional details on the explanations and verifications in our revised paper. We sincerely hope that you could take our discussion above into consideration and reevaluate our result. We appreciate your time and effort in reviewing our paper.
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 $\mathbb{E}_{\mu}$.

(i) and (iii): Given the calculation of derivatives we have presented above, it suffices to show that $\mathrm{diag}\Big([B\_{kj}(z\_j)]\_{1\leq k\leq D}\Big)$ and $\mathrm{diag}\Big([Z^T  B\_{kj}( z\_j)]\_{1\leq k\leq D}\Big)\frac{\partial\mathrm{softmax}(A\_i)}{\partial A\_i}Vz\_j$ are both locally Lipschitz continuous with respect to $Z$ and $\theta.$

Since each of $\mathrm{diag}\Big([B\_{kj}(z\_j)]\_{1\leq k\leq D}\Big),$ $\mathrm{diag}\Big([Z^T  B\_{kj}( z\_j)]\_{1\leq k\leq D}\Big),$ $\frac{\partial\mathrm{softmax}(A\_i)}{\partial A\_i}$ and $Vz\_j$ is obviously bounded by an increasing function of $N,D,\\|\theta\\|, K_T,L_T,$ it suffices to show that each of them is locally Lipschitz continuity with respect to both $Z$ and $\theta.$ This 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 $\sqrt{\sum\_{1\leq i,j\leq 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 $\theta$ and $Z.$ Again each component of $J\_{ij}$ that depends on $Z$ or $\theta,$ i.e. 
$Z,Q_i,W,P\_{ij},$ is bounded by an increasing function of  $N,D, K_P,L_T,K_T,\\|\theta\\|,$ and is locally Lipschitz continuous given sufficient smoothness. Hence, (ii) and (iv) also hold.

The proof for the HuberizedReLU MLP encoder is similar to the above and is not conceptually complicated, if not easier. We omit the details here, and save the space for a more detailed discussion about Assumption 4, as we expect that you may be more skeptical about the proof regarding the “universal kernel” assumption in Assumption 4.

---

For Assumption 4, we consider the pair $(g,\alpha) = (h,w),$ and the partition $\alpha = (\alpha\_1,\alpha\_2)$ with $\alpha\_1 = W\_2,$ $\alpha\_2 = W\_1$. We also let a compact set $\mathcal{K} = \\{ W_1: \\| W_1 \\| \leq 1 \\}$. Then Assumption 4(i) on the partial $1$-homogeneity property straightforwardly holds: 

$h(T,W\_1,c\cdot W\_2) = c\cdot W\_2\mathrm{HuberizedReLU}(W\_1H)= c\cdot h(T,W\_1,W\_2).$

Regarding Assumption 4(ii) on the universal kernel property, we first note that according to the choice $(g,\alpha) = (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 [1,2,3] 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 approximators.

Therefore, we know that the function class 

$ \mathrm{span} \\{  W\_2\mathrm{ReLU}^2(W\_1T):  W_1 \in \mathbb{R}^{\mathrm{dim}(W_1)}, W_2 \in \mathbb{R}^{\mathrm{dim}(W_2)}  \\} $ 

is dense in $\mathcal{C}(\\|T\\|\_{2-\mathrm{col}}\leq B,\mathbb{R}^{D\times(N+1)})$. Moreover, by the definition of HyberizedReLU, for any $B>0$ and any $\hat{W}_1$, $\hat{W}_2$, there exist small constant $c$ such that $c\cdot \hat{W}_1 \in \mathcal{K}$, $c\cdot \\|\hat{W}_1\\| \leq B^{-1}$, and

$ c^{-2} \cdot \hat{W}\_2\mathrm{HyberizedReLU}(c\cdot \hat{W}\_1T) =  c^{-2} \cdot \hat{W}\_2\mathrm{ReLU}^2(c\cdot \hat{W}\_1T) = c^2\cdot c^{-2} \cdot \hat{W}\_2\mathrm{ReLU}^2(\hat{W}\_1T) = \hat{W}\_2\mathrm{ReLU}^2(\hat{W}\_1T)$,

where the second equation follows by the positive $2$-homogeneity of $\mathrm{ReLU}^2$ activation. This implies that 

$  \\{  W\_2\mathrm{ReLU}^2(W\_1T):  W_1 \in \mathbb{R}^{\mathrm{dim}(W_1)}, W_2 \in \mathbb{R}^{\mathrm{dim}(W_2)}  \\}  \subseteq \\{W\_2\mathrm{HyberizedReLU}(W\_1T): W_2 \in \mathbb{R}^{\mathrm{dim}(W_2)}\times \mathcal{K}\\} $.

Therefore, we conclude that $ \mathrm{span} \\{W\_2\mathrm{HyberizedReLU}(W\_1T): W_2 \in \mathbb{R}^{\mathrm{dim}(W_2)}\times \mathcal{K}\\}$ is dense in $\mathcal{C}(\\|T\\|\_{2-\mathrm{col}}\leq B,\mathbb{R}^{D\times(N+1)})$. This finishes the validation of Assumption 4.

---

[1] Funahashi, Ken-Ichi (January 1989). "On the approximate realization of continuous mappings by neural networks". Neural Networks.

[2] Cybenko, G. (1989). "Approximation by superpositions of a sigmoidal function". Mathematics of Control, Signals, and Systems.

[3] Pinkus, Allan (January 1999). "Approximation theory of the MLP model in neural networks". Acta Numerica.
Proof. We first introduce some notations. We remind the readers that for a matrix $\mathbf{A}$, we denote by $\\| \mathbf{A} \\|\_2$ the spectral norm of $\mathbf{A}.$ We also denote by $\\| \cdot \\|\_{1-\mathrm{col}}$ the maximum $\ell\_1$-norm across all columns of a matrix. 

Denote $Z=(z\_1,\dots,z\_{N+1})\in\mathbb{R}^{D\times N+1}.$ Then the function $f$ can be rewritten as

$f(Z,\theta)=VZ\mathrm{softmax}(Z^TWZ)=(f(Z,\theta)\_{:,i})\_{1\leq i\leq N+1}$,

where

$f(Z,\theta)\_{:,i}=\sum\_{j=1}^{N+1}P\_{ij}z\_j$ and $P\_{i,:}=\mathrm{softmax}(Z^TWz\_i).$ Next, we calculate the derivatives of $f(Z,\theta)\_{:,i}$ with respect to $Z$ and $\theta$ as follows:


For $Z$: the Jacobian $J\in\mathbb{R}^{(N+1)D\times(N+1)D}$ is $J=(J\_{ij})\_{1\leq i,j\leq N + 1},$ where $J\_{ij}=\frac{\partial f\_{:,i}}{x\_j}\in\mathbb{R}^{D\times D}.$ After calculation, we obtain

$J\_{ij}=ZQ\_i\big[E\_{ji}Z^TW+Z^TW^T\delta\_{ij}\big]+P\_{ij}I,$

where $Q\_i:=\mathrm{diag}(P\_{i:})-P^T\_{i:}P\_{i:},$ $E\_{ij}$ is the matrix with zeros everywhere except one the $(i,j)$-th entry, and $\delta\_{ij}$ is the Kronecker delta.


For $\theta$: Define $A\_i=Z^TWz\_i.$ After calculation, we have

$\nabla\_{\mathrm{vec}[V]} f(Z,\theta)\_{:,i}=\sum\_{j=1}^{N+1}P\_{ij}\mathrm{diag}\Big([B\_{kj}(z\_j)]\_{1\leq k\leq D}\Big),$

$\nabla\_{\mathrm{vec}[W]} f(Z,\theta)\_{:,i}=\sum\_{j=1}^{N+1}P\_{ij} \Big(\mathrm{diag}\Big([z_l^T  B\_{kj}( z\_j)]\_{1\leq k\leq D}\Big)[\frac{\partial\mathrm{softmax}(A\_i)}{\partial A\_i}]\_l\Big)_{1\leq l \leq N+1} Vz\_j,$

where $B\_{kj}(z)$ is the $D\times D$ matrix with zeros everywhere except $z\_j$ for the $k$-th row.

---

We then verify the assumptions one by one. For Assumption 2 (i), we have


$ \\| f(T,\theta) \\|\_{2-\mathrm{col}}  = \\| VT \mathrm{softmax} (T^T W T ) \\|\_{2-\mathrm{col}} \leq \\| V\\|\_2 \cdot \\| T \\|\_2 \cdot \\| \mathrm{softmax} (T^T W T ) \\|\_{2-\mathrm{col}} \leq \\| \theta \\|\_2 \cdot \\| T \\|\_{2-\mathrm{col}}  \cdot \\| \mathrm{softmax} (T^T W T ) \\|\_{1-\mathrm{col}} \leq \\| \theta \\|\_2 \cdot \\| T \\|\_{2-\mathrm{col}},$

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

For $h$ in (Eq 2), we have 

$ \\| h(T,w) \\|\_{2-\mathrm{col}} = W\_2\mathrm{HuberizedReLU}(W\_1T) \\|\_{2-\mathrm{col}} \leq \\| W\_2 \\|\_2 \cdot \\|\mathrm{HuberizedReLU}(W\_1T) \\|\_{2-\mathrm{col}}  \leq  \\| W\_2 \\|\_2 \cdot \\|W\_1T \\|\_{2-\mathrm{col}} \leq 2 \cdot \\| w \\|\_2^2 \cdot \\|T \\|\_{2-\mathrm{col}}$,

where the second inequality follows by the property of HuberizedReLU that $|\mathrm{HuberizedReLU}(x)| \leq |x|.$ This demonstrates that Assumption 2 (i) with $K=1$ holds for $h$ in (Eq 2) as well.



For Assumption2 (ii), we have

$\\| \nabla\_{\mathrm{vec}[V]} f(T,\theta)\_{:,i} \\|\_{2} \leq \sum\_{j=1}\^{N+1} P\_{ij} \\| \mathrm{diag}\Big([B\_{kj}(T\_j)]\_{1\leq k\leq D}\Big) \\|\_{2} \leq \sum\_{j=1}^{N+1} P\_{ij} \\| T\_j \\|\_{2}
\leq \\|T\\|\_{2-\mathrm{col}}.
$

Similarly, we have 

$\\| \nabla\_{\mathrm{vec}[W]} f(T,\theta)\_{:,i} \\|\_{2} \leq \sum\_{j=1}^{N+1} P\_{ij} \Big\\| \Big(\mathrm{diag}\Big([z_l^T  B\_{kj}( z\_j)]\_{1\leq k\leq D}\Big)[\frac{\partial\mathrm{softmax}(A\_i)}{\partial A\_i}]\_l\Big)_{1\leq l \leq N+1} VT\_j \Big\\|\_{2}
\leq \sum\_{j=1}^{N+1} P\_{ij} \sum\_{l=1}^{N+1}\\|T\_l^T T\_j \\|\_{2}\cdot  (\frac{\partial\mathrm{softmax}(A\_i)}{\partial A\_i})\_l \cdot \\|V\\|\_{2} \cdot \\| T\_j \\|\_{2}
\leq \\|T\\|\_{2-\mathrm{col}}^3 \\|\theta\\|\_{2}.$

Combining the two equations above gives Assumption2 (ii) with $\phi\_P( \\|T\\|\_{2-\mathrm{col}} ) = \\|T\\|\_{2-\mathrm{col}} + \\|T\\|\_{2-\mathrm{col}}^3$.

For Assumption2 (iii), we have

$\\| J \\|\_{2}\leq \sqrt{\sum\_{1\leq i,j\leq N+1}\\|J_{ij}\\|^2\_2}\leq (N +1) \max\_{1\leq i,j\leq N+1} \\|J_{ij}\\|\_2.$

For any $1\leq i,j\leq N + 1,$ we have

$\\|J_{ij}\\|\_2\leq P_{ij} + \\|T\\|\_2 \\|Q\_i\\|\_2 \\|E\_{ji}T^TW+T^TW^T\delta\_{ij}\\|\_2
\leq 1 + 2 \\|T\\|\_2^2 \\|W\\|\_2
\leq 1 + 2 \\|T\\|\_F^2 \\|\theta\\|\_2.
$

Hence, we have

$\\| J \\|\_{2}\leq 2N \\|T\\|\_F^2\cdot (1 + \\|\theta\\|\_2).$

The above equation demonstrates that for $f$, Assumption2 (iii) holds with $\phi_T ( N, D, \\| T \\|\_F ) = 2N \\|T\\|\_F^2$.

We have verified Assumption 2 for the attention layer encoder $f$. The verification for $h$ is similar and easier. We omit the derivation for $h$ here to shorten our response, but we will make sure to include the full verification in the revised version of the paper.
Thank you very much for acknowledging that the asymptotic assumptions on $L$ and $M$ are common in the mean-field analysis literature. We also appreciate your suggestion on removing $C\_{\lambda}.$ We will follow your suggestion, choose $\lambda \le (2C\_2)^{-1}\epsilon,$ and adjust the proofs accordingly.

Regarding your concern about the concrete application to the Transformer architecture, we would like to clarify that verifying Assumptions 2-4 (please note that Assumption 1 is irrelevant to the Transformer model, and is only a fairly mild assumption on the data) for concrete examples of Transformer architectures with smooth activations is fairly intuitive and the proof is mainly based on a series of tedious calculations. Below, we give a concrete proposition and its brief proof as follows.


---

Consider  

$f(Z,\theta)=VZ \mathrm{softmax} (Z^T W Z ) $ $\qquad $ (Eq 1)

with the collection of parameters $\theta = \mathrm{vec}[V,W],$ where $\mathrm{softmax}$ denotes the column-wise softmax function. Moreover, consider

$h(Z,w)=W\_2\mathrm{HuberizedReLU}(W\_1H)$  $\qquad $ (Eq 2)

with the collection of parameters $w= \mathrm{vec}[W\_1,W\_2].$ Here, $\mathrm{HuberizedReLU}$ denotes the entry-wise HuberizedReLU activation function defined as 

$\mathrm{HuberizedReLU}(x) = \left\\{
\begin{aligned}
&0, &&\mathrm{if} \ z\leq 0;\\\\
&z^2/2, &&\mathrm{if} \ z\in [0,1];\\\\
&z - 1 / 2, &&\mathrm{if} \ z\geq 1.
\end{aligned}
\right. $

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.** Consider the Transformer model defined by equations (2.1), (2.2) and (2.4), with $f(Z,\theta)$ and $h(Z,w)$ defined in (Eq 1) and (Eq 2) respectively. Then Assumptions 2-4 all hold.

In our following comments, we will present a proof sketch of the proposition above. Due to the large character count of equations, we split the proof into several separate replies. We apologize for the long comments. To simplify our response, 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. We feel that the fact that $h(Z,w)$ satisfies Assumptions 2 and 3 is relatively more intuitive, especially given the proofs for $f(Z,\theta)$. We will also omit some calculation details to keep our response reasonably simple. We will make sure that all the omitted details will be added to the revised version of our paper.
I thank the authors for their answers. 

About Q2/A2: I do not think that an experiment on the MNIST dataset and just claiming global convergence would be a valuable addition to illustrate the claims in the paper. Studying the size of parameters to obtain this global convergence would be more interesting.

In any case, having read the answers to Q3 and Q4, I reckon this work appears as an improvement over the current literature and could motivate some further progress. This motivates me to increase my score.
The reviewer appreciates the authors' further response and acknowledges the effort made in addressing the concerns raised. Here are the reviewer's further comments:

1. The reviewer agrees that the term "major revision" used in the previous feedback regarding Corollary 4.1 was not accurate. For the sake of correctness, removing the lines related to the dependencies of $L$ and $M$ should be sufficient. When the reviewer suggested a revision, it was more about improving the clarity of the statement. For instance, as also noted by reviewer mtEE, the authors have treated many quantities as constants even though they still have dependencies on the problem parameters. The reviewer believes that this may have contributed to the inability to identify the incorrect conclusion in the initial paper. Another minor revision could involve removing the $C_{\lambda}$ constraint and instead simply choosing $\lambda \le (2C_2)^{-1}\epsilon$, and then adjusting the proofs on line 901 with the correct constants.

2. The reviewer also acknowledges that many current works in the mean-field regime exhibit asymptotic behavior. The reviewer recognizes that this work represents state-of-the-art results in this challenging area. However, the asymptotic nature of the results, particularly regarding the width and depth, still limits their practical applicability.

3. As noted in the previous response, the concern regarding the concrete application to the Transformer architecture remains the key issue that leads to a reject score. The reviewer appreciates the authors' willingness to include concrete examples of Transformer architectures and to address the case of ReLU activation. The reviewer agrees that these additions seem feasible, and if incorporated, the reviewer would be willing to raise the score and accept the paper. However, while these additions make sense intuitively and the authors seem confident in their ability to implement them, the reviewer remains skeptical that these theorems can be adequately developed and integrated into the camera-ready version, given the current form of the paper and the volume of work required. Nonetheless, including these components would make the work more convincing regarding the convergence story for Transformer networks.
We appreciate your careful feedback. We understand the concerns you have raised regarding Corollary 4.1, the asymptotic nature of our results, and the applicability to Transformer models. However, we would like to argue that these weaknesses you have mentioned are essentially the weaknesses of almost all existing mean-field studies of deep neural networks. In fact, our paper contributes to pushing the common assumptions in the mean-field literature towards more practical scenarios. Therefore, we respectfully disagree with your comment that our work is incomplete.

Corollary 4.1 should not need major revision, as its mathematical formulation and technical proof are accurate. To ensure rigor, it suffices to delete the sentence in line 302 "Here, $L_0$ scales as $\Omega(\epsilon^{-1})$, and $K_0$ as $\Omega((1+\delta)\epsilon^{-2})$", and remove the paragraph around lines 304-306.

Regarding your concern that the dependence on training time $\tau$, the depth $L$ and the width $M$ is asymptotic, we would like to emphasize again that our result remains, to our best knowledge, the state-of-the-art global convergence guarantees for large-scale Transformers. Moreover, we would also like to point out that our work has two main results: (i) connections between the discrete Transformers and their continuous limit, and (ii) global convergence guarantees of Transformers. We feel that it is unfair to deny all of our contributions based solely on that our global convergence guarantees for discrete Transformers require asymptotically large $\tau$, $L$ and $M$. 

Regarding your concern about the applicability to Transformers, we honestly believe that the Transformer architecture we consider is already among the ones considered in theoretical works that are closest to practical Transformer architectures. 

We will add concrete examples of Transformer architectures and rigorously prove that Assumptions 1-4 hold for these architectures in the camera-ready version. As clarified, these assumptions easily hold for activation functions that are sufficiently smooth. We will also clarify the case of the ReLU activation function, noting that it may not generally satisfy these assumptions, but it is possible for Assumptions 1-4 to hold when the data follow certain favorable distributions. 

Once again, we appreciate your detailed and comprehensive feedback. We sincerely hope that you can consider re-evaluating our contributions. Thank you for your time and attention.

Best regards,

Authors
Given the above points, the reviewer considers the paper to be an interesting attempt at tackling important questions. The paper introduces valuable techniques for addressing problems related to Transformer networks and has the potential to become a strong contribution. However, in its current form, the paper is incomplete and would benefit greatly from revision and resubmission to a future venue. The reviewer suggests the following improvements:
1. Please rewrite Corollary 4.1 and other theorems, ensuring that the dependencies for all constants are clearly specified. 
2. Prove that assumptions 1-4 are satisfied for Transformer networks in the form of (2.3).  It would be helpful if some parameters in the assumptions, including $K, K_T, K_P, \Phi_T, \Phi_P, \Phi_{PP}, \Phi_{TP}, \Phi_{TT}$ could be represented in more explicit forms with clear bounds.
3. If the authors intend to include ReLU activation beyond smooth activations, they should demonstrate that ReLU activation also satisfies Assumptions 1-4 and Proposition 3.1.
4. For the assumptions used in Theorem 4.1, even though it may be difficult to prove when these assumptions hold, it would be beneficial to provide heuristic examples to help readers understand when and why the assumptions are valid. For instance, the example provided by Reviewer accV in question 1 (Q3 in the rebuttal) offers a good illustration that this bounded assumption may not generally hold. Therefore, the authors need to make a stronger effort to justify the assumptions. In the current response to that question, phrases like "remains possible" are not convincing enough.

Again, the reviewer understands that theoretical work in this area is challenging and that compromises are often necessary. For example, the weak convergence assumption is difficult to avoid in the mean-field regime, and previous work on ResNet models [Ding et al. 2021] has also shared the same asymptotic nature as in this paper. This is why the reviewer initially leaned toward accepting the paper, hoping that some of these concerns could be addressed. However, as outlined above, the reviewer does consider the work to be incomplete and believes that the paper does not meet the standard for a conference paper due to its lack of direct application to Transformer models, its weak (asymptotic) results, and its reliance on assumptions that are difficult to verify. The reviewer apologizes for the change in rating, noting that this adjustment reflects a reassessment of the initial rating rather than a direct response to the rebuttal.

[Ding, Zhiyan, et al. "Overparameterization of deep ResNet: zero loss and mean-field analysis." Journal of Machine Learning Research 23.48 (2022): 1-65.]
The reviewer first thanks the authors for their detailed and thoughtful response. The reviewer also appreciates the authors' honesty and transparency in the rebuttal. Unfortunately, the reviewer does not feel that the concerns have been adequately addressed. As a result, the reviewer has decided to change the rating. The reviewer understands how frustrating this can be, especially given the thorough rebuttal provided by the authors in an effort to address all the questions raised. The reviewer struggled with the rating during the initial review phase and provided the rating with the hope that some of these concerns would be resolved. However, after careful and extensive consideration, the reviewer believes that the paper remains incomplete in its current form.

Before delving into specific concerns, the reviewer would like to acknowledge the strengths of the paper, as noted in the initial review. The paper is well-organized and clearly written, and it tackles a challenging and important problem: the theoretical understanding of training Transformer models. The reviewer recognizes that theoretical work, particularly in complex areas like Transformer networks, is inherently difficult. The reviewer is open to compromises, such as accepting reasonable assumptions that may be difficult to verify or acknowledging results with suboptimal dependence on problem parameters, including asymptotic results.

However, based on the authors' response, the reviewer believes the following deficiencies are significant:
1. Corollary 4.1 is incorrect and requires major revision.
2. The dependence on training time $\tau$, the depth $L$ and the width $M$ is purely asymptotic. 
3. The paper lacks clear results for the Transformer network (e.g., in the form of (2.3)) or for functions with ReLU activation.

First of all, having a misleading or incorrect conclusion for one of the main claims of the paper is a serious issue, as it undermines the credibility of the rest of the work. This problem also leads to the asymptotic nature of the results, which significantly diminishes the importance of the paper’s contributions from the reviewer’s perspective. To clarify for those who may be interested, the reviewer believes that the logic behind Theorem 4.1 and Corollary 4.1 is as follows: A specific training time is required to ensure that the limiting dynamics reach a small error $Q(\rho^{(\tau_0)})$, and $L$ and $M$ must be chosen large enough so that the discrete training dynamics remain close to the limiting dynamics within the time $\tau_0$. The lack of control over $\tau_0$ to achieve a small loss leads to the asymptotic nature of the time (with additional weak convergence assumption), depth, and width.

These first two deficiencies alone prevent the reviewer from giving the paper a high score. However, what is most concerning upon reevaluating the paper is the third point. The paper claims to provide convergence results for large-scale Transformers, yet no specific Transformer model is present in any of the main theorems. A Transformer should be defined in the form of (2.4) or at least in the form of (2.3), rather than as a concatenation of general functions $f$ and $h$. In the rebuttal, the authors used many "believe"s including the "belief" that (2.3) would satisfy Assumptions 1-4 and the "belief" that the results apply to the ReLU function with its Fréchet derivative. While the reviewer agrees with these beliefs on an intuitive level, the absence of a theorem that includes at least the form in (2.3) or the ReLU function makes these "beliefs" shaky and the work incomplete. It is somewhat surprising that the authors claim to achieve results for Transformers, yet provide no specific results related to the attention mechanism (2.4) (or its reparameterized form (2.3)) and the ReLU function. This presentation is therefore also misleading, making it difficult to connect the main results in the paper with a Transformer network.  Given the rebuttal, it seems that Assumptions 1-4 for (2.3) and ReLU activation for Proposition 3.1 are checkable, so including these components would form a more complete story for the convergence of Transformers in the asymptotic regime.
First of all, I would like to thank the authors for their thorough rebuttal.

I am satisfied with the answers and I would like to keep my score, conditionally on the authors addressing Q1 in the revised version.
Thank you for the detailed reply. I will maintain my score as I still feel the technical issue of Q1 ends up weakening the relevance of the paper, however I personally do think the approach and ideas are very interesting.