Conceptio › Archive › arXiv CS

arXiv CSopen access

On the global convergence of gradient descent for wide shallow models with bounded nonlinearities

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

neural-networks

machine learning, deep learning, neural networks

On the global convergence of gradient descent for wide shallow models with bounded nonlinearities

arXiv:2605.10775v1 [math.OC] 11 May 2026

Romain Petit CNRS and ENS, PSL Université [email protected]

Clarice Poon Mathematics Institute, University of Warwick [email protected]

Gabriel Peyré CNRS and ENS, PSL Université [email protected] May 11, 2026 Abstract A surprising phenomenon in the training of neural networks is the ability of gradient descent to find global minimizers of the training loss despite its non-convexity. Following earlier works, we investigate this behavior for wide shallow networks. Existing results essentially cover the case of ReLU activations and the case of sigmoid activations with scalar output weights. We study a large class of models that includes multi-head attention layers and two-layer sigmoid networks with vector output weights. Building upon [CB18], we prove that all non-global minimizers of the training loss are unstable under gradient descent dynamics. Thus, when the initial distribution of the parameters has full support (which includes the popular Gaussian case), and in the many hidden neurons or attention heads limit, continuoustime gradient descent can only converge to global minimizers. Establishing the instability of non-global minimizers corresponds to the construction of an “escaping active set” – we complete the proof of [CB18] to construct this set for models with bounded nonlinearities and scalar output weights. We also extend this construction to new cases for models with vector output weights. Finally, we show the well-posedness and the stability with respect to discretization of the mean field training dynamic for sub-Gaussian initializations.

Introduction

The training of neural networks leads to the minimization of objective functions that are highly nonconvex. As a result, first order optimization algorithms such as gradient descent could fail to reach an approximate minimizer and converge to non-optimal critical points. In certain settings, it has been observed that this behavior is not typical and that gradient descent finds global minimizers of the training loss despite its non-convexity. Among the various explanations that have been proposed, a line of works has focused on the study of the over-parameterized regime. The landmark results of [CB18] (see also [MMN18, RV18]) show that, if the model is sufficiently over-parameterized and the initialization densely covers the parameter space, then (continuous-time) gradient descent can only converge to global minimizers of the training loss. In these works, the over-parameterized models under study are typically fully connected neural networks with a single hidden layer and a large number of hidden neurons. Existing results essentially cover sigmoid activations with scalar output weights and ReLU activations. Besides networks with sigmoid activations and vector output weights, other popular architectures are not covered by these results. This is in particular the case of attention layers, which have become ubiquitous in recent years. In this work, we make a step towards bridging this gap by investigating global convergence results for a large class of over-parameterized models.

1.1

Main problem

Pm We study the training of parameterized models of the form (1/m) i=1 Φ(wi , θi ), where Φ is linear in the w variable. This setting covers two main examples of interest: transformers with a single multi-head attention layer and two-layer fully connected neural networks. In the first case, the w and θ variables respectively correspond to the value matrix and to the query and key matrices. In the second case, they respectively correspond to the parameters of the output layer and to the parameters of the hidden layer. We analyze the training of such models with continuous-time gradient descent (or gradient flow). To be more precise, we analyze the gradient flow of the function Fm : Ωm → R defined by Pm Fm ((wi , θi )1⩽i⩽m ) = R((1/m) i=1 Φ(wi , θi )), (1) where Ω = Rdw × Rdθ and R is a given loss function. The objective Fm is highly non-convex, and the iterates of gradient descent could get trapped by spurious local minima. Yet, in the large m (or over-parameterized ) regime, gradient descent often converges to global minimizers of Fm . The main objective of this work is to understand this phenomenon. Following [CB18], the key object of our analysis is the dynamic one obtains by encoding the parameters of the model as a probability distribution over the parameter space Ω. This evolution turns out to be a Wasserstein gradient flow (see [San17] for an introduction to this topic) of a function F defined on the set P(Ω) of probability distributions on Ω. By studying the long-time limit of this evolution, we aim to find conditions under which the gradient flow of Fm converges to a global minimizer.

1.2

Related work

Scaling limits. It is by now well understood that non-convexity is not necessarily an obstacle when training highly over-parameterized neural networks, including deep ones. Under suitable scaling regimes, local and sometimes global convergence can be established. Different scalings lead to distinct limiting descriptions of training dynamics, with the two most studied cases being the Neural Tangent Kernel (NTK) regime and the mean-field or “feature learning” regime. These asymptotic viewpoints clarify why optimization succeeds in over-parameterized models, although they may not always align with practical behavior [COB19, LZB20, YH21]. Wide networks and NTK. For nonlinear multilayer perceptrons, a standard scaling is to initialize √ weights with variance proportional to 1/ m, where m is the width of a hidden layer. In this setting, when the width grows, hidden units move little during training, and in the infinite-width limit, they remain fixed. The resulting dynamics correspond to kernel regression with the Neural Tangent Kernel [JGH18]. This framework provides a route to proving global convergence in sufficiently wide networks, with refinements depending on data separability and network depth. For two-layer ReLU networks, convergence has been analyzed under the NTK scaling [LY17]. For deeper models, general convergence results have been established [AZLS19, DLL+ 19, LZB20, ZCZG20, LXS+ 19, CC+ 20, Ngu21]. This analysis can be refined for residual networks, illustrating how skip connections stabilize training [HM17, BHL18, ZLG19, B+ 22]. Feature learning regime and mean-field models. An alternative scaling is to initialize weights with variance 1/m. In this case, hidden units evolve significantly during training, leading to the so-called meanfield or feature learning regime. Here, the network can be described by a probability measure over neurons, which adapts during training. Unlike NTK, this regime captures representation learning and explains the ability of neural networks to exploit low-dimensional structures in data [Bac17, GMMM20]. It also connects neural networks to functional analysis, where associated function spaces can be defined [EW22]. The meanfield scaling has been rigorously formalized as the limit of interacting particle systems, yielding a partial differential equation (PDE) of McKean–Vlasov type, which is a gradient flow for the Wasserstein metric [CB18, MMN18, RV18, Woj20, NP23]. This PDE viewpoint allows one to prove global convergence results when the support of the initial distribution is sufficiently broad [CB18]. The main goal of our work is to

extend the PDE framework (without diffusion and noise injection) and convergence analysis to a broader class of architectures—so as to accommodate attention heads and vector output weights for bounded activation functions—as well as to more general initialization schemes, including non–compactly supported distributions such as Gaussians and, more generally, sub-Gaussian distributions. Training attention and transformers. The convergence analysis of transformers and attention mechanisms has followed parallel lines in the NTK regime. [KVPF20] reformulated attention as a linear-time mechanism using kernelization, allowing convergence to be analyzed with tools from fast autoregressive training. [CLD+ 21] introduced the Performer architecture, which uses random feature approximations to make attention scalable while maintaining convergence properties, and [PHA+ 21] further developed random feature attention with theoretical guarantees. The feature learning behavior of attention under head scaling remains poorly understood. One focus of this article is to address this gap by uncovering and analyzing the mean-field structure of wide attention layers.

1.3

Contributions

Our main contribution, which is the object of Section 3, is the proof that, under suitable assumptions, the Wasserstein gradient flow of F can only converge to measures satisfying a certain condition. We establish this result for a large class of models, including multi-head attention layers and two-layer fully connected networks with vector output weights and sigmoid activations. Section 4 is dedicated to the proof that these two examples fit in our framework. In [CB18, Appendix C], a closely related condition was shown to be satisfied only at global minimizers of F for sigmoid networks with scalar output weights, hence providing a global convergence result for the Wasserstein gradient flow. We provide a corrected version of this proof, which contained an erroneous argument. We also investigate the extension of this result to sigmoid networks with vector output weights and attention layers in Section 3.3, and conjecture that it holds under mild assumptions. As the Wasserstein gradient flow of F is a key object in our analysis, we also study its well-posedness and stability with respect to discretization in Section 2. We improve previous results by proving its existence and uniqueness for all times, as well as a stability estimate, for a larger class of initializations. Our main assumption is that the initial measure is sub-Gaussian (instead of compactly supported for previous works), which includes the popular Gaussian case. In the case of two-layer fully connected networks, in addition to sigmoid activations, our assumptions also cover the popular GELU and SiLU / swish activations. Regularization. An important difference with respect to [CB18] is that they allow an optional regularization term in the objective. Our choice to cover the unregularized case is motivated by two reasons. First, the potential non-smoothness of the regularizer leads to several additional technical difficulties. Second, in the case of attention layers and fully connected networks with sigmoid activations, the analysis conducted in [CB18] would only allow us to treat the case of an ℓ1 regularization of the weights, which does not seem popular. Practitioners rather use weight decay, which corresponds to an ℓ2 squared regularization. Parameterization of the attention. Transformers have a complex architecture that relies on a number of crucial operations, such as layer normalization, skip connections, and variants of the classical softmax selfattention. In this work, we focus on a simplified task, that is, the training of a single softmax self-attention layer with a large number of attention heads. Moreover, we study the training of a parameterization that slightly differs from the most standard one. Namely, if (Q, K, V ) are the query, key, and value matrices of an attention head, its output only depends on V and the matrix A = K T Q. Optimizing over both factors (Q, K) of the product leads to several difficulties, including non-Lipschitz gradients. For this reason, we directly optimize over the variables (A, V ) and leave the investigation of the usual parameterization to future works.

Mean-field training dynamic

In this section, we introduce the evolution one obtains by encoding the parameters of the model as a probability measure on the parameter space. We prove the well-posedness and stability with respect to discretization of this mean-field training dynamic for a broad class of initializations, namely, sub-Gaussian distributions.

2.1

Setting

In all the following, we rely on the following assumption. As shown in Section 4 below, it covers multihead attention layers, as well as fully connected networks with a single hidden layer, provided the activation function is differentiable with a bounded and Lipschitz differential. This includes the popular sigmoid, GELU, and SiLU / swish activations (see [RZL17, EUD18, HG23] for more details on these recently proposed activation functions). Assumption 1. F is a separable Hilbert space and Ω ⊂ Rd is the closure of a convex open set. The loss R : F → R+ is differentiable with a differential that is Lipschitz on bounded sets and bounded on sublevel sets of R. The mapping Φ : Ω → F is differentiable, and its differential has the following properties. (i) Linear growth: ∥dΦ(u)∥L(Rd ;F ) ≲ 1 + |u|. (ii) Local Lipschitz continuity: for every r > 0 and u1 , u2 ∈ B(0, r) it holds ∥dΦ(u1 ) − dΦ(u2 )∥L(Rd ;F ) ≲ (1 + r)|u1 − u2 |.

Gradient flows. Gradient flows, which correspond to the vanishing step size limit of gradient descent, are curves of steepest descent. In the case of the objective function Fm , a gradient flow is a C 1 path t ∈ R+ 7→ (wi (t), θi (t))1⩽i⩽m such that (wi′ (t), θi′ (t))1⩽i⩽m = −∇Fm ((wi (t), θi (t))1⩽i⩽m ) for every t ⩾ 0. Under Assumption 1, the objective Fm is of class C 1 and lower bounded. As a consequence, standard results (see for instance [San17]) imply that, for every initialization (wi (0), θi (0))1⩽i⩽m ∈ Ωm , there exists a unique gradient flow of Fm . Lifting to the space of probability measures.

Following [CB18], we define the function

