Global Convergence in Training Large-Scale Transformers
Authors: Cheng Gao
Yuan Cao
Zihao Li
Yihan He
Mengdi Wang
Han Liu
Jason Matthew Klusowski
Jianqing Fan
Venue: NeurIPS 2024 poster (2024)
Keywords: Transformer
gradient flow
mean-field analysis
global convergence
Primary area: optimization_for_deep_networks
TL;DR: We prove that large-scale Transformers closely approximate their mean-field limit and that gradient flow converges to the global optimization minimum.

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

Global Convergence in Training Large-Scale Transformers
Cheng Gao
Yuan Cao²* Zihao Li
Yihan He²
Mengdi Wang!
Han Liu
Jason M. Klusowski"
Jianqing Fan
'Princeton University
The University of Hong Kong Northwestern University
{chenggao, zihaoli, yahan yihan, he mengdiv,jason. htt..ajjasso....ossskka.... (gfaan)pprinnettt edu yuancao@hku. hk hanliu@horthhwettern. edu
Abstract
Despite the widespread success of Transformers across various domains, their optimization guarantees in large-scale model settings not well-understood. This paper rigorously analyzes the convergence properties of gradient flow 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 the Wasserstein gradient flow, which is represented by a partial differential equation. Then, we demonstrate that the gradient flow reaches global minimum consistent with the PDE solution when the weight decay regularization parameter is sufficiently small, Our analysis is based series of novel mean-field techniques that adapt Transformers. Compared with existing tools for deep networks [47] that demand homogeneity and global Lipschitz smoothness, we utilize refined analysis assuming only partial homogeneity and local Lipschitz smoothness. These new techniques are of independent interest.
1 Introduction
Transformers have revolutionized the field of deep learning since their introduction in 666 These models are distinguished by their immense scales, often comprising billions of parameters to achieve state-of-the-arr performance. Notably, this massive parameterization enables them to excel in variety of domains, notably in natural language processing [2] 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 overpa- rameterized neural networks [51] 16 28 14
35 25 76.. In particular, several works 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 mean-field analysis. However, these results are established based on certain homogeneity and/or global Lipschitz smoothness properties of the
47 22 23
Equal Contribution.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

neural network, which are not applicable to Transformer models. Therefore, it remains an open question how gradient-based methods can effectively train large-scale Transformers.
1.1 Our contribution
In this work, we bridge the gap between Transformer theory and practice by demonstrating the global convergence of Transformer training optimization via gradient flow in large-scale model regime. We analyze the mean-field limit of the Transformer model, which is characterized by the distribution model parameters, shifting the focus from parameter space to distributional dynamics in the Wasserstein metric [16]. This approach yiclds two key theorems:
We show the closeness between practical discrete Transformers trained by gradient flow and continuous Transformers whose parameter distribution follows partial differential equation the Wasserstein gradient flow (Theorem (3.1). Our result demonstrates that large-scale discrete Transformers can be approximated by its mean-field limit and the approximation error can be expressed in terms of the width and depth of the Transformer models.
ii This approximation facilitates our analysis of the global convergence (Theorem [4.1 of discrete Transformer models. By leveraging the universal approximation capabilities of either the self-attention or feed-forward layers, we demonstrate that a basic gradient flow method can reliably find global optimum, despite the highly non-convex landscape of the training objective.
We also highlight our novel contributions Transformer theory through the development of these two core results:
The assumption on activation regularity conditions (Assumption) 2 is less stringent compared to those usually found in studies of two-layer neural networks [6 28 14 deep ResNet networks [22 23 47 In particular, many existing approximation guarantees reply on a Lipschitz continuity property of the network gradients, which limits the mean-filed study to neural networks with smooth activation functions. In comparison, our analysis relaxes this assumption and only requires local Lipschitz continuity of the gradient in expectation. This relaxation broadens the applicability of our approach and ensures that our result can cover more practical Transformer architectures.
ii Our model differs from the ResNet models in only a single identical encoder within each evolutionary block. Unlike the typical theoretical configurations, our model employs two distinct encoders ʃ 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 ʃ and h, providing rigorous validation of concepts proposed in (77 and 667 from a new perspective.
23 131, as those models incorporate
iii. Our global convergence guarantee for training Transformer models is also broadly applica- ble: our assumption (Assumption 4) ensures global convergence by relying the universal approximation capabilities of either the self-attention or the feed-forward encoder. Addi- tionally, we incorporate more flexible framework by adopting partial I-homogeneity for only subset of the parameters, in contrast to the full parameter homogeneity required in studies such as (7). 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 for a detailed discussion.
Notations. For any E Rd, dim(a) refers to its dimension d. For any B E Rdxd its trace is denoted by Tr(B). For any positive integer n, Let [m] = (1,2...., n]. Let 02 denote the d-dimension vector of all zeros. Let W (H, v) denote the Wasserstein-p distance between two probability measures U, E P(Rdd) for p 2 1. For matrix A (a1, .2. ... an), define its vectorization version as vec[A] : (a .T. TT Let 5(-) denote the Dirac mass and 1 be the indicator function. Let supp(-) denote the support of any distribution. Let ||-I| = I1-1/2 denote the /2 norm and |lll denote the maximum norm. For any subsets D1, D2 in Euclidean space, define C(D,, D2) as the collection of functions that map D) to D2 and are continuous over Dj. Define the Bounded Lipschitz norm for any measure μ E ((R || BL ;= s fdp f Rd + R. suplfis 1, f is 1 Lipschitz).

---

2
Transformer model
In this section, we describe our deep Transformer model with each data input as a sequence, and the gradient flow algorithm used for training.
2.1 Data setting
In our paper, the data input is both general and straightforward: an input sequence H E Dx(+1) consisting of N tokens, each with dimension D. We consider the setting where each input sequence H is associated with a label (H)) E IR, where y ()) is the target function we aim to learn. Furthermore, we assume that each instance H i.i.d. drawn from population distribution fa.
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] T0, 75]. where models are capable of making accurate predictions examples from the same pool. For clarity, consider the input sequence H IR Dx(++1) formatted as follows:
data when prompted with training
T2
EN
EN++ 0
== fl..,h2...... hhg
hN+1l=
y1
1/2
YN
iid.
P.
9N+1 = y(H).
P1
P2
PN
PN+1
Here, (EN)) are the input vectors, each associated with corresponding label {yr WEEN)- The last token, 2N+1 the test input for which a prediction made. The third row contains the customized and fixed positional encoding vectors {pe heEN, which typically include ones, zeros, and indicators denoting the token for prediction. The label for the query point NN+1 is then given by UN+1 (H) our terminology. operates in a zero-shot fashion, without any updates to the model's s parameters, highlighting unique and powerful capability of these systems to adapt and generalize based on the provided context alone. In [6), the authors demonstrate that fixed Transformers can approximate in-context penalized generalized linear regression to any desired degree.
2.2 Model
where each Transformer block consists of two distinct layers: self-attention mechanism layer and 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 the width across all blocks for simplicity. The formulation for matrix input Z DN++1) and given residual step size 1j 0 is follows: Each residual self-attention layer is represented by
We follow a common configuration of Transformer architectures 6
38 40 48
73
M
Ating, .2........................ BM (2,1) = 7 M
1(2,0,).
(2.1)
J=1
and each residual feed-forward neural network layer defined by
Ad
MLP an WF MA (Z, 2 + n
hZ, w3)
(2.2)
for parameter vectors 0 and w in the Fuclidean space. The encoders for the self-attention and feed-forward layers are denoted f Dx(++11 + [g Dx(++1 and h Dx + R Dx1+11) respectively. The self-attention encoder formulation, commonly adopting multiplicative or can be exemplified by
dot-produe approach as detailed in
66 73
f(z,0) WoWy (WxZ)
where Wv, WK. Wq E D and Wo E Dxs This formulation be reparametrized
1(2,0) = Vo z W2)
(2.3)
3

---

where V.W& W' E RDDD, f = vec[V. The activation dλ typically uses column-wiss softmax, but component wise ReLU is also viable, as in (6. For the ferd-forward layer, an example of the encoder is h( Z, w) = W2 2), as detailed in 6 38 73 where w vec[W; W2] and the activation ØM is component- wise ReI.U. Alternatively, setting h = O results in a Transformer block that comprises only the self-attention layer, referred to as "attention-only" Transformers, as discussed in [6 46 (6).
Next, we analyze Transformer network composed of L Transformer blocks, referring to L as the depth of the model. In this paper, we introduce an additional term, 7/ in (2.1) and (.2 to simulate the model's evolution in a residual manner. We set the step size 17 as At/2, where At = 1/ L. As increases, At approaches zero, allowing Transformer blocks incrementally contribute to the model's 's overall progression. The structure of the network is then defined as follows:
Te, L A1/2)
- BI M (T (H, A1/2)
(2.4)
(T(H, At) = MLPPum WA, 10, MM (To H A1/2),t/22)
for each = 0, At..... (L-1)44 with (H, = H. We abbreviate the subscript = 0, At...., (L 1)At by and = 1,2..... M by for simplicity. Here, θ = {0r.j. Wt, Ht denotes all parameters in the Transformer model.
Throughout this paper, we treat D and N as bounded finite values, while M and L, are treated as diverging, aligning with the setting of large-scale Transformers.
2.3
Gradient flow
For the la regularization with À > O, we consider training the constructed Transformer model using the following X-regularized risk objective:
M (110,,112+ |(w))1(),
( = (())) + 2MT. Σ>
(2.5)
5=1
with the population squared risk function defined as
R = -E 12 [Rcad[f((H.1)) y v(H)) )
1
we will show that Iz-rgullarization on the parameter norms is essential for the well- posedness of the (Wasserstein) gradient flow to control parameter growth under our mild assumptions, even with a very small À >> O. Similar strategies that consider necessary l2 regularization are employed in [2 and [70]. Then, drawing on the methodologies in [6] 32, 46. our model processes the final output through 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 thcoretical results.
In Section 3.2,
To minimize the objective function 2.5 we implement the standard gradient flow method as follows:
Step 1. Initially, for each t 0, At, ................................ (L- 1)At, sample M particles el) wh with E [M] independently from (0, (l), where pn is pre-defined distribution with bounded support.
(0)
Step 2. Then, we update all parameters 0/7) w/. in the set (r)) = {013. w/ J&s using gradient flow (scaled by ML), which is defined as follows:
(7)
aop3) = - 0, 1Q(0 dr
dw (T)
-M-V w.., ( (Q
(2.6)
dr
Define the function R(H10)) 3 (Read] (HH, (H))) and the partial derivative (H,
aR H; o0/oTo(THH... (H,t) for each t - 0, At/2 At, At...., (I - 1) (L -L- 1/2) Refer to Appendix
C.4 for the explicit formula of Per (H, t). Using the chain rule, we derive the explicit form of the
gradient flow as follows:
(0(T) 1.
dw (T)
- ((()))))(+))
(2.7)
dr
dr

---

where
1 G,(,,0,t) = - [o (f (1) Po +t/2)) + 20,
Gw 0,1) - -E VoTT( w/)), (H 1 40) + Aw
fort = 0, At... - 1)At.
3 Approximation by the mean-field limit
In this section, we present rigorous approximation result that bridges Transformer models (..4) with their mean-field limit as continuous Transformers. Thus, the width M and depth L in proposed model are treated as discretization of this continuous limit in the parameter space.
3.1 Assumptions
In addition, we introduce the norm I1-1l--on the maximum l2 norm across all columns of a matrix. We proceed under several mild assumptions related to the data distribution and the encoders f and Assumption 1 (Data regularity). There exists some universal constant > such that, for any H E supp(p), we have H y(H) < B. In addition, universal constant Ky > 0 ensures that y( (H) Ky-liipchhtz continuous for |I-I|r over H E supp(μs).
Remark Assumption is irrelevant to the Transformer model, and is only fairly mild assumption on data.
Assumption 2 (Transformer partiele growth bound). We assume that the gradient of f(T, 0) and h(T, w) exists. Furthermore, we have
(T(T,))ll2 S 22 ||0|| + [101
i. For every E [N + 1). we have o (T. S or ( + [10]).
ii. | vece ((T 0)]l < (T N, D, [1 + |1||| 1013))
for some continuous, monotonically increasing functions Φp, Φr for every coordinate, and universal constant K > O. Similarly, if we replace f with h and with w, the same conditions apply.
Remark
There are three key observations for Assumption 2 Firstly, it incorporates the I1-1/2-col norm, which is particularly useful for handling sequential inputs where each column represents token. Secondly, as we consider higher-order multiplications between data and parameters, this assumption accommodates a broader range of self-attention encoders, such as the one in (2.3) with softmax or ReLU activation (where the derivative defined as ReLUU (xx) = 1(x > 0}). Lastly, particularly interesting and frontier question is identifying the function dry, and we have listed related literature in Appendix B
Assumption 3 (Locally Lipschitz continuous gradient expectation). Besides Assumption 2 for functtion Ti = Ti (H) and Tr = T2(H), for every
any Ter O and any L.p-Lipschizz continuous functions i [N 1}, we have
E. Vef 0).,
Kr. LT) supl|T] 2-ol2- Vef(Tz, 0),alz S (PT(1)))
n. E, V vec vec - [eee hh
ΦTr (N, D D,supl [r, Kr, LT) 0'll
H1
8i. E VV. 0).,- Vef 0 < opp
sup Ti clo- LT) 0'1l.
#
iv. vec/ Vec - vf 0)]12 < GTT (N, D, Kr. ||0|| LT) supl Tallr
#
for KT max (sup # [| 2-oo,, supH (2l)2-col), ) Kp = max [100], [[0'1), and some continuous functions OPT, @TP, prp, ØTT that are monotonically increasing for every coordinate. Similarly, if replace ff with h and 0 with w, the same conditions apply.
5

---

Remark Assumption 3 states that functions are locally Lipschitz continuous in expectation, suitable for encoders that utilize ReLU functions and have second-ordes derivatives almost everywhere. This assumption is naturally satisfied if the activation has a locally Lipschitz continuous gradient.
Define 2 the set of probability measures endowed with the Wasserstein-2 distance, where the Lipschitz continuity with respect to the depth holds, i.c. there exists some universal constant Cp >> 0 1' for any t, E [0, 1].
such that llp(-;)) - p(, BBL. S C,
Choice of po Suppose Pa E P2 satisfies that for any E [o, 1], the support of .,( is contained within the set {(O, w) |101 ||||2 < I22} for universal constant R. Additionally, for each t [o, 1], it holds that So.. Pu (0, w. t)d(0, w) 1. This condition suits common bounded support distributions, and a natural choice uniform distribution across a disk with radius for each t E [0, 1].
We would like to clarify that verifying Assumptions and for concrete examples of Transformer architectures with smooth activation functions is fairly intuitive, and the proof is mainly based on series of tedious calculations. We give concrete proposition with its brief proof in Appendix G
3.2 Continuous Transformer and Wasserstein gradient flow
Drawing inspiration from 47 and 677. which suggest that deep residual networks behave like ensembles of residual networks locally, we apply similar manipulation to formulate the continuous version of (2.4). Consider the following continuous version (H, t) E R Dx(N+11) governed by the following continuous ODE that averages the two encoders:
| F(Tp(H,t),O (H, w) p(0, w. t)d((, w), THH H (3.1)
0,m
2
In 33.1 each encoder f or h is conceptualized as particle, and we consider the distribution of these particles denoted as p(A, w, t). For any E p2 that have bounded support, the well well-posedsess of (H, 1) that satisfies the Transformer ODE (3.1) shown in Proposition C.1 Transitioning to the framework with continuous Transformers, our objective shifts to minimizing the la risk function with regularization on the second moment p as follows:
A 5. N,w 2
/
Q(p) R(p)
(110112 + 11w112)((0 w. t)d(O, w)dtt
(3.2)
with
R(p) = F Read [T, -H)
(3.3)
Define (H, Dx the partial derivative of R(p) relative to (H, t) at a local query point H, as the solution derived in Appendix C.4 using the classical adjoint sensitivity method [58]
(H, [)] = (H) exp (f.
DN+d+11
Using this, we can compute the functional derivative to p as follows:
8Q 8p (0, w, EE (Tr [TT [TT
M(T, ( 0) (H, w)
[Ir
Po(H (110 ( (3.4)
2
The following Proposition claims that 4Q is indeed the derivative with respect to p (specifically, the 8p
Fréchet derivative (30) for the functional Q(p).
Proposition 3.1 (Functional derivative to p). Under Assumptions Mand? 2 for any pair v,v E P2
that have bounded supports, have
5Q Q(p + m(p 0)) = Q(p)) + +1
0) o(()),
5p
5Q where 80
is defined in
3.4 and
8Q
p) = So J(B
5Q
(v - p)d(0, w)dt IR
Spp
dp

---

Now, we are in a position to display the gradient flow of in the Wasserstein metric [16], given by Specifically, we study the following partial differential
McKean- Vlasov type equation 37 54
equation of the distribution p(T) w,t):
8Q (0, (0,w,1) div(o.w) (prr) ((..u) (0,w) 8p pmpitt)
dr
(3.5)
= dive (dive(p
1) div ( ((),,,
where p!0) Po. div is the divergence operator, and the gradient functions are defined as
1 Ey Gy(0,p,t) 2 p. 2 [ (MT,, w) (, + Au.
VOTr [T
(FT,, (.) + A0,
Propositions D and 3.2 provide the well-posedness of both gradient flow and Wasserstein gradient flow respectively. In both propositions, a À > 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 Similar adjustments to regularize the risk function also noted in (33]
Proposition 3.2 (Existence and uniqueness of Wasserstein gradient flow). Under Assumptions 7 and 2 there exists unique solution (pl)))r20 E pp x with p!0) = Po for (3.5) Additionally, for any have
10112 + /|l||2 S R2} x (0, where R2 = (R 1) exp(CRT)= 1 for some constant CR that only depends on N, D, X and the parameters of the assumptions.
plr) has a bounded support {0, w
i. (100|++||w|\)) (t) w, t)d(0, w)dt S Ab. where Ao = 1 exp (1+
R + R2))2).
ii.. [w, (0, w. t)d(G, w) 1 for any [0,
3.3 Approximation of large-scale Transformer
In this section, we discuss the general results associated with approximating our discrete Transformer model to its mean-fiee field limit. First, we highlight that the minimization of the risk function with discretization, whether or not regularization included, closely approximates the minimal risk achievable by continuous models.
Proposition 3.3 (Global minimum approximation of discretization.. Under Assumptions and we define par the set of distributions in p2 concentrated on {(0, w) -|01[ + || < <r 1]. for any r > 0. Then there exists a constant C dependent on N, D,r and the parameters of the assumptions such that
inf R(p) + +L + log( L (g((++)) 1
inf (())
pEpL.
inf 2()) inf Q(p) + C(1 + (1 V
log L + 1) M
pEP2.
Proposition 3.3 specifies that the distributions under consideration must have bounded support. While 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 (355). The proof with detailed explanation of the techniques used in Theorem is provided in Appendix D
Theorem 3.1 (Gradient flow approximation of discretization). Define the empirical distribu-
tion (()) ML E, 500() w .t) for any 2 0. Under Assumptions I
3
7

---

have that (+(0 weakly converges to (r'0))>2 almost surely along any sequence such that + oc, M /log - sc. Moreover, for any fixed > 0 and any ó > with probability at least 3exppp with respect to the parameter initialization ( we have
su Read | (H,t) Read pto) (H,1)l C(T.11+ + A+ I++1
it. 1 R(pb)))( c(1-1
/+too(L+1)) M
iii supseta-)(((00") - Q(p())) -(())))))))(())) C(L 6+log(L+1)
M
for some constant C that depends on on N, D, and the parameters of the assumptions.
Theorem 3.1 significantly advances our understanding by controlling the difference regarding both the Transformer output, the risk function, and the regularized risk function. It's noted that the difference bound in the model's approximation may increase, possibly exponentially [22] 23 51. as the time horizon extends. As argued in [51]. such behavior may be inherent to the systems being modeled.
Additionally, the technical uniqueness and innovation of this theorem contrast sharply with previous results from overparametrized ResNet models. Our analysis distinguishes itself in two ways. First, our discrete Transformer model (2.4) uniquely splits the averaged encoder (f h)/2 2 into two distinct blocks with encoders f and h. Second, we demonstrate uniform error control over any finite time interval [0, r]. enabling continuous monitoring of maximum error across the gradient flow's 's trajectory. In contrast, models prior studies such as 22 restricts the error analysis to a specific E [0,
4 Global convergence of gradient flow
In this section, we explore the optimization problem for gradient flow in the context of the discrete Transformer model, focusing our general global convergence results.
4.1 An additional assumption
To ensure the global convergence of gradient flow for our discrete Transformer model, we introduce the following assumption. While influenced by the work in [16 uniquely tailored to the context of Transformers:
22
23 47 our assumption is
Assumption 4. There exists a pair (g. a)) {U. 0), (h, w)} with partition (a1, Or2) such that
(Partial L-homogeneity) for any e Px(++1) and c E ER, we have g(T, ca1,a2) = (f,
ii (Universal kernel) a compact K. (ddim(as) ensures that the span of {l( :
(Rdim(an) x {] dense C(I S B, B.R.D. DV(+1)) for any B
We emphasize that the universal kernel property, as discussed in [52] closely relates to the universal approximation abilities. Under our assumption, 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
The universal kernel property of the feed-forward layer encoder 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 The interconnection between approximation abilities and memorization capacities established in (8). Notably, [49] investigated the expressive capabilities of one single multi-head softmax self-attention layer, thereby potentially validating our assumptions.
across multiple studies (27 31 38 39
63 73..
Finally, we posit that the universal kernel applies to 02 within compact set, as the function's scale can be moderated by the homogeneous part aq. In scenarios where d2 and I. are absent, our

---

assumption simplifies to that in [4]] characterized by complete homogeneity. Conversely, in the absence of the aj component, our framework aligns with [23 which necessitates a more stringent support condition for K,, as detailed later in Theorem 4.1
4.2 Global convergence result
In this section, we establish the convergence properties of the optimization task for discrete Trans- formers through gradient flow dynamics.
Theorem 4.1 (Global convergence up to A). Suppose that Assumptions 7 hold, and the Wasserstein
gradient flow ) weakly converges to some Poo E Pp. for some universal constant R the following two conditions hold:
i. (2())) is concentrated on {0, w. |10112 + ||01|12 S 12 x [0, 1] when T is sufficiently large.
i. If Assumption 4 holds with (g. a) = (f. we assume there exists at E [0, 1] such that the connected set supp(psc (., )) xxxx fwo..for some We E Rdim(w) and D c (ddi((1)
= 1/R and (0, : ||0 Roo) that separates {0, |10
ii' If Assumption holds with (g, = (h, w). assume there exists l E [0, such that the connected set supp(p (-,t") x x(o0 K. x D. for some 0g E Rdin(() and 'D C (ddimwe)) that separates (wi ||l1 || = Roo J and {wi ||will R.o
Then, for any e > O, there exists some TO > O such that
0 log 111 11)) +V M
sup T2.00
se+ set
+ C2A
with probability least 1- 1-3 exxp(-)) with respect the parameter initialization (0) for any Here, C₁ is some constant dependent only on N,D, V,D,oo, X and the parameters of the assumptions, while C2 depends only on N, Ros and the parameters of the assumptions.
Theorem 4.1 depicts the behavior of the risk function ((((()) as the training duration 7 is sufficiently large. Specifically, ?(((T)) asymptotically approaches zero both L and M/ L + oc, with an additional term that scales with X. 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 À > the risk approximates zero, demonstrating global convergence the minimum of R.
In addition, Theorem 4.1 posits some additional assumptions: the weak convergence of p!T), the long- time uniform boundedness, and the separation property for an with the support expansion of o2 K. Similar assumptions are made the literature of deep model optimization theory 47]. While these types of assumptions are typically challenging to justify, provide high-level justifications for them in Appendix C.5) deferring detailed verification to future research.
We then present corollary that directly follows from Theorem 4.1
Corollary 4.1. Continuing with the notations and assumptions from Theorem 4.1 suppose À S Cxe for some universal constant Cx > O. Then, for any 8 >> O, constants To. Lo, Ko can be found such that:
sup T2%o. Mooooo
(( < C
C2CA)E
(4.1)
with probability at least - exx((-8) 8) with respect the parameter initialization (0), Notably, if Cx E (2C2)-1. the upper bound in (4.1) is less than or equal to e.
Corollary 4.1 claims that with a fixed 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 guarantees for Transformers.
O, for any 6 > O, we can achieve an order of e-close
kind and lays the groundwork for future theoretical optimization

---

5 Proof ideas of main theorems
Given the technical nature of this paper, this section presents the key ideas behind the proof of our main novel results, along with outline of the proof preparation.
This convergence 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 s inequality, which allows us to bound output differences over time.
Idea for Theorem 3.1
Second, weak convergence of the empirical distribution process relies on optimal transport theory and stability results for Wasserstein gradient flows 44, 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 detailed illustration, including a description of each main step.
6p ps + ensuring remains constant if the derivative with respect to 3 is constant over region. Next, we derive the key bound for (o))), which is proportional to X. by analyzing the functional energy Q's landscape through its derivatives.
Idea for Theorem 4.1
We first establish the continuity of the functional gradient iQ
Finally, we show that the finite-time risk can approach this bound. Achieving c-level loss requires Q(p(t)) I for some large To- Applying Theorem 3.1 we show (TT) becomes sufficiently small, and since Q(p(r)) is non-inereasing, remains small for 2 To- See Appendix E for 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.I 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 ? 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.id. settings, these lemmas first establish 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.I D.3 which are essential for proving the approximation bound in Theorem 3.1.
and 3 that fit
plays key role by demonstrating that for any p. there is always a nearly descent direction around p for Q(p), implying that all local minima are nearly global. This motivates
In Theorem 4.1 Lemma E
further landscape analysis for bounding op Pso in the main theorem.
nQ
6 Conclusion
We conclude by summarizing our key contributions and suggesting future research directions. This paper establishes the global convergence of large-scale Transformer models through gradient flow dynamics, providing a thorough theoretical foundation. Our analysis, focused on the mean-field limit with infinite width and depth, shifts optimization from parameter space to distributional probability measures. We present two main theorems: one confirming the close approximation between discrete and continuous gradient flows, and another demonstrating global convergence, highlighting that basic optimization methods can successfully navigate complex landscapes to find optimal solutions. The techniques and results from this study lay the groundwork for further exploration into Transformer optimization. Future work could explore direct gradient descent with specific focus on step sizes, and expand on the in-context learning approximation capabilities of Transformers, as initiated by [6]. Additionally, it's crucial to rigorously assess under what conditions can self-attention layers serve as universal kernels to enhance our theoretical understanding, and to determine the generalization ervor bounds of Transformers trained on finite samples. These directions promise deepen the theoretical and practical insights into Transformer models.
10

---

Acknowledgments and Disclosure of Funding
We thank the anonymous reviewers for their helpful comments. Yuan Cao is partially supported by NSFC 12301657 and Hong Kong RGC-ECS 27308624. Mengdi Wang acknowledges the support by NSF IIS-2107304, NSF UPS-2312093, ONR 1006977 and Genmab. Han Liu's 's research is partially supported by the NIH R01LM01372201 Jason M. Klusowski supported in part by the National Science Foundation through CAREER DMS-2239448, DMS-2054808 and HDR TRIPODS CCF-1934924, Jianqing Fan's research partially supported by NSF grants DMS-2210833, DMS-2053832, and ONR grant N00014-22-1-2340.
References
[1] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference Learning Representations, 2022 2022.
[2] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overpa- rameterized neural networks, going beyond two layers. In Advances Neural Information Processing Systems, 2019.
[3] Zeyuan Allen Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference Machine Learning, pages 242-252, 2019.
[4] Luigi Ambrosio. Nicola Gigli, and Giuseppe Savaré. In metric spaces and in the space of probability measures. In Gradient Flows, 2005.
[5] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine line-grained analysis of op- timization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332, 2019.
[6] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Workshop on Efficient Systems for Foundation Models @ ICML2023, 2023.
[7] Raphaël Barboni, Gabriel Peyré, and François-Xave Xavier Vialard. On global convergence of resnets: From finite to infinite width using linear parameterization. Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[8] James Bernhard. Alternatives to the scaled dot product for attention in the transformer neural network architecture, 2023.
[9] Blake Bordelon, Hamza Tahir Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head transformer dynamics, 2024.
[10] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Berbert-Vosss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020.
[11] Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems, 2019.
[12] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations, Advances in neural information processing systems, 31, 2018.

---

[13] Yihang Chen, Fanghui Liu, Yiping Lu, Grigorios Chrysos, and Volkan Cevher. Generalization of scaled deep resnets in the Learning Representations, 2024.
regime. In The Twelfth International Conference on
[14] Zixiang Chen, Yuan Cao, Quanquan Gu, and Tong Zhang. A generalized neural tangent kernel analysis for two-layer neural networks. arXiv: Learning, 2020.
[15] Jingpu Cheng, Qianxiao Li, Ting Lin, and Zuowei Shen. Interpolation, approximation and controllability of deep neural networks. arXiv preprint arXiv:2309.06015, 2023.
[16) Lénaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS' 18, page 3040-3050, Red Hook,
NY, USA, 2018. Curran Associates Inc.
[17] Lénaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. In Advances in Neural Information Processing Systems, 2019.
[18] George V. Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of Control, Signals and Systems, 2:303-314 14, 1989.
[19] George Dasoulas, Kevin Scaman, and Aladin Virmaux. Lipschitz normalization for self- attention layers with application to graph neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 2456 2466. PMLR, 18-24 Jul 2021.
[20] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xiao Wang. Carlos Riquelme Ruiz, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar. Sjoerd Van Steenkiste, Gamaleldin Fathy Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Collier, Alexey A. Gritsenko, Vighnesh Birodkar, Cristina Nader Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetic, Dustin Tran, Thomas Kipf, Mario Lucic, Xiaohua Zhai, Daniel Keysers, Jeremiah J. Harmsen, and Neil Houlsby. Scaling vision transformers to 22 billion parameters. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, edi- tors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 7480-7512. PMLR, 23 Jul 2023.
[21] Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In North American Chapter of the Association for Computational Linguistics, 2019.
[22] Zhiyan Ding, Shi Chen, Qin Li, and Stephen Wright. On the global convergence of gradient descent for multi-layer resnets in the mean-field regime, 2021.
[23] Zhiyan Ding, Shi Chen, Qin Li, and Stephen Wright. Overparameterization of deep resnet:
Zero loss and mean-field analysis. Journal of Machine Learning Research, 23:48-1 1, 2022.
[24] Alexey Dosovitskiy. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 19999, 2020.
[25] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pages 1675-1685, 2019.
[26] Weinan E, Jiequn Han, and Qianxiao Li. A mean-field optimal control formulation of deep learning. Research in the Mathematical Sciences, 6(1):1-41, 2019.
[27] Benjamin L. Edelman, Surbhi Goel, Sham M. Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms, 2022.
[28] Cong Fang, Hanze Dong, and Tong Zhang. Over parameterized two-level neural networks can learn near optimal feature representations. ArXiv, 15s/1900.1150S 2019.
12

---

[29] Cong Fang, Yihong Gu, Weizhong Zhang, and Tong Zhang. Convex formulation of overparam- eterized deep neural networks. IEEE Transactions on Information Theory, 68(8):5340-55 5352. 2022.
[30] Bela A. Frigyik, Santosh Srivastava, and Maya R. Gupta. Introduction to functional derivatives. UWEE Tech Report 2008-0000 University of Washington Department of Electrical Engineering. 2008.
[31] Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? a study through the random features lens. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[32] Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? case study on learning with representations. In The Twelfth International Conference on Learning Representations, 2024.
[33] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05229, 2023.
[34] Ken ichi Funahashi. On the approximate realization of continuous mappings by neural networks. Neural Networks, 2:183-192, 1989.
[35] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, pages 8571-8580, 2018.
[36] Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho. and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 37822-37836. Curran Associates, Inc., 2022.
[37] Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker- planck equation. SIAM Journal on Mathematical Analysis, 29, 04 2000.
[38] Tokio Kajitsuka and Issei Sato. Are transformers with one layer self- attention using low-rank weight matrices universal approximators?, 2024.
[39] Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention, 2021.
[40] Junghwan Kim, Michelle Kim, and Barzan Mozafari. Provable memorization capacity of transformers. In The Eleventh International Conference on Learning Representations, 2023.
[41] Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. In Forty-first International Conference on Machine Learning, 2024.
[42] Qianxiao Li, Long Chen, Cheng Tai, and Weinan E. Maximum principle based algorithms for deep learning. Journal of Machine Learning Research, 18(165):1-29, 2018
[43] Qianxiao Li and Shuji Hao. An optimal control approach to deep learning and applications to
neural networks. In International Conference on Machine Learning, pages 2985 2994. PMLR, 2018.
[44] Qianxiao Li, Ting Lin, and Zuowei Shen. Deep learning via dynamical systems: An ap- proximation perspective. Journal of the European Mathematical Society, 27151116,,,179,, 2022.
[45] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In International Conference on Machine Learning, pages 19689-19729. PMLR, 2023.
[46] Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. In The Twelfth International Conference on Learning Representations, 2024.
13

---

[47] Yiping Lu, Chao Ma. Yulong Lu, Jianfeng Lu, and Lexing Ying. A mean field analysis of deep ResNet and beyond: Towards provably optimization via overparameterization from depth. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 6426 6436. PMLR, 13-18 Jul 2020.
[48] Thang Luong, Hieu Pham, and Christopher D. Manning, Effective approaches to attention- based neural machine translation. In Lluís Märquez, Chris ceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pages 1412-1421, Lishon, Portugal, September 2015. Association for Computational Linguistics.
and Jian Su, editors, Pro-
[49] Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi- head attention in transformers. In The Twelfth International Conference on Learning Represen- tations, 2024.
[50] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-fin free bounds and kernel limit. In Conference on Learning Theory, 2019.
[51] Song Mei. Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7655- E7671, 2018.
[52] Charles A. Micchelli, Yuesheng Xu, and Haizhang Zhang. Universal kernels. Mach. Learn. Res., 7:2651-2667, 2006.
[53] Atsushi Nitanda and Taiji Suzuki. Stochastic particle gradient descent for infinite ensembles, 2017.
[54] Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Particle dual averaging: optimization of mean field neural network with global convergence rate analysis Journal of Statistical Mechanics: Theory and Experiment, 2022(1 1:1144010, nov 2022.
[55] OpenAL. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023.
[56] Felix Otto. The geometry of dissipative evolution equations: The porous medium equation. Communications in Partial Differential Equations, 16(1-2)111-177, 2001.
[57] Allan Pinkus. Approximation theory of the mlp model in neural networks. Acta Numerica, 8:143 195, 1999.
[58] L. S. Pontryagin, V. G. Boltyanskii, R. V. Gamkrelidze, and E. F. Mishechenko. The mathemati- cal theory of optimal processes. Zamm- zeim--eitschriff Fur Angewandte Mathematik Und Mechanik, 43:514-515, 1963.
[59] H. Risken. The Fokker-Planck Equation: Methods of Solution and Applications. Springer, 1996.
[60] Filippo Santambrogio. Optimal transport for applied mathematicianss Birkäuser, NY, 55(88- 63):94, 2015.
[61] Filippo Santambrogio. (Euclidean, metric, and Wasserstein) gradient flows: overview. Bulletin of Mathematical Sciences, 7:87-154, 2016.
[62] Alain-Sol Sznitman. Topics in propagation of chaos. In Ecole d'Eté de Probabilités de Saint-Flour XIX-1989, pages 165-251. Springer, 1991.
[63] Shokichi Takakura and Taiji Suzuki. Approximation and estimation ability of transformers for sequence-to-sequene functions with infinite dimensional input. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, edi- tors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 33416-33447. PMLR, 23-29 Jul 2023.
[64] Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao and Che Zheng Synthesizer: Rethinking self-attention for transformer models. In International Conference on Machine Learning, 2020.
14

---

[65] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal. Eric Hambro, Faisal Azhar, Aurelien Rodriguez. Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. ArXiv, awwpaaa2..//3330 13971, 2023.
[66] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[67] Andreas Veit, Michael Wilber, and Serge Belongie. Residual networks behave like ensembles of relatively shallow networks. Advances in neural information processing systems, 29, 2016.
[68] C. Villani. Optimal Transport, Old and New, volume 338 of Grundlehren der Mathematischen Wissenschaften. Ppringer-Verlage 2008.
[69] James Vuckovie, Aristide Baratin, and Rémi Tachet des Combes. mathematical theory of attention. ArXiv, abs/2007.022 02876, 2020.
[70] Colin Wei, J. Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization neural nets v.s. their induced kernel. In Neural Information Processing Systems, 2018.
[71] Weinan. A proposal machine learning via dynamical systems. Communications in Mathematics and Statistics, 1(5):1111, 2017.
[72] E Weinan, Chao Ma, and Lei Wu. Machine learning from continuous viewpoint, i. Science China Mathematics, 63:2233 2266, 2019.
[73] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Ku- mar. Are transformers universal approximators of sequence-to-sequeeee functions? ArXiv, abs/1912..0007, 2019.
[74] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021.
[75] Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in- context. In Models, 2023.
of Few-shot and Zero-shot Learning Large Foundation
[76] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep ReLUU networks. Machine Learning, Oct 2019.
15

---

A Overview of Appendix
The appendix is organized as follows:
Appendix
B Additional related works are discussed.
Appendix
In Appendix additional notations and preliminary details introduced.
In Appendix we show the proof Proposition C.I concerning the existence and
uniqueness of the continuous Transformer ODE (3.1)
In Appendix useful lemmas for the main proofs detailed.
In Appendix
C.4 the explicit formulas for Pp H, and pee H, are explored via the
adjoint sensitivity method.
In Appendix C.5
high-level explanations are provided to substantiate the assumptions
made in Theorem
Appendix
Includes proofs of main results from Section
Appendix Includes proofs of main results from Section
Appendix Includes proofs of all auxiliary technical results mentioned in Appendix C
Appendix G
The verification of Assumptions 2 for a concrete example of Transformer
architectures is provided.
Appendix H We provide simple experiment results and discuss the widths and depths of Vision Transformers that guarantee convergence, achieving near-zero training loss and 100% training accuracy on the CIFAR-10 dataset.
B Additional related work
Theory of Transformers. Some very recent works have studied theoretical properties of Transformer models from different aspects. [75 33 studied the in-context learning guarantees for single-layer Transformers perform linear regression predictions after being trained with linear regression example tasks. 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] 455 investigated how single-layer Transformers can be trained learn simple image models and topic models respectively.
The most closely related work ours is which the only study aware of that addresses the general and universal in-context learning capability large-scale Transformers through optimization dynamics. It shows that two-layer MLP followed by 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 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 heads increase to infinity.
In addition, we have noticed a lot of theoretical interest in identifying the optimal choice of ør in Assumption 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 [F), where poly(-) denotes a polynomial function. In the context of l2 self-attention, [39] find Φr to be /N log N/D, which notably does not depend on || and [69] demonstrate that for 4 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 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 functional of the distribution of network parameters. A separate line of works also established the global convergence guarantees for training overparameterized neural networks the "neural I_ where the gradient descent training iterates
in the mean-field regime [16, 51 50 70
tangent kernel" regime [35] 3
25
76 17
16

---

asymptotically equivalent to the training iterates of kernel regression based on the neural tangent kernel.
Connection between ordinary differential equation models and infinite-dppth ResNets. Our work is also closely related the recent literature aiming understand ResNets by analyzing their
connections to ordinary differential equations [71 I2 43, 42 72 26
22
47
23
44
7
15 [3]
Specifically, [2 [5 studied the approximation flow-bate bow-based networks
discrete networks.
43
42
72
26
7 studied the optimization of the infinite-depth and infinite-width
ResNets.
13 studied the generalization properties of the ResNet trained in the mean-fiet regime.
C
Proof setup
C.1
Additional technical notations
Define
9(T,3) f(I.0) + LT,,
Thus, 5Q could be expressed as
8Q
(B.t)= Tr (T((O(T, Do 2
(C.1)
6p
Additionally, we can combine Gy with Gn. and Gj with GA to reformulate
G(3,p,t) = E,, V TT (ITT po(H, + XB,
(C.2)
and
0(3,0,t)=B, VOTT (VTT ( w2 Po (H, At)
(H, 0)/2
( HHH -11)1 AB. (C.3)
Remark 1. To facilitate the proof, we restate Assumptions 2 andd 3or for g(T, Under Assumption the gradient of g(T, respect to T and B exists. Additionally, have Under Assumption 2
(199(T, B)1l2-oo1 I [|] + ||3|| + 11
it. For every E [N + 1], we have [T, 3).,ll2 I (1 + [|]]))
ii.. V veclggee B)]ll2 S (T D. [1] + ||||| + 1181²²)
For any LT > O and any Lr- Lipscchhtt continuous functions Ti = (H) and T3 (H), for every E [N + 1), we have
Under Assumption 3
Va B) - Veg 2l22 s @PT([ KT, LT) supl[Ti - T2
i.. E, V vo
(T) 3)
V
111112 S ΦTP(N, D,supp|Ti Kp, LT) l'll
as.. E, Veg B).. Veg(Ti, B*) ) < opP (K sup |Ti o2o- L ) 8'11,
iv. E2 V (er 8)] vecll B)]ll S ΦTT(N, D, Kr, ||0|| L supl|Ti T2lle
H
Verifyingg all these results above only needs the basic triangle inequality of general norms, omit the trivial proof. We will apply them directly throughoul the proofs of results. Additionally, we omit writing Lo for simplicity in the proof, as all functions applied Assumption 3 will be Lipschitz continuous with some universally bounded Lipschitz constant.
Next, we introduce some additional technical notations. Denote the identity matrix with d-dimension as Ld. Define the sample space Ω == Rdim@ x [0, 1]. and P(2)) the probability measure space
17

---

defined on S2 For any θ = {Be. Ie and ê = (bshhaaatee,)eemM define
d(O, 6) E, ش | - Be, Define the local risk function as
1 Rc v(H))
R(:: p)
2
Define the nested family of compact subsets (P+)r>0 as
P = {3: || < r} x [0, 1], vr > O.
For any V E P(2) and 2 1, define the b distance |lp - 1/1p as
|| p vlle (Lo |l00) (o(())-v(xx)) [ -(x)
1/8
Specifically, when 1, we have
Wi(p,v) = sup { S. f(p f(p-
- v)ddtt - (i0 0]
sup L///v / |llle- v|dBdt (i,,,,i = 0}
S S
1112
for any p, E p2 concentrated on Pr. For simplicity, any H discussed throughout this paper is assumed lie within supp(u).
C.2 Transformer ODE existence and uniqueness
In this section, establish the existence and uniqueness of the solution (H, to the ODE presented in (3.1 for any H. given that p E P2 is concentrated on bounded support, specifically, P. for some r This following proposition forms the cornerstone of the subsequent technical analyses:
Proposition C.1 (Existence and uniqueness of Transformer ODE). Under Assumptions 7 and 2 for any p E 22 that has a bounded support, there exists a unique solution of (3.1 on E 00, 1] that is Lipschitz continuous with respect to (H,t)
Initially, we demonstrate that the integral So P(3, t)d3 is bounded. According to the definition p2, it follows that
(C.4)
1f8 p(B.tt) / p(B,t)dBls Cplt 1l
for any t, E 1]. Integrating C.4 over l2 [0, obtains
/ p(3, t)d =1 -
p(3,t)d3 S.' p(B,t')dBdt
<1 + Cp / 1t -t'''dtt
(C.5)
<1 Cp/2.
For the remainder of the technical proof, we will employ CC.5 without additional elaboration.
Proof of Proposition C.I Step I: Create a small neighboring area with local Lipschitz continuity
Consider any vector 3 such that || r. Define F(T,t) (B) For T within the rectangle (T || max E d), where 8 > 0 bounded, both || and ||7||] are also bounded. Given Assumption (i. gT, (8) is universally bounded by some constant Ks,-, Moreover,
18

---

under Assumption 2 (ii) and (ii), (T, is Lipschitz continuity with some constant Lø,r. Hence, within the rectangle [T': IT H max I 63 x [0, 1] the following properties hold:
(F) P(T,12)) max Ka,r, I le 41) (2:,42)l(B1 Cp max Karr 1a,r314 12/- (C.6)
which indicates that F(T,t) is continuous with respect within the rectangle {T': || H Ilmax S
6} x [].
Moreover, within the bounded region {T' ||T - HI| max S 8}. Assumption (iii) ensures that IV vederr) 3)lll bounded. Consequently, g 9) is Lipschitz-contnnnee with respect to T for |I-11 F. Denote this Lipschitz constant by a/r,, Therefore, for any T, T" (T: |T - H mmmmx S 8)
and E [0, 1]. it follows that
() t)ls LT TFE p(3. (1, Cp/2) -le
(C.7)
which deduces the Lipschitz continuity of F(T, 6) with respect to T for ||-I|
Step II: Show that the maximal existence interval is infinite by repeatedly using the Picard- Lindelöf Theorem
Invoking the Picard- Theorem, there exists some > 0 such that the initial value problem
(H,,)) F(T,)) T(H,0) H
has unique solution on & E [0,e]. Given that this claim holds for any H, the standard ODE Extensibility Theorem guarantees a continuation of T(t) to maximal interval of existence, denoted as [0, max]
Assume by contradiction that mmx < see that
From and Assumption 2 i), for any E [0, [max]
d dt (H. I 2212-
11 1
())
K(1 + +1 -
(C.8)
Therefore, by the Grönwall's inequality, we have
[T, 122- I [T,(H,0))| exp( L' / K(1
exp
Cp/2 + r mmxx)
This presents contradiction to the notion that fmax 1. This is because, By reapplying the local Picard-Lindelle Theorem using the state T(H, fmax) as the new initial condition, we can extend the interval of existence beyond fmax- Consequently, must conclude that f and the existence and uniqueness follows.
Step III: Show that the Lipschitz continuity with respect to (H,
In the final part of our proof, demonstrate that (), Lipschitz continuous with respect to (H, t) for H E suрр(μ) and any t E [0, 1]. Given that T,(H, t) is universally bounded within H supр(μ) and any E [0, 1], we only need focus on establishing its Lipschitz continuity with respect to H and separately. The Lipschitz continuity with respect to is derived from
S [ Alle (l. BUA.:A/AA..............ECCCCCCCA
(1+C,2))CCLLH))))()
(C.9)
19

---

for any f2 E [0, 1]. Given Assumption 2 (ii), we have
(H, - -T ,1) - S. -)) 3)lpp),) sdddds
So (N ((( +
To(H' (3)
for any H. H' E supp(u). Define LH := D, (ww.Beeo(((((///2r+++))) (C.10) Utilizing Grönwall's inequality, we establish:
tH) - (H 1f S H H'll exp(Le [ p(B, p(B,s) = H'IlF
given that (-, 0) serves the identity mapping. Consequently, (H, demonstrates Lipschitz continuity with respect to H E supp(μr.
C.3 Useful technical lemmas
Lemma C.1 (Continuous Transformer output bound). Under Assumption 2 for any distribution
E ((2) where fo 1g11 fgl18112(33 ()dBdi S A² for some constant A > O and for any L E [0, 11.
have
(H, 1)ll |1 ( + A A2)).
Lemma C.2 (Continuous Transformer difference bound). Under Assumption 2 and given H, for any P. E p that satisfy fo Sg [ t)ddct I A² and have bounded supports P. for some
constants A, r > 0, we have that
-T, S C..)
sup tE[0.1]
Here, the universal constant C, only depends on N, D,r, A, and the parameters of the assumptions. Lemma C.3 (Continuous Transformer gradient component bound). Under Assumption 7 and 2 for
any p P()) where supp(p) C Pr and fo 111181) ((3, < A?, we have
sup [I, t), 8) S VN 1KBeep (K + 11 (8.0)P,
(I,, S + exp (OT(N, (N
4+) A +A4)).
sup (8,0)P,
8Q (3,t) s V - 1K + (3) +r +3) (e Bexp(K(1 + A+ (2))
sup (8,0)P,
8p
exp (Pro(NN 1KB K(1 + + 41) + A A²) + n,
Lemma C.4 (Discrete Transformer bound). Under Assumptions for any + where ML Ee 1/ < A2 for some universal constant A 0 and any At/2, Al, ...
1/2) L, we have
[H, S ||l |2-col exp 4
Lemma C.5 (Discrete Transformer difference bound). Under Assumption andd for any H, let
{Br, 1 Ô = {Be, ///441141))) such that max [ 1l. [B $ r for any = 0, 0..... A)A j = M. Then, we have that
(H,t) T (H, S Cr
ML d(),)).
1
Here the universal constant C, only depends N, D, and the parameters of the assumptions.
20

---

Lemma C.6 (Discrete Transformer gradient component bound), Under Assumption and 2 for any O such that sup Bt, S r and MT. - 1
1|:1| _ AP, we have, for any t = 0, ................................ ,L.
11,
sup max (0,))P.
sup sup max [ |1tp:///00...0.................0 | s +1), (EN++1)
[P - (B + Br) exp D. (N 4A))))
sup (3.1)EP,
where BT = B (ex-Bxexxx++++>>00 + A2)).
Lemma C.7 (Norm average concentration). Under Assumption 2 consider a parameter setting
/ MM] i..d drawn from ( where p E pp² is concentrated on P. and satisfies Ig p(3, t)d3 1 for every E [0, 1]. Then, with probability at least 1 exp(-
with respect to the parameter initialization (0), we have
ΣE 1 L. 11919000,0 J L- 0 + log(L
MT.
1=1
for any 8 > O. 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 A1, A2,..., AL and B1, B2..... B1, satisfy the following conditions:
1. For each I 1,..., T,, the norms of the matrices are bounded as ||l S + az |Bill S + bi, where a,, by
2. The product of the increments for each matrix is bounded by I 1 max b1) S C for some constant C
Under these conditions, it holds that
L
L
L
III A
-
II
Bill S
C
|
Bill.
C.4 Solution of adjoint ODE
In this section, we define the partial derivative
PH t) = aR(H;p) DDD(()+) OT, (H,
without specifying its explicit formula. Denote the derivative of 1H, to (H,,) (after vec- torization) by the Jacobian (H, t) E +1NN+)) and assume that J H, t) exists for any t e [0, 1]. Then (58) shows that J (H,, satisfies the adjoint equation of the ODE.
JAH - J, (V veeery veel
(H,
(C.11)
for any / E [0, 1]. By applying the chain rule and exchanging the order of the derivative and integral, we have, for any t [0, 1]. that
dR(H;p) Ovec 1)]
vec vec (H,
(Tp 1)] Ovec (H,
0R(H: p)
Övee[T,
veol_
(T, 1)
(voc(T,(HH())
=veclp, T)] ()(Htt)
21

---

Hence, by taking the derivative with respect to t, we obtain that
(o, -voolppp V
with the solution
veclp, vec exp (f.") Vveeryeedd(((144.))))9000..800 (C.12)
On the other hand, we have
ÖR(H:
Pp(H, = (H, (Read[T,(H,1) ( Bread:
(C.13)
where Eread is x (N + 1) zero matrix except at the (d + 1, + 1)-th entry. Moreover, from C..22 and C.13
that
veclp, (H, t)] (H. exp (f. Voooop
( ()) s)d (C.14)
DN+4+1;
Additionally, we could explicitly derive the formula for pe a H:)) for the discrete
(T.(H.))
Transformer: By applying the chain rule multiple times across each layer with the encoder either f or h, for any = 0, At.... ........... (L. 1) 1, we have
vee rpe H, t)] Reead[Te 1)] v(H)
(Tainwoe(r) + (At/2)M 11-1//110110//// JE[M]
II
II (s-t)/As At41)//44((( (1(+1)(()))/)))) JE(M]
[dimnvec/2] + (At/2)M
V (ee A1/2),
DN (C.15)
and
veclpe H,, + At/2)] (Read (H, 1)] v(H))
II
(Jainvee(r) + (At/2)M
vorly (e
(e- (0-t/A442((0--)/4 jE[M]
II (G-t)AAt -t))/A)) JE[M]
[dimvee|22] + (A(/2))
Σ V Veed A1/2),
wnill) DN
(C.16)
C.5 Explanation for assumptions made in Theorem 4.1
The justification of the two assumptions outlined in Theorem 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 X, which penalizes the magnitude of the parameter norms, implicitly promotes solutions that confined 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. is naturally satisfied as long as the origin Odim@+dimw remains interior point of supp(pec (-)t).. This condition is relatively mild and
22

---

is generally satisfied. The challenge arises in verifying that supp(Poc t)) for the ag component extends encompass the entire space K. While direct confirmation elusive, 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 cireumstances, providing a basis for this assumption.
D Proofs of main results in Section 3
D.1 Proof of Theorem 3.1
This convergence is detailed two parts. First, the finite time result, stated in points (i)-(ii), 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 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 any specific time, along with the difference regarding different time for the same dynamic. By applying probability union bound on a dense set of L points, we can extend this 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- mean-field limit.
Preparatory Step: Nonlinear dynamics
We first define some auxiliary quantities and differential equations that are useful for the proof. For any gradient flow parameter setting ( (3()).,) Je.j. from its definition 2.7 we could rewrite the dynamics as
3 = 3(0) - 1/" (((())) @(e), t)ds
(D.1)
for any gradient flow time > O, depth index 0, At...., (L 1) and width index 1...., .M. For simplicity, any mentioned constant only depends on N, D, T, À and the parameters of the assumptions, and we abbreviate the subscript t 0, At.... ..... (L- 1) j = I....,,M M by t.j throughout the proof.
Inspired by the "propagation of chaos" idea 62 could define the "nonlinear dynamies" with the
same initialization setting (+)) = { hags i.e.
{
3((7 310)
So G( p(0)
(()),t)ds. t)ds,
(D.2)
3 3(0)
for any f,j. Here, (220 the solution 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 (35 its existence
and uniqueness are guaranteed by Proposition 3.2
Observing that htr are independent due to the dynamics only involving (p(*)) s20 with i.i.d
initialization over Po. we can consider (3(-) Jes as i.i.d. samples drawn from (p(*) J... In addition, from Propositions 3.2 and D.I for any t.j we have max(]];9t.ll] 1|3 < Rz, where R2 is defined as in these propositions and does not depend on M and L.
Preparatory Step: Bound the gradient difference regarding parameter settings
As the second preparatory step for the proof of Theorem 3.T that will be helpful:
we present the following three lemmas
23

---

Lemma D.1 (Continuous gradient difference bound). Suppose Assumptions 13 hold. we we have
foror P.V. E p concentrated Pr for some > O, and B, 3 such that maxx{118]] || S then
G(B, (1at))-G(B,,,, v, 1) - Co( exp(CoW1(p. v)) - -1+ (1+ 1)11B - 31).
for any t E [0. Here, the constant CG only depends D. and the parameters of the assumptions.
Lemma D.2 (Discrete gradient difference bound), Under Assumption 3 for any (BL.)h//ar-e((().eet Q = any = O..... 1) j -....,, M. Then have that
such that max (1 ll. (13 S r for
||ō(8, 0,t) G(3 ((3,6,e)1 t) S CG
d(o, d(0, +) -81)
ML
for any 8,3 E {8: : ||8|| rr and t = 0, At.... (L 1) Here, the constant Cg only depends on N, D,r and the parameters of the assumptions and d(O), () is defined as E. EgE= [B, Br,
Lemma D.3 (Oracle gradient approximation with discretization). Under Assumptions
suppose
that the parameter setting @ is i.i.d. drawn from // [/ for some p pa
concentrated on P, and satisfies that Io ((3,t)ds = for any t E [0, 1]. Then with probability at least 1 4cxp(-5) with respect to the parameter initialization ( we have
[((,,0,)) - (03,0,,))-(3,,,0 SL- + + log (L + 1)
M
+ log + 1)
[G((,A,t) G(3,p,t)) S L-1
M
0, -((,A))) ST.
for any 3 E {B : ||8||| S r), 1 = 0, (I- 11),,, 1 and any 8 > O. Here, S hides the dependencies on N,D, 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 W2 continuity of parameter (sample) distributions
Our analysis commences with the bound for 0 < 81 < $2 T, we have
M ΣΣ 0(0,
(2(0()) 2( <
1 ML
2lea)
(82 81)
ML
where each particle at time $1 is paired with its position at time 82, leveraging the Jensen's 's inequality. Recalling the identity
)Q())))
M
ds ML Σ>Σnn
(G( g(+),)]("
shown in Proposition D.I it follows that
Wa(p(or), p(oz)) (82 - (101110))))))) < AB.
where the last inequality uses (D30 Since AB 1 we see that (2(t))),, C(1 +X(2222 1 2/2 for some constant dependent on the parameters listed in the result. Similarly, we have
< B 36m0l| s (82- ST ['' G e)|l
(1( (0)) <
dsdtds & (s2- s1)((p)) & (11=
where B(or) (B, 1) is embedded with its future position at Blea), The last step is feasible by setting C large enough, noticing that ((o)) S X2 To summarize, we have
max { (W2(ptn)) ((00))
(e( , S C(1 A)
(D.3)
24

---

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 A (()) == sup,/ supt,,11333
31-21
I| for any 8 E [0, Taking the
difference of [D.I and D D we obtain that
| - - 11 -
[ 000) C(3 t) ds
- l
[ 000) t) 6 60, 1) ds
+ /
|0 3 t) G t) ds
SCo L''
1 10 ( + (1 + 2)1[3(() - B(ol1) do + L. |000 1) G p()) .0) ds
(M..
SCo L'' (11-000 (MI ( (1 + \)]|3()) B(011) ds + sup
| 60 1) G(B plo) .1)
sE(U,-]
(D.4) where the final inequality stems from Lemma D.2, employing a constant Cc dependent on the parameters listed in the result. By taking the supremacy over 1,j in (.4). and considering
ML (( ((0) - supr, [1 - || for any s 2 0, we derive
1
sup supll - s Cc(2+X) I supll - || ds sup tj sE(U,r]
|0003,000 ;0, 1100))((0)))0 COG 20,)))00)))
(.j
t.j
Further supremacy taken over E [0, r] yields:
A(T) < (()) + Co(2 + A) 1-' A(s)ds,
(D.5)
where we define A(7) == supt,j (0. G( ((0) () G(Al (000).pp),.0 |0 for simplicity of notation. Apply the Grönwall's inequality D.5) yields
( S exp CCa(2 + A)T) ((7). )
(D.6)
It remains to bound ((T) to bound A(T). It's worth noting that by Lemmas D.I and D.2 for any
81, 82 E T] and 4,J, we have
( (loz) t) G( p00) rt - | ((m)),)) -( pbm) .t)
< ((3 O(0)) ()- G(((5)) @(0)), .1) + c( p(m)), -G((())) (6m)),0)
Cos (01) (+))) ) - 1) + Co
d (())))) ((ez)) + Ca(1 + 2)) - A3''
ML
Sexp(CaCVV3 81) 1 + (1 + s) s B - 3(52)
t,j
Sv32 51 + sup ((.11) 3( ||
(D.7)
for some constant Ca dependent on the parameters listed in the result. Moreover, for any 1,j, we have
| - ((22) || S [
|G ((0, S v82 -
| |GGB 1000 & (1+x² - 81 81
(D.8)
where the universal boundedness of [|G
t.j (() (++), t) A36 can be readily derived from the third result of Lemma C.3 alongside Assumption 2 (ii). Therefore, we obtain
(( 3(0z), t) - G(6(3) ((0)) b) ( 3(x)), .) - G(863) ((0)), t)|l > (1+A+22-29) 81
(D.9)
25

---

for any $1, $2 10.
Define on I for 1, ...... T22. Leveraging Lemma D.3 and applying the union bound over n E [1,2] (updating the the lemma to 5 + 2 2og L), we obtain, with a probability of at least 1- -xxp(-00 with respect to the parameter initialization (
G (003),00)0)) +(1+++X
X2 >
V
8 + log T. + log(( + 1)
sup (E(T))nejdR
M
8 + log L
ST. +
M
Furthermore, leveraging (.9). we deduce
sup sup
G ((0) G G ,(0),))
Ssup sup
[(3 0 G( ,001))
tj sE[0,r]
( sE(Tainc
+ sup
(1+ + / - $1
log I. + 1)
SL
M
Returning (D.6), we establish that with probability of at least exp(-8) with respect the parameter initialization (10).
+ log(L + 1)
sup supll -
<E(0.r] 6
M
We denote the event where the above inequality holds as E,, thus we have P(B,) - exp(-6). -eexp((0))
Step III: Prove the finite time results
Now, we are poised to demonstrate the results in Theorem 3.1 that concern supremacy over s [0, The verification of Lemmas C.4 reveals the existence of universal constant B, := B exx((K( R, + R²)) such that
max (IT ( (H, 12ll (H,
(H, ( (H,0)l2-col) S Br
for any H and E [0, 1]. Utilizing Lemma D.6 and applying the union bound over E [12], we observe
8 + log(L + 1)
li (H - (pi (H, L- +
sup (E(Tn)nE(A))
Additionally, note that for any $1, $2 [0, r]. derive from Lemmas
C.2 and C.5 that
[ tH p (H, f) - T (H, T (p (H, t)
<I Boo (H,t) T Tou (H, + /1T (H.t) (H.)) T T (H, t)ll
(D.10)
Suppl 3(1.) 3(020)1 Wi(por) ((o))) 1,
Sv82 - 51
where the last inequality is derived utilizing
D.3 and D.8 Consequently, have
sup sE(0,7]
18 (H, Tocn (H, MT
sup SE NE(Tn)neL2)
| (H, Tycs) (H,0)l
sup
V82 81
[82 1ST/L2
-L- V 1 + [ 1l)
M
(D.11)
with probability at by - exp(epp) -). We denote the event where the above inequality holds as Ez. thus have P(E2) 2 - exp( -)). Considering that (B(5) Jas could be regarded as i..d. samples
26

---

drawn from fpl7),, h.j. employing a similar method with the concentration guarantee from Lemma
C.7
we can readily deduce the existence of an event Eg with P(Ex) 2 -exp(-o) such that, under
Fgg, we have
(111) 1. + V 8 + log(L+
sup +E[O,T]
ML
M
j=1
Now, let's analyze the scenario under the probability event E1 Eg n Eg with P(E, n E2 Es) 2
1 - exxp(-8), Lemma C.5 demonstrates that
sup RRead To H, 8)] Read T
To (H, t)]s
sup (H, - TBw (Hrt)le sE(0,r]
=E[o,]]
(D.12)
X sup supllet. 360 L- + <E[o.r] 1.j This further implies, by [.II), that
log + 1)
M
sup Read (H,0) [)] Read Read[Tpr)) (1, (H,
=E(0.r]
$ sup Read (oo (H, t)] Read Tgo. +)]++ sup [ (H, (D.13) =E(0,r]
sE(0,r]
ST. /
15 + log( T. + 1)
M
Since
max p(r) (H, 1o [|T (H,0022 i (H, (2-0o) is universally bounded, (D13 immediately indicates that
8 + log(L + 1)
sup Reead [)] [el]]-Rad[TTpt. (H, t)][ L --1 sE(0,7]
sup .E(0r]] and
M
(D.14)
M Σ
1 / |311 (3)) ()33)
sup -((()))(((()))) sup ((()())) - R((p())l+ =E(0,-] + log(L + +1
12
sup sE(U.7]
ML
+E[0.r]
j=1
S6
M
(D.15)
Step IV:
Prove the weakly convergence
For the remainder of the proof, adopt similar approach in the proof of Theorem 2.6 in [16] We denote á as PM., for any given M and L. It's essential to note that we treat PMM. as probability measures in this step of the proof.
Recalling D.3). for any 81, $2 E [0, T]. we have
W2(P (...)) S C(1+ -
for some constant C' dependent on the parameters listed in the result. We observe that the family =(x) of curves (s + PM,L) is equicontinuous in W2 on [0, uniformly in M, L. Additionally, the family (PM.L) lies within W2 ball, thus weakly precompact. As the weak topology is weaker than the topology induced by W2, according to the Arzelà-Ascoli theorem, along any sequence where L and log L/M + oo, we can identify subsequence that converges weakly to certain process (1=))) E p2 x I, concentrated focus on this subsequence, still denoted as (ÂM,L.)M,..
PR. at all times. In the subsequent analysis, solely
For any E [0. 1). let's define the sequence MM,, of momentum fields, which is a vector- valued measure on 7] I2, denoted by EM,L = G (()) ()) (3, t)ds. We also define E = G(3, (0), ()) t)ss.
27

---

Considering that both PM.L and v are concentrated on PR., we also have uniform convergence in the Bounded Lipschitz metric. Hence, for any bounded and Lipschitz function : [0, 7] x Rdiin + Rdimm it holds
| v" IIBL. +
uniformly among E [0, r] along the sequence. Note that
Jo L. 1.
p. [(E.L E)| s 511 max J. L' / ( M. G(3, - (1 (3,
+ [ L / 8. (p(0) - -3) (3,1)dBddds
J. / [G 0 ) - ( pM..p 1) (1)) 1)dBdids
' L'' / G(3, ()) C(B, t)] p()) ())dddds
+| +1lyp|lma.
+ sup
11 v/ VBL
sc [0,]
[1 - v² B +
(G(3,p)),)) G(3, ((33,p),)))G(((,())
sup
sup
sE(U,T]
ME
5 + log( [. +1)
SL
sup
v* vBILBI
M
sE(O,T]
(D.16)
for some constant Cg dependent on the parameters listed in the result and with probability at least 1 exp(p0) ())) with respect to the parameter initialization ( for any > O. Here, the third inequality of 6.16 utilizes the third result of Lemma D.3 and the fourth inequality uses the second result of Lemma D.3 following similar process in Step II to achieve supremacy over E ().
From D.16 we infer that So So fg d(EM,L E) 0 almost surely along the sequence.
Hence, EM.L converges weakly to E almost surely along the sequence, and the particle gradient flow for (z(7)))20 almost surely satisfies (2.7) on [0, r] for any arbitrarily given >> 0. According to the Fokker- Plankr- equation without noise involved [59]. conclude that (v(T))z20 almost surely
satisfies (3.5). Consequently, the uniqueness stated in Proposition 3.2 ensures that (u(T))
(plT)) 20 almost surely.
(r2
D.2 Proof of Proposition 3.1
Proof. Suppose that the Fréchet derivative
8R indeed exists, we establish öp
8Q
8p (0, w, t) = 8R 0, w. t) + 2 (1101 + ||w]]
óp
Therefore, it suffices to show that the Fréchet derivative of T. with respect to p is
8R (3. = E, (TTT, ())
(D.17)
5p
Denote Py p ryv p). We provide the following lemma to bound To, -H, by expanding the first-order derivative as follows
Lemma D.4 (First-order derivative of Transformer output). Under Assumption for any H and P.V pa that have bounded supports, we have
vec[T,,(HH, ( (H. -n S. / exp ( VwwerryeelllllTccc ())3))))
vec[g(T, (H, ()) (3, o()).
(D.18)
where py := = 1/22
28

---

Given (.188 we observe from the solution of Pp CC.33 and C. 14 that
(H Read (p, (H, 1) (H,
=vec Pp H, 1)] vec[Tp, (H,1) 1))
=n //' I. vocip, vocl)))))T t), B)] (B..t))d o(n)
(D.19)
!/' /, (I P.HH, 11,, t)dddt +
Hence, by applying (D.19) to the risk function, we obtain
R(p,) R(p) 2 F (Read (H, 1)] u( v(H) (Read[T, v(HF) I
E [Read[T,(H,, 1)] (H) Read (H,, Tp, (H,, Read[T, Tr,, (H, (Read [T,, (H,1) 1H)) 88R V- /- p) + o(T).
n 8p
which indicates D.17 and concludes the proof.
D.3 Proof of well poledposs of Wasserstein gradient flow
Following similar idea as Proposition 2.5 of [16]. we leverage the general theory of Wasser- stein gradient flow developed in 4 Define the functional family (p) as
Proof.
Qr(p) - ((()) oc
p(P:)=1,
otherwise.
For any 0, let's consider any admissible transport Y paxx concentrated on P,.. By definition, both of its marginals, denoted by P1 and p2, are concentrated on P,.. We define the transport cost for y
cp(r) = / 11
for p 2 1. Additionally, we denote the transport interpolation as ph ((1 (1--ap) a.pp #T. Our proof consists of several steps outlined below.
Step I: Show that Q, is proper and continuous for W2 on its closed domain:
Note that the parameters r, D,N, A remain fixed throughout this proof step, so we hide the constant
dependencies on them. For any (B,t) E Pr, we have (r(6(s,))) = [[Redd H + Atg(H, H, B)]
( A 111B112 < This indicates that Q, is proper. Moreover, for any p, v p whose bounded support belong to Pr. Lemma C. ensures that
[T,, +T (H, = 0(1),
and Lemma C.2 guarantees that
[T,, (HH, ((W
for any H. Therefore, have
R(v) R(p) 1 Fr [ReaRRdTTT - v(FF)) (H). -())) [(H))
=Eyl (Read[T, - -v(H) Read [T, (H, -H,(H,1))
(D.20)
+ RRead[T,(H,)) TH,H, (T 1,HH, O(W2(p,v))
29

---

Furthermore, since both and have bounded support. || is Lipschitz continuous with respect to (3, t). Therefore, by the Seattorovic-RRubinstte Theorem (see Theorem 5.10 of [68] for example). we have
A / 2 3
1181
- v)dBdu = OWWi (p, = O(W((W v)).
(D.21)
Combining (D200p and (D.2T), we obtain that Q(x) - Q(v) = O(W2(p,))) Therefore, Qr
continuous for W2 on its closed domain.
Step II: Show that + ((P))/O(()) is differentiablee and has a Lipschitz continuous derivative Let's denote h((a) == Qr(p%). Lemma C.3 ensures that for any E p2 with bounded support belonging to Pr, we have bounded ||
AQ ]| on Pr. Therefore, (p(r) is differentiable with
respect to t, and the derivative reads
iQ
h((a)= = (
PaPa
Pla)
0p
de
p=pă
iQ
(11 ((1-a)(B1,th) + ((22, 12) ((),))))))))))))))))))) ((2,1)))))) (7((3) (82, b2))
6p
p=
(D.22)
Then, it suffices to show that h/ (a) is Lipschitz continuous. To accomplish this, we first propose the following lemma for later use:
Lemma D.5 (Locally Lipschitz of p for the gradient). Under Assumptions 7 and for any p.v E pa concentrated on Py. there exists some constant L, depending on r, N, and parameters of the assumptions such that
sup (H, (H, S Lrlpp ul11. eE(0.1]
8Q 0p
iQ
sup (B.1)EP,
(3,t)
(3. S L, vl11-
6p
Returning to the lemma proof, for 01 E [0, 1]. by the triangle inequality have [N'(a)) (((22) E J1 + J2 where
J1
/
d
SQ
+ (1 (B1, ( (82, 12) (1)
8p
p=pā.
1
d 8Q
(1-aa)())t)+a((bbb)) (B), (2), dr ())
+ar
0p
p=pó2
8Q
8Q
(8,t) / (((1,.t) - (B2 (82,t2)
< sup
(g, f)
(B.1)P.,
0p p
ōp
<LrlIpa, - -Pa, C1(7)
< L. (7,Ca()))ln 02l
SL,C2(7))a1 a2
(D.23)
by Lemma D.5 The final inequality of (D.23 applies Hölder's inequality to obtain that C?(x) s (7()). Furthermore,
5Q
Jz :
- |
((1-aa)(311))
t1) + (82,
8p
p=pb2
f
3Q
az (2) -1)
8p
p=pda
< sup
8Q
(8,t)
lai a2 (B2)) (10))=tt))(((,b))]((((((.,.,,,),
8p p=pde
(B.t)EP.
<
<T,C%(7))an 02l
(D.24)
30

---

where L, == sup (B,,)) 11 p=p2, (3,t)ll < oo from Lemma C.3 leads us to the result that h' (n) is Lipschitz continuous.
Combining
D.23 and D.24
Step III: Show the well-posedness of Wasserstein gradient flow at some finite time
We follow similar approach the proof of Proposition 2.5 in 16]. Since N'(a)) is Å (7)
Lipschitz continuous with respect to a for some λ, the well-posedness of the Wasserstein gradient flow for Qr with the velocity field constrained on P. corollary of Theorem 11.2.2 of [4] Specifically, there exists a unique curve (pr (TT) )r20 continuous in p2 such that:
dp. dr
= div (of
where
G(
(8,1) E Pr, otherwise.
of
(3,t)
0,
for pf -aee Given the initialization Po concentrated on PR. for any R, the unique pf exhibits a first exit time denoted as
Tr = inf > 0: pf (P,) <}}.
By defining this exit time, for any F > r and T E [0, Tr], we observe uy (3), = G(3,pt .) and (B, G(3, pf t). Due uniqueness, we infer (T) (7 on [0, Considering pf as the solution to (3.5) we establish the existence and uniqueness of the Wasserstein gradient flow for Q over [0, Tr]
Step IV: Show the well-posedness of Wasserstein gradient flow at all time
To establish the Wasserstein gradient flow's definition for 2 0, it's necessary to demonstrate that lim Tr oc. For any r >> R, according to the energy identity Theorem 11.2.1 of [4), on [0, ty], we observe that + Q(p(T)) is non-increasing. Specifically, this represents
aQ(p(t)) dr
dQ L.' /
dp p=p(+) div Glvs .t))
!.' / (G(3, div (1)
(D.25)
/ / 13 S
Therefore, for any E 10, r.), utilizing Lemma C.1 we have
((p(1)) s Q(po) -E,, [1 Read][T (H, 1)] v(H1) 2 91+
2 // /
A
1
XR?
2 B, (H, 1)l12-col + B) 1+ +
(D.26)
XR²
SE,KK(TTT (H, 1) -0c BB1
2
2
AR22
<B² + B² exp [[(1 [K[1 R
2
31

---

Thus, we have Jo [1118)12 s (0(6()) - R² + x
x 2B2²
2B2 exp [(1 + R
R²)) Ag. According Assumption (ii) and Lemma C.3 for any (8,1) E Pr. we have
N+1
(1,0) - ا8 - - -
V (,, )
< sup N+[++1]
I Vagg(T 1), B). Σ ( t)ll
< N +1 sup
(T (H, ().t),))... IIp, (H,
iE[N+1]
SvN Tor(( 2- (H,
+ [|
<Cn(1 + (13)),
(D.27)
where CR
Vop exp [[1 As + A + B ex + Ao
A)))) exp ( D, D. exp K (1 + Ao + A1)) + A + 43)) A Applying (D.77) to the
gradient flow equation
d367)
= 16 (7) (3,t), =3
dr
d (3.),)0) S ( >811 S CR(1+ I). This 1361
for r 2 0, we obtain
dr
indicates
| S (13) + 1) exp(CRTT) 1 (R + exp(CCrT) -1
(D.28)
by the Grönwall's inequality. Therefore, for any 7 > 0, pLT) is concentrated on P(R+1) exp(x))))) implying that for r > (R + 1) exp(CRIT), we have T, > T. Hence, we conclude lim Tr = OC,
establishing the existence of unique Wasserstein gradient flow from (3.5) over 0.
Eventually, we establish the three properties listed in Proposition 3.2 for p(T). By 3.5) for any E [0, 1], have
| (0) (00 ((0 ds
=1 1' 0.
indicated by the Divergence Theorem as pl+) has bounded support. Next, for any 2 0. (D28
shows that (()) is concentrated on PR.. Moreover, (D255) now holds for any > 0, implying
So 1111811 (r) (B, s AB for any 0.
D.4 Proof of poseddossss of gradient flow
Proposition D. (Existence and uniqueness of gradient flow). Under Assumptions 7 3 for any
initialization of ()))) i.i.d. drawn from {Pu(O, w()))... there exists unique solution ( for
2.7 Additionally, for any т 2 O, we have
i. (T) has a bounded support, meaning sup 10 + 1103 - R.
i. E.E,, 1 | 17 A3.
Here, R2 and An are defined as in Proposition 3.2
Proof. The local Lipschitz continuity established in Lemma D.2 directly implies the continuity of Gg(Bi, O, t) with respect to O. Since {ML AML-Gg(BL., θ h serves the gradient of (()), it follows that (()) is continuously differentiable, indicating the local semiconvexity of
(()). Specifically, for any θ, there exists some such that Q(0) E, (l)0r,, (1
32

---

| is convex within small neighborhood of θ The existence and uniqueness of a gradient flow over the maximal interval [0, Faax] is standard result (see Section 2.1 of [T]).
For any T E [0, [max]] it holds that
((110)) 2 ((010)) ((0())) -((1))))(((1))
/ Σ
M
dQ(e)
.L dQ(e) dBe, (e=(())
dT
dāe,
(=
M
ML Σ Σ
M
ΣΣ
So
00)) [I ,0)d))))
MLT
(D.29)
where the last inequality follows from Jensen's 's inequality. (D299 establishes that (()(+)) is both upper and lower bounded, and ((r) exhibits bounded curve length over any the time interval [0, Tmax]] By compactness, if Tmax is finite, then (7mmxx) exists and thus must exist beyond Tmax. which leads contradiction. Therefore, Tmax oc, and the well-posedness of the gradient flow for
2 0 consequently follows. Additionally, (D.29) shows that for any T 2 0,
M
Σ
Σ j=1
18, 12 s 22 2--------0 (Q((()) s 22 ( =x-'B (Readd (H, 1)] - -((()) '
MT.
1
ΣΣ
MT.
<x Ex [(] () + B) R²
52X IEL (H, 1) cool + B²] + R²
SR² + + (2B²
(2B² + 2B exp
[ + R²)
=4
(D.30)
The last inequality of (D.30) follows from Lemma C.4 thereby showing that NEΣ Σ, E, 113 A for any T 2 0.
M
As the final part of our proof, we demonstrate that the norm of any entry of e is bounded at any given time + 2 0. Note that
(G, f) A31
V (fo JOOHH, /2+ ( A/2), w):s 1 (o (H, At/2),, /2
N+1
N+1
< sup (HV, ve(++1]
Σ
(1/. w).ll
[p, 1 11/2)),11//
SVN+I +I sup
1E[++1]
V
+ 1
max (H, (H,t
<Cn(1 + [[31))
(D.31)
Applying (D311) to the gradient flow
d3(7)
[.j
-
dr
33

---

for 2 O, we have 4110 (-G(of .e ,6,1)))))))11
[|0(0!7 0i,, -)
- A3 Cn(1+ ().
|
This indicates
| S (1131 +
-
Rr.
(D.32)
D.5 Proof of Proposition 3.3
Before commencing the proof, we introduce two proxy Transformer procedures in addition to Te and
Tp The first proxy, denoted as Te. involves moving the layers with the encoder h slightly forward by A1/2 in the depth index. This adjustment results in a discrete Transformer with only L layers, where each layer has a step size of A/ and an encoder of (ʃ + h)/2, represented by Specifically, Te can be written as
M
M
Te + At) (To(H,) AtM
Σ
t), +
Σ ==1
h(To (H, mnr))
f=1
(D.33)
M (H
-H,(H, ALM
Σ
The second proxy, denoted as Tp. extends the width to infinity by letting M + co, effectively replacing the average with integral:
+HH = T,HH, J(TTAU.,))))
(D.34)
We let all four Transformers share the same initial state (H, (H (H,
(,H 0) = H.
We first present the following lemma, considering parameters i.i.d. drawn from some distribution E p2 with bounded support:
Lemma D.6 (Oraele approximation of discretization). Under Assumptions 7 andd 2 suppose that the parameter setting θ is i.i.d. drawn from [p(0, w i)),, for some E p2 concentrated on {(0, w): |10]22 + |! < r²) x [0,1] and satisfies that So p(0, w, t)d(0, w) = for any t E [0, 1].
Then with probability at least 1 - exx(--) with respect to the parameter initialization (0), we have
log( (
l HHtt (T T,(H,t) S LL-1
M
for any H, 1- Al,... ... (L (1-1) and any ō > Here, S hides the dependencies on N D.r and the parameters of the assumptions.
3.3 Since p E p22r has bounded support for any p. there exists some p² p2r such that R(p" inf per R(p). According to Lemma D.6 we can find specific O such that
Proof of Proposition
(H, - (H, S <C(
log +
M
where C' depends N, D,r, and the parameters of the assumptions. Moreover, from Lemma D.6 ensure that each entry Br.j of θ satisfies || S r. Verification of Lemmas C. and C.4 on T and
TO respectively leads to their uniform boundedness, i.e., sup, T,H,, S and sup, Te H t) S 1. Therefore, we have
((()) R(pp)((()))))))))))))) ) SE, Read T. (H, - Te (H, R)]l][ (H + (HH + Hy
(og(++))
S--1
M
Here, S hides the dependencies on N, D, r and the parameters of the assumptions. The result then follows.
34

---

The proof for the energy functional Q (and Q) follows a similar approach. There exists some p pptr such that Q(p") peppe pa.r Q(p). From Lemmas D.6 and C.7. such that
we can find specific +
[Ta(H, (H, f) S c(( + V log(L-1) )
MatE E. 113 - [3313)1
C(L log +1]
by setting C' large enough. Verification of Lemmas their uniform boundedness. Hence, we have
C.I and C.4 on T, and To respectively leads to
M
((00)-Q(")) (())) ( R(p))
Σ -
1 112 1110
log L + 1)
ML
M
2=1
The result thus follows.
E Proofs of main results in Section 4
For simplicity, we assume that Assumption holds with (g, = (f, 0). The proof for the case of (g, a) (h, w) is symmetric, involving simple substitution of f with f and 0 with w.
E.1
Proofs of Theorem 4.1 and Corollary 4.1
Our proof of Theorem consists of three parts, each focusing bounding differences related to the energy functional Q or Q
The first step establishes the continuity of the functional gradient #Q P=o This ensures that if the derivative with respect 3 for the functional gradient is constant region, then the functional gradient remains constant within that region.
The second step provides the key bound for ((pc), which is proportional to x. This involves detailed analysis of Q's s' landscape by bounding its derivatives.
After obtaining the bound for ((ss)), the final steps are to show that the finite-time risk can approach this bound. Achieving a loss as small as requires (p((To)) S e for some sufficiently large To- We then apply Theorem 3.1 with constant dependency Te. show that Q(To) becomes sufficiently small. Since (()(r) non-increasing. Q(T) remains small for all r 2 To-
Preparatory Step: Landscape analysis
First, the following lemma suggests that as long as the risk R(p) remains positive, descent direction Q(p) can constructed at any depth index, provided that À sufficiently small. This implies that by adjusting X, one can influence the gradient flow to effectively reduce Qpp).
Lemma E.1 (Landscape of Q(p)). Suppose that Assumptions hold. For any E pp concentrated
on P. with some r >> O, any wo E Rdimm, and any t" E [0. 1 such that Sq ((B,tt) 2 1/2, there exists P(Rdim0) such that
for any 3 E supp(v), have 1/B, S || B, it for any (01, (01,02,0 w). we have f₂ K. and wo-
ii.. Ig 6p
8Q (B,1) (2()) ((113) -()))) S Ci C2R(p)
Here, Br, C1, C2 are constants that depends on N, d, and the parameters of the assumptions.
Given the t" and wo specified in the theorem, Lemma E. indicates that there exists some
E P((Buime, (0, BR Bamee, (0, (0.1/BRa)) x K. x {wol)
such that JA 8Q Ie- (3,1* ((()) P (31 dB S C,A - C2R(p), where B C1, C2 are con- 6p stants dependent on N, d, Rog and the parameters of the assumptions.
35

---

In addition, for any p E p?? we define the following two functional derivatives:
iQ1 (0, w, - E,, (5.[Tr(( (..TTT(((TTALU,))... p.(H,,))] + 1l01k
Öp
8Qh
(0, w. EE [E,,TT(((T,((H)))))) (p)HH. + A1ww112
6p
It is obvious that
6Q
1
+ 1/2
5p
8p
Proof of Theorem 4.7 Step I: Show that 6Q 6p Pm (3,t) is continuous with respect to (3,t)
In Step of the proof of Lemma E.1 we establish that Poo E p2, with bounded support PR implies Poz (H, t) C,-Lipschitz continuous, where Cp constant dependent solely on N, D, Roo. and the parameters in our assumptions.
Next, we would like to show that 8Q (8, t) continuous with respect to (3,t) E R x [0, 1].
6p pro
Let's focus on the region (3,1) E P. for any r > Rec, so that Pro is also concentrated on P.. It's
noteworthy that for any bounded support (8,() E P,, Lemma C.1 and Assumption 2 (i) ensure the universal boundedness of gTT, (H, t), 3). and Lemma C.3 ensures the universal boundedness of PHHH, t) for any H E supp(u), with the constants depending solely on N, D. r. and the parameters of the assumptions.
Combining the Lipschitz continuity of Pp (H, t) and Tp (H.t) t) with respect to (H, t). as shown in Proposition C.2 along with the Lipschitz continuity of 9(T, 3) with respect to (T, B) when || is universally bounded (as guaranteed by Assumption 2 (ii) and (iii), and their universal boundedness,
we derive that (T([sT(T... 1)) (H,
Ca-LLipschizz continuous for ||-1/2 with respect (8,1) E P. for some constant CG that depends only N, D.r and the parameters of the assumptions. Since the Lipschitz constant Cg is independent of the choice of H, we see that
Tr (T (1) Poz tH) is uniformly continuous across all Hl supp(u) with respect to (3,t) Pr. Consequently, we have that
5Q 5p
(3, t) p0= (By.t))===E<TTT (1
Pon (H, (H.
A + 1101
is continuous with respect to (B,t) E P,.. Since the choice of r is arbitrary, conclude that 6Q
öp Pmu (3,t) f) is continuous with respect to (3,t) E Rdivvp j0, 1).
Step II: Show that Q(pr)) S X with further landscape analysis
In the first part of the proof. we will adopt a similar approach to Theorem 3.9 of [7 to demonstrate that the stationary point of the Wasserstein gradient flow, denoted Pac. satisfies (p S X. It's worth noting that in [47], the authors assume and conclude R(pco) = (p)))) = 0, but this claim relies 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 Pso of the Wasserstein gradient flow 3.5 i.e.
dp(3,t) diva pV
8Q
dr
5p
must satisfy V
6Q
Bp P= 0 almost everywhere over supp(poo). This further indicates that
V sQ 5p Ipw. (01, f22 w, t*) = 0 almost everywhere over supp(p t")). The fact p(,,t") is con-
nected set, coupled with the continuity of the Frechét differential 42 low with respect to 3, implies sQ that, 8p lew (01, f22 w, t*) = C for some constant C' over (3, E supp(p(::t").
Given the separation assumption on the support of (,o(;,"", we ensure that for any (01,02) E Baamo, (0, BR.)) Baino (0, 1 1/BRw)) x K,, there exists IR, 1/Ro BR S Io/5 Rx BR_
36

---

such that (c81, B2. wo. t²) supp(pac). Combined with Assumption (i), which implies the 1-
homogeneity of f(T, 0, 82) with respect 8, we have
8Q1
8Qs
(c01, wo, 1*)
(01,02, wo, (-) ((c
6p Pso
6p
Pro
(E.1)
hQh
hQh
(c01, 02 (0)
(01, 02, wo, t*).
8p
8p
Pro
Pso
Hence, given that
5Q V
Pro (,,t²) almost everywhere over supp(poc (,) it also holds that 6p
(2,00)
8Q
02,
02 wo, t-) V (0)00)2)0)
8Q,
(01, 02 wo t²) + V (00102200) 8p
8Qh
(01, 02 wg wg. (t),tt)) /2
0p
8p
= ( |cl-1 A01, Odim0 Ow)
which implies
8Q (0,,0,, wa, wo, 12) S BR BR Roood.
(E.2)
8p Pr
Given the condition that ||(c01, 82 wo, wp, 1²) (01, 02 up, wb, t²)ll S R2 BR and recalling that öp Pw (01, 02 w, t²) = across (3,t*) E supppp(..tt) ,²)))
E.2 further indicates that
8Q
(01. 02 wo, wo, c/ (Re + R BR ).
(E.3)
öp pmo
for any (01, 02) E Baame (0, BR.....))) )/ Batmos (0, (0,1/BR.)) x K. Hence, by Lemma B.1 we have
8Q
C₁ - C2R(poo) 2
(B,t*)(((6)
-
6p Poc
/ 6p
8Q
(3,1)
()((3) (20(311))) 13
(E.4)
2 / (R's Bi + Roc BR 1> ) (v(3) + p (c0
22 RT BR R BR..! )
C1+2(R2 B
+R BR X, and
Therefore, we have R(poo) S
Q(px) S R(pso) + ( / [1811 p(3, t)d3dt <
C1 2 BR
R B + RZ,) B2x) (E.5)
C2
which completes the first part of our proof.
Step III: Bound the difference between ((((T)) and Q(pcc) when 7 is large
Proposition 3.2 establishes that the second moment for p(T) is uniformly bounded across all 7 2 0:
S. 1101 (B, < Ag.
where Ao is defined as in Proposition 3.2 Therefore, the weak convergence of probability measures (p(T))z2 equivalent to the convergence in the Wasserstein-2 distance, i.e.
lim (z(p(T), , = 0.
(E.6)
When sufficiently large, pfr) concentrates on PR.-- Therefore, according to Lemma C.2 there exists a constant Co depending solely on N, D, Roo, and the parameters of the assumptions that
|| (H, To ()lt CoWa(p!7) poo) poo)
(E.7)
37

---

for any H and t E [0, when r is sufficiently large. Note that Lemma C.1 shows that
max p(r) (22,))l2-ool) pT (H2,))ll-coll I exp((((1 + Roo + R2o)) Br
for any H and t E [0, 1]. Thus, from E.7 we have
10(ps0) -()
>/ 1 [Read][Tpre p(r) [H 1T.1)|+TTppom (H. 1)I++ PT (H, - Tosrs (H, 1)1l
+ A
/ [1191P6SS (B,,
X Rox Wi ,oo)
<Br + B)COW2(D()) (Po) (oo)) 2
<((Br B)Co + A We ,oo):
(E.8)
where the second inequality incorporates the Kantorovich-Rubinstenn Theorem (see Theorem 5.10 of [68], for example) and the -L<ssLLiphhhtt continuity of ||s|| over the region (3,t) E PR
Combining equations (F.6 and (E.8) we deduce that for any > 0, there exists some To > 0 such
that (px))) Q(p ]|S
Step IV: Complete the proof by bounding the difference between ((0(r)) and ((p())) when
is large
The final step can be seen as direct corollary the approximation result in Theorem 3.1 According to Theorem 3.1 there exists constant C₁ dependent on N, D, To. X, and the parameters specified in the assumptions, such that
(((( - Q(p (T) Ci Ci(
( log(L
+ +++++)1) +1)
6
M with probability at least - 3 exx((-0 with respect to the parameter initialization ( for any > 0. Combining the outcomes from the preceding steps, we obtain
8 + log + +1) O2A.
SE
(L-1
(E.9)
Note that
M
(
ML dQ(())
M (00
((0())) = ΣΣ<
ML Σ
C s0,
dr
dBrj
dBes
5=1
(=0(()
(0=(())
j=1
the sequence (²(((T))) ))
-20 is non-decreasing. Hence, for any T 2 TO.
?( < (((())) s @(((m)) Se+ c
8 loggL 1)
+ C2A,
M
which completes the proof, recalling that Ci depends only on N. D, To, À and the parameters of the assumptions, and C2 depends only on N, D, Roc, and the parameters the assumptions.
Proof of Corollary 4.1 Given the choice of By Theorem 4.1 there exists some To such that
8 bog(L+1) + C2OsA
sup 5c/2 (1.1+ T2
M
V8+ log To + 1))
$11/4+ C2CAc+ C2CA)c+ Ci (I- Ci
M The result holds by setting and M/log sufficiently large to ensure that
C. c1 (1-
8 log L + 1)
(ci(L-1
2(1 + 8) log + 1)) E
M
M
38

---

Proofs of auxiliary results
F1 Proof of Lemma D.1
Proof. Lemma C.I
confirms that (1)l and (H, 2-oo are bounded uniformly by
Br exp( K r))). Considering the definition CC.2
it suffices to demonstrate that
V ((T,, (.((H.))))))))((.)) VeTT p (I
E
S + J2.
where
J1
V V (H. (H, (3) (.(H, ~ 118 - 311
J2
E,y V (OTT, (H, (H, +)) 2 exp(CoWi(p,u)) -1
Here, the symbol X hides dependencies on N, D, r, and the parameters of the assumptions. To bound JI, consider that
J1 < sup E [e[++]]
-TT,
Σ Il (H, t)
By ()) - 1), (H..
+1 sup
iE[N-1]
(F.1)
$ sup B (TT, -H),)),,)))) ie[N+1]
<OPT(T, (T, + o 311
SWi(p, + ||3 - B||-
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 (H, t) T f) by W1 (p, v) is achieved with Lemma
On the other hand, to bound J2, we have
Ja SVN+1 +1
sup By (H, (). (H, Pp(H, iE[N+1]
(F2)
(I,, (ll (H, pp (H, 1)
(H, - - (H,,
In F2 the third inequality relies on Assumption 2 (i). Consequently, to establish J2 A24 exp exp(Co (p, it is adequate demonstrate that (H, (H), S 11 12 2A24 Wi(p v), where
exp (f,
1 Read [Read,,, 1H, (H, 1)]l
// = =[Read[T, (H, 1)] x(H)
exp ( V vor (3) - exp ( V B))()) s))
From Lemma is trivial that (H, - T,HH [T,(H,1) (HH, x
Wi(p, v). Thus, I < V)( given the boundedness of Veee (H, t), 3)]] as provided
39

---

in Assumption 2 (iii). To bound 1z. we have
exp S..
J,
exp (f. v Troerryeegg(II s), 3))(( ss)ddss)
exp (.. / V
Vesr77redoo((1EE.333303.66..33 [.') voedrry B)]p(3 s)dBds) Ldinver(T)
Sexp (f. / |
vedti Bi]l M3, s)d
Sexp
(orr(N, Der ((T, -1
Sexp
(N D, Br, r) D, /N + (1-++ (1 -1
(F3)
for some constant C, dependent on the parameters listed in the result setting. Here, the first inequality
in F33) stems from [Read[T, 1)] - y {H)] Br + B, the second inequality is ensured by
the boundedness of V veccr] (T, t), 30]l as stated in Assumption 2 (iii. and the fourth
inequality is provided by Assumption ( (iv). The last inequality in (F.3) arises from Assumption 2
(iii) and Lemma C.2 By combining Equation F3) with the bounds of J1 and 71, deduce that
J1 + J2 S (xp(CaWi(p..))
1+ || 3|| for some constant CG dependent on the parameters listed in the result, thereby completing the proof.
F.2 Proof of Lemma D.2
Proof.
demonstrates that |iTo t) and ligy (H, 2)112-col are bounded by Br B (exp(K(1+A++ A for any H and i E [0, 1]. We begin by bounding
Lemma C.4
-(3,),,) - (1
T (6(H,1+41/22)
E 1,110011000000000000
E Tr (T M(T (H At/2), PO
2. TT pe (H. Pg (H,1))
To demonstrate that
Co (( θ). it suffices to show
J =
En -
60(11, 8 44/2) (1/1))+)1/2)11
< Co ML d(O,)),
(F.4)
and
Jz = Ep (HT,) T
(De (H, (T (t/ - +
< CG ML
d(o,0),
(F.5)
as the other part for h and w follows a similar proof approach.
To bound J1, by Assumption
have
N+1
J1
Vo I(T -(f 0).
li (H, At/2).
N+1
R
sup
Σ
lipe H + A1/2).all
iEiN+1
(F.6
3 S0r(r, BT) (111,,,)1220000 V N + De (H, At/2)]lF
Spr(1. Br) Br (,, A1/2)] ML ((,
ML d().
40

---

where the fourth inequality utilizes Lemma C.5
and the last inequality uses Lemma C.6
To bound J2, by Assumption 2 (ii), have
J2 SVN + 11,
sup iE[N++]
| (),
(), 0).s
[l A1/2) PH 41/2)]lr)
(FE7)
=VN 16 BT)( + F |1 H A1/2) Ps HH, A1/2)]lr]
Hence, it suffices to show that (H, + A1/2) - (H, A1/2)l S ML d(O, θ) to establish J2 < पI. d(e, O). Recalling the formula Pe (C.16), we have [1,, At/2) Py (11, At/2) F S I + I2, where
4 -Reed (o T (H,
II
(Tainve((T) (At/2)M
V voe [e
(0-t)/A-2E((1-00 JE[M]
M (Laimvoc'/7 + (At/2)M
II
Vwerrycch//bb///ss4--M///011l DNN+4+1.
V
(a-1) (1 (1tt/Ad) jE[M]
d(e,)) G)
ML
II
Idimvoe[TT) + (At/2)M
Σ
(o-1)At+2el 1/(1//1/22(((1//A [)/AA)] JE[M]
M
II
(Aiimvec/27 + (At/2)M
(s-1) /) (1- jE(M)
d(e, θ)
ML
exp (1A1/2)
Σ
M-1
V Veee)
+(Al/2)
Σ
M
V Veedty (H w.)))
(-1)/A+HH((
d(e,0) exp D,
ML
1
d(
ML
41

---

and
12 =|Readd (3 (HH, + v(H)
M
II JE[M]
(lamveetrr + (A1/2)M
V Veer)vec((76(4,3)00...]
==1
II
[Tdimvec|T] + (Ar/2)M
Vo w..)l)
(s (5-0/At+1efe [ jE[M]
II (1 (1----)1111)))11111 jE[M]
[diimvor/Ty + (At/2)M
V (e)
M [dttnveciyy + (
II
vec A/2), Wniill) )
(5- (1 -t)Ar] .E(M]
j
<(B + Br)
cxp (1At/2) (s (s-t)/AHtE((111 f)/A
Σ
M
max V (va (4
+(A1/2)
Σ
M
Σ max
V /eeer//e///aaaa.....////////////
(s-t)/ (i-t1/A+1(((--)/A))
M
(Att/2)
Σ
M-
(ee M
V veo[T] vec ( To.))
(a- (/-t/AAC A1 (1-))
JE(M]
M
M
veerr M
M-
5=1
<B + Br) exp
ML d(0.
"MI ())
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 I2 Z Mr. d(e, +). By combining the bounds of (F5) holds, thereby establishing |ā @, G, s
In and T2, we observe that Equation
Ca Certzzl d( θ) for some CG dependent on N, d, r, and the parameters of the assumptions.
It remains to prove that ||ā 0, G(, 6, t)ll < CG(1 + x) - B||. Note that
(G G(, <\|B 31|
+ { V (1 ( At/2)]
Do
Ex
h (H, (MT
- At/2), be(HH)))}}
Therefore, we only need to show that
E, VoTT (01r (
T To(H,,t+5//2)
< Coll0 Ō|,
and
E
CC
=|| -|1,
42

---

Here, we only establish the inequality 3|1. to obtain |G -(3,), I C](1+
above for and as the proof of the other inequality follows similar pattern. Note that by Assumption 3 (iii), we have
F(a 'o(H, + A
||
By VoTT
< sup Ey
(H, Vo (1I, - S(T | (H, At/2)]l
ve[N+1]
(ppp(r, (r, 0|| V ( 1t/2)lF
-110 - O1|
where the last inequality applies Lemma
C.6
Therefore, we conclude that (C(8,6,0
6, s Ca(1 + A)IB - B||. completing the proof.
F33 Proof of Lemma D33
Proof. Lemmas C.I and C.4 establish that max
[T, f)ll2-ool I By ==
B Bexp( (1 + ) According to Lemma D.6 there exists an event with P(E) 2 1- exp(-0) such that under E, we have
log(L L + 1)
lTe T,, SL
(F.8)
M
1. Following the same proof procedure as in Lemma D.6. with replaced by P. and bounding only J1 and J3 in the proof (as there is no need to utilize Hoeffding's 's inequality to bridge the difference due to finite width M), we could obtain the bound
for any H and
0, (L.
1)41,1 (T.
ITee (H, T2(H, t)le & LL
(F.9)
We present the proof only for the case involving p. The bounding of []G θ, - P. f) can be derived analogously by substituting P with ñ using Equation (F9), and skipping the process of bound- ing [ D2ll where D₁ and D2 will be defined later. The bounding of ||G(B, p, t) - G(B, p. t)ll can be straightforwardly achieved by combining the results obtained from the other two cases.
By the definitions of the gradients in Equations (G(3,,,t) - -(B,,,,)-G(B.)). p, t) < (,(0,,,t)) - Gy((, p.t), t)
we observe that + (w, O,t) - Gn(w, p.
and
t)|l
We will focus on showing that 2(,,, ()) i SL-1 8+ as the other part of the proof follows a similar approach.
M
Let's define the following quantities
A,, Ag, B1, B2, C1, C2:
A == V (H,
A2 = Vovee (Te (T) (H,,),,)
B1 Read To Te (H, vHH), B2 Read (H, y(H)
M
C (1 02/4/ JE[M]
II
(Talnved/11 + (AI/2)M /veepve(/(((((4.))0.....
(s
II
[Iaimvoc(T] + (At/2)M v
vo. At/2), At/2), wasll)} DN+d+1:
/ At St///A-eeliie (0---e((--)/1 jE[M]
C2: = :=pppp ( /
V
1)d3ai
DN+4+1;
The universal boundedness of |A] and || |2l implied by Assumption 2 (ii), while the universal
boundedness of B1| and [B is implied by Assumption |
Additionally, ||C|] and ||Call can be
43

---

bounded via Assumption 2 (iii,, and one can refer to the proofs of Lemmas
C.3 and C.6 for detailed
explanations.
From
C.14 and
C.16 we could rewrite J = - as
J || A2
( A2) [||[[] B2)
| | -2)1
A - 42 B2 || + || - Be + || | C21
A1 A2l + +1 B2ll + ||Cr Call.
We claim that to obtain the result, it suffices to show that
i. |A1 A2ll - 1 /0tlg(L++) under event E.
ii [B B2ll - ---114 +
6+1og(L+1) under event E.
iii. There exists some event Eg with P(E)) 2 1 - -exp(-o) such that |C -2l1 S --1 V 6+1og((+1) under E n Ez.
This is because if we can establish the above statements, then under the event EnE, with P(ENE)) 2
Given the similarity in proof for |2 0.t) (w. + we deduce that with probability at least
1 exp( 8), we obtain J 2 GG 1)) S I.-1 + 8+ (L
1-4 4 exp( 5) with respect the parameter initialization (0), we have [ G(3, x
-- + (+ogg(++1) The remainder of the proof focuses bounding the quantities in statements (ii-).
Proof for statement (i):
By Assumption
(iv). we have |A1
Azll T-1 + (Ing([+))) under event
Orr D. E.
T,(H,()]
Proof for statement (ii): Under event E, it obvious see | - Ball S
1-1 + (+1)
44

---

Proof for statement (iii): We further define the following quantities
M
D1 =
(Jasnvedtt1 + (At/2)M
V vedf21 ved/zy /eeppee//((6((Hs))....].
-0- ()/4 jE[M]
[Tāimve((T] + (At/2)M
Σ
vec|Ty voH At/2), wrs'l)
(e / At- (e(((-()/a)) JE[M]
[diinvee((] + (At/2) / Vverrveee//(((4sss)0))44554
D2 =
II
(s- (-t)/At- At-2//t-(((1 [)/A(]
[Tainve((]] (At/2) / ( ( At/2),, w (1
II
(e-t)/Al-EEl)
(A1/2) / / /reeppee///(oooHH.)))))0003303
D3 = exp
Σ
(5-t) (5-1/AA+2E(( (1-1)A1]
(A1/2) / Vverryce(((To0111:==//11
Σ
(4-0)ALHHE(11-)//)
D4 = exp (f.)) 1 vec[i vec[g(T, 1), (),
Note that |[0 - C2ll S D1 D2ll + D2 - Dzll + || - D1ll- Assumption (iii) indicates that for any 0, At (L 1) and j = M, we have
+ At/2), us 3 S B, for some constant B, dependent on the parameters listed in the result. This implies that each column of V Ve norm upper bounded by B, as well. Applying Hoeffding's 's inequality to each column of vecejry (To (, s), 0s.)] and V veel/yyveehhi ( + At/2), waj)l. and subsequently cal- culating the union bound across all columns yields:
V ( s), (z. VIV
[..]] Ve vec[ + At/2), ws,j)l has l2
M
P
M
V
vH 0,..))
(H
2 V(N [(N+1)De)
z²
<2(N
+
D)D exp(
M)
233
(F.10)
and
P(
M
V
(o wxyl (eveerrvoo(((((((TT.)) ()..,t)).s))))))) (s)s))((()))
(N 1)D2)
z²
<2(N 11D exp(
2B² .M)
(F.11)
for any 2 > For (F.10) and (III). we further the consider the union bound across all B,y 2M + log(2(N + 1)DL)), which implies that with probability at least 1 - exp( 55) with respect to the parameter initialization (0), we have
* = 0, Al, Al,..., (L 1) and
M
I|M
V
and
||M
veefo vech haaa::/.)....llll/ooooeeeohhaaa () (w..)/))44
45

---

bounded by B, (N + + 1)DL)) I CyV&tloot for any = At.t.. ... (L 1)A. Here, CJ is some constant that only depends on N, D,D and the parame-
M
ters of the assumptions, Denote this probability event by Eg. and we have P(B)) 2 1 - exxp(-)).
Under Ez, by Lemma C.8
have
1|D1 D2ll
M
S 2L
1
Σ
M
V vod Ve /ttp:///////aal/////eee//////t// Vwwernyecc/(iilllsss)))......
(s-1) (11////122((11/////11111 (1-1)/4j
M (e) w.s)l (eeeealll/eeeeeeettteeee Vreegvee(M(76(HHs+ (12))
(/(-)///+(((111 (1-))AA
jml
+ logoo + +1
SCJV
M
(F.12)
For any
0, At.
(L 1)4, define
Aa.j (At/2) / Vverryoll//o((445)))))88883
and
Baj = / V |
Since Assumption (iii) indicates that max || |A. 2 At L-1, we have
[ccp(AA..) - As,i S || 1?
|exp (B,j) - B.ill - |.,sll?
Applying Lemma C.8 once more, we have
DD DDll <
Σ
|
Σ
L & |Barll
(F.13)
(+=)/A44E((LL0)))
(e-t/Art1E(1 (11//
Since Assumption
(iii) ensures the boundedness of |aDa|, we have
DD3 - D4ll S
exp
Σ
(At/2) [reeppell//(c((Hss))))0d...
(G-))AA+22(p [t/AA
(A1/2) / V Vwerrpechh(///HHH:==)///
Σ
(a- -)- A110//4441((11//0
L.' | ved veclg (), ())p(3, ()dddt)) - -1/-
(F.14)
Therefore, to show that
[ Dall S T.-1 +
A+ M
suffices to show that
Σ (At/2) / (Ho.
Jz4
(8- (6-1)/A44((11-)//
Σ
(At/2) / Vwee At/2), w)]p(3,s)dp
(s- (==t)/A4tT((111)//
[.¹ Voo t), (Bll(33))dddd IL- /
8 + log +1)
M
By Assumption 3 (iv) and Lemma D.6 we have
+ log L + 1)
Vrecfzy (e Ve (1 --
M
46

---

and
veolT At/2), w)] Vveetrjoc(hh(Tpp s), wi)l (H,, A1/2) (T 8)1F
Te H, 21/2) To H ))(( + [T, ()lH,
8 + log L + 1)
<L
L |
h h(Te (H, 5), wa,
M
8 6+log(L ( +11
SL- +
M
where the last inequality employs Assumption (i). Therefore, we conclude that
8 log L + +1)
+ (((A//2) / / (ee.eeeee// 1). 0))p(3, ()8B
|13 Dall E Jo4 SL-1 +
M
(A1/2) / (F) s)d3
Σ
(0- (4-0)AAt1[(11-0/))
Σ (At/2) / Wwnerryee(hh(,,,444)),.... [
(a- (e-t)/A--E(1111///
veol 3))p(3 1)d3dt
8 log( I + 1)
sup
Veoor vec (H (1)))))))) - -aa.3)))vvvellll (H, (2)
M
81 11-821 SAt
5 + log [. 1) M
SL
sup
[Tp(H, s1) (2H
SAt
SL-1 Log(L+))
M
(F.15)
where the third inequality uses Assumption
3 (iv), and the last inequality relies the Lipschitz
continuity as demonstrated in Proposition C.I Combining F.12). (F.13) and (F.15) yields ||C] -
Call < D1 D2ll D2 Dall + D
Da $ I-1 + +oog(L+1 L+)
47

---

E4 Proof of Lemma D.4
Proof. Define () = So (1,0017(H,,,,)))008.3.0 for any P. E P2. From Taylor's
expansion, we have
voclT, (H,t) (,) Fp(Tp,: FF (T,::)) Fp(Tpn::)
V (H, t)d3)
(vee(TT, (H,t) (H, (H,,)))))
veolg(T (H. +
V (F (vec(T (H,t) (H,))))
(H,, ()(3
"
Voee17/eecgg(T, 3)]
p)(3)) -))(v=)((B)))d)) (vec[Zz,, (H,t) (H, + +(())
V (vec|TA,, (H,t) (H)
n 100 (3) o(1).
of which the last equality holds as Lemma C.2 shows that [T (H, T,(H, O(W2(Pm p)) O(n), where we hide the constant dependence on B, K, N.r. Therefore, we have
d d exp (-f² (weepvellg((.((4.))))))) (vec))) (H,, (o(H, {
di
eexp
(-- (- / Veclg 3). s)d3)
vvoofTp, (H, (o(H.t)) t)] (1 (. (vec/Te, (H,t) (H) }
=exp (- /
veolT vec
s) vo) (3,,t) + o())
(F.16)
which leads to [D.18)
F.5
Proof Lemma D.5
Proof. Fix (3,1) P,.. Lemma C.2
implies that
1H) [Roodd(T,HH,) Read 1)15 C,Wi(p, s (-(r+))l-=vl) v/1
for some constant C, dependent on the parameters listed the result. Our goal to regulate the difference between P. (H, and p. (H,,t) to control the the propagation of (H, Pr -)]-
48

---

Note that by CC.4 and Assumption 2
dt (H, (H t)lle - (H, - (2( F)
vHcclp /. V t), (3),
-veolp / Voeet) (8) 1)d3
S1po(H,1)) (1( / / Vewe 8]] ((,1)d33
+H,, |
- --)
Sknu (IIo(H)) - - p(3, (p. ((.,))l
SLn,1 (I),( - - p(3,t)dB L.2lp((( v1,, v(.,))112)
(F.17)
where
= 1
and
Lr,2 = BB+ (BBBBex(((L+rrtr)) exp (1N
The third inequality of (F.17) uses Lemma
obtain
[TT(H, S VN +
S VN (Bep)((((1++++++
and
[1,(,, S +
V IBNxDeep((((( (1++++2)))+
with Assumption 2 (ii) to bound the norm the Jacobian matrix with Ly, The last inequality of
employs Lemma C.3 to bound [H, with L Applying the Grönwall's inequality to we obtain
(H,()) - PH <0,(r+ + 1)++11)xppppp 1)exp(L,,))le - ull2 + /'
J. / 0(3, () v(:4)lldd <C,(r + e)exp((LL))) Lr, )lpp w/l + Lr, bn exp(((r)) S (1((-0)) -5:1)-(((- ()
=(C, + Ln exp( 1) 2/l
(F.18)
Since
I. ((t,t)))((.)))dde 10(3) ((B,t)ddtt 11p 2/11.
Thus, we complete the proof of the first result.
By Lemma C.3 under Assumption
I we have
(H, < T.. == (B+Bexx(K exp (Ør(N. D, /N +KK (
In addition, by Lemma
under Assumption 2 we have
(I 8)ll) s vN (Hl 8)1l2-o & L. == KB (Bexp((
49

---

Therefore, for the gradient function
XQ by Lemma C33 we have
5p
8Q 6p
8Q
(3.t) öp
(3,t) = (TO (o
-... (H, Pv (H(,
+ B
- T
T + L,,11 2 u/12
S v +
22o- + Lna(C, br 2111.
(F.19)
Hence, it suffices to show that for any Hl such that || ||2col I B, have (l (H, (H, 8)
g(T, H. t), B)l12-col I L 2/11 for some L,,5 > 0 in order to obtain the second result of this
lemma. By 2 (iii), we see that
9(1 o D, max | 34 (H,
<or(N, D, VN +10001/000X0|++++++++++++++++++0
SØT(N, D, vN IK +4,.://001..00/0(((11++++++1+++0
=/11,
(F.20)
where the second inequality again uses Lemma proof the second result.
C.2
Combining
F.19
and F.20 completes the
F.6 Proof of Lemma D.6
Proof. Denote the empirical distribution of Te and Tp by
ML Σ M
8(Be.: p= H (((t)p())t)
respectively. It's straightforward verify that p and 5 meet the conditions outlined in Lemma C.I. and Tp = To and Ta Tp- Hence, Lemmas C.1 and C.4 indicate that max (I)Te (H, I () and t [0, 1]. We then define By == B exp( (1 +r+)))).
S B Bexp( ((1 + for any H
The following decomposition equation holds:
[1, [1,,1, (HH (Te (11, t) (1,, [1,(1, T, (H,
J
J1
8
(F.1))
Our proof will bound J1, J2 and J3, possibly in probabilistic manner, to obtain the desired result. Bounding J1: Note that according Assumption (i),
M
f(Te t), 0r.
<1106(44))))-cco (Te, 0c.
M
(4A6 >
()l + 110 + (10
1=1
<B, (KA(1(((AA))))+++++++
<Br(1 +0 (K/2)(())(
50

---

Denote (1(1+ + by Br. Combining the two equations in 2.4 gives us To(H, + At) = MLP ww....,.,W... (Attne....... &MM
M
1 (T, (1/2) (A1/2)
Σ
I. (Te 21/2),
wrs) wr.))
==1
M +To(H,))+((112))))) Σ
M Σ
1 (To(H,1)
==1
+(A1/2)M Σ (H, j=1
(F.22)
Then, from the formula of (.22 and Assumption 2 (ii). we see that for any t = 0, ,4, 1)4,
[ At) Te 1 At)lE
<HITo(H, (H( (1)(/2))) Σ WN + 1|0, ||0ts]1))
M
+(At/2)M Σ
Φt (N,D, 1B7) + || + /| 113)
M
(H, + (A1/2)M
Σ
(H, - (H,,t))
<1TT((H,t) - AtM Σ
DN VN ||3,,|| ||8t,s 12))
M
+VX-119/((A//((0++++))))))) Σ
DN 1/ (1we,ol3)
<H (H. (N 1Br
+ NNN- (AAA0)0+0)0)///+)))++++++++++
(F.23)
Therefore, applying F.23 we deduce that for any 0, At.
(L - 1) 1, we have:
t) 2201 I V IB /))(+ exp(or (N, D. vN + IBr (1 +r
At
(F.24)
where C1 CL==AAAABBBBBK11+-7))ppppll..
IB))(1(I (1+
r))44
Bounding J2: For any = 0, At..... (L - 1)4, e IN + 1),5 E [M]. we have
(H, :ell s Br. Hence, by the Hoeffding's 's inequality, for any 2 > 0 we have
P(IM -1 Σ (),,),)),B))) Bes)) (c) / ((TT((H,)))) (.) 2) < 2 exp(
2²
M).
2B2
By the union bound over i IN + 1] and t 0. ..... (L-1)4, the above inequality implies
M ((sup J Σ
- (H, Br.a) - / ((T
2) 2(NN+))Lexp)
M).
2B
1=1
(F.25)
We let = = BT 2M 16
log((N + 1).J). Then, F25 turns into
M
(isup|/M- Σ
(TT((H,))) 8. 2311))))) Br 2r- log( (N + 1)L)) s exp(-5))
1=1
(F.26)
51

---

Denote the event such that
sup M Σ (T,,H))))))
(a00::/))8aa0a110000000000000010 I BrV2M log( + 1]]
by Es. F226 directly indicates P(Es) 2 exx((=).
Suppose that the high probability event Es occurs. Let's denote Br 2M by Bg for brevity. From Assumption 2 (ii), it follows that for any = 0, At.... (L 1)A.
2M 1- log((N + 1)L))
ITe (H, At) - THH, + ST - (H,,
M (l)) (.(
+AIIIM Σ
j=1
+ / (H,
(To(H.t) TH AtVN 1BB+
At S [Ilo(Te(H.. 11 81.))l
- (1 +
+AIV 1Bs.
(F:27)
Repeatedly applying Equation F.27 yields
VN + TBs (N D, N (1 +r+ exp(or(N, D, (Br) (r++++
(HH S
=C2VM-1(1 log((N + 1)..)..
(F.28)
for some constant C2 > 0 dependent on the parameters listed in the result.
Bounding J3: It's worth noting that the convergence proof with a convergence rate of O( = (())) for (H,,t), t), the first-order Euler method for (H, t), is non-standard. This departure from convention arises because do not assume the boundedness of the second-order derivative dHI,(H.))
dea
instead relying on the continuity of p(:,t) t) with respect to the depth index t. In this proof, O(.) hides dependencies N, D, and the parameters of the assumptions.
From the definition of T, and Tp. we have
|T, At) - (,,(H+++0))
- -
+ |To At) - -T,(H, - At / [ 00C.001.,,,,,00000003300
+ At - (.0.:/.......-000.......).111111 (100))13
(F.29)
52

---

To bound I1, we use
3.1 to get
++1r
1 51
(T,, (1)(HH.)))))(((1)33333 - Al 91
At sup +E|IIAA 1/,
((I,(H, 03,03144
<At sup +E(12+A1) 1
(F.30)
1f ((TT((H,), -(3))))31+
At sup SE(LI+A
<At sup
8 (H, 1), (, ()
1 (TT -(s1))) (E(L)AA
+ At
sup
Given that Proposition C.I establishes the Lipschitz continuity of (H, with respect to under the condition that P has a bounded support, we can conclude:
sup S (,,++1
s), 8) (TT, B))l S C3.1lt s|s C3.
(F.31)
for some constant C3.1 dependent on the parameters listed in the result. Furthermore, Lemma C.I and Proposition C. demonstrate that (H, t), 3) is both bounded and Lipschitz continuous with respect to 3. Thus, we have
sup sup 1/ʃ (CTT, p(31t))d31r== - sE(,,tAr
+E(t,t-AN) 11
((3,t))) I sup Cs.elp((, -p)s)))))a xE(t,t+At)
C03,CC,Lt
(F.32)
for some constant C3.2 > 0 dependent on the parameters listed in the result. Substituting (F.1 and F.32 into F.30), we find that there exists a constant C3,5 dependent on the parameters listed in the result such that I I At?.
Additionally, Assumption
2 (iii) implies that
12 < / -(T
(F.33)
SoT(N, D,
B) (H,,
Therefore, by bounding I, + I2, we obtain the following inequality (T)
B) S At (
(F.34)
which implies, after being used multiple times, that
(H, (H, s Ca (1 C,
& C3L-1
(F.35)
C2.5At
for any - 0, .......1 1)411 1 and some constant C3 dependent on the parameters listed in the result. Combining (F.24), (F28). and F335 yields the desired result.
E.7 Proof Lemma
Proof. Let ((() be the measure induced by TH(H,t) with H μ, and ç be ūa(t). By verifying Lemma C.1 we establish that [|T, t)22 < Br B Bexp(A +r x²)) for any H and
53

---

t [0, 1]. Consequently, supp(p(t)) C {T: ||T S Br} for any t E [0, 1]. The remainder of the proof involves four steps:
Step Show that Pp ()( 1) is Lipschitz continuous with respect to (H.t)
In the proof of this proposition, when referring to the Lipschitz continuity of a function, imply its Lipschitz continuity for !/ within the support of μ. Recall that we have shown in Lemma C.3 that Po( t) is universally bounded and Lipschitz continuous for ||-I| and any E [O, 1].
From the formula of Pp in (C.14). for any H, H', we have
1ppp (H,)) 1) Pp (H', 1) voolpp H, 8)] vec (H' :))]
Read Read[Tp(H, (Tp(H, ( -p - exp (f, / V veeclgy (), (3)
+
[Read[Tp(H) 1)] v(H))
exp (f.
exp (f. voodt veclo(TT s))
()),
(F.36)
we observe that (, Lipschitz continuous. Since v--) is K2 -ipsscic continuous, as given in Assumption (iii, we obtain that I X || H' |[.. Moreover, from Assumption and Lemma C.I. we see that /a is universally bounded. Therefore, to show that (H I) Lipschitz continuous for |F- suffices to demonstrate that I2 S and I4 S || H'
From Proposition C.1
From Assumption 2 (iii), we have
12 exp (..
V vee[T Veeet veclg(T, s)) p(9, s)dBds) S exp (OT(N, 1 (1)
Thus, dividing I by the uniformly bounded part Ig. we obtain
1 $ exp (f, (mmmm(vecii
(vecig(T, s),
V veder vocl(T B31 p(B3, s)d Ittivoolrr]
Sexp (f," 1 [11 veet21 (veclg(T, 5), 3)] vedgg B)) p(3, s)ddds)
sexp (OTr(N, VN
[Tp t) -T, (H(H', 1,
sup
(E(0.1]
(F.37)
where the final inequality holds by Assumption (iv).
By the Lipschitz continuity of (H,,)) for ||:1|r as stated in Proposition C. suprejo.yylIT,
we have
TT, ,)]l X || H' Hence, utilizing the universal boundedness of the last equation in (F.77). we derive 14 S || H'lle.
Considering all assertions regarding Jr, I,, Is and Ia. we conclude that PH is Cp-Lipschitz continuous with respect to H for some constant Cp dependent the parameters listed in the result that sufficiently large. The Lipschitz continuity of p,,, t) with respect to could be easily derived from the boundedness of the Jacobian matrix, as asserted in Assumption 2 (iii). Moreover, since pp (H, is universally bounded shown in Lemma C.3 we conclude that P, HT, () is Lipschitz continuous with respect (H, () for some universal Lipschitz constant Cn dependent the parameters listed in the result that is sufficiently large.
Step II: Prepare bounds related to Pa for later
C.14 implies that Pp solves the adjoint equation
-veclp, (H, / V
veeli, (H,
(F.3)
54

---

with
I// V
()d3 - (1+Cp/2)0T(N, D, 1 = Co
implied by Assumption 2 (iii). Therefore, the Grönwall's inequality directly indicates that
E2 (1,, O) 2 (1p 1)I 2 exp(_2Coo),, 1)] =(H)) 2 2exp(-2Co)((o))
E2 Ip H, 1) I exp(2Cot)E,.llp H,, 1)I cx 1)] y ( - cxx((2))
for any E [0, 1].
Step III: Construct the descent direction
By the pell-posednees of the ODE solution to (3.1] shown in Proposition + the solution map (H, invertible. Hence, for any p E pa, there exists continuous inverse map T, such that 7, (H,) for any H and t [0. Let's define the following function F(T) to approximate
/ (TT,9))(81))) h(T, wo)
F(T) = -o
The function F(T), arising from the composition of pp())) and T.). exhibits continuity T' E supp(ā) owing to the continuous nature of tp(-t) and TT- (-). Therefore, Assumption 4 (ii) could be applied to F. Since [(, 0) is universal kernel constrained on 0 E RRdime
K. [52). and
F(T) is continuous with respect to T, there exists sequence {(cK, )=22 C. (R pddim K)) such that
20
F(T)
Σ (T
s//3
(F.39)
max
and any T such that || 2-col I Br. Notably, since J(T, 0 05) -f(,, - 0%), we could assume without loss of generality that c for any k. Furthermore, there exists constant ke such that
given
F(T) Σ CT((T
S 2e/3.
(F.40)
We define C(c) = Σ 1 Ck, and vB3 E (R as the probability distribution such that, given D and for any k 0, 3 has probability ch /C((e) of being (2C(e)0), 02 0). Then, (F.40) transforms into
F(T) - / ((T, ())
S 2c/3
(F.41)
max
where
== =p))===p((
g(T, ()
From
F.4
claim that there exists some R(c) such that
v{B 1/B(e) < ||3|| S R(e))) 2 1/2.
1 gT, 3)B(3)d3 I E/3..
We are now the position to define the descent direction v. By defining v E ((R as the measure obtained by truncating any part outside 1/R(c) || R(c) from 7, and scaling the measure function by 1/0(18: 1/R(e) < |||||| < R(e)}) $ 2, we can establish that
F(T) / (T, ()
S.
(F.42)
for any T such that ||Tll-coll & Br. A straightforward deduction from F.42 is that for any T such that ||T||2-08 S. Br, we have [TT) Sg ((T, <N(N + 1)D. It is clear that has
55

---

a bounded support as {5 1/ I/)(c) & ||]| I R(e)). We will determine the value of e later, ensuring based only on N, D, r, and the parameters of the assumptions.
Step IV: Upper bound Jo o (6, t*) (=))))(())))) ((())) d3 to complete the proof
p(
Utilizing the gradient definition in 3.4 and So (T,, ())((1)))3)))) F(T) + Po(T we obtain
/ 8Q 6p
(8, (1B)
/ Tr (T, po(Hl, (1)) + A |181| ((()) [10)((((0) (e((((())p())) -(())t))))a 2
E
SBr-ao Tr g(T, (5(B)
((6) ((3)t)) PT (T),
> (R(c)+ +r)
Er (F(T) ( ((T)))) (07
ET-N(H) Tr TT (T), (T) (TT + (R(c) +r)
For p p2 concentrated on P.,. observe
R(p) - 1 Eul 2
(H)))) 1
(Br B)
Hence, bound J1, we have
J1 1BT~200) a(T,
(T))
S SV + +1)DeBBTTae Ip (T), ll
S(N + 1)De ET (BA- (T)) (T), 1/2
(F.4)
=(N + 1)DE (H,))) 1/2 dt
<Vē(N + +))D
S(N + 1)(B++++ B) + ((pp) CCyR(p)c,
where C3s (N + L)(Br B) 11DD
To bound J2, we have
J2 - EE (T), ()
EElllo (H,1)
EE exp(eepppp)pppeeeee (-2 P
(F.4)
22 R())
Combining (F:43) and (F44)), by choosing expp(-2Co)Ca, which only depends on N, D, and the parameters of the assumptions, we have
1 5Q (8, -((()) dB 6p
+ 2 [R( exp(-2C6)/C3) +r).
Setting B, R(} (xpccx(--Ca)/C3)) [R(} (20(l(2C)))))))))) +r) /2, and C2 (exp(-220) completes the proof.
56

---

F.8 Proof of Lemma C.1
Proof. By applying the Cauchy-Schward inequality, we trivially obtain So (31 A. Thus, leveraging (3. and Assumption 2 (1), we can infer
d dt (H, 20o & 11T, 1)
9 t)d
8
/
- 9
K(1 + |1
(F.45)
Therefore, by the Grönwall's inequality, we have
(T, ~ (H, cxp( 1.¹ / K(1 +1+112++ t)d3dt)
H 22
[Xp12- exp(K((1 A A?)).
F9 Proof of Lemma
Proof. As per Lemma C.I the boundedness of [H, 2) and [TT, (H, t)l12-oo1 is established
by a constant C := H col (x + A A))) for all E [0, 1]. Consequently, from (3.1
this implies
T,, 5
(1,,6)))22-o0
11
1/² | (1),
(F.46)
$ /n
5(1+
for any f1,t2 E Therefore, [/L///////////////CCC/CL+CCC/_// t for I1-12-coll and thus vN + + + Lipsschhtt with respect to for ||-]]r- Note that by (.TI.
with respect to
A(H,) 1) = ()) ()),
-11 '.
(H,
-11 -1l/" (H),(H.)),,)p(3 (), 1. J o(T.((H.)) (,)
(F.4))
[1,(T,,H,,,,, - (TT 5), sdd3ds
+ 11 J. (CTLLH a) u) s)ddds||
We then bound J1 and Jg using the following two lemmas separately. Firstly, since ||Tplle 1 VN + TC and |r & + we have by Assumption that
(I, (H, 8)
sup
V Vecl (T, (H,
IT PTTSS
(F.48)
57

---

Therefore, by F.48 we have
/. A(H, 8)ds.
J S N
(F.49)
Secondly, we aim to bound the integral J2 given Assumption 2 and (H, Conl E 1].
Again by Assumption 2 we have
N+1 IV (H,
| g s), B)ll
S VN 10p(C)(1+t) (F.50)
VX +1+))++++)))))
(F.51)
Since (H, 8) VN= + (2)(-+C,,2)K((++++(22 -iipschtzz with respect for |llle (as shown in (H46)) ) by (F51), , we obtain that (H, 8), 8), /N + C(1+CC
D, /N + (1+ + 2 --iisshhtt with respect s. Thus, (H,, is C'-Lipschitz with respect (s, B), (5/22)CCrrll D, /N [C)(1 r +2) /N + [Cp(C) r). This indicates, by the Kantorovich-Rubingg
where C' VN + 1(1
Theorem (see Theorem of 68 for example), that
(), J2 =111 J. ((T,((H,), (p-)) (3,))ddddl1r S Cp,
(F.52)
Define C max (T, NN +C)( + r)) By combining (.49 and F have
A(1, C* J. A(H.s)
(F:53)
Applying the Grönwall's incquality then shows
Specifically, we have AHH, S C exp(C*t)Wi(n.v))
(F:54)
S exp((")Wi(Aue)
F.10 Proof of Lemma
Proof. Let's consider a fixed (3,t) pair within P... Given Lemma we establish
(), S B Bexp(K(11 + A A²)). Consequently, under Assumption we deduce
(H, - /N + (H, 9)
S N + (1
(F.55)
s VN + 1K (
On the other hand, from (C.14. we have
(H, t)ll s (T, 1)) y (H)) ||0pp////p//// ( ( s), sdddss)
< (B+ (e A+ A2)) exp (f." / (ewerree 8) ((B.s)dBl
< (B + Bexp ((1 + A+ A3))) exp (f. / / |eveerrjeely((T, (( x)dBds)
< (B + Bexp K (1 A+ A3)) exp (f() /
+ [|(|() )p(3, s)d3ds
< (B + B Bexp ( A+ A3)) exp
D. IK (1) A+ A3))
(F.56)
Combining
F.55
and F.56 we could obtain
8Q 8p
(B, t)
SE, ll (T, 11
+ 1KB + A+
A 4 A+ 12)) + A222
exp DV,
58

---

E11
Proof of Lemma C.4
Proof. By employing the Cauchy-Schwarz inequality, we readily observe that 31. E, ΣM - |1 2 we ascertain that for any t 0, Sl,..., ,.,, (T.
A. Consequently, leveraging (2.4) and Assumption
1)At obtain:
[H, A1/2)]l2-col S (To 2ML
Σ
K(1 + 10, + 110,,1333
M
[ 2t)ll--col I [H, + 21/2)]l
Σ
K(1 + || |1 + -/ 113] }.
2ML
j=1
(F.57)
Therefore, by applying (157 multiple times, we obtain that for any = 0, (1.- 1,
M
M
|iTe ()l2-ooo S1T 0)1
K(1 + |0 + +10+
Σ
K - | || + +/u,,, 112)
2ML
2ML
j=1
K ML
H || 2-col exp
Σ
+ |1 + 113,ol1))
j=1
[(1+ + A+ A?)).
H 2-col exp
F.12 Proof of Lemma C.5
Proof. Lemma C.4 shows that [ 22 and (H,) are bounded by Br B (I + A A))) for any H and t E [0. From 2.4 for any = 0, O...., (L 1) At. from Assumption 2 (ii) and (iii) we have
(H,1+ A1/2)|I
A1/2
-Tg
-
M
M
<H - -AV + +rM Σ 10. Ōeill
(F.58)
(A./200NN. +r (1 T H,, 1)ll
(HH -
( Σ 100. 0..
for some constant C1 and C2 depending only N, D, and assumptions. Similarly, we have [TT At) - TT (H, At) |r (H) + At/2) - T3 (H, (1 Cr
H Σ ||w. Wz
(F.59)
Combining (F38 and F.59 we derive
M
[ (1, <
Σ
18r. Broll
(F.60)
where C3 a constant depending solely on N, D, r, and the parameters of the assumptions. Iterating (.60)) multiple times yields
(H,t) (H t)lr I exp(Ca) ML d(e Õ)
59

---

for any t = 0,
E.13 Proof of Lemma C.6
Proof. By verifying that Q satisfies the conditions outlined in Lemma C.4 we establish (H, t)ll I B exxp + A + A2)). The first two results stem from Assumption ii) and (ii), with recognition that ||]'| vN + 2 for any T E Dx(+1)) As for the third result, consider 0, At, (L- - 1), 1, where
max Pal H, f) IPe (H, t At/2)lla)
=maxf]| [ee [ t)]]l,
l/voolfa (H, [
At/2)]II)
S|Read H, 1)] v(H
II
Idimo + (A//////e//vve///0((((....
1etp///-e1(11--111 JE[M]
II
Idima +
,t/22), ws,
(s- (--)AAr+1e(1 A// JE[M]
Read [Te 1)] (H)
II
(1+ (1+ (At/2) (retz/vell/(T((l.,))....
(x-0//:+1E[[11))) ]E[M/]
(F.61)
+ (A1/2)
V vec A1/2), wans
<(B + Br)
M
exp ((At/2)
E
M
V veo [c
(1-0//44E((111)// [//
5=1
M
+(At/2)
Σ
M-1 1
Σ
V Veo voclh (H), w),
j=1
M ΣΣ
<(B + ex) ( (N
(1 + ||31| + 118113)
MT.
<(B + ΦT(NN, D,
+ A A
where the first inequality arises from the fact that the matrix 2-norm is greater equal than the norm of any of its columns, and the fourth inequality follows from Assumption 2 (iii)
E114 Proof of Lemma
Proof. Note that [| S r for any denoted as fgl1 (1, t)dB. Applying Hoeffding's inequality yields, for any > and t
(I- 1) and j E [M] with its expectation
Al, (I 1)
2²
P(M
11 8
1111112 2) < 2 exp(
M).
2p
By applying the union bound over t. 0, At.... (L 1) the inequality above implies
M
M
P(
Σ >
Brsll 1
I Σ/ / |111) 2) <Pisup| -1 Σ
11
/ p111)P((,,)))32 2)
ML
1=1
j=1
2²
<2L exp( M).
2,
(F.62)
60

---

In addition, we have
L Σ/ p(), / [911 ((3,
|t-e|5 , |1811 t)d3 [1311"((3,sdd)
sup /
21 sup [p((;)) p(;,
|0-s/SAr
SP'
(F.63)
Combining (F.62 and (F.3)) and setting = - rv - + log((L)) completes the proof.
F.15
Proof of Lemma C.8
Proof. The proof will be trivial by noting the equality
L
L
II
A
II
B
(a(A: Bi)
II A.). a=1+1
1=1
1=1
Hence, we have
L
II 1=1
A
Bd
Ball ||A,
Bell
II e=1+1
A.ll S
Σ
A:
Bill-
1=1
1=1
s=1
1=1
G Assumption verification for a concrete example
In this section, we consider
f(z,0) V Zsootmaax2 (Z" WZ)
(G.1)
with the collection of parameters vec.V W], where softmax denotes the column- wise softmax function. Morcover, consider
h(Z, w) = Ww.ubberizelRlLLUUWWW
(G.2)
with the collection of parameters vcc W2]. Here, HuberizedReLU denotes the entry-wise
HuberizedReLU activation function defined as
0, 2²/2, - 1/2,
S 0;
HuberizooReLU(()
2 E [0, 1]: if z 2 1.
Then, consider Transformer model defined equations 2.1 2.2 and (2.4) in the paper, where the functions f and are specified above. We suppose that this Transformer model is applied learning task with data that satisfies Assumption [ We have the following proposition.
Proposition G.I. Consider the Transformer model defined by equations 22.1
2.2 and (2.4 with f f) and h( Z, w) defined in (G.I) and (G.2) respectively. Then Assumptions all hold.
Proof. We omit the detailed derivations for the function A(Z, w). which corresponds to the MLP part, our verification of Assumptions and h(Z, w) satisfies Assumptions 2 and is relatively more intuitive, especially given the proofs for f(2, 0).
Denote Z (21,..., (N++1) E R D-N-1 Then the function can be rewritten as
f(Z,0) 7aoot 7 TTT) ((.
where N+1 P,y and P.. = softmax (2
(TT W2,). Next, we calculate the derivatives
of f(Z,0). with respect to Z and 0 as follows:
61

---

For Z: the Jacobian J E ((N+ Js J== (Jo)1s(JEN+1 where J.
(Z.0) E
8ej
DxD, After calculation, we obtain
J1p EZ W + Z W Piy,,
where Q: := diag((:))) P. P.:: Fag the matrix with zeros everywhere except one the (f, entry, and dig is the Kronecker delta (1 i = j. 0 otherwise).
For 0:: Define A, = Z T
W 2,. After calculation, we have
OP:
0Ai
Piçlóst Pu),
Zzi
(G.3)
AAu
Thus, by the chain rule, we have
N+1 Vw. (2, Σ ZriZziPiy(Oas - Pa)
(G.4)
Moreover, we have
N+1
V (2,
Σ j=1
Po(e:.....2.)T
(G.5)
where (2) ....,. 2
contains D copies of zj. We then verify the assumptions one by one.
For Assumption 2 we have
|TT'oomma WT) 22-
V|[ ||7 ||softma(TT' 12)]l2-an 2--ool
<|00|| [T]]2-ea) ||softmax(TT WT)
< |T]|2-col
where the second-to-the last inequality follows by the fact that 22-norm can be upper bounded by the fioorm, and the last inequality follows by the fact that each column of the softmax output has an fi-norm equaling one. Therefore, the first condition in Assumption with K = is verified for the function f in I.
For h in G.2 we have [[<[T, w) 2oooll=
col = [yytmbbrrieedkeLL[(([W]]] S
W2ll2 [[HuberizeeRRLLL(WWW11]]
W2ll2 [[
| 11:1 1|]|2
where the second inequality follows by the property of HuberizedReLU that [HuberizodReI.(()))) |æ]. This demonstrates that Assumption with = 1 holds for h in G.2 as well.
For Assumption
2
(ii), G.5
leads to
N+1 Σ
N+1
vee| (VT
P. T.
1 < Σ f=1
Ti S [T| 2-col-
1=1
Moreover,
G.4 leads to
N+1
5 Σ
Z - Pa)
IEkIED #=1
N+1
2 Σ
Σ
Σ
max 15
ISkIED
<2
max
T |V2
18
ISkIDD
<2 2
62

---

Combining the two equations above gives Assumption 2 (ii) with (P((TT]22cc)) = 1|21 2|[T||2-00
For Assumption (iii), we have
| -
Σ
1 (N +1)
max
||5g|12..
1SAJSN+1
ISAJEN+1
For any S i S No + 1, we have
MM5ll2
+ 11770//11+ W +1 W Totyll22/VV
<|V1 [ |+ 11+
~ 21
<(1 2(1+2177)) + |01| + 11013)
Hence, we have
| S 2NITT (N 1)(1 21 ||0||2 + (101)))
The above equation demonstrates that for Assumption 2 (iii) holds with Φr D, [7]rr) (N + 11 + 21 We have verified Assumption 2 verification for is similar and easier.
for the attention layer encoder f. The
Given that we are currently considering the example where the encoder employs smooth univariate activation function, we can prove stronger results by removing the expectation Ep.
Next, we verify Assumption 3
(i) and (ii)): Given the calculation of derivatives in G.4 and G.5 show that
we have presented above, we first
P.j softmax 7 W
is locally Lipschitz continuous with respect to and 0. By that
G.3 and the chain rule, we can derive
N+1
ap.
P. = Σ Zaa Zi:Pis
:P));
[[[Eaa] W + zJ W TTy).
f=1
and the local Lipschitz continuity is then obvious given the boundedness of V V Pi and B2g
0R; with
respect to necessary parameters.
As we prove that Piş 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 6, and obviously bounded by an increasing function of N, D. ||@]|, Kr, Lr. 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 || - -- is bounded
by v Eisissn4 1-] - L3 11. it suffices to show that Jij is locally Lipschitz continuous with respect both / and Z. Again each component of Jij that depends on or 0, i.e. Z, Q. W, P.ys bounded by an increasing function of N, D, Kp, LT, Kr, ||0]|, and is locally Lipschitz continuous given sufficient smoothness. Hence, (ii) and (iv) also hold.
For Assumption 4 we consider the pair (g, a) = (h, w). and the partition (a) 0g) with 01 W2, a2 WI.
i. We also let compact the partial 1 h-homogeneity property straightforwardly holds:
(W, || < 1}. Then Assumption (i) on
(T, Wi.c W2) c W.HHubeizeeRRLL((WW.. c h(T, W1, W2).
Regarding Assumption (ii) on the universal kernel property, first note that according to the choice (g. a) = (h, w) this assumption is purely an assumption on the MLP part of the Transformer. Here 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 [8 [8 57 for more details), we know that two-layer fully-connected
63

---

networks with non-polynomial activation functions and without any constraints on its parameters are universal approximates.
Therefore, know that the function class span (oan(WW Wi gdii(W)), W2 E [ddm((2)) } is dense in C(I < B, D. Moreover, by the definition of Huberize- dReLUU for any > and any W1. W2, there exist small constant e such that c W1 K. c. [ 11 S B B= and
c² 2 We WyHbberizedRelUfe WIT) = (e (e-WT) c c W
-W>ReLU3
where the second equation follows by the positive 2-homogeneity of ReLU activation. This implies that
(W2ReLU' (WiT) Wi Rdim(WW)) W2 E Reim(Wa)) S (WaHyyerivveRRaaUUWW.1 [
W2 Rddim(W)) x xXC].
Therefore, we conclude that span Wy..HberzeRRRLLUUWWW (([ & B, (x(+)))) This finishes the validation of Assumption
W2 E Rdim(Wa) KC} is dense in
H Experiments
As discussed in Sections and 4 mean-field approximation results and global convergence results are asymptotic guarantees requiring exponentially large number of heads M and number of layers L. Such results are due to the nature of mean-fiell type analysis. In practice, we frequently observe that global convergence can be achieved by Transformer models of reasonable sizes. In this section, run simple experiments training Vision Transformers (ViT) 24 the CIFAR-10 datasets to demonstrate global convergence in practical applications.
We train Vision Transformers with different numbers of heads and layers. In all our experiments, we split each CIFAR-I0 image into four patches and then pass the patches into Vision Transformer models. We keep the dimension of each attention head to be 128. The output of each self self-attention layer is passed through a
feedforward component with 128 hidden neurons and GeLU activation. Both the self-attention and feedforward components include skip connections. We implement dropout in the self-attention layers as well as the feedforward layers with a dropout probability of 0.1. The model is attached to a linear classifier.
In all experiments, we train the VIT models using Adam for 200 epochs with a mini-batee size 512. We set the initial learning rate to be le 4, and implement a cosine annealing learning rate schedule. We do not use any data augmentation or explicit regularization techniques, so that global convergence for large enough models implies close-to-zero training loss and close to 100% training accuracy.
In the first set of experiments, we fix the depth of the VIT to 6 layers (i.e., there six self-attention feedforward component). We train such Vision
layers, each followed by hinle-hidden--aye Transformers with the numbers of heads per layer ranging from 4 40, and record the training loss and training accuracy throughout training. The results are given in Figure [ Based on the results, it clear that for ViT models with more than 20 heads can achieve close-to-zerd training loss and close to 100% training accuracy, demonstrating global convergence on the CIFAR- 10 training data.
In the second set of experiments, we fix number of heads in the ViT model per layer to 8. We train such Vision Transformers with depths ranging from 2 to 20, and record the training loss and training accuracy throughout training. The results are given in Figure Again, the results indicate that ViT models with more than 16 layers can achieve close-to-zero training loss and close to 100% training accuracy on the CIFAR-I0 10 dataset, implying global convergence.
We note that all these experiments are conducted on standard GPU card. We can observe clear global convergence when the Vision Transformer is sufficiently wide deep, but still within reasonable scales. This indicates that, although our theoretical guarantees require extremely large numbers of heads and layers due to the limitations of the mean-field technical tools, global convergence can be achieved by Transformers of reasonable sizes in practice.
64

---

Reads, layers
12 heads, 6 leyers
90
layers
80
24 heads, 6 leyers
Ursiing 65
70
layons
heads, layers
se heads, n layers
layers
heads,
training
12 hewds, layers
16 heads, 6 laayers
heads, layars
layers
28 heads, 6 layers
30
layors
40 heads. layers
20
23
50
100
125
150
175
200
25
50
100
125
350
115
200
spccn
(a) Training loss
(b) Training accuracy
Figure 1: Training loss and training accuracy of Vision Transformers with different numbers of heads. (a) gives the curves of training loss, while (b) gives the curves of training accuracy.
hhaads
100
2.
hhadds layer
Sheods h lyyes
90
heads,
layors
heads, 12 layers
traitin bos
heeds,
training accurssy
layers
8 heads, layers
2.0
heads, layers
heads, 18 layars
60
layers
layers
50
heeds, layers
heads, layors
8 heeds,, 14 layers heads,
heads, layers heads, 20 layers
23
50
100
125
150
175
200
25
50
100
125
:5a
115
200
opoch
opcch
(a) Training loss
(b) Training accuracy
Figure 2: Training loss and training accuracy of Vision Transformers with different depths. (a) gives the curves of training loss, while (b) gives the curves of training accuracy.
65

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: The abstract and introduction accurately summarize the main contributions and scope of the paper (Section L.I.. clearly outlining the theoretical advancements while aligning with the stated assumptions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We operate under noiseless label settings, noisy conditions typically preclude achieving zero risk. The limitations Assumption and the assumptions in Theorem 4.] are discussed in the respective sections. Although some assumptions currently untestable, we provide high-level justifications for their validity.
jugh-leveljusúfficttonn
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
66

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All necessary assumptions (Assumptions 14 are stated the main content prior presenting the theoretical results, with complete and correct proofs included in the supplementary material.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes)
Justification:
We clearly disclose all the information for reproducing experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g.. in the of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution a new model (e.g.. a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
67

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification:
This paper includes only simple experiments in the supplemental material to verify the main theoretical results. Therefore, the access to the data and code is not central to the paper's contribution.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (kttpss//nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ ltttlssccc/ubblc/guiiceee Setguddee/oodeeaeeeee wwaibbmssssooooooo ssionol for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: This paper specifies all necessary details understand the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification:
This paper reports appropriate information about the statistical significance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
68

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors 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 out of range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
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 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 estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: The research presented in this paper fully complies with the NeurIPS Code Ethics. All ethical guidelines pertinent the theoretical nature of the work have been adhered to rigorously.
Guidelines:
The answer NA means that the authors have 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., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
69

---

The paper does not discuss societal impacts as it strictly addresses theoretical optimization guarantees for Transformer models, focusing solely on providing rigorous proofs without exploring practical applications.
Justification:
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: The paper is theoretical and does not involve the release of data or models. Therefore, there no risks for misuse dual-use associated with this research.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging. and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The public CIFAR-10 datasets are explicitly mentioned and properly respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
70

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperssi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should 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 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper is theoretical and does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper is theoretical and does not involve crowdsourcing nor research with human subjects.
71

---

Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
72


Reviews:
[{"review_id": "RLdgldlxvs", "reviewer": "Reviewer_iBME", "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.", "questions": "1. Please address the above concern for Corollary 4.1 carefully. The answer may greatly affect the final rating from the reviewer.\n\n2. Questions related to assumption 1-4:\n    1. For the purpose of confirmation, have you shown in which cases Assumption 1-4 are met for the Transformer model? My current understanding is that they would be met when you reparameterize $f$ as in (2.3) but not when you keep the W_V, W_K, W_Q and W_O matrices (at least Assumption 2 won't be satisfied in this case). Is that correct? The review is not against the reparameterization and considers this is still a Transformer model, but want to understand the difference and the claim of the paper more carefully. \n    2. As the focus of the paper is on the Transformer model, the main results are based on the structure defined in (2.1) (2.2), (2.4) and Assumptions 1-4 over f and h. The review wonders if it is possible to provide a separate corollary for a commonly used form that satisfy Assumption 1-4, such as (2.3), making it clearer on an example where minimum assumptions are used. Or in other words, are there technical difficulties in showing Assumption 1-4 for (2.3)? Is the universal approximation capabilities the key obstacle?", "limitations": "The authors have adequately addressed the limitation of the paper.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Strength:\n\n1. The paper studies the Transformer model in the context of in-context learning. Although the mean-field approach has been applied to infinite width and depth ResNet models before, the Transformer model is currently the most widely used and thus worthy of detailed study. This paper extends the analysis by incorporating two distinct encoders, filling the gap in understanding the Transformer model in infinite limits.\n\n2. The paper introduces new techniques by extending previous approaches used for ResNet models. The authors refine the analysis by assuming only partial homogeneity and local Lipschitz smoothness, which is a key extension. The idea of considering the continuous dynamics as a function of the average behavior of f and h is new and interesting.\n\n3. The final result has a mild requirement for the dependence of depth and width, where one would only need depth $L = \\Theta(\\epsilon^{-1})$ and the number of Transformer block $M = \\Theta(\\epsilon^{-2} \\log (\\epsilon^{-1}))$ to achieve a loss as small as $\\epsilon$.  (However, the reviewer has some questions about the result, so this strength can still be questionable.)\n\n4. The paper is well-organized and clearly written. The authors effectively communicate the ideas, contributions, approaches, and limitations of their work, making it accessible and comprehensible to the reader.", "weaknesses": "## Potential fallacy in the claim of Corollary 4.1:\n\n1. The reviewer is convinced by most results presented in the paper, but has concerns about the claim in Corollary 4.1. The main issue arises from the fact that the authors treat many objects as constants, even though they clearly depend on some key quantities. In short, the reviewer thinks that $C_1$ in Theorem 4.1 depends significantly on $\\epsilon$, which makes the claim in Corollary 4.1 questionable.  \nThe concern stems from the observation that $C$ in Theorem 3.1 may depend exponentially on $\\tau$. Following the proof of Theorem 4.1, to achieve a loss as small as $\\epsilon$, one would need $Q(\\rho^{(\\tau_0)}) \\le \\epsilon$. This introduces a non-trivial dependence of $\\tau_0$ on $\\epsilon$. Given that $C$ may depend exponentially on $\\tau_0$, $C_1$ in Theorem 4.1 would thus depend exponentially on $\\tau_0$, leading to a complex dependence on $\\epsilon$. This dependence cannot be ignored, and treating $C_1$ as a simple constant to achieve Corollary 4.1 appears incorrect. (The reviewer believes this is how Corollary 4.1 is currently derived.)  \nBased on the above, the reviewer thinks that \n    1. The current approach in proving Corollary 4.1 is flawed.\n    2. A naive application of Theorem 4.1 might only lead to asymptotic bounds for $L$ and $M$, instead of the linear and quadratic terms presented nicely in the current version of Corollary 4.1.\n\n## Limitations of settings and results\n\nThough the reviewer agrees that studying and achieving a good understanding of the case where both width and depth go to infinity is valuable, there are some limitations in the results presented in the paper.\n\n1. The main result, Theorem 4.1, depends on Assumptions 1-4. Although the reviewer considers Assumptions 1-4 to be reasonable and likely to hold under certain settings, the authors do not make it clear when these assumptions would hold. Consequently, the setting of the Transformer model might differ from classical Transformer models (e.g., it might only hold under the reparameterized case (2.3) instead of keeping all value, key, query, and output matrices). The reviewer will address this point more in the Questions section.\n\n2. The main result, Theorem 4.1, relies heavily on a few additional assumptions that are not justified in the paper:\n    1. the Wasserstein gradient flow weakly converges to some distribution $\\rho_{\\infty}$\n    2. the uniform boundedness of $\\rho^{(\\tau)}$ for large time $\\tau$\n    3. the separation property for $\\alpha_1$ with the support expansion of $\\alpha_2$ to $K$.\n\nThe reviewer acknowledges that some of these assumptions are also adopted in other mean-field approaches and appreciates the authors' efforts to justify them intuitively in Section C.5. However, they are not examined carefully and may harm the validity of the statement. For example, for the second assumption, the reviewer agrees that a large regularization $\\lambda$ can implicitly bound the norm of the solution. However, this might require a large $\\lambda$, which could make the results less meaningful. Notice that the current bound shown in Proposition 3.2 implies that $R_{\\tau}$ may depend exponentially on $\\tau$, thus it is not clear that a small $\\lambda$ can result in a bounded solution. And based on Theorem 4.1, a small $\\lambda$ would be required to achieve a small loss.\n\n3. The result does not provide a non-asymptotic bound for the required training time. In other words, there is no guarantee that $\\hat{Q}(\\Theta^{(\\tau)})$ will be small for any specific $\\tau$. This is because there is no control over the convergence rate of $\\rho^{(\\tau)}$ to $\\rho^{\\infty}$, which would require a stronger assumption beyond the current weak convergence assumption of $\\rho^{(\\tau)}$.\n\n4. The loss and training dynamics are only considered for the population risk under $\\mu$. There is no consideration of the effect of an empirical dataset. This is a clear limitation, but the reviewer can accept this as an initial step toward studying the Transformer model. \n\n## Minor:\n\n1.  Some parts of the statements are not well organized. For example, Assumption 2 is an assumption about norm bounds of function $f$ and its gradient. The line \"define ReLU'(x) = ...\" has no direct connection with the assumption. The reviewer thinks that the authors are trying to argue that by using ReLU'(x) as the gradient of the ReLU function, $f$ defined in (2.3) would satisfy Assumption 2. If this is the case, the reviewer suggests the authors separate the statement into 1) what the assumption of $f$ is. and 2) when the assumption would hold. Another minor point is that the Fréchet derivative still acts differently from the gradient, so some careful treatment might be needed if the authors claim that the result in Theorem 4.1 applies to Transformers with ReLU activation. \n\nTypos:\n1. Line 111, it should be \"h(Z, w) = W_2 \\sigma_M(W_1 Z)\" instead of $W_1 H$. \n2. Line 156, in Assumption 2, $\\Psi_j$ in the description should be $\\Phi_P$. \n3. Line 86: \"~~the~~ our deep Transformer model\"\n4. Line 901, For the line that bounds $R(\\Theta)$, there is no $C_{\\lambda}$ in the first line, and the transition from the first line to the second line is currently shown as $1/2 \\epsilon + ... \\le 1/4 \\epsilon + ...$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your constructive and helpful comments and suggestions. We address your concerns as follows.\n\n**Q1**: Corollary 4.1 only implies asymptotic bounds for $L$ and $M$, instead of the linear and quadratic terms in the current version of Corollary 4.1.\n\n**A1**: Thank you for the crucial clarification on this point. You are correct, and we will revise Corollary 4.1 and remove the comment on how $L_0$ and $K_0$ depend on $\\epsilon$. Our result only ensures that the risk can be bounded when ​​$L$ and $M$ are sufficiently large. \n\nOur work does not aim to provide a tight bound on $L$ and $M$ with respect to $\\tau_0$ or measure the exact reduction in risk. Instead, our goal is to take the first theoretical step in demonstrating that large-scale Transformers can achieve global convergence via gradient descent-based optimization. Even 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.\n\n**Q2**: Assumptions 1-4 may rely heavily on the reparameterized form of transformer as in (2.3). The reviewer is not against the reparameterization and considers this is still a Transformer model, but wants to understand the difference and the claim of the paper more carefully. Additionally, is it possible to provide a separate corollary for a commonly used form that satisfies Assumption 1-4, such as (2.3)?\n\n**A2**: Thank you for your suggestion. We indeed mainly focus on the reparameterized form of transformer in (2.3). However, given the theoretical nature of our work, we believe that the form of transformer in (2.3) covered in our paper is already pretty close to the practice, compared with most existing theoretical analyses of transformers.\n\nWe believe that Assumptions 1-4 are satisfied when we consider:\n- The form of (2.3) for the attention layer, with $\\sigma_A$ being the column-wise softmax.\n- $W_2\\sigma_M(W_1H)$ for the MLP layer with the ReLU activation or smooth activations with Lipschitz derivatives.\n\nThe first two assumptions are easier to verify by definition. Assumption 3 is straightforward to verify when all activation functions are smooth. For ReLU activation, Assumption 3 can still hold when the data are generated from a smooth distribution – all the properties in Assumptions 3 are in the form of expectations over the data distribution, and such an expectation can have a smoothing effect.\n\nFor Assumption 4, the partial 1-homogeneity is clearly met in our paper given the partially linear structure of Transformers. Furthermore, since the universal kernel property applies to either the attention or the MLP layer, we can choose the MLP layer with the ReLU activation, whose universal approximation ability has been discussed in [12].\n\nWe will add more discussions in the revision.\n\n**Q3**: Theorem 4.1 relies heavily on a few additional assumptions that are not examined carefully and may harm the validity of the statement. It is not clear that a small $\\lambda$ can result in a bounded solution. And based on Theorem 4.1, a small $\\lambda$ would be required to achieve a small loss.\n\n**A3**: Thank you for acknowledging the common adoption of these assumptions. We agree that a small $\\lambda$ may not always result in a bounded solution. We treat the assumption of “the uniform boundedness of $\\rho^{(\\tau)}$ for a long time $\\tau$” as a data-related assumption. Our heuristic thinking is that if the true data distribution $\\mu$ is simple, and the true label generation process $\\rho^*$ is bounded, then a bounded solution suffices to minimize the loss for such a simple learning task.\nLikewise, the first and third assumptions are also assumptions on the learning task regarding $\\mu$ and $\\rho^*$. We believe they should be satisfied when the task is sufficiently simple. However, like most other papers considering mean-field analysis across multiple network layers, we acknowledge that it is challenging to construct a general class that could be verified to meet these assumptions.\n\n**Q4**: The result does not provide a non-asymptotic bound for the required training time.\n\n**A4**: Thank you for your comment. Similar to our response in Q1, we believe that even establishing an asymptotic convergence bound is a novel and significant finding given the technical complexities of the Transformer models. Establishing more concrete convergence rates is an important future work direction.\n\n**Q5**: The loss and training dynamics are only considered for the population risk under $\\mu$, but the reviewer can accept this as an initial step toward studying the Transformer model.\n\n**A5**: We thank the reviewer for the careful reading and understanding. We consider the extension to finite sample results with explicit generalization error as the next step of our research.\n\n**Q6**: In Assumption 2, the line \"define ReLU'(x) = ...\" has no direct connection with the assumption. The authors should separate the statement into \n1) what the assumption of $f$ is.\n2) when the assumption would hold. The Fréchet derivative still acts differently from the gradient, so some careful treatment might be needed.\n\n**A6**: We thank the reviewer for the careful reading and the excellent suggestions. We indeed aimed to claim that by using $\\text{ReLU}'(x)$ as the gradient of the ReLU function, $f$ defined in (2.3) would satisfy Assumption 2. We will revise the statement following your suggestions.\n\nWe believe the formula of the Fréchet derivative holds under ReLU activation (Equation (3.4) and Proposition 3.1), as the Fréchet derivative is more akin to the weak derivative and milder than the standard definition of the derivative. We will add a detailed discussion in the revision.\n\n**Q7**: Some typos.\n\n**A7**: We thank the reviewer again for pointing out the typos, especially with the careful reading of the proof. Line 901 indeed contains an error in the constant term. We will fix them all in the revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GttjalahUV", "reviewer": "Reviewer_oct4", "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.\nThen, they show that the gradient flow in this limit is well posed, by using the notion of Wasserstein gradient flows since the mean field model is defined on space of measures. They also show consistency with the original discrete transformer model.\nThen, they study the global convergence properties of the gradient flow: they prove a sort of global convergence result, which requires some hypotheses.", "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?\n\nThe justification of the hypothesis on the separability property in C5 is difficult to understand. Can you elaborate more on this assumption and try to report precisely on the progress of the literature on this assumption?\n\nI do think it is not fair to claim global convergence for such kind of \"if\" theorem with a hypothesis that is uncheckable. I suggest the authors to rephrase their contributions.", "limitations": "Not applicable.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The writing of the paper is quite clear.\n\nProving mean field convergence results has been studied a lot in the literature on theoretical deep learning. The result of the authors is interesting because they deal with the transformer architecture which ends up being a sum of two maps in the mean field representation. The authors can prove similar results in this context, which are new to the best of my knowledge of the literature.", "weaknesses": "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. \n\nThe novelty of the results is the fact that transformer architecture is taken into account in the mean-field limit. However, the stated results do not show to what extent this architecture helps in obtaining the results. In fact, the transformer architecture appears more as something that hinders the standard proof on global convergence. \nThe paper does not shed light on the interest/peculiarity of the transformer architecture.\nI would have appreciated if some simple experiments could be done to assess the range of validity of the results and in particular testing the importance of the transformer architecture.\n\nThe hypothesis in section 4 to obtain a global convergence result seems not checkable in practice. In my opinion, it is not a global convergence result but rather an \"if\" theorem. Although I know that this assumption has been put forward in other previous papers, it is still an very demanding hypothesis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable suggestions. In the following, we give point-by-point responses to your questions.\n\n**Q1**: One of the main weaknesses for a submission to NeurIPS is the format. The main paper is devoted to a (well done) explanation of main results, but insufficient explanations of the proof ideas are given in the main paper. \n\n**A1**: Thank you for your comment. Since our work is novel in the context of Transformer models, we believe comprehensive explanations of the main results are necessary to clearly convey the key theoretical messages to readers. Consequently, there is limited space for including the concrete ideas of proofs due to their highly technical nature.\n\nWe have made efforts to address the lack of proof sketches. In the appendix, we included proof ideas for the most important results before presenting the rigorous proofs. For example, we discussed the idea of obtaining the adjoint ODE solution of $p_\\rho$ in Appendix C.4, and the proof idea of Theorem 3.1 at the beginning of Appendix D. For the proofs of the main results, we also explicitly listed the goal of each proof step to enhance understanding.\n\nWe agree that it would be helpful to explain the ideas of proof in the main paper. If this paper is accepted, we will revise the paper structure and also utilize the one additional page for the camera-ready version to add a proof sketch.\n\n**Q2**: The novelty of the results is to take the transformer architecture into account in the mean-field limit. However, the results do not show how this architecture helps in obtaining the results. Can you extend your results to more general architecture than transformers? In addition, some simple experiments could be helpful to assess the range of validity of the results and in particular testing the importance of the transformer architecture.\n\n**A2**: We thank the reviewer for the insightful question. The purpose of our paper is to enable mean-field analysis on Transformers, so that  global convergence guarantees can be established for Transformers. The complexity of the Transformer structure makes achieving such convergence results particularly challenging, and existing mean-field tools cannot be directly applied. Therefore, our paper's contribution lies in overcoming these difficulties. \n\nWhile our focus is on Transformers, our result can indeed cover (or be easily extended to cover) more general architectures. Importantly, Assumption 4 only requires partial 1-homogeneity. This can potentially enable mean-field analysis of very general architectures and activation functions\n\nRegarding experiments, in practice, we observe that Vision Transformers can fit simple datasets such as MNIST and achieve a near-zero training loss, indicating convergence to a global minimum. We will include experimental results on validating the mean-field limit of transformers in the camera-ready.\n\n**Q3**: Assumptions in Theorem 4.1 are not checkable in practice. It is rather an \"if\" Theorem. Although these assumptions appear in other previous papers, they are still demanding.\n\n**A3**: We thank the reviewer for the question and the acknowledgement about the common issue in literature. Indeed, many other papers in the mean-field literature make global convergence claims using similar assumptions. Although our paper retains these untestable assumptions to prove global convergence, we emphasize that we have already made significant contributions towards for practical assumptions by:\n1. Weakening these “if” assumptions compared to [41], as we do not require full homogeneity.\n2. Weakening Assumptions 1-3 compared to [19, 20, 41], whose assumptions cannot be extended to Transformer structures.\nTherefore, while validity concerns about the “if” Theorem still exist, we believe that weakening these assumptions to a milder version is a significant step towards the ultimate goal of achieving “checkable” assumptions.\n\n**Q4**: Discussions about the hypothesis on the separability property in Appendix C.5 are difficult to understand. Can you elaborate more, and report precisely on the progress of the literature on this assumption?\n\n**A4**: The first assumption in Theorem 4.1 requires that the parameter distribution $\\rho^{(\\tau)}$ is concentrated in a bounded region across all time. Though this assumption is uncheckable, the introduction of the regularization parameter $\\lambda$ penalizes the parameter norms, which may lead to a $\\rho^{(\\tau)}$ with compact support. Additionally, this assumption is more likely to hold on simpler learning tasks: if the true label generation process is defined by a $\\rho^*$ with compact support, then a bounded solution suffices to minimize the loss for such a simple learning task.\n\nThe second assumption concerns the separation property. It requires that the support of the convergence point $\\rho_\\infty$ “separates” a small and a big sphere at any local point $\\theta_0 (w_0)$ regarding the 1-homogeneous parameter part, and the support always “spans the full disk” $\\mathcal{K}$ used for universal approximation. A mild case that satisfies the “separation part” of this assumption is that the original point $0_{\\text{dim}\\theta+\\text{dim}w}$ is an interior point of the support for some $t^*$. This condition is relatively mild and is generally satisfied, though no paper has rigorously shown it when considering deep networks where $L$ tends to infinity. The challenge arises more in verifying the “spanning full disk” part, i.e., $\\text{supp}(\\rho_\\infty(\\cdot,t))$ extends to encompass the entire space $\\mathcal{K}$. While no paper can prove this property for the limit $\\rho_\\infty$, [20] with a similar theoretical setting shows that if the initial parameter distribution $\\rho_0(\\cdot,t^*)$ spans $\\mathcal{K}$, then $\\rho^{(\\tau)}(\\cdot,t^*)$ spans $\\mathcal{K}$, i.e., this expansive support property is maintained at any finite time. We will revise Appendix C.5 with clearer statements and more detailed discussions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Kd84uEHxSg", "reviewer": "Reviewer_eZmg", "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.", "questions": "N/A", "limitations": "N/A", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This is a well-written paper and it seems the results are strong and clean. However, I have very little background in this area and cannot verify the correctness of the statements. I hope the AC can find another qualified reviewer and delete this review if possible.", "weaknesses": "N/A", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your recognition of our paper. Your comments provide valuable guidance on our presentation. In this work, we aim to take the initial steps towards studying the theoretical optimization guarantees of Transformer models and, for the first time, prove the global convergence property via gradient descent-based methods.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NuailT48uF", "reviewer": "Reviewer_accV", "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$.", "questions": "On the technical side: \n1. While I understand the need of assumption i) in theorem 4.1, it is not obvious to me how stringent it is. Already for an extremely simple problem like $Q(\\rho) = \\int \\rho \\log(\\rho) + \\lambda \\int x^2 d\\rho$ the minimizer is a Gaussian, that has unbounded support. \tGiven the last example, the heuristic justification given by the authors seems to not be sufficient. Do the authors see a way to transfer this condition on the growth rate of $R(\\Theta)$ at infinity?\n2. The immediate extension of the deterministic case would be to consider the full Fokker-Plank, as it may emerge as a model of stochastic optimization. Already in the simplest case, the PDE under consideration would be $\\partial_\\tau \\rho^{(\\tau)} = div(\\rho^{(\\tau)} \\nabla U) + \\sigma \\Delta \\rho^{(\\tau)}$. The stationary distribution in this case would be $\\rho \\propto e^{-U/\\sigma}$, that has unbounded support for $U$ defined on all parameter space. This suggest that already in the simplest case (just pure isotropic diffusion added), the extension of the results may not be straightforward. I would appreciate the authors' comment on this, and I would like to know if they see an easy way to weaken that hypothesis as it's something one has no control on.\n3. While mentioned in the proofs, it is not clear in the main manuscript what $\\delta$ is. For example, in theorem 4.1 the statement is \"with high probability\", but it is not mentioned on what. I would add two words in the theorems/lemmas that are in high probability to specify with respect to what. \n4. It sounds strange to me that theorem 3.1 gives a bound on a compact set, while theorem 4.1 on an unbonded time horizon, with what appears to be exactly the same constant. Is the constant in theorem 4.1 really the same of the one in 3.1 or it includes other terms?", "limitations": "The authors correctly assessed the limitations in the manuscript, I have no further suggestions.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "The tackled problem is extremely relevant and the derived results are sound. Given the size of the work, the authors did a good job overall in organizing the contents.", "weaknesses": "1. The content is intense and the presentation doesn't help a novel reader. I believe polishing of the manuscript would improve the quality. \nMore concretely, I would add a small paragraph in the main manuscript summarizing the main ideas behind the novelties. In the appendix \t\t\tinstead it is really easy to get lost in a long proof, loosing track of the main goal. I suggest to include a small paragraph at the beginning of \t\teach section with the main steps required for the final proof, to guide the reader and make even the technical part more accessible \t\t\t(something on the line of what the authors did in D.3 would be perfect). \n\n2. While this did not and will not impact my score, for manuscripts of this size I believe at least a really small experimental section (even in the appendix) to test the claims would make the results more sound.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your supportive comments, which are greatly helpful for us to improve our work. We address your questions as follows.\n\n**Q1**: The content is intense and the reviewer suggests polishing the manuscript. Adding a summary of the main ideas in the manuscript and brief introductory paragraphs in each appendix section are helpful.\n\n**A1**: We thank the reviewer for this great suggestion. We agree that it would be helpful to briefly summarize the main ideas in the manuscript and the main proof ideas in the appendix. (If we can efficiently utilize the additional page for the camera-ready version, we may even add a brief proof sketch if space permits.) We will revise the paper following your suggestions.\n\n**Q2**: The reviewer believes that at least a really small experimental section (even in the appendix) to test the claims would make the results more sound.\n\n**A2**: In practice, we observe that using vision Transformers to fit simple datasets such as MNIST can easily achieve near-zero training loss, indicating convergence to a global minima. This observation reinforces our confidence that our results are valid and that our Transformer structure can minimize the training loss to zero as $L$ and $M$ increase. We will also conduct experiments investigating the relation between discrete transformers and their continuous limits and include the results in the revision.\n\n**Q3**: Assumption i) in Theorem 4.1 has ambiguity on how stringent it is without sufficient heuristic justification. A related example: the minimizer of $Q(\\rho)=\\int \\rho\\log(\\rho) +\\lambda x^2d\\rho$ is a Gaussian that has unbounded support. Do the authors see a way to transfer this condition to the growth rate at infinity? The immediate extension of the deterministic case would be to consider the full Fokker-Planck, whose stationary point of the PDE has unbounded support. \n\n**A3**: We thank the reviewer for the insightful comments about the extra entropy and diffusion terms. Since these two questions are related, we have combined them into one and provided a comprehensive response.\n\nYes, several papers [13,24,45] in the mean-field analysis literature study the same terms you mentioned, considering noisy gradient descent (NGD). By introducing an additional random Gaussian noise term $N(0,\\lambda_2)$ into the gradient flow/descent formula, the PDE evolution of $\\rho^{(\\tau)}(\\theta,w)$ becomes a diffusion process with an additional term $\\nabla^2\\rho^{(\\tau)}(\\theta,w)$ based on the Fokker-Planck equation. Consequently, the stationary point of the PDE is the local minimum of the regularized risk function with a regularization term $\\int \\rho\\log(\\rho) +\\lambda x^2d\\rho$. In this case, we note that the minimizer of $\\tilde{Q}(\\rho)$ will not have compact support. In comparison, (noiseless) gradient flow minimizes the risk function with $\\int \\lambda x^2d\\rho$ regularization, and it remains possible that the minimizer still has compact support.\n\nWhile the extension to NGD is very interesting, after careful consideration, we believe this extension does not fit within the framework and assumptions of our paper. The key reason is that our Assumptions 1-3 are much milder than the corresponding assumptions in [13,24,45], as the constant terms in our assumptions also depend on the parameter norms to accommodate Transformer structures. Given the dependence on parameter norms (e.g., $(1+||\\theta||)$ in Assumption 2(ii)), we must ensure that the parameter distribution is always bounded within any infinite time $s \\in [0,\\tau]$ to apply these assumptions. If we consider a simpler structure like residual networks with stronger assumptions, then the extension you mentioned could perfectly fit or even enhance the theoretical proof.\n\nLastly, we hope to clarify the first assumption in Theorem 4.1 to alleviate concerns about its stringency. Although this assumption is uncheckable from the given assumptions, the weight decay regularization can penalize the parameter norms, potentially leading to a $\\rho^{(\\tau)}$ with compact support. Additionally, this assumption is more likely to hold for simpler learning tasks: if the true label generation process is defined by a $\\rho^*$ with compact support, then a bounded solution suffices to minimize the loss for such a simple learning task.\n\n**Q4**: It is not clear in the main manuscript what $\\delta$ is. The reviewer would add two words in the Theorems/lemmas that are in high probability to specify with respect to what in Theorem 4.1.\n\n**A4**: We thank the reviewer for the careful reading and the great suggestion. The probability is with respect to the parameter initialization $\\Theta^{(0)}=\\\\{\\theta^{(0)}\\_{t,j},w^{(0)}\\_{t,j}\\\\}_{t,j}$. We will clarify it in the revision.\n\n**Q5**: It sounds strange 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 does it include other terms?\n\n**A5**: We thank the reviewer for the question and careful reading. The constant in Theorem 4.1 is indeed the same as that in Theorem 3.1, but with $\\tau$ is fixed as $\\tau_0$. To clarify how we obtain the constant in Theorem 4.1, we first select a large time horizon $\\tau_0$ and apply Theorem 3.1 with respect to this specific choice to bound $|\\widehat{R}^{(\\tau_0)} - R(\\rho^{(\\tau_0)})|$. The results follow if we can show $\\sup_{\\tau \\geq \\tau_0} R(\\rho^{(\\tau)})$ is asymptotically smaller than $\\epsilon + \\lambda$. This is why $C_1$ is dependent on $\\tau_0$ in Theorem 4.1. Further illustration can be found in the proof steps of Theorem 4.1 (pages 34-36), where $\\tau $ being large means exactly $\\tau \\geq \\tau_0$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M1BmXlP0HC", "reviewer": "Reviewer_mtEE", "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.", "questions": "* Can the dependency on time horizon be alleviated with stronger model assumptions?\n* Is the $L^{-1}+\\sqrt{\\log L/M}$ rate tight/expected? (see Weakness)\n* The widths of both the feedforward layer and the attention layer are both set equal to $M$, however the number of heads cannot typically be expected to be very large compared to the width of feedforward layers. Does the analysis generalize to when they are different?\n* Some recent papers [1,2] have also studied mean-field limits of Transformers from different perspectives, it would be nice to add a comparison in the related works section (although the latter is a contemporary work).\n\n[1] https://openreview.net/forum?id=xm2lU7tteQ\n\n[2] https://arxiv.org/abs/2405.15712", "limitations": "See Weaknesses.", "rating": 5, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 2, "strengths": "* The paper is a nontrivial extension of existing mean-field analyses of two-layer neural networks and ResNets to Transformers, which is challenging due to the existence of both feedforward and attention layers, the latter of which is new in the literature.\n* The assumptions for the networks and input data are quite general and can encompass e.g. the softmax mechanism and in-context learning settings.\n* The analyses and techniques are sufficiently novel and detailed and serve as an initial characterization of the dynamics of large scale Transformers, in particular the global convergence analysis of Section 4.", "weaknesses": "* The obtained error bounds (Theorem 3.1) are so large that they likely have little practical relevance beyond the initial moments of training. In particular, the hidden constants are in the worst case super-super-exponential with respect to the time horizon. Specifically, the bound is exponential in $\\phi_T(N,D,C\\times B_\\tau)$ where $B_\\tau$ is exponential in $R_\\tau$, which in turn is exponential in $\\tau$. While bounds exponentially diverging in time are common in the literature (which is a priori expected from an ODE discretization argument) as mentioned in  the paper, the dependency in Theorem 3.1 is much worse and some effort needs to be made to at least justify this. For example, one exponentation seems to be removable if e.g. $\\phi_T$ is bounded. Currently the dependency in the hidden constants are not made obvious without going through the proofs in the Appendix, which I feel should be addressed more transparently.\n* The convergence analysis (Theorem 4.1) also suffers from the same problem. The analysis requires a time horizon $\\tau_0$ large enough so that $W_2(\\rho^{(\\tau)}, \\rho_\\infty)$ is exponentially small (w.r.t. $R_\\infty$) compared to the desired error $\\epsilon$. This horizon is then fed into the approximation bound of Theorem 3.1, again yielding very large constants unless convergence is achieved in the very early stages of training. (This issue seems less critical since it only affects the $C_1$ term, although the $C_2\\lambda$ term is still super-exponential in $R_\\infty$.)\n* The recurring rate $L^{-1}+\\sqrt{\\log L/M}$ is not shown to be tight, leading me to further question the utility of the provided bounds. Why is the scaling for $L,M$ different? How does the rate compare with the analysis for ResNets or other ODE systems with 'width' & 'depth' dimensions?\n* A nitpick: the terms 'universal constant' or 'universal bound' are used quite frequently, however they should be reserved for constants that do not depend on *any* problem parameters.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments and suggestions. They are extremely helpful in improving our work and presentation. Below, we address your concerns point by point.\n\n**Q1**: The error bounds in Theorem 3.1 are so large they may be impractical beyond initial training moments, with hidden constants possibly being super-super-exponential over time. The hidden constants' dependency is unclear without reviewing the Appendix proofs and should be addressed more transparently. Additionally, can stronger model assumptions alleviate the time horizon dependency? One exponentiation might be removable if $\\phi_T$ is bounded.\n\n**A1**: We thank the reviewer for this insightful suggestion and careful clarification on this point. We agree that the hidden constants in Theorem 3.1 have sensitive dependencies on the time horizon, which is partially due to the structural complexities of the Transformers. Since Assumptions 2-3 are mild, allowing additional parameter norm factors in the bound, it is inevitable that the constant dependency would accumulate, thus easily causing super-exponential constants. Our goal is to take the first theoretical step in analyzing Transformers using mean-field tools. We can always make $L$ and $M$ sufficiently large to ensure vanishing approximation error. Though the result becomes asymptotic, our work lays the groundwork for future mean-field theory for Transformers.\n\nYes, we could alleviate the dependency with stronger model assumptions, as $\\phi_T$ may not depend on the Transformer output magnitude. We have noticed a lot of interest in identifying the optimal choice of $\\phi_T$, i.e., the Lipschitz constant of the Jacobian matrix of the self-attention term. This is a particularly challenging frontier question. For instance, [R1] suggests that $\\phi_T$ can be bounded by $\\sqrt{N/D} + \\mathrm{poly}(\\|T\\|_F)$, where $\\mathrm{poly}(\\cdot)$ denotes a polynomial function. In the context of $l_2$ self-attention, [R2] finds $\\phi_T$ to be $\\sqrt{N \\log N / D}$, which notably does not depend on $\\|T\\|_F$, and [R3] demonstrates that for $l_1$ distance metrics in attention layers, $\\phi_T$ could be $\\sqrt{D \\log N}$. We will include this illustration if our submission is accepted.\n\n[R1] Dasoulas, G., Scaman, K., & Virmaux, A. (2021). Lipschitz normalization for self-attention layers with application to graph neural networks. ICML\n\n[R2] Kim, H., Papamakarios, G., & Mnih, A. (2021). The lipschitz constant of self-attention. ICML.\n\n[R3] Vuckovic, J., Baratin, A., & Combes, R. T. D. (2020). A mathematical theory of attention. arXiv preprint arXiv:2007.02876.\n\n**Q2**: The convergence analysis in Theorem 4.1 also has issues with large constants. It requires a large time horizon $\\tau_0$ to make $W_2(\\rho^{(\\tau)},\\rho_\\infty)$ exponentially small relative to the error $\\epsilon$. This leads to large constants in Theorem 3.1's approximation bound unless early-stage convergence occurs. While this mainly affects the $C_1$ term, $C_2$ remains super-exponential in $R_\\infty$.\n\n**A2**: We thank the reviewer for the crucial clarification on this point. We agree that $C_1$ heavily depends on the choice of $\\tau_0$, and $\\tau_0$ heavily depends on the prefixed $\\epsilon$. Nevertheless, we can always make $L$ and $M$ sufficiently large (independent of the $\\tau_0$ choice) to ensure the risk is asymptotically bounded. We will focus more on refining these constants in future work after current initial steps towards extending the architecture into the realm of Transformers.\n\n**Q3**: The rate $L^{-1} + \\sqrt{\\log L/M}$ is not shown to be tight. Why is the scaling for $L$ and $M$ different? How does this rate compare to the analysis for ResNets or other ODE systems with 'width' and 'depth' dimensions?\n\n**A3**: The different scaling for $L$ and $M$ also appears in the analysis for ResNets (see [20] Theorem 9), where $L$ depends linearly and $M$ depends quadratically on $\\epsilon$. We think such rates are indeed expected. Intuitively, the different scaling is from the source of the approximation difference: $L^{-1}$ arises from the discretization of the ODE into $L$ small steps, and $M^{-1/2}$ comes from the implicit use of the Hoeffding’s Inequality considering the average of the outputs of $M$ nodes/heads. Furthermore, since our result in Theorem 3.1 considers the maximum difference across the entire time interval, we add an additional $\\log L$ term as we apply the probability union bounds.\n\n**Q4**: The terms 'universal constant' or 'universal bound' are used quite frequently, however they should be reserved for constants that do not depend on any problem parameters.\n\n**A4**: We thank the reviewer for the careful reading. We will fix this issue in the revised version.\n\n**Q5**: The widths of both the feedforward layer and the attention layer are both set equal to, however the number of heads cannot typically be expected to be very large compared to the width of feedforward layers. Does the analysis generalize to when they are different?\n\n**A5**: We thank the reviewer for the great question. The results can be easily extended when the widths of the feedforward layer ($M_1$) and the attention layer ($M_2$) both go to infinity, and we conjecture that the main results still hold by replacing $M$ with $\\min\\{M_1, M_2\\}$. However, by the nature of the mean-field analysis, we do need $M_2$ to be sufficiently large to achieve the theoretical convergence. The shift from discrete parameters to parameter distributions requires a large number of heads to provide a close approximation for the \"average\".\n\n**Q6**: Two recent papers referenced by the reviewer have also studied mean-field limits of Transformers from different perspectives. It would be nice to add a comparison.\n\n**A6**: We thank the reviewer for pointing out the related works. We will cite them and add discussions in the revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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`.
Decision: Accept (poster)