F : P(Ω) → R ∪ {+∞} ( R R R( Ω Φdµ) if Ω ∥Φ(w, θ)∥F dµ(w, θ) < +∞, µ 7→ +∞ otherwise, Pm and notice that the mapping Fm defined in (1) satisfies Fm ((wi , θi )1⩽i⩽m ) = F ((1/m) i=1 δ(wi ,θi ) ). In other words, the restriction of F to empirical distributions coincides with Fm . The main advantage of this lifting is that it allows to consider models associated with more general parameter distributions, such as distributions with an infinite support. Wasserstein gradient flow. In [CB18, Proposition B.1], it is shown that, if t 7→ (wi (t), θi (t))1⩽i⩽m is a Pm gradient flow of Fm , then t 7→ µt := (1/m) i=1 δ(wi (t),θi (t) is a Wasserstein gradient flow of F , that is to say a curve of steepest descent for the geometry induced by the Wasserstein distance on P(Ω). To be more precise, it holds ∂t µt = −div(µt vt ) with vt (u) = −∇F ′ (µt )(u) (2)

for every t ⩾ 0, where F ′ (ν) : Ω → R denotes the function representing the differential of F at some ν ∈ P(Ω). We refer the reader to Section B for the precise meaning of (2) and to [San17] for more details on Wasserstein gradient flows. When the evolution (2) is initialized with a more complicated distribution, proving its long-time existence and stability is more involved. This is the object of the following subsection.

2.2

Well-posedness

Under Assumption 1, we are able to show that the mean field training dynamic is well-posed for every sub-Gaussian intialization. Theorem 1. Under Assumption 1, if µ0 ∈ P(Ω) is sub-Gaussian, then there exists a unique Wasserstein gradient flow (µt )t⩾0 of F starting from µ0 . Moreover, if µm,0 ∈ P(Ω) converges to µ0 for the W2 distance when m → +∞, then (µm,t )t⩾0 converges as m → +∞ to (µt )t⩾0 . Moreover, there exists C > 0 such that W2 (µt , µm,t ) ⩽ eCt W2 (µ0 , µm,0 ) for every t ⩾ 0. The proof of Theorem 1 is given in Section B. Our proof strategy and our assumptions on the initialization are discussed below. Assumption on the initialization. We stress that our condition on the initialization is essentially weaker than that of [CB18, Proposition 2.5], which requires the w-marginal of the initial measure to have compact support. In particular, our condition covers the important case of Gaussian initializations. In Section 3 below, we focus on the case where Ω = Rdw × Rdθ and Φ : (w, θ) 7→ ϕ(θ)w where ϕ(θ) is a linear mapping from Rdw to F . The w variable corresponds to the parameters of the output layer (in the case of fully connected networks) or to the value matrix (in the case of attention layers). The θ variable corresponds to the parameters of the hidden layer (for fully connected networks) or to the query and key matrices (for attention layers). For attention layers or fully connected networks with sigmoid activations, we could deal separately with these two sets of variables and exploit the boundedness of ϕ to remove the sub-Gaussian assumption on the θ-marginal. However, this leads to more involved computations. For the sake of clarity, we choose to deal with both sets of variables jointly, at the price of requiring that the full distribution is sub-Gaussian. Proof strategy. In [CB18], the analog of Theorem 1 (which is Proposition 2.5) is proved by relying on the theory of gradient flows in metric spaces developed in [AGS09]. In contrast, our proof, which allows to remove the compact support assumption, only uses elementary tools and is based on a fixed point argument. We stress that this is made possible because we do not consider a possibly non-smooth regularization term in the objective, contrary to [CB18]. Our proof strategy is borrowed from the literature on interacting particle systems (see also [CACP25] for an application to the analysis of infinitely deep transformers) and uses Dobrushin-type estimates (see for instance, [Dob79, CCR11]). However, the control we have on the velocity field is weaker than in these works. We circumvent this issue by exploiting the sub-Gaussian assumption on the initialization and by controlling the growth of the tails along the evolution.

Global convergence

In this section, we consider the case Ω = Rdw × Rdθ and Φ : (w, θ) 7→ ϕ(θ)w where ϕ(θ) is a linear mapping from Rdw → F . We focus on the case where ϕ is bounded, which covers two-layer fully connected networks with sigmoid activations and multi-head attention layers. We investigate the limit points of the Wasserstein gradient flow of F and wish to find conditions under which they are global minimizers. Existing results. In [CB18], two different cases are considered depending on the homogeneity of the mapping Φ. The first corresponds to settings where ϕ is positively 1-homogeneous, meaning that ϕ(λθ) = λϕ(θ) for every λ > 0. In this case, Φ is positively 2-homogeneous. This essentially models two-layer networks

with ReLU activations. The second case corresponds to the situation where ϕ is bounded, and Φ is hence only partially 1-homogeneous. This models two-layer networks with sigmoid activations. In [CB18, Theorems 3.4 and 3.5], it is proved that, in both cases and under suitable assumptions on the initial distribution, the Wasserstein gradient flow of F can only converge to global minimizers. This result is established for every dw ⩾ 1 in the 2-homogeneous case but only for dw = 1 in the partially 1-homogeneous case. Section outline. In this section, we investigate the extension of the global convergence result mentioned above in the partially 1-homogeneous setting when dw > 1. In Section 3.1, we leverage the tools introduced in [CB18, Appendix C] to prove that under suitable assumptions, the Wasserstein gradient flow of F can only converge to measures satisfying a certain condition. A proof that this condition holds only at global minimizers was proposed in [CB18, Proposition C.4]. However, a part of this proof relies on an erroneous argument. In Section 3.2 and Section C, we provide a corrected version of this proof and show that the result holds under slightly revised assumptions. Finally, in Section 3.3, we discuss the extension of the result when dw > 1.

3.1

Limit points of the training dynamics

We first isolate an abstract instability mechanism for ruling out non-optimal limit points of the training dynamic. The idea is simple: if the vector field associated with a candidate limit point µ contains a region of the parameter space where output weights must grow at a uniform rate, then a trajectory cannot converge to µ in W2 , because a positive amount of mass would escape to infinity. Throughout this subsection, we denote by Cb1,1 (Rdθ ; Rdw ) the space of bounded differentiable maps whose differential is bounded and Lipschitz. In the case of bounded nonlinearities, the relevant perturbation norm is ∥h∥X := ∥h∥∞ + supr⩾0 ∥rJh (r·)∥∞ , where Jh denotes the Jacobian matrix of h. Definition 1 (Escaping active set). Let g ∈ Cb1,1 (Rdθ ; Rdw ). We say that g admits an escaping active set and write E(g) if there exist constants ε > 0, η > 0, and a Borel set A ⊂ Rdw × Rdθ containing a nonempty open set such that the following holds. For every absolutely continuous curve (gt )t⩾0 ⊂ Cb1,1 (Rdθ ; Rdw ) satisfying sup ∥gt − g∥X ⩽ ε,

(3)

t⩾0

and for every initial condition (w0 , θ0 ) ∈ A, the solution of ( ẇt = −gt (θt ), θ̇t = −Jgt (θt )T wt

(4)

satisfies d 1 |wt |2 ⩾ η dt 2

for all t ⩾ 0.

Remark 1. If g ≡ 0, then g does not admit an escaping active set. Indeed, taking gt = g for all t, every solution of the ordinary differential equation is constant. Thus the escaping active set property is a genuine non-stationarity condition. Assumption 2. Let F be a separable Hilbert space and let Ω = Rdw ×Rdθ . Assume that the loss R : F → R+ is convex and differentiable, with differential Lipschitz on bounded sets and bounded on sublevel sets of R. Assume also that Φ : Ω → F has the form Φ(w, θ) = ϕ(θ)w, where ϕ(θ) ∈ L(Rdw ; F) is bounded and differentiable, with bounded Lipschitz differential. R Under Assumption 2, the first variation of F (µ) = R Ω Φ dµ is given by R F ′ (µ)(w, θ) = ⟨R′ Ω Φ dµ , ϕ(θ)w⟩F . 6

We define gµ (θ) := ϕ(θ)∗ R′

R Ω

Φ dµ ,

so that F ′ (µ)(w, θ) = ⟨gµ (θ), w⟩. We are now able to state the following result, which restricts the potential limit points of the training dynamic. Theorem 2 (Escaping active sets cannot occur at W2 -limits). Assume that Assumption 2 holds. Let (µt )t⩾0 be a Wasserstein gradient flow of F , and assume that µ0 has full support. If W2 (µt , µ) → 0, then gµ does not admit an escaping active set. Proof. We argue by contradiction and assume gµ admits an escaping active set. We have R R ∥gµt − gµ ∥C 1 ⩽ ∥ϕ∥C 1 ∥R′ ( Ω Φdµt ) − R′ ( Ω Φdµ)∥F . Using the fact that R′ is Lipschitz on bounded sets together with Lemma 10 and the convergence of (µt )t⩾0 towards µ for the W2 distance, we obtain the existence of t0 ⩾ 0 such that ∥gµt − gµ ∥∞ + ∥Jgµt − Jgµ ∥∞ ⩽ ε for every t ⩾ t0 . Using again that W2 (µt , µ) → 0 when t → +∞ and that (w, θ) 7→ |w| is Lipschitz, we obtain: Z Z lim |w|dµt (w, θ) = |w|dµ(w, θ). (5) t→+∞

Rdw ×Rdθ

But since E(gµ ) holds, we also have Z Z |w|dµt (w, θ) = Rdw ×Rdθ

Rdw ×Rdθ

Z |wt |dµt0 (wt0 , θt0 ) ⩾

|wt |dµt0 (wt0 , θt0 ) A

⩾ η(t − t0 )µt0 (A). Using the results of Section B, we know that µt0 is the pushforward of µ0 by Xt0 , where X is the flow of the velocity field v : (t, w, θ) 7→ −∇F ′ (µt )(w, θ). We also have that Xt is a homeomorphism for every t ⩾ 0. Since µ0 has full support, the same holds for µt0 , which implies µt0 (A) > 0. We therefore obtain that the left-hand side of (5) equals +∞, which yields a contradiction, as its right-hand side is finite. Arguing as in [CB18, Appendix C.1], one can show that µ ∈ P(Ω) is a global minimizer of F if and only if F ′ (µ) is identically zero, which is equivalent to gµ being identically zero. As a result, we immediately obtain the following consequence for the global convergence of the training dynamic. Corollary 1 (Global convergence criterion). Assume that Assumption 2 holds. Further assume that every µ ∈ P(Ω) satisfying gµ ̸≡ 0 has the property that gµ admits an escaping active set. Let (µt )t⩾0 be a Wasserstein gradient flow of F with full-support initial condition µ0 . If µt → µ in W2 , then µ is a global minimizer of F .

3.2

Scalar linear parameters

In this section we provide a proof that, when dw = 1, if g is not identically zero and satisfies additional regularity properties, then E(g) holds. Since these regularity properties are satisfied when g = gµ for some µ in the domain of F , Corollary 1 shows that the training dynamic can only converge to global minimizers. Our corrected proof of [CB18, Proposition C.4] involves very slightly different assumptions (see (ii) − (iii) below), which are satisfied in the examples we consider in Section 4. If a C 1 function f is defined on Rd or the sphere Sd−1 for some d, we say that a value η in the range of f is regular if the differential of f does not vanish on f −1 ({η}), and we say it is singular otherwise. Proposition 1. Assume that dw = 1 and g ∈ Cb1,1 (Rdθ ; Rdw ) is not identically zero. Then g admits an escaping active set holds provided all the following conditions are satisfied: (i) the set of singular values of g has Lebesgue measure zero in R, 7

(ii) the mapping θ 7→ g(rθ) converges uniformly on Sdθ −1 as r → +∞ towards some mapping g∞ , (iii) the mapping θ 7→ r∇g(rθ) converges uniformly on Sdθ −1 as r → +∞, (iv) the set of singular values of g∞ has Lebesgue measure zero in R. Proof. If g is constant then the result is easier to prove. Indeed, let us assume that it is everywhere equal to −η with η > 0 (the other case can be treated in a similar way). Then ẇt = −gt (θt ) ⩾ −g(θt ) − ε = η − ε. As a consequence E(g) holds with e.g. ε = η/2 and A = R∗+ × Rdθ . If g is not constant then its range contains a non-empty open interval. Let us assume this interval contains negative values (the other case can be treated in a similar way). Using (i), we obtain that the set of negative regular values of g has positive Lebesgue measure. Now, let us also assume that there exists η > 0 such that −η is a regular value of g and K = {θ ∈ Rdθ | g(θ) ⩽ −η} is such that ∂K is bounded. The other case is more involved and its treatment is postponed to Section C. Since −η is a regular value of g, we have that ∇g(θ) ̸= 0 for every θ ∈ ∂K. Since ∂K is bounded by compactness there exists β > 0 such that inf θ∈∂K |∇g(θ)| ⩾ β. Using this, we obtain: d g(θt ) = −wt ⟨∇g(θt ), ∇gt (θt )⟩ ⩽ −wt (β 2 − ∥∇g∥∞ ε) < 0, dt provided wt > 0, θt ∈ ∂K and ε < β 2 /∥∇g∥∞ . This shows that (wt , θt ) ∈ R∗+ × K for every t as soon as (w0 , θ0 ) ∈ R∗+ × K. Since ẇt = −gt (θt ) ⩾ −g(θt ) − ε ⩾ η − ε ⩾ η/2 as soon as θt ∈ K and provided ε ⩽ η/2, the result is proved. Corollary 2. Assume that Assumption 2 holds and that the following holds for every ν ∈ P(Ω) in the domain of F : (i) the set of singular values of gν has Lebesgue measure zero in R, (ii) the mapping θ 7→ gν (rθ) converges uniformly on Sdθ −1 as r → +∞ towards some mapping gν,∞ , (iii) the mapping θ 7→ r∇gν (rθ) converges uniformly on Sdθ −1 as r → +∞, (iv) the set of singular values of gν,∞ has Lebesgue measure zero in R. Then if (µt )t⩾0 is a Wasserstein gradient flow of F such that µ0 has full support and (µt )t⩾0 converges for W2 towards some measure µ, we have that µ is a global minimizer of F .

3.3

Vector linear parameters

As explained above, the case dw > 1 is not covered in [CB18, Theorem 3.5]. For attention layers, imposing dw = 1 amounts to dealing with scalar tokens, which is not realistic. As a result, we would like to extend the global convergence theory beyond the case of scalar output parameters. In this setting, the dynamic of the linear parameters w ∈ Rdw becomes significantly more complex: they can change direction in addition to growing or shrinking in norm. In Section 3.3.1, we explain why a natural generalization of the proof of Proposition 1 might fail. In Section 3.3.2, we discuss an alternative construction of the set A. It allows to prove an analog of Proposition 1 when dθ = 1 and dw ⩾ 1 under reasonable assumptions. However, these assumptions do not seem realistic when dθ > 1. We conjecture that E(g) holds when dw > 1 provided g is not identically zero and satisfy some regularity assumptions similar to those of Proposition 1. Proving this would show that the mean field training dynamic can only converge to global minimizers of the training loss, even in the case of sigmoid networks with vector output parameters and multi-head attention layers.

3.3.1

First construction

A natural idea could be to define A = {(w, θ) ∈ Rdw × Rdθ | (w/|w|, θ) ∈ K}, where K = {(v, θ) ∈ Sdw −1 × Rdθ | ⟨g(θ), v⟩ ⩽ −η} and −η is a regular value of the mapping h : (v, θ) ∈ Sdw −1 × Rdθ 7→ ⟨g(θ), v⟩. If (vt , θt ) ∈ K then it holds d |wt | = −⟨vt , gt (θt )⟩ ⩾ −[⟨g(θt ), vt ⟩ − ε] ⩾ η/2 dt provided ε is small enough. We also have 1 d ⟨g(θt ), vt ⟩ ⩽ −|wt |[|Jg (θt )T vt |2 − ∥Jg ∥∞ ε] − [|proj{vt }⊥ (g(θt ))|2 − ∥g∥∞ ε]. dt |wt | We know that |∇h(v, θ)|2 = |proj{v}⊥ (g(θ))|2 +|Jg (θ)T v|2 is non-zero on ∂K. We would like to guarantee that d dt ⟨g(θt ), vt ⟩ < 0 if (vt , θt ) ∈ ∂K, which would yield that (vt , θt ) ∈ K for every t ⩾ 0 as soon as (v0 , θ0 ) ∈ K. d The lower bound on dt |wt | would finally allow us to conclude. The problem is that if |Jg (θt )T vt | is very d small for some (vt , θt ) ∈ ∂K and |wt | is very large, we cannot say anything about the sign of dt ⟨g(θt ), vt ⟩ (ensuring wt is not too small is possible so that this problem does not appear for the other term). As a result, it is unclear how to generalize the proof of Proposition 1. 3.3.2

Second construction

Another possibility is to construct the set A using a fixed vector v ∈ Sdw −1 . In Section D, we prove the following result. Proposition 2. Let g ∈ Cb1,1 (Rdθ ; Rdw ) be non identically zero and let v ∈ Sdw −1 and η > 0 be such that −η is a regular value of θ 7→ ⟨g(θ), v⟩. Let K = {θ ∈ Rdθ | ⟨g(θ), v⟩ ⩽ −η}. Provided ∂K is bounded and |proj{v}⊥ (Jg (θ)Jg (θ)T v)| |⟨g(θ), v⟩| < inf , T θ∈K |g(θ)| |Jg (θ)Jg (θ) v| θ∈∂K sup

(6)

there exists δ ∈ (0, 1) such that the set A defined by A = {(w, θ) ∈ Rdw × Rdθ | θ ∈ K, ⟨v, w⟩/|w| ⩾ δ} is stable: for every initial condition (w0 , θ0 ) ∈ A, the solution of (4) satisfies (wt , θt ) ∈ A for every t ⩾ 0. d In this case, by construction, we have a positive lower bound on dt |wt | on A, which shows that E(g) holds. Condition (6) is difficult to analyze. However, assuming that h : θ 7→ (1/2)|g(θ)|2 has a non-degenerate local maximizer θ∗ , one can perform a local analysis of (6) with v = −g(θ∗ )/|g(θ∗ )| and −η close to ⟨g(θ∗ ), v⟩. This leads us to the following sufficient condition ensuring E(g) holds, which is always satisfied when dθ = 1.

Proposition 3. Let g ∈ Cb1,1 (Rdθ ; Rdw ). Assume that h : θ ∈ Rdθ 7→ (1/2)|g(θ)|2 has a local maximizer θ∗ at which the Hessian is negative definite. Then E(g) holds provided sup ⟨Hx,x⟩=1

|Jx| <

|Hx|2 with J = Jg (θ∗ ) and H = −Hg (θ∗ )[g(θ∗ )], ⟨Hx,x⟩=1 |JHx| inf

which is always the case when dθ = 1.

Examples

In this section, we show that the two examples mentioned in the introduction fit into the abstract frameworks of Sections 2 and 3, so that Theorems 1 and 2 both apply. 9

Model space and loss function. We consider two settings, namely regression with the square loss and din dout multi-category classification with the cross-entropy loss. Given a data ) over R distribution ρ ∈ P(R × R pairs of features and labels, we minimize the expected risk R : f 7→ Rdin ×Rdout ℓ(f (x), y)dρ(x, y), where ℓ is P either the square loss ℓ : (z, y) 7→ (1/2)|z −y|2 or the cross-entropy loss ℓ : (z, y) 7→ −⟨z, y⟩ +log( 1⩽i⩽k ezi ). We denote by ρx and ρy the marginals of ρ on Rdin and Rdout , and ρy|x the conditional distribution of labels given a feature x. In the case of multi-category classification, we ask that ρy is concentrated on the canonical basis Bk of Rk , where k = dout is the number of classes and the i-th element of Bk is associated to the i-th class. Lemma 1 below, whose proof is given in Section A, shows that if F = L2 (ρx ; Rdout ) then in both cases R fits in our abstract framework. Lemma 1. In the following two cases: (i) ℓ is the square loss and ρy has finite second order moments, (ii) ℓ is the cross-entropy loss and ρy is concentrated on Bdout , it holds that R is convex and differentiable on F = L2 (ρx ; Rdout ). Moreover, its differential is bounded on sublevel sets and Lipschitz.

4.1

Two-layer fully connected networks

The mapping implemented P by a two-layer fully connected neural network with m hidden neurons and m activation function σ is (1/m) i=1 Φ(wi , θi ), where Φ(w, θ) : x 7→ σ(⟨θ, x⟩)w, dθ = din and dw = dout . Lemma 2 below states that Assumptions 1 and 2 hold under suitable assumptions on the activation function and the data distribution. As a consequence, we obtain that Theorem 1 applies to a large choice of activation functions, including sigmoid, GELU, and SiLU / swish activations, and that Theorem 2 applies to sigmoid activations. Lemma 2. If σ is differentiable with a bounded Lipschitz differential and ρx has finite moments up to order 4, then Assumption 1 holds. In addition, provided σ is bounded, then Assumption 2 holds. The proof of Lemma 2 follows by arguing as in [CB18, Lemma D.2], which deals with sigmoid activations. A closer inspection reveals that the bounds only depend on the supremum norm and the Lipschitz constant of σ ′ , so that the result also applies to GELU and SiLU / swish activations. We end this section with the following result, whose proof is postponed to Section A.1. In particular, it implies that when dw = 1, assumptions (i) − (iii) of Corollary 2 are satisfied. Indeed, in this case, the mapping gν writes R R gν (θ) = Rdin R′ ( Ω Φdν)(x)σ(⟨θ, x⟩)ρx (x)dx R and R′ ( Ω Φdν) is bounded and continuous. As explained in [CB18, Section 3.2 and Appendix D.3], condition (iv) is hard to check but seems purely technical, as known counterexamples involve artificial constructions. Lemma 3. Let σ : s 7→ 1/(1 + e−s ) be the sigmoid function. If ρx has a bounded continuous density with respect to the Lebesgue measure on Rdin , has finite mooments up to order 2dθ + 2 and satisfies |ρx (x)| ⩽ C(1 + |x|)−p for some C > 0 and p > din , then for every bounded continuous function f : Rdin → R the following holds: R 1. the mapping gf : θ ∈ Rdθ 7→ Rdin f (x)σ(⟨θ, x⟩)ρx (x)dx is of class C dθ , R 2. the mapping θ 7→ gf (rθ) converges uniformly on Sdθ −1 towards θ 7→ {⟨θ·x⟩⩾0} f (x)ρx (x)dx as r → +∞, R 3. the mapping θ 7→ r∇gf (rθ) converges uniformly on Sdθ −1 towards θ 7→ {⟨θ,x⟩=0} f (x)ρx (x)xdHdin −1 (x) as r → +∞.

4.2

Multi-head attention layers

We consider a multi-head attention layer taking as input contexts of size n composed of tokens in Rd . The number of attention heads is m and the query, key and value matrices associated to the P i-th head are m (Qi , Ki , Vi ) ∈ (Rk×d )3 . In a next-token prediction setting, the implemented predictor is (1/m) i=1 Vi ψ(Ai ) where Ai = KiT Qi and ψ(A) : Rn×d → Rd is defined by ψ(A) : X = xT1 · · · xTn

7→

n X

n X e⟨Axn ,xi ⟩ x = σi (XAxn )xi = X T σ(XAxn ), i n P ⟨Ax ,x ⟩ i=1 i=1 e n j j=1

Pn where σ : z ∈ Rn 7→ (ezi /( j=1 ezj ))ni=1 ∈ Rn is the softmax function. One can see that this setting is covered by our abstract framework by taking dw = kd, dθ = d2 , din = nd and dout = k with w being the vectorized version of V and θ the vectorized version of K T Q. The mapping ϕ is then defined by the relation ϕ(θ)w = V ψ(A). For this reason, properties of ψ (boundedness, regularity, etc.) directly transfer to ϕ. As ψ has a simpler expression, we use it rather than ϕ in our computations. The following lemma, whose proof is postponed to Section A, shows that, under a finite moment assumption on the distribution of input contexts ρX ∈ P(Rn×d ), the mapping ϕ is sufficiently regular, so that, by Lemma 5, Assumptions 1 and 2 hold and Theorems 1 and 2 apply. R Lemma 4. If Rn×d max1⩽i⩽n |xi |6+4p dρX (X) < +∞ with p ⩾ 1, then ϕ is bounded and of class C p and its differential is bounded and Lipschitz. 4.2.1

Uniform convergence towards hardmax attention

The aim of this subsection is to obtain an analog of (ii) − (iii) from Lemma 3 in the case of attention. This property amounts to the uniform C 1 convergence when A becomes large of the parameters-to-predictor mapping associated to the softmax attention towards the one associated to the hardmax attention. Below, we prove the uniform convergence of the mapping itself. To our knowledge, this result is new and of independent interest. We conjecture that the uniform convergence of the differential holds, but leave the proof of this result for future works. def. def. Given z ∈ Rn , we define zmax = max1⩽i⩽n zi and Σ(z) = {i ∈ {1, ..., n} | zi = zmax }. We also define σr : z 7→ σ(rz) and σ∞ : z 7→ limr→+∞ σ(rz) = (1/#Σ(z))1Σ(z) ; where 1Σ(z) is the n-dimensional vector 2 with i-th entry 1 if zi ∈ Σ(z) (that is zi = zmax ) and 0 otherwise. Abusing notation, we denote by Sd −1 the set of d × d matrices with ∥A∥ = 1, where ∥ · ∥ is the Frobenius norm. Finally, we define ψr and ψ∞ on 2 Sd −1 by ψr (A) = ψ(rA) and ψ∞ (A)(X) = lim ψ(rA)(X) = r→+∞

n X

[σ∞ (XAxn )]i xi =

i=1

1 #Σ(XAxn )

xi .

i∈Σ(XAxn )

We stress that, when Σ(XAxn ) = {i} for some i ∈ {1, ..., n} (that is to say, when i is the only index for which ⟨Axn , xj ⟩ is maximal), then ψ∞ (A)(X) = xi . The following proposition, whose proof is postponed to Section A.2, states the uniform convergence of softmax attention towards hardmax attention in the large A regime. Proposition 4.R Assume that ρX has a bounded continuous density with respect to the Lebesgue measure on Rn×d and that Rn×d max1⩽i⩽n |xi |2 dρX (X) < +∞. Then we have: lim

r→+∞

sup

∥ψ(rA) − ψ∞ (A)∥F = 0.

A∈Sd2 −1

R If Rn×d max1⩽i⩽n |xi |3 dρX (X) < +∞, one can show that for every bounded and continuous function R f : Rn×d → Rd the mapping gf : A 7→ Rn×d ⟨f (X), ψ(A)(X)⟩ρX (X)dX is of class C 1 and n X n X

Z ∇gf (A) =

Rn×d i=1 j=1 j̸=i

σi (XAxn )σj (XAxn )ρX (X)⟨f (X), xi ⟩(xi − xj )xTn dX.

When r → +∞, we conjecture that, under suitable assumptions on ρX , the mapping A 7→ r∇gf (rA) 2 converges uniformly on Sd −1 towards n X n Z X 1 ⟨f (X), xi ⟩ρX (X)(xi − xj )xTn dHnd−1 (X), (7) A 7→ α (A, X) i=1 j=1 Γij (A) ij j̸=i

where Γij (A) = {X ∈ Rn×d | ⟨Axn , xi − xj ⟩ = 0 and ⟨Axn , xi − xk ⟩ ⩾ 0 for all k ∈ {1, ..., n}} and p 2 T 2   p2|Axn | + |A (xi − xj )| αij (A, X) = |∇X ⟨Axn , xi − xj ⟩| = |Axn |2 + |AT (xn − xj ) + Axn |2  p |Axn |2 + |AT (xi − xn ) − Axn |2

if i ̸= n and j ̸= n, if i = n, if j = n.

(8)

The main difficulty towards a rigorous proof of this result is that the integrand in Equation (7) is singular when Axn = 0 and AT (xi − xj ) = 0. As a result, one has to ensure that ρX vanishes sufficiently fast around 2 these points to ensure some uniform integrability in A ∈ Sd −1 .

Conclusion

In this work, we studied the global convergence of gradient descent for wide shallow models with bounded nonlinearities. In the case of models with vector output weights, we reduced this question to a blow-up condition for a family of ordinary differential equations. The main question that we left open is the complete proof that this condition holds under mild assumptions. There are several natural perspectives to this work. Concerning the training of attention layers, other variants of attention such as masked attention could be considered. An interesting but challenging question is to study the traditional (Q, K, V ) parameterization, instead of optimizing on (A, V ) where A = K T Q. One could also investigate whether the abstract framework in which our results are proven allows to handle contexts with a variable (and possibly arbitrarily large) size. For training transformers, the Adam algorithm seems to significantly outperform gradient descent. Studying its global convergence properties is a highly interesting challenge. Weight decay (which amounts to a squared ℓ2 regularization on the weights) is also popular and is not covered by our analysis. Finally, activation functions that are smooth approximations of ReLU (such as SiLU / swish and GELU) have become increasingly popular in modern architectures. As they are neither positively 1-homogeneous nor bounded, they are not covered by our global convergence result or that of [CB18]. Investigating whether an analogous property can be proven in this case is particularly interesting.

Acknowledgments The work of G. Peyré was supported by the French government under the management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute). The works of G. Peyré and R. Petit were supported by the European Research Council (ERC project WOLF). 12

References [AGS09]

Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows: In Metric Spaces and in the Space of Probability Measures. Birkhäuser Basel, September 2009.

[AZLS19]

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

[B+ 22]

Bubacarr Bah et al. Learning deep linear neural networks: Riemannian gradient flows and convergence to global minimizers. Information and Inference: A Journal of the IMA, 11(1):307– 353, 2022.

[Bac17]

Francis Bach. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research, 18(1):629–681, 2017.

[BHL18]

Peter Bartlett, David Helmbold, and Philip Long. Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks. In International Conference on Machine Learning (ICML), pages 521–530. PMLR, 2018.

[CACP25]

Valérie Castin, Pierre Ablin, José Antonio Carrillo, and Gabriel Peyré. A Unified Perspective on the Dynamics of Deep Transformers, January 2025.

[CB18]

Lénaı̈c Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. In Advances in Neural Information Processing Systems (NeurIPS), 2018.

[CC+ 20]

Zixiang Chen, Yuan Cao, et al. How much over-parameterization is sufficient to learn deep relu networks? In International Conference on Learning Representations (ICLR), 2020.

[CCR11]

J. A. Cañizo, J. A. Carrillo, and J. Rosado. A well-posedness theory in measures for some kinetic models of collective motion. Mathematical Models and Methods in Applied Sciences, 21(03):515–539, March 2011.

[CDGS23]

Guillaume Carlier, Arnaud Dupuy, Alfred Galichon, and Yifei Sun. SISTA: Learning Optimal Transport Costs under Sparsity Constraints. Communications on Pure and Applied Mathematics, 76(9):1659–1677, 2023.

[CLD+ 21]

Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations (ICLR), 2021.

[COB19]

Lenaı̈c Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems (NeurIPS), pages 2937–2947, 2019.

[DLL+ 19]

Simon S. Du, Jason D. Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Learning Representations (ICLR), 2019.

[Dob79]

R. L. Dobrushin. Vlasov equations. Functional Analysis and Its Applications, 13(2):115–123, April 1979.

[EUD18]

Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3–11, November 2018.

[Eva15]

Lawrence C. Evans. Measure Theory and Fine Properties of Functions, Revised Edition. Chapman and Hall/CRC, Boca Raton, 2015.

[EW22]

Weinan E and Stephan Wojtowytsch. Representation formulas and pointwise properties for barron functions. Calculus of Variations and Partial Differential Equations, 61(2):46, 2022.

[GMMM20] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pages 14820–14830, 2020. [HG23]

Dan Hendrycks and Kevin Gimpel. Gaussian Error Linear Units (GELUs), June 2023.

[HM17]

Moritz Hardt and Tengyu Ma. Identity matters in deep learning. In International Conference on Learning Representations (ICLR), 2017.

[JGH18]

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2018.

[KVPF20]

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156–5165. PMLR, 2020.

[LXS+ 19]

Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha SohlDickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Advances in Neural Information Processing Systems (NeurIPS), 2019.

[LY17]

Yuanzhi Li and Yang Yuan. Convergence analysis of two-layer neural networks with relu activation. In Advances in Neural Information Processing Systems (NeurIPS), volume 30, pages 597–607, 2017.

[LZB20]

Chaoyue Liu, Libin Zhu, and Mikhail Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, 2020.

[Mag12]

Francesco Maggi. Sets of Finite Perimeter and Geometric Variational Problems: An Introduction to Geometric Measure Theory. Cambridge Studies in Advanced Mathematics. Cambridge University Press, Cambridge, 2012.

[MMN18]

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):E7665– E7671, August 2018.

[Ngu21]

Quynh Nguyen. On the proof of global convergence of gradient descent for deep relu networks with linear widths. In International Conference on Machine Learning (ICML), pages 8056–8062. PMLR, 2021.

[NP23]

Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. Mathematical Statistics and Learning, 6(3):201–357, 2023.

[PHA+ 21]

Hancheng Peng, Ramin Hasani, Alexander Amini, Daniela Rus, and Thomas Serre. Random feature attention. In ICLR, 2021.

[RV18]

Grant Rotskoff and Eric Vanden-Eijnden. Parameters as interacting particles: Long time convergence and asymptotic error scaling of neural networks. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. 14

[RZL17]

Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for Activation Functions, October 2017.

[San17]

Filippo Santambrogio. {Euclidean, metric, and Wasserstein} gradient flows: An overview. Bulletin of Mathematical Sciences, 7(1):87–154, April 2017.

[Ver26]

Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science. Cambridge University Press, 2nd edition, March 2026.

[Woj20]

Stephan Wojtowytsch. On the convergence of gradient descent training for two-layer relunetworks in the mean field regime, 2020.

[YH21]

Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning (ICML), pages 11727–11737. PMLR, 2021.

[ZCZG20]

Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes overparameterized deep relu networks. Machine Learning, 109:467–492, 2020.

[ZLG19]

Difan Zou, Philip M. Long, and Quanquan Gu. On the global convergence of training deep linear resnets. In International Conference on Learning Representations (ICLR), 2019.

Auxiliary results

We begin with the proof that, when Φ has the specific form used in Section 3, then Assumption 1 holds under a regularity assumption on ϕ. Lemma 5. Assume that Ω = Rdw × Rdθ and that Φ : Ω → F has the form Φ(w, θ) = ϕ(θ)w where ϕ(θ) ∈ L(Rdw ; F). If ϕ is differentiable with a bounded and Lipschitz differential, then Assumption 1 holds. Proof. We have that dΦ(w, θ) · (w′ , θ′ ) = ϕ(θ)w′ + [dϕ(θ) · θ′ ]w. As a result, we obtain that ∥dΦ(w, θ)∥ ⩽ (∥ϕ(θ)∥2 + |w|2 ∥dϕ(θ)∥2 )1/2 . Using that dϕ is bounded and Lipschitz, we obtain that ∥dΦ(u)∥ ≲ 1 + |u|. Now, defining ∆ = [dΦ(w1 , θ1 ) − dΦ(w2 , θ2 )] · (w′ , θ′ ) we also have ∆ = [ϕ(θ1 ) − ϕ(θ2 )]w′ + [(dϕ(θ1 ) − dϕ(θ2 )) · θ′ ]w1 + [dϕ(θ2 ) · θ′ ](w1 − w2 ). As a consequence, we obtain ∥dΦ(w1 , θ1 ) − dΦ(w2 , θ2 )∥ ⩽

∥dϕ∥2∞ |θ1 − θ2 |2 + (Lip(dϕ)|θ1 − θ2 ||w1 | + ∥dϕ∥∞ |w1 − w2 |)2

and a direct computation allows to conclude. Now, we proceed with the proof of Lemma 1, which states that the square loss and the cross-entropy loss are both compatible with Assumptions 1 and 2. Proof of Lemma 1. We argue as in the proof of [CB18, Lemma D.1]. First, a direction computation shows that, in both cases, the integral in the expression of R(f ) is finite for every f ∈ F. The convexity of R directly follows from the fact that, for every y ∈ Rdout , the mapping z 7→ ℓ(z, y) is convex. In both cases, z 7→ ℓ(z, y) is differentiable for every y and there exists L > 0 such that, for every (z1 , z2 , y) ∈ (Rdout )3 , it holds |∇z ℓ(z1 , y) − ∇z ℓ(z2 , y)| ⩽ L|z1 − z2 |. To see this, notice that ∇z ℓ(z, y) = z − y for the square loss so that we can take L = 1. For the cross-entropy loss, it holds ∇z ℓ(z, y) = −y + σ(z), where σ denotes the softmax function. By Lemma 8, we can take L = 2. Now, let us prove that the differential of R is given by R dR(f ) · h = Rdin ×Rdout ⟨∇z ℓ(f (x), y), h(x)⟩dρ(x, y). We have def.

∆(h) = |R(f + h) − R(f ) − dR(f ) · h| R = | Rdin ×Rdout [ℓ(f (x) + h(x), y) − ℓ(f (x), y) − ⟨∇z ℓ(f (x), y), h(x)⟩]dρ(x, y)|. Using that |ℓ(z + z ′ , y) − ℓ(z, y) − ∇z ℓ(z, y) · z ′ | ⩽ L|z ′ |2 /2, we obtain ∆(h) ⩽ (L/2)∥h∥2F . A similar computation also shows that dR is L-Lipschitz. Finally, we claim that |∇z ℓ(z, y)|2 ⩽ 2ℓ(z, y). This is immediate for the square loss. For the cross-entropy loss, since ρy is concentrated on Bk with k = dout , it is enough to show that property when y ∈ Bk , say y = eq . We use |a − 1|2 ⩽ − log(a) if a ∈ (0, 1) to obtain 2

ezq

|σ(z) − y| = Pk

j=1 e

zj j=1 e

i̸=q

−1

j=1 e

ezq

+ Pk

zj j=1 e

⩽ −2 log(σq (z)) = 2ℓ(z, y). 16

ezi

X i̸=q

ezq

= Pk

+ 2

ezq

⩽ Pk

−1

ezi

zj 2

−1

To conclude, we notice that |dR(f ).h| ⩽

Rdin ×Rdout

⩽ ∥h∥F

|∇z ℓ(f (x), y)||h(x)|dρ(x, y)

Rdin ×Rdout

|∇z ℓ(f (x), y)|2 dρ(x, y)

2R(f ),

so that dR is bounded on sublevel sets of R.

A.1

Fully connected networks

The two properties below are stated (without proof) in [CB18, Appendix D.3]. We provide a proof of both results for the sake of completeness. Lemma 6. Assume that ρx has a bounded continuous density with respect to the Lebesgue measure. If σ is the sigmoid function then when r → +∞ the mapping ϕr : θ ∈ Rdθ 7→ ϕ(rθ) ∈ F converges uniformly on Sdθ −1 towards ϕ∞ : θ 7→ [x 7→ 1{θ·x⩾0} ]. Proof. We fix η, R > 0 and split Rdin in three sets Eηi (i = 1, 2, 3) defined by Eη1 = {x ∈ Rdin | ⟨θ, x⟩ ⩾ η}, Eη2 = {x ∈ Rdin | ⟨θ, x⟩ ⩽ −η} and Eη3 = Rdin \ (Eη1 ∪ Eη2 ). If x ∈ Eη1 then |σ(rθ · x) − 1{θ·x⩾0} | = |1/(1 + e−rθ·x ) − 1| ⩽ e−rη . If x ∈ Eη2 we also have |σ(rθ · x) − 1{θ·x⩾0} | = 1/(1 + e−rθ·x ) ⩽ e−rη . As a consequence, we have ∥ϕr (θ) − ϕ∞ (θ)∥2F =

|σ(rθ · x) − 1{θ·x⩾0} |2 dρx (x) R + E 3 |σ(rθ · x) − 1{θ·x⩾0} |2 dρx (x)

Rdin −2rη

⩽e

−2rη

⩽e

+ 4ρx (Eη3 )

⩽ e−2rη + 4ρx (Eη3 ∩ Bdin (0, R)) + 4ρx (Rdin \ Bdin (0, R)). Now, we notice that since |θ| = 1, up to a rotation the set Eη3 ∩ Bdin (0, R) is included in the cylinder (−η, η) × Bdin −1 (0, R). As a consequence, we have ρx (Eη3 ∩ Bdin (0, R)) ⩽ 2η∥ρx ∥∞ Ldin −1 (Bdin −1 (0, R)). √ Taking for instance η = 1/ r, we obtain lim sup sup ∥ϕr (θ) − ϕ∞ (θ)∥F ⩽ 4ρx (Rdin \ Bdin (0, R)). r→+∞ θ∈Sdθ −1

As this holds for every R > 0 and the right-hand side converges to 0 as R → +∞, we obtain the result. In [CB18, Appendix D.3], the result below is stated under a finite moment assumption on ρx . However, the integrability of x 7→ f (x)ρx (x)x with respect to the surface measure on the hyperplane {⟨θ, x⟩} cannot be guaranteed without additional assumptions on ρx . For this reason, we introduce an additional polynomial decay assumption under which the result is indeed true. Lemma 7. Assume ρx has a finite moment of order 1, has a bounded continuous density with respect to the Lebesgue measure and |ρx (x)| ⩽ C(1 + |x|)−p for some C >R0 and p > din . Then for every continuous and bounded function f : Rdin → R, the mapping gf : θ ∈ RRdθ 7→ Rdin f (x)σ(⟨θ, x⟩)dρx (x) is of class C 1 and θ 7→ r∇gf (rθ) converges uniformly on Sdθ −1 towards θ 7→ {⟨θ,x⟩=0} f (x)ρx (x)xdHdin −1 (x). 17

Proof. The gradient of the integrand with respect to θ is f (x)σ ′ (⟨θ, z⟩)ρx (x)x, whose norm is bounded by ∥f ∥∞ ∥σ ′ ∥∞ |x|ρx (x). Since ρx has a finite moment of order 1, this last function is integrable, which shows that gf is of class C 1 and that for every θ ∈ Rdθ it holds Z f (x)σ ′ (⟨θ, x⟩)xρx (x)dx. ∇gf (θ) = Rdin dθ −1

For every r ⩾ 0 and θ ∈ S , the co-area formula [Eva15, Theorem 3.11] yields that the restriction of x 7→ f (x)ρx (x)x to {⟨θ, x⟩ = s} is Hdin −1 -integrable for almost every s ∈ R and Z Z +∞ rσ ′ (rs) f (x)ρx (x)xdHdin −1 (x) ds. r∇gf (rθ) = −∞

{⟨θ,x⟩=s}

dθ −1

For every s ∈ R and θ ∈ S , since p > din we obtain Z Z din −1 |f (x)||ρx (x)||x|dH (x) ⩽ ∥f ∥∞ {⟨θ,x⟩=s}

|ρx (x)||x|dHdin −1 (x)

{⟨θ,x⟩=s}

1 dHdin −1 (x) p−1 (1 + |x|) {⟨θ,x⟩=s} Z 1 ⩽ C∥f ∥∞ dHdin −1 (y) (1 + |y + se1 |)p−1 {y1 =0} Z 1 ⩽ C∥f ∥∞ dHdin −1 (y) p−1 (1 + |y|) {y1 =0} ⩽ C∥f ∥∞

< +∞, where we have used the change of variable y = Rθ−1 x − se1 with R Rθ a rotation with Rθ e1 = θ. Applying the dominated convergence theorem, we also obtain that (θ, s) 7→ {⟨θ,x⟩=s} f (x)ρx (x)xdHdin −1 (x) is continuous, and hence uniformly continuous on Sdθ −1 × [−1, 1]. Since Z +∞ Z 1/√r ′ rσ (rs)ds = 1 and rσ ′ (rs)ds −→ 1, √ −∞

r→+∞

−1/ r

the uniform boundedness and continuity obtained above allow to conclude.

A.2

Attention layers

Pn In all this subsection, we denote by σ : z ∈ Rn 7→ (ezi /( j=1 ezj ))ni=1 ∈ Rn the softmax function. We start with the proof that, under a finite moment assumption on the distribution of the input contexts ρX , the mapping ϕ is sufficiently regular. Proof of Lemma 4. As ψ(A)(X) is a convex combination of the (xi )1⩽i⩽n , using Jensen’s inequality we obtain n X e⟨Axn ,xi ⟩ 2 |ψ(A)(X)|2 ⩽ max |xi |2 . n P ⟨Ax ,x ⟩ |xi | ⩽ 1⩽i⩽n i=1 e n j j=1 def.

As a result, provided Rn×d max1⩽i⩽n |xi | dρX (X) < +∞, we have that ψ(A) ∈ H with H = L2 (ρX ; Rd ) for every A ∈ Rd×d and that supA ∥ψ(A)∥H < +∞. Pn Now, we notice that ψ(A)(X) = X T σ(XAxn ) where σ : z ∈ Rn 7→ (ezi /( j=1 ezj ))ni=1 ∈ Rn is the softmax function. We claim that ψ is differentiable with differential dψ given by R

[dψ(A).B](X) = X T (dσ(XAxn ) · (XBxn )) n n X X = σi (XAxn ) ⟨Bxn , xi ⟩ − σj (XAxn )⟨Bxn , xj ⟩ xi . i=1

j=1

We have def.

∆(B)2 = ∥ψ(A + B) − ψ(A) − dψ(A) · B∥2H Z = |X T (σ(X(A + B)xn ) − σ(XAxn ) − dσ(XAxn ) · (XBxn ))|2 dρX (X) Rn×d

Z ⩽

max |xi |2

" X n

1⩽i⩽n

Rn×d

#2 |σi (X(A + B)xn ) − σi (XAxn ) − dσi (XAxn ) · (XBxn )|

Now, using that σi (z + h) − σi (z) − dσi (z) · h = Lemma 8 below, we obtain ∆(B)2 ⩽

dρX (X).

i=1

R1 0

(1 − t)d2 σi (z + th) · (h, h)dt for every z, h ∈ Rn and

Z Rn×d

Z ⩽9 Rn×d

1⩽i⩽n

X 1⩽i⩽n

This shows that ϕ is differentiable with differential dϕ. A similar computation shows that dϕ is bounded and Lipschitz. Lemma 8. For every z, h ∈ Rn , it holds |dσ(z) · h|1 ⩽ 2|h|∞ and |d2 σ(z) · (h, h)|1 ⩽ 6|h|2∞ . Proof. We have dσ(z) · h = σ(z) ⊙ h − ⟨σ(z), h⟩σ(z). Using the fact that the components of σ(z) are positive and sum to one, we obtain that |dσ(z) · h|1 =

n X

σi (z)|hi − ⟨σ(z), h⟩| ⩽ max |hi − ⟨σ(z), h⟩| ⩽ 2|h|∞ . 1⩽i⩽n

i=1

Now, for every 1 ⩽ i ⩽ n, we have that d2 σi (z) · (h, h) = σi (z) (hi − ⟨h, σ(z)⟩)2 − ⟨σ(z), h ⊙ h⟩ + ⟨σ(z), h⟩2 . Using again that the components of σ(z) are positive and sum to one, we obtain that |d2 σ(z) · (h, h)|1 ⩽ max |hi − ⟨σ(z), h⟩|2 + |⟨σ(z), h ⊙ h⟩| + ⟨σ(z), h⟩2 1⩽i⩽n

= max ⟨σ(z), h − hi 1⟩2 + |⟨σ(z), h ⊙ h⟩| + ⟨σ(z), h⟩2 1⩽i⩽n

⩽ max |hi − hj |2 + 2|h|2∞ 1⩽i,j⩽n

⩽ 6|h|2∞ .

We now turn to the proof of Proposition 4. Our strategy is to restrict our study to contexts whose tokens all lie in a given ball by using that ρX has finite second order moments and the rest of the integrand is bounded. Then, given η > 0, we can directly bound the integral on the region of the input space where Σ(z) = {i} for some i and zmax −zj = zi −zj ⩾ η for every j ̸= i, where z = (⟨Axn , xi ⟩)1⩽i⩽n . The remaining region is contained in the region where there exists i ̸= j such that |zi − zj | ⩽ η. This is a region of small width around the zero set of some non-trivial polynomial, and we are able to show that its measure tends 2 to zero when η → 0 uniformly in A ∈ Sd −1 .

Proof of Proposition 4. In all the following, we set z = (⟨Axn , xi ⟩)1⩽i⩽n ∈ Rn . We notice that ! X X 1 σr,i (z) xi − xj |ψr (A)(X) − ψ∞ (A)(X)| = #Σ(z) i∈Σ(z) / j∈Σ(z) X ⩽ 2 max |xj | σr,i (z). 1⩽j⩽n

i∈Σ(z) /

Given R > 0 and η > 0, we split the input space Rn×d in three regions R1 = {X ∈ Rn×d | ∀i ∈ {1, ..., n}, |xi | > R}, R2 = {X ∈ Rn×d \ R1 | Σ(z) is a singleton and zmax − zi ⩾ η for every i ∈ / Σ(z)}, R3 = Rn×d \ (R1 ∪ R2 ). P Pn Bound on R1 . Since i∈Σ(z) σr,i (z) ⩽ i=1 σr,i (z) = 1, we have / Z Z 2 2 |ψr (A)(X) − ψ∞ (A)(X)| dρX (X) ⩽ 4 max |xi | dρX (X). R1

1⩽i⩽n

By our finite moment assumption on ρX , this quantity therefore goes to 0 as R → +∞. Bound on R2 . If Σ(z) is a singleton and zmax − zi ⩾ η for every i ∈ / Σ(z), then for every i ∈ / Σ(z), it holds P X e−r(zmax −zi ) 1 1 i∈Σ(z) / P . = ⩽ σr,i (z) = Pn rη −r(z −z ) −r(z −z ) −1 max j max i 1 + e /(n − 1) 1 + ( i∈Σ(z) e ) / j=1 e i∈Σ(z) /

As a consequence, we obtain Z |ψr (A)(X) − ψ∞ (A)(X)|2 dρX (X) ⩽ R2

4R2 (1 + erη /(n − 1))2

Bound on R3 . We notice that R3 is included in the set R′3 = {X ∈ Rn×d \ R1 | ∃i ̸= j s.t. |zi − zj | ⩽ η}. As ∥A∥ = 1, A is nonzero and its kernel is a linear subspace with co-dimension at least 1, so that the set R′3 ∩ {Axn = 0} has measure zero. If |zi − zj | ⩽ η and Axn ̸= 0, then up to a rotation xi − xj is contained in the cylinder Bd−1 (0, R) × (−η/|Axn |, η/|Axn |), where Bd−1 (0, R) denotes the (d − 1)-dimensional ball of radius R. Let us first assume that i ̸= n and j ̸= n. d d As a consequence of the above, denoting by R L the2dLebesgue measure on R and defining ε = η/|Axn |, we nd ′ d n−3 obtain that L (R3 ) ⩽ [L (Bd (0, R))] L (Exn )dxn where Bd (0,R) Exn = {(y1 , y2 ) ∈ (Bd (0, R))2 | y1 − y2 ∈ Bd−1 (0, R) × (−ε, −ε)}. Now, we observe that L2d (Exn ) ⩽ min([Ld (Bd (0, R))]2 , L2d (Ex′ n )) where Ex′ n = {(y1 , y1 + y2 ) | y1 ∈ Bd (0, R), y2 ∈ Bd−1 (0, R) × (−ε, −ε)}. Since the map f : (x, z) 7→ (x, x + z) is injective and its Jacobian determinant is equal to one, the area formula [Mag12, Theorem 8.1] yields L2d (Ex′ n ) = 2εLd (Bd (0, R))Ld−1 (Bd−1 (0, R)). Since ∥A∥ = 1, the √ largest singular value of A is at least 1/ d. Applying a rotation to obtain that the first coordinate of the integration variable is along the right singular vector associated to the largest singular value, we obtain Z Z 2d 2d−1 L (Exn )dxn ≲ R min(R, η/|y1 |)dy. Bd (0,R)

Bd (0,R)

Final step. Collecting all three terms, we see that our final bound does not depend on A. Taking (for √ instance) η = 1/ r, we see that the integral on R2 goes to 0 as r → 0, and, by an application of the dominated convergence theorem, the same holds for the integral on R3 . As a result, we obtain that the limit is smaller than the integral on R1 . As this holds for every choice of R > 0 and this quantity goes to 0 when R → 0, we obtain the result. Finally, to deal with the case where i or j equals n, without loss of generality, we can assume i = n and j = 1. Arguing as above, we obtain Z nd ′ n−2 L (R3 ) ⩽ [L(Bd (0, R)] Ld (Exn )dxn Bd (0,R)

where Exn = {x1 ∈ Bd (0, R) | x1 − xn ∈ Bd−1 (0, R) × (−ε, ε)}. Arguing in a similar fashion as above finally allows to conclude.

Proof of Theorem 1 To prove Theorem 1, we show the existence and uniqueness of a weak solution to the continuity equation ∂t µt = −div(µt vt ) with vt (u) = −∇F ′ (µt )(u)

by a fixed point argument. These solutions should be understood in the weak sense, that is to say, for every test function φ : (0, +∞) × Ω → R, it holds Z +∞ Z [∂t φt (u) + ⟨∇φt (u), vt (u)⟩]dµt (u)dt = 0. 0

Ω

In Section B.1, we introduce some useful notions that are used in the proof. The core of the proof is given in Section B.2, while the technical lemmas on which it relies are proved in Section B.3.

B.1

Preliminaries

R If µ ∈ P(Ω), we denote by m2 (µ) = Ω |u|2 dµ(u) its second order moment. We denote by P2 (Ω) the set of probability measures on Ω that have finite second order moments. We recall that (P2 (Ω), W2 ) is a complete metric space. Sub-Gaussian norm.

We define the sub-Gaussian norm of a probability measure µ ∈ P(Ω) by ∥µ∥ψ2 = inf{c > 0 |

R Ω

e|u| /c dµ(u) ⩽ 2}

and say that it is sub-Gaussian if ∥µ∥ψ2 < +∞. We stress that, when d ⩾ 2, this definition differs from the usual definition of sub-Gaussianity for random vectors, which can be found for instance in [Ver26, Section 3.4]. Namely, if X is a random vector with law µ, then with our definition ∥µ∥ψ2 = inf{c > 0 | E[|X|2 /c2 ] ⩽ 2}, whereas the sub-Gaussian norm of X is usually defined as sup|v|⩽1 inf{c > 0 | E[(v T X)2 /c2 ] ⩽ 2}. Our choice is motivated by the following lemma, which shows that with our definition, the pushforward operation by a map with linear growth preserves sub-Gaussianity. Lemma 9. If µ ∈ P(Ω) is sub-Gaussian and T : Ω → Ω satisfies |T (u)| ⩽ C(1 + |u|) for some C > 0, we have p ∥T# µ∥ψ2 ⩽ 2C max ∥µ∥ψ2 , 1/ ln(2) . (9)

Proof. Let X be a random vector with law µ. As a result, |X| is sub-Gaussian with sub-Gaussian norm 2 2 ∥µ∥ψ2 . By [Ver26, Prop. 2.6.1], for every t > 0, we have that P(|X| ⩾ t) ⩽ 2e−t /∥µ∥ψ2 . If t ⩾ 2C, then |T (x)| ⩾ t implies |x| ⩾ 1, which yields |T (x)| ⩽ 2C|x|. As a result, we obtain 2

P(|T (X)| ⩾ t) ⩽ P(2C|X| ⩾ t) ⩽ 2e−t /(2C∥µ∥ψ2 ) . p If t ⩽ 2C, we use the trivial bound P(|T (X)| ⩾ t) ⩽ 1. Taking K = 2C max(∥µ∥ψ2 , 1/ ln(2)), we see that 2 2 P(|T (X)| ⩾ t) ⩽ e−t /K for every t > 0. Applying the same proposition as above, we obtain the result. Continuity of the predictor for W2 . The following lemma shows that the measure-to-predictor mapping R µ 7→ Ω Φdµ is continuous with respect to the W2 distance. Lemma 10. Under Assumption 1, for every µ1 , µ2 ∈ P2 (Ω), it holds R R ∥ Ω Φdµ1 − Ω Φdµ2 ∥F ≲ (1 + m2 (µ1 ) + m2 (µ2 ))W2 (µ1 , µ2 ). Proof. Considering Π an optimal transport plan (for the quadratic cost) between µ1 and µ2 , we obtain R R R ∥ Ω Φdµ1 − Ω Φdµ2 ∥ = ∥ Ω×Ω [Φ(u1 ) − Φ(u2 )]dΠ(u1 , u2 )∥ R ⩽ Ω×Ω ∥Φ(u1 ) − Φ(u2 )∥dΠ(u1 , u2 ) R ≲ Ω×Ω (1 + max(|u1 |, |u2 |))|u1 − u2 |dΠ(u1 , u2 ) qR qR 2 dΠ(u , u ) (1 + max(|u |, |u |)) |u1 − u2 |2 dΠ(u1 , u2 ) ⩽ 1 2 1 2 Ω×Ω Ω×Ω qR = (1 + max(|u1 |, |u2 |))2 dΠ(u1 , u2 )W2 (µ1 , µ2 ) Ω×Ω ≲ (1 + m2 (µ1 ) + m2 (µ2 ))W2 (µ1 , µ2 ).

B.2

Core of the proof

In all the following, we work under Assumption 1. We set an initial condition µ̄0 ∈ P2 (Ω) and α > 0 such that F (µ̄0 ) ⩽ α. Truncation of the objective. We claim that there exists a twice differentiable function ξ such that ξ(x) = x if x ∈ [0, α], g is (strictly) increasing on [α, 2α], ξ(x) = 2α if x ⩾ 2α, ∥ξ ′ ∥∞ ⩽ 3/2 and ∥ξ ′′ ∥∞ ⩽ 4/α. One can check that a possible choice for ξ|[α,2α] is ξ(x) = x +

(x − α)(2α − x)(1 − cos(π(x − α)/α)) . 2α

We define R̃ = ξ ◦ R and F̃ = ξ ◦ F . We notice that F̃ = F on {F ⩽ α} = {F̃ ⩽ α}. As a consequence, for every T > 0, a curve (µt )0⩽t⩽T with µ0 = µ̄0 is a Wasserstein gradient flow for F if and only if it is a gradient flow for F̃ . This trick is standard in the literature on optimization (see for instance [CDGS23, Equation 35]). Construction of the flow map. Let T > 0 and µ ∈ C([0, T ]; P2 (Ω)). The continuity of µ implies that t 7→ m2 (µt ) is continuous, so that m2 (µ) = sup0⩽t⩽T m2 (µt ) is finite. There exists a unique solution to ( ṙ(t) = −∇F̃ ′ (µt )(r(t)), 0 ⩽ t ⩽ T, r(0) = u. Indeed, Lemma 11 below shows that the velocity v(t, u) = −∇F̃ ′ (µ)(u) is continuous, locally Lipschitz in space, and has linear growth in space (uniformly in time), so that traditional Cauchy theory allows to conclude. We define the associated flow map Xt [µ] : u 7→ r(t). 22

Fixed point argument on [0, T ].

We set an initial condition µ̄0 ∈ P2 (Ω) and define the map

F : C([0, T ]; P2 (Ω)) → C([0, T ]; P2 (Ω)) µ 7→ (Xt [µ]# µ̄0 )0⩽t⩽T . Our aim is to find a complete metric space X ⊂ C([0, T ]; P2 (Ω)) such that F (X ) ⊂ X and F is a contraction on X when T is sufficiently small. This will allow us to apply a fixed point argument which will yield the existence and uniqueness of the solution on [0, T ]. In Section B.3 below, we show that there exists a continuous function f : (R+ )2 → R+ that is non-decreasing in each variable and such that, if µ̄0 is sub-Gaussian and X = {µ ∈ C([0, T ]; P2 (Ω)) | ∀t ∈ [0, T ], ∥µt ∥ψ2 ⩽ f (t, ∥µ̄0 ∥ψ2 )} (10) then F(X ) ⊂ X . We also show that, equipped with the distance W2,T : (µ, ν) 7→ sup0⩽t⩽T W2 (µt , νt ), the space X is complete. Finally, we prove that there exists a positive continuous function g : R+ × R+ → R+ satisfying g(T, K) → 0 as T → 0 for every K > 0 and W2,T (F(µ1 ), F(µ2 )) ⩽ g(T, K0 )W2,T (µ1 , µ2 ) for every K0 such that ∥µ̄0 ∥ψ2 ⩽ K0 . As a result, for every sub-Gaussian initial condition µ̄0 , there exists T > 0 such that F is a contraction and hence has a unique fixed point. This shows the existence and uniqueness of the solution on [0, T ]. Construction of a global solution. Now, we can iterate the previous argument by applying the result again with the updated initial condition ΦT [µ]# µ̄0 , which satisfies F (ΦT [µ]# µ̄0 ) ⩽ F (µ̄0 ) ⩽ α, so that α does not need to be updated. But since the sub-Gaussian norm of the initial condition increases at each step, the time interval that we can add at each step decreases. Assume by contradiction that this procedure fails to guarantee existence beyond some time Tlim > 0. Reasoning as in Lemma 13, we have that ∥µt ∥ψ2 ⩽ f (t, ∥µ̄0 ∥ψ2 ) for every t ∈ [0, Tlim ). Let T ′ > 0 be such that g(T ′ , f (Tlim , ∥µ̄0 ∥ψ2 )) < 1. Since f is nondecreasing in its first variable, we have that ∥µTlim −T ′ /2 ∥ψ2 ⩽ f (Tlim − T ′ /2, ∥µ̄0 ∥ψ2 ) ⩽ f (Tlim , ∥µ̄0 ∥ψ2 ) and we can hence apply the previous argument to extend µ|[0,Tlim −T ′ /2] to [0, Tlim + T ′ /2], which gives a contradiction.

B.3

Technical lemmas

By Assumption 1, R′ is bounded on sublevel sets of R. There hence exists a constant C > 0 depending only on α such that ∥R̃′ (h)∥F ⩽ (3/2)∥R′ (h)∥F ⩽ C for every h ∈ F such that R(h) ⩽ 2α. Since R̃′ (h) = 0 if R(h) > 2α, the previous bound in fact holds for every h ∈ F . Lemma 11. Let α > 0, T > 0 and µ ∈ C([0, T ]; P2 (Ω)). Defining v : [0, T ] × Ω → Rd by v(t, u) = −∇F̃ ′ (µt )(u), we have that v is continuous, locally Lipschitz in space, and has linear growth in space (uniformly in time). R Proof. By Lemma 10, Assumption 1 and the fact that ξ ′ is continuous, we have that t 7→ R̃′ ( Ω Φdµt ) and dΦ are continuous. Using the continuity of h1R, h2 7→ ⟨h1 , h2 ⟩F , we obtain the continuity of v. Now, it holds |v(t, u1 ) − v(t, u2 )| ⩽ ∥R̃′ ( Ω Φdµt )∥F ∥dΦ(u1 ) − dΦ(u2 )∥. Using Assumption 1 (which guarantees that dΦ is locally R Lipschitz), we obtain that v is locally Lipschitz in space. Finally, using that ∥R̃′ ( Ω Φdµt )∥F ⩽ C and the fact that Assumption 1 guarantees that dΦ has linear growth, the result is proved. Lemma 12. If T > 0 and µ̄0 ∈ P(Ω) is sub-Gaussian, there exists a constant C(T ) > 0 depending on α and T such that, for every µ ∈ X , it holds |Xt [µ](u)| ⩽ C(T )(1 + |u|) for every t ∈ [0, T ] and u ∈ Ω. Rt Proof. Using Assumption 1 and Xt [µ](u) = u + 0 v[µt ](Xs [µ](u))ds, we obtain: Rt R |Xt [µ](u)| ⩽ |u| + 0 ∥R̃′ ( Ω Φdµs )∥F ∥dΦ(Xs [µ](u)∥L(Rd ;F ) ds Z t Z t ⩽ |u| + C ′ (1 + |Xs [µ](u)|)ds = |u| + C ′ t + C ′ |Xs [µ](u)|ds, 0

for some constant C ′ > 0. Applying Grönwall’s inequality yields the result. Using Lemmas 9 and 12, we immediately obtain the following result, which shows that F (X ) ⊂ X . Lemma 13. There exists a continuous function f : (R+ )2 → R+ that is non-decreasing in each variable and such that, for every µ ∈ X , defining ν = F(µ), we have ∥νt ∥ψ2 ⩽ f (t, ∥µ̄0 ∥ψ2 ). Lemma 14. The space (X , W2,T ) with X defined by (10) and f as in Lemma 13 is a complete metric space. Proof. Let (µn )n⩾0 be a Cauchy sequence in (X , W2,T ). As X is a subset of the complete metric space C([0, T ]; P2 (Ω)), we know that (µn )n⩾0 converges for W2,T towards µ∗ ∈ C([0, T ]; P2 (Ω)). Now, for every c > 0, the mapping u 7→ exp(|u|2 /c2 ) is continuous (hence lower semi-continuous) and positive (hence bounded from below). Since (µnt )n⩾0 converges weakly towards µ∗t for every t ∈ [0, T ], we obtain by the Portmanteau theorem that R R exp(|u|2 /c2 )dµ∗t (u) ⩽ lim inf n→+∞ Ω exp(|u|2 /c2 )dµnt (u) (11) Ω R for every c > 0. Taking c = f (t, ∥µ̄0 ∥ψ2 ) yields Ω exp(|u|2 /c2 )dµ∗t (u) ⩽ 2 and hence ∥µ∗t ∥ψ2 ⩽ c = f (t, ∥µ̄0 ∥ψ2 ). This shows that µ∗ ∈ X and concludes the proof. Lemma 15. There exists a positive continuous function g : R+ ×R+ → R+ satisfying g(T, K) → 0 as T → 0 for every K > 0 such that, for every µ1 , µ2 ∈ X , we have W2,T (F(µ1 ), F(µ2 )) ⩽ g(T, K0 )W2,T (µ1 , µ2 ) for every K0 such that ∥µ̄0 ∥ψ2 ⩽ K0 . Proof. Let Xt1 and Xt2 be the flow maps associated to µ1 and µ2 and K0 such that ∥µ̄0 ∥ψ2 ⩽ K0 . Using the Rt fact that Xti (u) = u + 0 v[µit ](Xsi (u))ds, we obtain that |Xt1 (u) − Xt2 (u)| ⩽

Z t

|v[µ1s ](Xs1 (u)) − v[µ2s ](Xs2 (u))|ds

Z t ⩽

[|v[µ1s ](Xs1 (u)) − v[µ1s ](Xs2 (u))| + |v[µ1s ](Xs2 (u)) − v[µ2s ](Xs2 (u))|]ds.

To bound the first term, we use the boundedness of R̃′ , the local Lipschitz bound for dΦ from Assumption 1 and the bound from Lemma 12 to obtain Z t Z t 1 1 1 2 |v[µs ](Xs (u)) − v[µs ](Xs (u))|ds ⩽ C(T )(1 + |u|) |Xs1 (u) − Xs2 (u)|ds. 0

To bound the second term, we first notice that, for every h1 , h2 ∈ F, we have: ∥R̃′ (h1 ) − R̃′ (h2 )∥F ⩽ ∥ξ ′ ∥∞ ∥R′ (h1 ) − R′ (h2 )∥F + ∥R′ (h2 )∥F |ξ ′ (R(h1 )) − ξ ′ (R(h2 ))|. Moreover, it holds R ∥ Ω Φdµis ∥F ≲ 1 + m2 (µis ) ≲ 1 + ∥µis ∥2ψ2 ⩽ 1 + f (T, K0 )2 ,

(12)

so that { Ω Φdµis , (i, s) ∈ {1, 2} × [0, T ]} is included in some bounded set BT,K0 . Since R′ is Lipschitz and hence bounded on BT,K0 , using the fact that ∥ξ ′′ ∥∞ ⩽ 4/α and Lemma 10, we obtain R

R R R R ∥R̃′ ( Ω Φdµ1s ) − R̃′ ( Ω Φdµ2s )∥F ⩽ C(T, K0 )∥ Ω Φdµ1s − Ω Φdµ2s ∥F ⩽ C(T, K0 )W2 (µ1s , µ2s ).

Finally, using the fact that ∥dΦ(Xs2 (u)∥ ≲ 1+|Xs2 (u)| ⩽ C(T )(1+|u|) (the last inequality being a consequence of Lemma 12), we obtain Z t

|v[µ1s ](Xs2 (u)) − v[µ2s ](Xs2 (u))|ds ⩽ C(T, K0 )(1 + |u|)

Z t

W2 (µ1s , µ2s )ds.

Combining the two bounds together and applying Grönwall’s inequality, we finally obtain Z t exp(C(T, K0 )(1 + |u|)(t − s))W2 (µ1s , µ2s )ds. |Xt1 (u) − Xt2 (u)| ⩽ C(T, K0 )(1 + |u|) 0

Now, using the transference plan (Xt1 × Xt2 )# µ̄0 , we obtain Z 2 1 2 W2 ((Xt )# µ̄0 , (Xt )# µ̄0 ) ⩽ |Xt1 (u) − Xt2 (u)|2 dµ̄0 (u) Ω

Z ⩽

C(T, K0 )2 (1 + |u|)2

Z t

2 eC(T,K0 )(1+|u|)(t−s) W2 (µ1s , µ2s ) dµ̄0 (u)

Ω 2

⩽ C(T, K0 ) W22 (µ1 , µ2 )

Z (1 + |u|) Ω

Z t

C(T,K0 )(1+|u|)(t−s)

2 ds dµ̄0 (u)

2 C(T,K0 )(1+|u|)t Z e −1 dµ̄0 (u) = C(T, K0 )2 W22 (µ1 , µ2 ) (1 + |u|)2 C(T, K0 )(1 + |u|) Ω Z h i2 eC(T,K0 )(1+|u|)t − 1 dµ̄0 (u). = W22 (µ1 , µ2 ) Ω

Using Cauchy-Schwarz inequality and the fact that

R Ω

exp(−|u|2 /K02 )dµ̄0 (u) ⩽ 2, we obtain (11) with 2

g(T, K) = 2 sup (eC(T,K)(1+|u|)T − 1)2 e−|u| /K , u∈Ω

which satisfies g(T, K) → 0 as T → 0 for every K > 0.

Proof of Proposition 1 with unbounded sublevel sets

The main issue when ∂K is not bounded is that |∇g(θt )| could become arbitrarily small when |θt | is large. This would prevent us from ensuring that θt cannot exit K. To circumvent this, it is suggested in [CB18] to leverage the convergence assumptions (ii) − (iii) and the Sard regularity of the limit g∞ in (iv). However, the claim, in the proof of [CB18, Proposition C.4], that inf θ∈∂K |∇g(θ)| > 0 provided −η is a regular value of both g and g∞ is incorrect. In the following, we show that, despite this issue, one can still find a positive lower bound for ẇt , which allows to correct [CB18, Proposition C.4] and prove Proposition 1.

C.1

Setting

Let g be as in Proposition 1 and assume it is not constant and its range contains a non-empty open interval included in R∗− . Further assume that {θ ∈ Rdθ | g(θ) = −η ′ } is unbounded for every η ′ > 0 such that −η ′ is a regular value of g (the other case has been treated in the proof of Proposition 1 above). Choice of a common regular value. First, let us show that there exists η > 0 such that −η is a regular value of g and g∞ . To this aim, we prove that, if a level set {θ ∈ Rdθ | g(θ) = −η ′ } is unbounded for some η ′ > 0, then −η ′ belongs to the range of g∞ . In this case, there exists (θn )n⩾0 such that g(θn ) = −η ′ for every n ⩾ 0 and |θn | → +∞. Let us define rn = |θn | and φn = θn /|θn | so that θn = rn φn . Since |φn | = 1

for every n ⩾ 0, up to the extraction of a subsequence, the sequence (φn )n⩾0 converges to φ ∈ Sdθ −1 as n → +∞. Given that |g∞ (φ) − (−η ′ )| = |g∞ (φ) − g(rn φn )| ⩽ |g∞ (φ) − g∞ (φn )| + |g∞ (φn ) − g(rn φn )| ⩽ |g∞ (φ) − g∞ (φn )| + supψ∈Sdθ −1 |g∞ (ψ) − g(rn ψ)|, the uniform convergence of (ii) (which also yields the continuity of g∞ ) gives that g∞ (φ) = −η ′ , so that −η ′ belongs to the range of g∞ . As a consequence, the set of η ′ > 0 such that −η ′ is a regular value of g, which has positive Lebesgue measure, is included in the range of g∞ . Since the set of critical values of g∞ has Lebesgue measure zero, we obtain the existence of η > 0 such that −η is a regular value of g and g∞ . In all the following, we fix such a η and define K = {θ ∈ Rdθ | g(θ) ⩽ −η}. Definition of the constants and choice of parameters.

For every r̄ > 0, we have βr̄ , β∞ > 0 with

βr̄ = inf {|∇g(θ)| | θ ∈ Rdθ , |θ| ⩽ r̄, g(θ) = −η} and β∞ = inf {|∇S g∞ (φ)| | φ ∈ Sdθ −1 , g∞ (φ) = −η}, def.

def.

where ∇S g∞ denotes the element of {φ}⊥ representing the differential of g∞ at φ ∈ Sdθ −1 . Since −η is a ′ regular value of g∞ and Sdθ −1 is compact, we also have the existence of γ∞ > 0 such that inf {|∇g∞ (φ)| | φ ∈ Sdθ −1 , g∞ (φ) ∈ [−η − γ∞ , −η + γ∞ ]} ⩾ β∞ /2

(13)

′ ′ , η/4). We define Cw = ∥g∥∞ + 1 and . We fix γ∞ = min(γ∞ for every γ∞ ⩽ γ∞

Cθ = max(∥∇g∥∞ , supr⩾0 ∥r∇g(r·)∥∞ ), this last quantity being finite by (iii). We fix α > 0 small enough to have 2(4 + Cθ )(α + Cw α2 /2) ⩽ 1 and define C1 = 9 + Cw (4 + Cθ )α2 , C2 = 2(4 + Cθ ) and r τ = max 1,

(14)

C1 η . C2

(15)

Finally, we define C2 3β 2 . c = min γ∞ , ∞ log 1 + α 32C2 C1 Using (ii) and (iii), we choose r̄ ⩾ 1 large enough to have c η , , γ∞ , sup |g(rφ) − g∞ (φ)| ⩽ min 4 4 φ∈Sdθ −1

(17)

sup |rproj{φ}⊥ (∇g(rφ)) − ∇S g∞ (φ)| ⩽ min

φ∈Sdθ −1

for every r ⩾ r̄. We also take

ε ⩽ min

2 β∞ η , ,1 , 16∥∇S g∞ ∥∞ 4

2 β∞ η 1 , , . 16∥∇S g∞ ∥∞ 4 r̄

(16)

(18)

(19)

C.2

Preliminary lemmas

In all the following, given θt ∈ Rdθ , we use the notation φt = θt /|θt | ∈ Sdθ −1 . Lemma 16. It holds |ẇt | ⩽ Cw ,

(20)

|θ̇t | ⩽ |wt |[ε + Cθ min(1, 1/|θt |)].

(21)

Proof. Using (3) and the boundedness of g, (19) and r̄ ⩾ 1, we obtain: |ẇt | = |gt (θt )| ⩽ |g(θt )| + |gt (θt ) − g(θt )| ⩽ ∥g∥∞ + ε ⩽ ∥g∥∞ + 1. In the same way, we obtain |θ̇t | = |wt ||∇gt (θt )| ⩽ |wt |[|∇g(θt )| + |∇gt (θt ) − ∇g(θt )|] ⩽ |wt |[∥∇g∥∞ + ε]. Finally, we can use ∇g(θt ) = (1/|θt |)|θt |∇g(|θt |φt ) to get the improved bound supr⩾0 ∥r∇g(r·)∥∞ , |∇g(θt )| ⩽ min ∥∇g∥∞ , |θt | which in turn yields (21). The following lemmas shows that g∞ (ϕt ) is strictly decreasing near the regular value −η, and hence, decreases by a constant in finite time under appropriate bounds on θt , wt . Lemma 17. If wt > 0 and θt ∈ [r̄, +∞) is such that g∞ (φt ) ∈ [−η − γ∞ , −η + γ∞ ], then d β 2 wt g∞ (φt ) ⩽ − ∞ < 0. dt 8 |θt |2

(22)

Proof. Using (3), (13), (18) and (19), we obtain d g∞ (φt ) = ⟨∇S g∞ (φt ), φ̇t ⟩ dt wt ⟨∇S g∞ (φt ), proj{φt }⊥ (∇gt (θt ))⟩ =− |θt | wt = − 2 |∇S g∞ (φt )|2 + ⟨∇S g∞ (φt ), |θt |proj{φt }⊥ (∇g(θt )) − ∇S g∞ (φt )⟩ |θt | + ⟨∇S g∞ (φt ), |θt |proj{φt }⊥ (∇gt (θt )) − |θt |proj{φt }⊥ (∇g(θt ))⟩ 2 2 wt β β∞ ⩽ − 2 ∞ − ∥∇S g∞ ∥∞ − ∥∇S g∞ ∥∞ ε |θt | 4 16∥∇S g∞ ∥∞ 2 β wt ⩽− ∞ . 8 |θt |2

Lemma 18. For every t0 ⩾ 0 with θt0 ∈ [2r̄, 3r̄], g∞ (φt0 ) ∈ [−η −γ∞ , −η +γ∞ ] and wt0 ⩾ τ r̄ (with τ defined in (15)), there exists t1 > t0 such that t 7→ g∞ (φt ) is decreasing and |θt | ∈ (r̄, 4r̄) on [t0 , t1 ]. Moreover, we have g∞ (φt1 ) ⩽ g∞ (φt0 ) − c with c defined in (16).

Proof. Without loss of generality, we can assume t0 = 0. Let δ = inf{t ⩾ 0 | θt = r̄ or θt = 4r̄}. We have δ > 0. For t ∈ [0, δ], by (19) it holds ε ⩽ 1/r̄ ⩽ 4/|θt |, so that (20) and (21) yield |θ̇t | ⩽ wt (ε + Cθ /|θt |) ⩽ (4 + Cθ )

w0 + Cw t wt ⩽ (4 + Cθ ) . |θt | |θt |

As a result, we obtain d |θt |2 ⩽ 2(4 + Cθ )(w0 + Cw t), dt which in turn yields ||θt |2 − |θt0 |2 | ⩽ 2(4 + Cθ )(w0 t + (Cw /2)t2 ). Provided t ⩽ t′1 with t′1 = αr̄2 /w0 and using (14), we obtain |θt | ∈ (r̄, 4r̄). As a result, δ ⩾ t′1 and the previous bound holds for every t ∈ [0, t′1 ]. By Lemma 17, we have that t 7→ g∞ (φt ) is decreasing on [0, t′1 ], or until the first time t′′1 for which g∞ (φt′′1 ) < −η − γ∞ . If t′′1 ⩾ t′1 then using (17), (22) and γ∞ ⩽ η/4 we obtain ẇt = −gt (θt ) ⩾ −g(θt ) − ε ⩾ −g∞ (φt ) − ε − sup ∥g(r·) − g∞ ∥∞ ⩾ η − γ∞ − ε − sup ∥g(r·) − g∞ ∥∞ ⩾ η/4, r⩾r̄

r⩾r̄

which in turn yields β2 g∞ (φt1 ) ⩽ g∞ (φ0 ) − ∞ 8

Z t′1

β2 ⩽ g∞ (φ0 ) − ∞ 8

Z t′1

0 ′

⩽ g∞ (φ0 ) −

2 Z t1 β∞

wt dt |θt |2 w0 + (η/4)t dt 2 |θ0 | + 2(4 + Cθ )(w0 t + (Cw /2)t2 ) w0 + (η/4)t dt. C1 r̄2 + C2 w0 t

Now, using that (b/d)x + [(ad − bc)/d ] log(|c + dx|) is a primitive of x 7→ (a + bx)/(c + dx), we also obtain t′1 Z t′1 η C2 w02 − (η/4)C1 r̄2 w0 + (η/4)t 2 dt = t+ log(C1 r̄ + C2 w0 t) C1 r̄2 + C2 w0 t 4C2 w0 C22 w02 0 t=0 C2 w0 t′1 C2 w02 − (η/4)C1 r̄2 η ′ log 1 + t + = 4C2 w0 1 C22 w02 C1 r̄2 1 C2 C1 ηr̄2 ⩾ log 1 + α 1− . C2 4C2 w02 C1 As a result, using w0 ⩾ τ r̄ and (15), we obtain the result with t1 = min(t′1 , t′′1 ).

C.3

Proof of Proposition 1

Assume that w0 ⩾ τ r̄ and θ0 ∈ K. In the following, we show that g(θt ) ⩽ −η/2 for every t ⩾ 0. To do this, we use the above lemmas to show that, if g(θt0 ) = −η for some t0 > 0, then there exists t1 > t0 such that g(θt1 ) < −η and g(θt ) ⩽ −η/2 on [t0 , t1 ]. We recall that, if g(θt ) ⩽ −η/2 for some t, then ẇt ⩾ η/4. Small regime. If |θt0 | ⩽ 2r̄, by compactness of ∂K ∩ B(0, 2r̄) we can conclude as in the case where ∂K is bounded that we cannot exit K. Medium regime. If |θt0 | ∈ [2r̄, 3r̄] and g∞ (φt0 ) ∈ [−η − γ∞ , −η + min(γ∞ , c/2)] (which, by (17), is in particular the case if g(θt0 ) = −η) then, since wt0 ⩾ τ r̄, we can apply Lemma 18 to obtain the existence of t1 > t0 such that t 7→ g∞ (φt ) is decreasing on [t0 , t1 ] and g∞ (φt1 ) ⩽ g∞ (φt0 ) − c. Using (17), we obtain g(θt1 ) ⩽ g∞ (φt1 ) + sup ∥g(r·) − g∞ ∥∞ r⩾r̄

⩽ g∞ (φt0 ) − c + sup ∥g(r·) − g∞ ∥∞ r⩾r̄

c c ⩽ −η + − c + < −η. 2 4 28

Moreover, for every t ∈ [t0 , t1 ], using (17) again, the following also holds: g(θt ) ⩽ g∞ (φt ) + sup ∥g(r·) − g∞ ∥∞ ⩽ g∞ (φt0 ) + sup ∥g(r·) − g∞ ∥∞ ⩽ −η + γ∞ + r⩾r̄

r⩾r̄

η η ⩽− . 4 2

Consequently, in the medium regime, θt can exit K but re-enters K in finite time with |θt | ⩾ r̄ and, in the meantime, we still have g(θt ) ⩽ −η/2. Large regime. If |θt0 | ⩾ 3r̄ and g(θt0 ) = −η, then reasoning as above we obtain g∞ (φt0 ) ⩽ −η + γ∞ , so that by Lemma 17 the quantity g∞ (φt ) decreases until one of the following happens: • g∞ (φt ) < −η − γ∞ , which yields g(θt ) < −η − γ∞ + supr⩾r̄ ∥g(r·) − g∞ ∥∞ ⩽ −η by (17). In the meantime reasoning as above we also had g(θt ) ⩽ −η/2. • |θt | = r̄. For this to happen we would have to enter the medium regime before. We would enter it with g∞ (φt ) ⩽ g∞ (φt0 ) ⩽ −η + supr⩾r̄ ∥g(r·) − g∞ ∥∞ ⩽ −η + min(γ∞ , c/2), so that the condition to ensure we re-enter K is satisfied. Conclusion.

From the reasoning above, we have that, taking n o A = (w0 , θ0 ) ∈ R × Rdθ | w0 ⩾ τ r̄, g(θ0 ) ⩽ −η ,

then |A| > 0 and, provided (w0 , θ0 ) ∈ A, it holds wt ⩾ τ r̄ and g(θt ) ⩽ −η/2 for every t ⩾ 0. This yields ẇt ⩾ η/4 and shows that E(g) holds.

Construction of a stable set when dw > 1 We begin with this section with the proof of Proposition 2.

Proof of Proposition 2. We define vt = wt /|wt |. Sublevel set condition. If ⟨v, vt ⟩ ⩾ δ and ⟨g(θt ), v⟩ = −η, we have that |⟨vt , v⊥ ⟩| ⩽ v⊥ ∈ {v}⊥ with |v⊥ | = 1. As a result, using that β := inf θ∈∂K |Jg (θ)T v| > 0, we obtain

√

1 − δ 2 for every

d ⟨g(θt ), v⟩ = ⟨Jg (θt )θ̇t , v⟩ dt = −⟨Jg (θt )Jgt (θt )T wt , v⟩ = −|wt |⟨Jg (θt )T v, Jgt (θt )T vt ⟩ = −|wt | ⟨Jg (θt )T v, Jg (θt )T vt ⟩ + ⟨Jg (θt )T v, (Jgt (θt ) − Jg (θt ))T vt ⟩ = −|wt | ⟨Jg (θt )T v, Jg (θt )T (⟨vt , v⟩v + proj{v}⊥ (vt ))⟩ + ⟨Jg (θt )T v, (Jgt (θt ) − Jg (θt ))T vt ⟩ = −|wt | ⟨vt , v⟩⟨Jg (θt )Jg (θt )T v, v⟩ + ⟨proj{v}⊥ (Jg (θt )Jg (θt )T v), vt ⟩ + ⟨Jg (θt )T v, (Jgt (θt ) − Jg (θt ))T vt ⟩ p ⩽ −|wt | δ⟨Jg (θt )Jg (θt )T v, v⟩ − 1 − δ 2 |proj{v}⊥ (Jg (θt )Jg (θt )T v)| − |Jg (θt )T v|∥Jg − Jgt ∥∞ |proj{v}⊥ (Jg (θt )Jg (θt )T v)| ⟨Jg (θt )Jg (θt )T v, v⟩ p T 2 − = −|wt ||Jg (θt )Jg (θt ) v| δ 1 − δ |Jg (θt )Jg (θt )T v| |Jg (θt )Jg (θt )T v| |Jg (θt )T v| − ∥Jg − Jgt ∥∞ |Jg (θt )Jg (θt )T v| |proj{v}⊥ (Jg (θt )Jg (θt )T v)| ⟨Jg (θt )Jg (θt )T v, v⟩ p ε 2 2 ⩽ −|wt |β δ − 1−δ − , |Jg (θt )Jg (θt )T v| |Jg (θt )Jg (θt )T v| β

where we have used that |Jg (θt )T v|2 = ⟨Jg (θt )Jg (θt )T v, v⟩ ⩽ |Jg (θt )Jg (θt )T v|. Using s |proj{v}⊥ (Jg (θt )Jg (θt )T v)| 2 ⟨Jg (θt )Jg (θt )T v, v⟩ = 1 − |Jg (θt )Jg (θt )T v| |Jg (θt )Jg (θt )T v| √ √ and the fact that x 7→ a 1 − x2 − 1 − a2 x is non-increasing on [0, 1] for every a ∈ [0, 1], we obtain i h p p d ⟨g(θt ), v⟩ ⩽ −|wt |β 2 δ 1 − γ ′2 − 1 − δ 2 γ ′ − ε/β dt with

|proj{v}⊥ (Jg (θ)Jg (θ)T v)| . |Jg (θ)Jg (θ)T v| θ∈∂K

γ ′ = sup def.

As a result, provided δ > γ ′ , we can take ε small enough to ensure that θt cannot leave K. Angle condition.

If ⟨g(θt ), v⟩ ⩽ −η and ⟨v, vt ⟩ = δ then

d 1 ⟨v, vt ⟩ = ⟨v, (I − vt vtT )ẇt ⟩ dt |wt | 1 = ⟨v − δvt , −gt (θt )⟩ |wt | 1 = ⟨(1 − δ 2 )v − δproj{v}⊥ (vt ), −gt (θt )⟩ |wt | 1 = [⟨(1 − δ 2 )v − δproj{v}⊥ (vt ), −g(θt )⟩ + ⟨(1 − δ 2 )v − δproj{v}⊥ (vt ), g(θt ) − gt (θt )⟩] |wt | i p 1 h ⩾ (1 − δ 2 )|⟨g(θt ), v⟩| − δ|proj{v}⊥ (vt )||proj{v}⊥ (g(θt ))| − 1 − δ 2 ∥g − gt ∥∞ |wt | √ i 1 − δ 2 hp 1 − δ 2 |⟨g(θt ), v⟩| − δ|proj{v}⊥ (g(θt ))| − ∥g − gt ∥∞ = |wt | √ |proj{v}⊥ (g(θt ))| ∥g − gt ∥∞ η 1 − δ2 p |⟨g(θt ), v⟩| ⩾ [ 1 − δ2 −δ − ] |wt | |g(θt )| |g(θt )| η s √ ε η 1 − δ2 p |⟨g(θt ), v⟩| |⟨g(θt ), v⟩|2 2 −δ 1− − ⩾ 1−δ |wt | |g(θt )| |g(θt )|2 η √ i p η 1 − δ 2 hp ⩾ 1 − δ 2 γ − δ 1 − γ 2 − ε/η , |wt | where def.

γ = inf

θ∈K

|⟨g(θ), v⟩| . |g(θ)|

Provided δ < γ, we can take ε small enough to ensure we always have ⟨v, vt ⟩ ⩾ δ.

D.1

Local analysis

In the following, we assume that h : θ 7→ (1/2)|g(θ)|2 has a non-degenerate local maximizer θ∗ ∈ Rdθ , that is to say ∇h(θ∗ ) = Jg (θ∗ )T g(θ∗ ) = 0, (23) ∇2 h(θ∗ ) = Jg (θ∗ )T Jg (θ∗ ) + Hg (θ∗ )[g(θ∗ )] ≺ 0,

where Jg (θ) is the Jacobian of g at θ ∈ Rdθ and Hg (θ)[u] =

" d w X

# ∂ij gk (θ)uk

k=1

1⩽i,j⩽dθ

for every u ∈ Rdw . By (23), we also obtain that θ∗ is a non-degenerate local minimizer of θ 7→ ⟨g(θ), v⟩ with v = −g(θ∗ )/|g(θ∗ )|. Since this function is of class C 2 , we obtain the existence of c, r > 0 such that |θ−θ∗ | ⩽ c(ηmax −η)1/2 for every θ ∈ B(θ∗ , r) such that ⟨g(θ), v⟩ ⩽ −η, where ηmax := |g(θ∗ )|. Taking K = {θ ∈ B(θ∗ , r) | ⟨g(θ), v⟩ ⩽ −η},1 we would like to obtain an expansion of the two quantities involved in (6) when η → ηmax . D.1.1

Right-hand side

Let us prove that θ∈K

inf

ηmax − η |⟨g(θ), v⟩| = 1 − c21 + o(ηmax − η), |g(θ)| ηmax

sup

|Jx| with J = Jg (θ∗ ) and H = −Hg (θ∗ )[g(θ∗ )].

(24)

where c1 =

⟨Hx,x⟩=1

Using Taylor’s theorem, we obtain for every θ ∈ Rdθ 1 ⟨g(θ), v⟩ − (−ηmax ) − ⟨Hg (θ∗ )[v](θ − θ∗ ), θ − θ∗ ⟩ = R1 (θ − θ∗ ) with 2

lim

θ→θ∗ |θ − θ∗ |2

R1 (θ − θ∗ ) = 0. (25)

This yields 1 ⟨Hg (θ∗ )[v](θ − θ∗ ), θ − θ∗ ⟩ ⩽ (ηmax − η) + |R1 (θ − θ∗ )|. 2 p As a result, θ′ := θ∗ + (θ − θ∗ )/ 1 + |R1 (θ − θ∗ )|/(ηmax − η) is such that 1 ⟨Hg (θ∗ )[v](θ′ − θ∗ ), θ′ − θ∗ ⟩ ⩽ ηmax − η 2 and

1 |R1 (θ − θ∗ )| ⩽ c(ηmax − η)1/2 |θ′ − θ| = |θ − θ∗ | 1 − p 2(ηmax − η) 1 + |R1 (θ − θ∗ )|/(ηmax − η)

provided θ ∈ K. Since |⟨g(θ), v⟩| 1 = 1 − 2 |Jg (θ∗ )(θ − θ∗ )|2 + R2 (θ − θ∗ ) with |g(θ)| 2ηmax

lim

θ→θ∗ |θ − θ∗ |2

R2 (θ − θ∗ ) = 0,

we obtain |⟨g(θ), v⟩| 1 = 1 − 2 |Jg (θ∗ )(θ − θ∗ )|2 + R2 (θ − θ∗ ) |g(θ)| 2ηmax 1 1 = 1 − 2 |Jg (θ∗ )(θ′ − θ∗ )|2 − 2 [|Jg (θ∗ )(θ − θ∗ )|2 − |Jg (θ∗ )(θ′ − θ∗ )|2 ] + R2 (θ − θ∗ ) 2ηmax 2ηmax 1 ∥Jg ∥2∞ ⩾ 1 − 2 |Jg (θ∗ )(θ′ − θ∗ )|2 − |θ − θ′ | |θ − θ∗ | + |θ′ − θ∗ | − |R2 (θ − θ∗ )|. 2 2ηmax 2ηmax 1 Although the two conditions above were obtained with K a sublevel set of θ 7→ ⟨g(θ), v⟩, the same reasoning carries out if we take K to be the intersection of a ball with this sublevel set.

for every θ ∈ K. Combining the above yields 1 |⟨g(θ), v⟩| 2 1 ⩾ inf 1 − 2 |Jg (θ∗ )(θ − θ∗ )| ⟨Hg (θ∗ )[v](θ − θ∗ ), θ − θ∗ ⟩ ⩽ ηmax − η + o(ηmax − η). inf θ∈K |g(θ)| 2ηmax 2 To obtain the reverse inequality, one can argue in the same way by observing that, if 1 ⟨Hg (θ∗ )[v](θ − θ∗ ), θ − θ∗ ⟩ ⩽ ηmax − η, 2 p then |θ − θ∗ | ⩽ 2ηmax (ηmax − η)/λmin (−Hg (θ∗ )[g(θ∗ )]). Now, either ⟨g(θ), v⟩ ⩽ −η and we are done, or ⟨g(θ), v⟩ > −η and by the intermediate value theorem there exists ρ(θ) ∈ [0, 1] such that ⟨g(θ∗ + ρ(θ)(θ − θ∗ )), v⟩ = −η. We obtain 1 −η = ⟨g(θ∗ + ρ(θ)(θ − θ∗ )), v⟩ = −ηmax + ρ(θ)2 ⟨Hg (θ∗ )[v](θ − θ∗ ), θ − θ∗ ⟩ + R1 (ρ(θ)(θ − θ∗ )), 2 which yields s

s " # |R1 (ρ(θ)(θ − θ∗ ))| |R1 (ρ(θ)(θ − θ∗ ))| 1− and |θ∗ + ρ(θ)(θ − θ∗ ) − θ| ⩽ |θ − θ∗ | 1 − 1 − . ηmax − η ηmax − η

ρ(θ) ⩾

Since ρ(θ) ⩽ 1, we obtain √

1 ηmax − η

sup |θ∗ + ρ(θ)(θ − θ∗ ) − θ| (1/2)⟨Hg (θ∗ )[v](θ − θ∗ ), θ − θ∗ ⟩ ⩽ ηmax − η

−→ 0.

η→ηmax

We can then conclude as above to finally obtain |⟨g(θ), v⟩| 1 2 1 inf = inf 1 − 2 |Jg (θ∗ )(θ − θ∗ )| ⟨Hg (θ∗ )[v](θ − θ∗ ), θ − θ∗ ⟩ ⩽ ηmax − η + o(ηmax − η) θ∈K |g(θ)| 2ηmax 2 1 2 1 = 1 − 2 sup |Jg (θ∗ )(θ − θ∗ )| ⟨Hg (θ∗ )[v](θ − θ∗ ), θ − θ∗ ⟩ ⩽ ηmax − η + o(ηmax − η) 2ηmax 2 2ηmax (ηmax − η) =1− sup |Jg (θ∗ )(θ − θ∗ )|2 ⟨−Hg (θ∗ )[g(θ∗ )](θ − θ∗ ), θ − θ∗ ⟩ ⩽ 1 + o(ηmax − η), 2 2ηmax which yields (24). D.1.2

Left-hand side

Let us prove that |proj{v}⊥ (Jg (θ)Jg (θ)T v)| ηmax − η = 1 − c22 + o(ηmax − η) T v| |J (θ)J (θ) ηmax g g θ∈∂K sup

where c2 =

|Hx|2 with J = Jg (θ∗ ) and H = −Hg (θ∗ )[g(θ∗ )]. ⟨Hx,x⟩=1 |JHx| inf

Reasoning as above and using (25) and |⟨Jg (θ)Jg (θ)T v, v⟩| |Hg (θ∗ )[v](θ − θ∗ )|2 = + R3 (θ − θ∗ ) with |proj{v}⊥ (Jg (θ)Jg (θ)T v)| |Jg (θ∗ )Hg (θ∗ )[v](θ − θ∗ )| 32

lim

θ→θ∗ |θ − θ∗ |

R3 (θ − θ∗ ) = 0,

we obtain

|proj{v}⊥ (Jg (θ)Jg (θ)T v)| 1 =p |Jg (θ)Jg (θ)T v| θ∈∂K 1 + (I + o((ηmax − η)1/2 ))2 sup

with |Hg (θ∗ )[v](θ − θ∗ )|2 1 I = inf ⟨Hg (θ∗ )[v](θ − θ∗ ), θ − θ∗ ⟩ = ηmax − η |Jg (θ∗ )Hg (θ∗ )[v](θ − θ∗ )| 2 s 2(ηmax − η) |Hg (θ∗ )[g(θ∗ )](θ − θ∗ )|2 = − ⟨Hg (θ∗ )[g(θ∗ )](θ − θ∗ ), θ − θ∗ ⟩ = 1 , inf ηmax |Jg (θ∗ )Hg (θ∗ )[g(θ∗ )](θ − θ∗ )| which finally yields the result.

Record · ID 175279 · SHA-256 14864d0811ceb6f0
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.