High-Dimensional Learning Dynamics of Attention-Indexed Models

High-Dimensional Learning Dynamics of Attention-Indexed Models Yizhou Xu1,2 , Margarita Sagitova2 , Lenka Zdeborová2 , and Florent Krzakala1 1 Information, Learning and Physics Laboratory, École Polytechnique Fédérale de Lausanne (EPFL)

arXiv:2609.03858v1 [cs.LG] 3 Sep 2026

2

Statistical Physics of Computation Laboratory, École Polytechnique Fédérale de Lausanne (EPFL)

Abstract Attention mechanisms are central to modern foundation models, yet their training dynamics remain poorly understood, especially when the attention matrices have extensive rank. In this work, we study attention-indexed models, a broad framework that can represent multi-layer and multi-head attention architectures. First, we show that, in a suitable high-dimensional limit, the population-loss landscape is characterized by a finite set of trace order parameters. In contrast, online stochastic gradient descent (SGD) is governed by an infinite hierarchy of matrix moments, which we show can be exponentially well-approximated by a finite truncated system. Second, this framework reveals that attention parameterization itself can act as an architectural implicit bias. Direct optimization of an attention matrix S ∈ Rd×d can remain trapped in an uninformative state. Tied attention (S = W W ⊺ ) induces an automatic symmetry-breaking mechanism and yields weak recovery in Θ(d2 log d) samples. For untied attention, S = U V ⊺ , we uncover a fast-slow mechanism: the pre-activation mean first evolves on a fast timescale, while the overlaps evolve on a slower one. Weak recovery on the Θ(d2 log d) scale occurs when the state selected by the fast dynamics breaks the initial symmetry.

1

Introduction

Attention mechanisms [36] are the central building blocks of modern foundation models, but their training dynamics remain poorly understood. This is particularly true in the high-dimensional feature-learning regime with extensive-rank attention matrices, whose ranks diverge as the embedding dimension grows. Understanding this regime raises two related but distinct questions. First, what is the correct macroscopic description of learning when the attention matrices have extensive rank? Second, once such a description is available, how does the parameterization of an attention matrix affect whether features can be learned? Existing theoretical work has made substantial progress on several aspects of attention. One line studies the mechanisms of in-context learning [15, 30, 23], while another analyzes training from scratch in simplified transformer models [37, 14, 16]. For feature-learning dynamics, however, most existing analyses either do not take a high-dimensional limit [31, 32, 17, 22, 37] or focus on finite-rank attention matrices, corresponding to sequence-indexed models [9, 8, 34, 11, 18, 1]. Such low-rank descriptions do not capture the extensiverank regime naturally associated with attention matrices in practice, whose rank scales with the embedding dimension. In this work, we study this regime through attention-indexed models, building on the formulation introduced in [7, 6]. The model class describes losses depending on collections of quadratic token interaction √1 xT Sij xk with the attention matrix Sij ∈ Rd×d of extensive rank. The formulation is broad: multi-layer and d i multi-head attention-only architectures can be expressed in terms of finitely many such quadratic forms; we give the construction in Section 2 and Appendix B. Previous analyses of attention-indexed models focused on static Bayes-optimal or empirical-risk-minimization [7, 6], with tractable analysis restricted to a single-layer, single-head tied setting. Dynamical analyses of extensive-rank matrices have so far been available mainly for substantially simpler quadratic models [20, 21, 38]. A first difficulty is that the static and dynamical descriptions have fundamentally different dimensionality. We prove that, as d → ∞, the population loss is determined by finitely many order parameters. At the level of 1

the loss landscape, the original Θ(d2 )-dimensional optimization problem therefore admits a finite-dimensional description. The training dynamics are richer. Unlike multi-indexed and sequence-indexed models whose dynamics can be tracked by a finite set of order parameters [13, 1, 4], for the tied and untied factorizations analyzed in this paper, online SGD generates an infinite hierarchy of joint matrix moments, and no fixed finite collection of moments is sufficient in general to close the dynamics. We nevertheless show that this hierarchy is tractable: its empirical trajectory converges to a deterministic infinite-dimensional ODE, and finite-order truncations approximate the limiting dynamics with an error that decays exponentially with the truncation degree. The second question is what this dynamical description reveals about parameterization. We compare three ways of representing an effective attention matrix: direct optimization of S, the tied factorization S = W W ⊺ , and the untied factorization S = U V ⊺ . Although these parameterizations can describe closely related predictors, they induce qualitatively different optimization geometries. This distinction becomes especially sharp when learning starts from an uninformative initialization. In high-dimensional inference, direct optimization can be bottlenecked by the information exponent [2]: when the first-order correlation vanishes, SGD over S can remain trapped on the d2 log d sample scale. Our analysis shows that factorized attention can alter this conclusion, but through different mechanisms in the tied and untied cases. For tied attention, S = W W ⊺ , the positive-semidefinite (PSD) factorization provides an automatic symmetry-breaking mechanism. Under a mild non-degeneracy condition, this mechanism yields weak recovery in Θ(d2 log d) samples. The untied factorization S = U V ⊺ behaves differently. We find a fast–slow learning mechanism. The pre-activation means √1d Tr(U V ⊺ ) evolve on a fast timescale, while the matrix moments evolve on a slower timescale. The subsequent weak-recovery behavior is therefore controlled by the state selected during the fast boundary layer. If this shift breaks the relevant symmetry, weak recovery occurs in Θ(d2 log d) samples; if the symmetry survives, the dynamics remains uninformative on this sample scale. These results give two complementary messages. At the methodological level, extensive-rank attention requires a new macroscopic dynamical description: a finite-dimensional loss landscape coexists with an intrinsically infinite-dimensional training dynamics. At the learning level, the same framework shows that attention parameterization acts as an architectural implicit bias: tied and untied factorizations can overcome—or fail to overcome—uninformative symmetries through mechanisms that are absent from direct optimization of the attention matrix. Our main contributions are therefore: • A high-dimensional dynamical framework for extensive-rank attention. We formulate a broad class of attention-indexed models and show that their population loss admits a finite-dimensional characterization. For the tied and untied parameterizations, we derive high-dimensional limits for online SGD. The resulting dynamics involve an infinite hierarchy of matrix moments, for which we establish well-posedness, convergence, and exponentially accurate truncations. • Tied attention induces automatic symmetry breaking and achieves weak recovery in Θ(d2 log d) samples. Untied attention learns through a fast–slow mechanism, which leads to a weak-recovery dichotomy: recovery on the Θ(d2 log d) scale occurs when the fast process breaks the relevant symmetry, whereas the dynamics stalls on this scale when it does not. The remainder of the paper is organized as follows. Section 2 introduces attention-indexed models and establishes the finite-dimensional characterization of their population-loss landscape. Section 3 studies tied attention, derives its infinite-dimensional macroscopic dynamics and truncation theory, and identifies the symmetry-breaking mechanism responsible for weak recovery. Section 4 turns to untied attention and derives its two-timescale dynamics and associated weak-recovery dichotomy. The appendices contain the proofs, additional examples of attention-indexed models, and the analysis of direct optimization over the attention matrix S.

2

Attention-indexed models and their population loss landscape

The first goal of this paper is to identify a macroscopic description that remains tractable when the ranks of the attention matrices diverge with the embedding dimension. Before turning to training dynamics, we 2

first ask a simpler static question: which quantities determine the population loss in the high-dimensional limit? We show in this section that, despite the high dimensionality of the underlying matrices, the answer is finite-dimensional. This finite-dimensional loss landscape will provide the potential that drives the dynamical equations in Sections 3 and 4. Let L = Θ(1) represent the sequence length (number of tokens), and K = Θ(1) denote the number of indices, such as attention heads or layers (K = 2 in [6] with one student and one teacher index). Let the data vectors x1 , . . . , xL ∈ Rd be jointly Gaussian with zero mean and covariance given by E[xia xjb ] = Cij δab , where k C ∈ RL×L is a PSD matrix. Let Sij ∈ Rd×d for i, j = 1, . . . , L and k = 1, . . . , K be a set of deterministic weight matrices. The assumption E[xia xjb ] = Cij δab is made without loss of generality, as any arbitrary covariance structure can be absorbed into the weight matrix S. We consider losses of the form ⎫L,L,K ⎞ T k ⎛⎧ ⎪ ⎪ ⎪ xi Sij xj ⎪ ⎟. ⎬ L ⎜⎨ √ ⎪ ⎪ d ⎪ ⎪ ⎝⎩ ⎭i,j,k=1 ⎠ under the following assumption. Assumption 1. In the limit as d → ∞ with K, L = Θ(1), the following conditions hold for all i, j, k, i′ , j ′ , k ′ : ∣∣S k ∣∣op

1. The matrices satisfy limd→∞ ∣∣Sijk ∣∣

ij F

= 0. ′

k k ) = µijk , limd→∞ d1 Tr(Sij (Sik′ j ′ )T ) = Ωijk,i′ j ′ k′ and 2. The scaled traces have finite limits: limd→∞ √1d Tr(Sij ′

k k limd→∞ d1 Tr(Sij Si′ j ′ ) = Ψijk,i′ j ′ k′ .

3. The loss L ∶ RL×L×K → R is continuous with at most polynomial growth. Remark. While Assumptions 1.2 and 1.3 are primarily regularity conditions, Assumption 1.1 identifies the extensive-rank regime of interest here. We argue that this assumption is highly relevant in practice. In standard multi-head attention, for example, the rank of an individual attention matrix can grow with the embedding dimension when the number of heads remains fixed. This marks a critical difference from existing literature on low-rank attention (sequence-indexed models) [9, 34, 1, 11]. k In this section, we do not impose any structural assumptions on the attention matrices Sij . We will specifically analyze both tied and untied attention mechanisms in subsequent sections. Remarkably, the attentionindexed model inherently covers standard multi-layer, multi-head attention networks (i.e., attention-only transformers).

Example (Multi-layer multi-head attention). Consider an M -layer, H-head attention network: H

(h)

L×d Xm = Xm−1 + ∑ softmax (A(h) , m (Xm−1 )) Xm−1 Vm−1 ∈ R

m = 1, . . . , M

(1)

h=1

where the pre-softmax attention matrix is given by: 1 (h) (h) T T L×L A(h) . m (Xm−1 ) = √ Xm−1 Km−1 (Qm−1 ) Xm−1 ∈ R d (h)

(h)

(h)

(2)

The matrices Km , Qm , and Vm represent the key, query, and value weights of the m−th layer, respectively. T ) for some activation function σ. The final output is y = σ ( √1d XM KM QTM XM Consider a teacher-student setup with input data X0 ∼ N (0, C ⊗Id ) and a loss function L(y, ŷ), where both teacher and student networks employ multi-layer multi-head architecture. This setup naturally falls within our general framework (see Appendix B for detailed derivations). Notably, we do not require the teacher weights to be Gaussian, rendering the existence of such a teacher model a relatively plausible assumption.

3

In Appendix B, we introduce several additional models that fit this framework, including auto-regressive sequence models, multiple-location regression models, multiple-step reasoning models and matrix denoising problems. k Although the matrices Sij contain O(d2 ) degrees of freedom, the loss does not retain this full complexity in the high-dimensional limit. The reason is that, under the extensive-rank condition, the finite collection of quadratic forms entering the loss becomes jointly Gaussian. Its limiting distribution is therefore completely specified by the first- and second-order trace statistics in Assumption 1. The following theorem makes this reduction precise. Theorem 1. Under Assumption 1, the joint distribution of the quadratic forms converges weakly to a multivariate Gaussian, and expectations converge as follows: ⎡ ⎧ T k ⎫L,L,K ⎤ ⎢ ⎛⎪ ⎞⎥⎥ ⎪ ⎢ ⎪ xi Sij xj ⎪ ⎟⎥ = E [L({Gijk }L,L,K ⎬ lim Ex ⎢L ⎜⎨ √ i,j,k=1 )] , ⎥ ⎢ ⎝⎪ ⎪ d→∞ d ⎪ ⎪ ⎠ ⎥ ⎢ ⎩ ⎭ i,j,k=1 ⎦ ⎣

(3)

where {Gijk }L,L,K i,j,k=1 is a multivariate Gaussian vector whose mean and covariance are given by: E[Gijk ] = Cij µijk , Cov(Gijk , Gi′ j ′ k′ ) = Cii′ Cjj ′ Ωijk,i′ j ′ k′ + Cij ′ Cji′ Ψijk,i′ j ′ k′ .

(4)

Theorem 1 gives the basic reduction underlying the rest of the paper. In the high-dimensional limit, the population loss depends only on the finite collection of order parameters {µ, Ω, Ψ}. Thus, although the parameter space grows with d, the limiting loss landscape is finite-dimensional. The Gaussian limit in Theorem 1 can be viewed as a consequence of the fourth-moment theorem for Wiener chaos [26]; for completeness, we provide a direct proof in Appendix C.1. A centered version needed for tied attention is given in Appendix C.2. Theorem 1 characterizes the loss at any fixed state. It also implies a stronger optimization-level statement: subject to an appropriate realizability condition, optimizing over the high-dimensional matrices is asymptotically equivalent to optimizing over their finite set of limiting order parameters. Corollary 1. Let K ⊆ {1, 2, . . . , K} denote the index set of the learnable weights, while the remaining weights k k k {Sij }k∉K are fixed and satisfy the condition limd→∞ ∥Sij ∥op /∥Sij ∥F = 0. Define the order parameter set Q as ⎫ ⎧ ⎪ ⎪ ′ 1 1 ⎪ ⎪ 1 k k k′ k ]} , { Tr[Sij ⎬. Si′ j ′ ], Tr[Sij (Sik′ j ′ )T ]} Q ∶= ⎨{ √ Tr[Sij ⎪ ′ ′ ′ d d i,j,k,i ,j ,k ⎪ d ⎪ ⎪ i,j,k ⎭ ⎩

(5)

Let Q denote the feasible domain of Q realizable by matrix sequences satisfying Assumption 1 in the limit k d → ∞, subject to the constraints imposed by the fixed weights {Sij }k∉K . Let R(Q) ∶= EG [L({Gijk }L,L,K i,j,k=1 )] denote the expected population loss in Theorem 1 specified by Q and the covariance C. Under the regularity assumption specified in Appendix C.3, ⎡ ⎧ T k ⎫L,L,K ⎤ ⎢ ⎛⎪ ⎞⎥⎥ ⎪ ⎢ ⎪ xi Sij xj ⎪ ⎟⎥ = inf R(Q). lim inf Ex ⎢L ⎜⎨ √ ⎬ k } ⎢ ⎝⎪ d→∞ {Sij d ⎪ k∈K ⎪ ⎠⎥⎥ Q∈Q ⎢ ⎪ ⎩ ⎭ i,j,k=1 ⎣ ⎦

(6)

Corollary 1, proven in Appendix C.3, makes the finite-dimensional nature of the limiting landscape explicit: The optimal asymptotic population loss is determined by an optimization over the finite-dimensional feasible set Q. Importantly, this does not imply that training can be described by gradient descent directly on Q. Whether the order parameters form a closed dynamical system depends on how the attention matrices are parameterized. Direct optimization over S does close at the level of finitely many first- and second-order quantities (Appendix E), whereas the tied and untied factorizations studied in Sections 3 and 4 generate additional matrix moments. This distinction has important consequences for weak recovery.

4

3

Online SGD of tied attention

We now turn to a question that cannot be answered from the loss landscape alone: does the parameterization of the attention matrix change the way SGD learns features? In this section, we first derive the macroscopic dynamics of tied attention and show that, despite the finite-dimensional population loss, the factorized flow requires an infinite hierarchy of matrix moments. We then use these dynamics to characterize weak recovery and contrast it with direct optimization over S.

3.1

Macroscopic dynamics

For tied attention, the trace of Sk = Wk Wk⊺ is typically of order d. The pre-activation d−1/2 x⊺i Sk xj therefore contains a diverging contribution. Following the centered formulation of Appendix C.2, we consider the regularized loss L,L,K ⎞ γ K ⎛ Tr[Wk WkT (xj xTi − Cij Id )] 2 √ } + L { ∑ ∣∣Wk ∣∣F , 2d ⎠ ⎝ d k=1 i,j,k=1 where the input data x1 , . . . , xL ∼ N (0, C ⊗ Id ) and γ = Θ(1) > 0 denotes the weight decay strength. For k = 1, . . . , K, let Wk ∈ Rd×dk and Sk = Wk Wk⊺ ⪰ 0. As before, let K ⊆ {1, . . . , K} denote the trainable indices. For k ∈ K, Wk evolves during training, whereas for k ∉ K it remains fixed and may, for example, represent a teacher matrix. At iteration n, online SGD uses a fresh sample x(n) ∼ N (0, C ⊗ Id ) and updates, for k ∈ K, (n+1)

Wk (n)

(n)

= Wk (n)

− αd ∇Wk L(G(n) ) −

αd γ (n) Wk , d

(7)

(n)

where we denote Gijk ∶= √1d Tr[Wk WkT (xj (xi )T − Cij Id )] and αd is the step size. (n)

Define Sk

(n)

(n)

= Wk (Wk )⊺ and, for any multi-index α = (k1 , . . . , kw ), µ(n) α =

1 (n) (n) (n) Tr [Sk1 Sk2 ⋯Skw ] . d

(8)

In particular, qkl = µ(k,l) denotes the second-order moment matrix. We introduce the continuous time t = nτd , τd = 4αd d and let µ̃(d) (t) denote the piecewise-constant interpolation of these moments: µ̃(d) (t) = µ(n) ,

for t ∈ [nτd , (n + 1)τd ).

(9)

Define the potential as the limiting population loss in Theorem 1: Φ(q) ∶= EG [L(G)] ,

q ∈ SK + ,

(10)

where G ∈ RL×L×K is a zero-mean Gaussian tensor with covariance Cov ((G)ijk , (G)i′ j ′ l ) = (Cii′ Cjj ′ + Cij ′ Cji′ )qkl . We define its matrix derivative as ∇Φ(q)1 for q ⪰ 0. Assumption 2. 1. The loss function L ∶ RL×L×K → R is four times continuously differentiable. L and its derivatives up to the fourth order have at most polynomial growth. √ (0) (0) 2. At initialization, the weights Wk satisfy ∥Wk ∥F = Θ( d). Additionally, there exists a constant (0) C0 > 0 independent of d such that sup1≤k≤K ∥Wk ∥op ≤ C0 almost surely. (d)

3. limd→∞ ∣µ̃α (0) − µ̄α (0)∣ = 0 in probability for every multi-index α. 4. The regularization is sufficiently large, i.e., γ > γ∗ (C0 , C, L, K, L) for some γ∗ independent of d. Remark. Assumption 2.4 is used only to guarantee uniform spectral control and prevent finite-time blow-up. Equivalently, one may impose an a priori uniform bound on the operator norms of the weights; numerically, our trajectories remain bounded even without weight decay. 1 More rigorously speaking, it should be defined as lim

ϵ→0+ ∇Φ(q + ϵI).

5

The key point is that the dimensionality of the loss and that of the dynamics are different. The potential Φ depends only on the finite-dimensional matrix q. However, the evolution of q is not closed: second moments depend on third moments, third moments on fourth moments, and so on. The tied factorization therefore converts a finite-dimensional loss landscape into an infinite-dimensional dynamical hierarchy. The following theorem shows that this hierarchy nevertheless has a deterministic and well-posed high-dimensional limit. Theorem 2. Under Assumption 2, the online SGD trajectory is governed by an infinite-dimensional ODE driven by the gradients of Φ: w dµ̄(k1 ,...,kw ) γ =− ∑ ∑[∇Φ(q̄)]ki l (µ̄(k1 ,...,ki−1 ,l,ki ,...,kw ) + µ̄(k1 ,...,ki ,l,ki+1 ,...,kw ) )− ∣α∣K µ̄(k1 ,...,kw ) , dt 2 i=1 l

(11)

ki ∈K

where ∣α∣K ≤ w is the number of trainable indices in the multi-index α of length w. (11) has a unique admissible solution and the empirical trajectory uniformly converges to it in probability. Specifically, if the width scales such that max1≤k≤K dk ≤ dcw for some fixed cw < +∞ and the learning rate scales such that αd ≤ c0 (d log d)−1 and αd = Ω(d−ι ) for some constants c0 > 0 small enough and ι > 0, then the empirical SGD trajectory satisfies: lim P

d→∞

⎞ ⎛ sup ∣µ̃(d) α (t) − µ̄α (t)∣ > ϵ = 0, ∀ϵ > 0 ⎠ ⎝t∈[0,T ]

(12)

for every multi-index α and for any fixed T > 0. Remark. We regard (11) as an ODE on the moment space ⋃s>0 Xs0 , where Xs0 ∶= {µ ∶ limw→∞ sup∣α∣=w ∣µswα ∣ = 0}. A trajectory µ on [0, T ] is called an admissible solution if q(µ(t)) ⪰ 0 for every t ∈ [0, T ] and there exist t 0 < sT < s′T < ∞ such that µ ∈ C([0, T ]; Xs0T ) and µ(t) = µ(0) + ∫0 V (µ(r)) dr holds as an identity in Xs0′ , T where V refers to the right side of (11). Theorem 2 gives the macroscopic dynamical counterpart of the static reduction in Section 2. The same finite-dimensional potential Φ(q) determines the coefficients of the flow, but it does not determine a closed ODE for q. Instead, the geometry of the factorization S = W W ⊺ propagates information through the entire hierarchy of joint moments. This infinite-dimensionality is not merely an artifact of our choice of coordinates. Appendix D.3 gives an exactly solvable quadratic example. In that example, no finite-dimensional ODE with a continuous initialization map can reproduce all possible trajectories. The proof, given in Appendix D.1, combines four ingredients: local well-posedness in a scale of weighted moment spaces, uniform spectral bounds obtained through a stopping-time argument and matrix Freedman inequalities, tightness and convergence of the empirical moment process, and global well-posedness of the resulting admissible infinite-dimensional trajectory. The infinite hierarchy is required for an exact description, but this does not make the theory numerically intractable. Higher-order moments can be truncated with exponentially small error. For M ≥ 2, we define the degree-M truncated trajectory µ(M ) (t) by d (M ) µ (t) = [V (µ(M ) (t))]α , dt α (M )

∣α∣ ≤ M,

(13) (M )

with the initial condition µα (0) = µ̄α (0) for ∣α∣ ≤ M and the boundary condition µα (t) ≡ 0 for ∣α∣ > M . Here V denotes the right side of (11)2 . The following corollary guarantees that µ(M ) (t) can approximate the true dynamics µ̄(t) to arbitrary precision. Corollary 2. Under the conditions of Theorem 2, there exist constants s0 > C∗2 and c < ∞ independent of M , α and T such that, for all sufficiently large M , the truncated ODE (13) has a unique solution and for every multi-index α, M +1 C2 ∣α∣ ) s0 . (14) sup ∣µ(M (t) − µ̄α (t)∣ ≤ c ( ∗ ) α s0 t∈[0,T ] 2 More specifically, V refers to the extended vector field defined in (94) because the truncated moment sequence might not satisfy q ⪰ 0.

6

Figure 1: Finite truncations accurately approximate the infinite moment dynamics. Comparison of online SGD with degree-M truncations of (13) for L = 2, softmax activation, and MSE loss. Increasing M improves the agreement with the empirical trajectory. We use d = 400, learning rate 0.01, batch size 512, and no weight decay. Corollary 2 suggests that while the limiting trajectory is intrinsically infinite-dimensional, any fixed collection of low-order observables can be approximated to arbitrary accuracy by increasing the truncation degree. In the exactly solvable quadratic setting (quadratic networks [5, 20, 21]) of Appendix D.3, these finite-order systems coincide with the Taylor expansion of the exact matrix flow. Corollary 2 is proven in Appendix D.2. As a side remark, Assumption 2.4 is used in its proof to obtain a uniform-in-time truncation error bound; without this condition, the constants in Corollary 2 would generally depend on T . Figure 1 verifies the same phenomenon beyond the exactly solvable setting. For L = 2, softmax activation, and MSE loss, increasing the truncation degree systematically improves agreement between the deterministic theory and online SGD.

3.2

Weak recovery and its sample complexity

We now use the macroscopic dynamics to address the main learning question of this section: can the tied factorization escape an initially uninformative state, and does its behavior differ from direct optimization over S? The relevant notion of information is not the raw overlap d1 Tr(S1 S2 ). Because tied matrices have a nonzero trace, this quantity can be O(1) even when the student contains no information about the teacher. We therefore remove the isotropic components and measure learning through the overlap of the traceless parts. (d) For clarity, consider one trainable student S1 = W1 W1T and one fixed teacher S2 . Define tk ∶= d1 Tr[Sk ], (d)

S̊k ∶= Sk − tk Id and the structural overlaps (d)

pkl ∶=

1 (d) (d) (d) Tr[S̊k S̊l ] = qkl − tk tl . d

(15)

Their deterministic limits are denoted by t̄k , q̄kl , and p̄kl . Thus p̄12 , rather than q̄12 , measures alignment between the student and teacher. We impose the following uninformative initialization. Assumption 3. As d → ∞, the initial student S1 (0) and the teacher S2 satisfy p̄12 (0) = 0,

1 Tr[S̊1 (0)2 S̊2 ] = 0, d→∞ d lim

1 Tr[S̊1 (0)S̊22 ] = 0. d→∞ d lim

(16)

Moreover, we assume that t̄1 (0) > 0 and p̄22 > 0. Assumption 3 places the student on an uninformative manifold: its centered overlap with the teacher, as well as the mixed third-order centered moments vanish asymptotically at initialization. An independently initialized Gaussian W1 (0) satisfies these conditions. 7

The tied parameterization nevertheless contains one additional macroscopic quantity that does not vanish: t̄1 (0) > 0 as a direct consequence of S1 = W1 W1⊺ ⪰ 0. Evaluating the moment dynamics (11) at the uninformative initialization gives dp̄12 ∣ = −2[∇Φ(q̄(0))]12 t̄1 (0)p̄22 . dt t=0

(17)

Hence, whenever [∇Φ(q̄(0))]12 is nonzero, the uninformative state has a nonzero velocity in the informative direction. This is the symmetry-breaking effect of the tied factorization. Corollary 3 turns this heuristic into a sample-complexity statement, proven in Appendix D.4. For a fixed recovery threshold c > 0, define the weak-recovery sample complexity by (d,n)

(d) Nwr (c) ∶= inf {n ≥ 0 ∶ ∣p12 (d,n)

where p12

∣ ≥ c} ,

(18)

denotes the structural overlap after n steps of online SGD (7).

Corollary 3. Under Assumption 3 and the conditions of Theorem 2, suppose additionally that [∇Φ(q̄(0))]12 ≠ 0,

(19)

Then there exist constants c > 0 and 0 < T− < T+ < ∞ independent of d, such that lim P (⌊

d→∞

T− d T+ d (d) ⌋ < Nwr (c) ≤ ⌈ ⌉) = 1. 4αd 4αd

(20)

1 In particular, choosing the largest learning-rate scaling αd = Θ ( d log ) covered by Theorem 2 gives the weak d recovery sample complexity (d) Nwr (c) = ΘP (d2 log d). (21)

Remark. The condition (19) is mild in a broad class of teacher-student losses. To see this, consider the MSE loss L(G1 , G2 ) = 12 ∥F (G1 ) − F (G2 )∥2 with the same activation F for student and teacher. Then 1 [∇Φ(q̄)]12 = − E [⟨∇F (G1 )B 1/2 , ∇F (G2 )B 1/2 ⟩F ] , 2

(22)

where B(ij),(i′ j ′ ) = Cii′ Cjj ′ + Cij ′ Cji′ . Suppose additionally that the initialization scale is chosen such that q11 (0) = q22 > 0. A Gaussian Hermite expansion then gives −[∇Φ(q̄(0))]12 =

1 ∣ν∣ 2 ∑ ρ ν!∥Hν ∥F , 2 ν

(23)

√ where {Hν } are the Hermite coefficients of ∇F ( q22 B 1/2 z)B 1/2 and ρ = √ t̄1 (0)t̄2 > 0. Hence [∇Φ(q̄(0))]12 < q̄11 (0)q̄22 √ 0 whenever ∇F ( q22 B 1/2 Z)B 1/2 is not constantly zero. The significance of Corollary 3 is not only the d2 log d sample scale, but the mechanism that produces it. The student begins with zero structural overlap with the teacher, yet the tied factorization makes this uninformative state non-invariant: the nonzero isotropic component t̄1 (0) couples to the teacher variance p̄22 through the cross-gradient of the population loss and immediately generates informative overlap. This behavior should be contrasted with direct optimization over S. As shown in Appendix E, the direct S-flow closes at the level of finitely many first- and second-order parameters. More importantly, when the relevant cross-gradients vanish throughout the uninformative manifold, that manifold is invariant under the limiting S-flow, and weak recovery does not occur on the d2 log d sample scale. The tied factorization therefore acts as an architectural implicit bias: it changes the geometry of the gradient flow so that an initialization that is structurally uninformative can nevertheless have a nonzero escape direction. Figure 2 illustrates this distinction. For a linear activation, both parameterizations can develop overlap. When we do online SGD directly over S on higher-order activations (h2 , h3 ), however, it remains uninformative on the displayed sample scale, whereas the tied flow escapes and learns the teacher. Corollary 3 identifies the mechanism governing weak recovery. Under an additional Polyak–Łojasiewicz condition, Appendix D.5 also gives a convergence rate for the subsequent strong-recovery phase. 8

Figure 2: The parameterization of the attention matrix qualitatively changes weak recovery. We compare online SGD on the tied factor W , with S = W W ⊺ (Left), and direct online SGD on S (Right), starting from uninformative initializations in the same teacher–student setting. For the linear activation both parameterizations recover the teacher on the displayed scale, whereas for the higher-order activations the direct S-flow remains close to the uninformative state while the tied W -flow develops substantial overlap. We use d = 400, learning rate 0.005, batch size 512, and no weight decay. The activations are chosen from: h1 (x) = x, 3 2 and h3 (x) = x √−3x , respectively (L = 1). The input data is N (0, Id ). h2 (x) = x√−1 2 6

4

Online SGD of untied attention

The symmetry breaking described in Section 3 relies crucially on the PSD geometry S = W W ⊺ . This raises the natural next question: what remains of the mechanism when S is not PSD? This section answers this question for the untied parameterization S = U V T . The mechanism is qualitatively different from the tied case. For untied attention, the pre-activation has a nontrivial mean mk = √1d Tr(Uk Vk⊺ ), and this mean evolves faster than the matrix moments. As a result, training separates into two stages. On the fast timescale, the mean variables evolve while the moments remain asymptotically frozen. The fast flow drives the trainable means toward a critical manifold of the population loss. Only afterwards, on a slower timescale, do the matrix moments undergo O(1) evolution, with the means remaining asymptotically slaved to this manifold. This separation has a direct consequence for feature learning. The fast phase does not itself create an O(1) teacher overlap, since the overlaps remain frozen on that timescale. Instead, it selects the effective state from which slow feature learning begins. Weak recovery is therefore controlled by the loss gradients evaluated after the fast relaxation. The fast flow can break a student-side symmetry by shifting the preactivation mean and then the slow dynamics escapes the uninformative manifold. If the relevant symmetry survives the fast relaxation, this recovery mechanism is absent. The remainder of this section makes this fast–slow picture rigorous. We first derive the two-timescale macroscopic limit and then use it to characterize a corresponding dichotomy for weak recovery.

4.1

Macroscopic dynamics

For each index k = 1, . . . , K, let Uk , Vk ∈ Rd×dk and define the associated matrices Mk ∶= Uk Vk⊺ , Ak ∶= Uk Uk⊺ and Bk ∶= Vk Vk⊺ . We consider the regularized loss L,L,K

⎞ γ ⎛ 1 2 2 L { √ xTi Uk VkT xj } + ∑ (∥Uk ∥F + ∥Vk ∥F ) , 2d k∈K ⎝ d i,j,k=1 ⎠

(24)

where x1 , . . . , xL ∼ N (0, C ⊗ Id ), and K ⊆ {1, . . . , K} denotes the set of trainable indices. At iteration n, a fresh sample x(n) ∼ N (0, C ⊗ Id ) is drawn, and for k ∈ K the online SGD updates are αd γ (n) U , d k αd γ (n) (n+1) (n) Vk = Vk − αd ∇Vk L(G(n) ) − V , d k (n+1)

Uk

(n)

= Uk

− αd ∇Uk L(G(n) ) −

9

(25)

(n)

(n)

(n)

(n) T

where Gijk ∶= √1d (xi )T Uk (Vk

(n)

) xj

and αd is the learning rate.

Order parameters. Unlike the centered tied model of Section 3, the untied pre-activations can have an (d) O(1) limiting mean. We therefore need to track separately mk ∶= √1d Tr(Mk ). The remaining macroscopic state is described by normalized joint moments of the matrices Mk , Mk⊺ , Ak , Bk . Introduce the alphabet Z ∶= {Mk , MkT , Ak , Bk ∶ 1 ≤ k ≤ K} and for a word α = (Z1 , . . . , Zw ) ∈ Z w , define µ(d) α ∶=

1 Tr[Z1 ⋯Zw ]. d

(26) (2)

(1)

B In particular, we define the traces tA k ∶= µ(Ak ) , tk ∶= µ(Bk ) and the overlaps qkl ∶= µ(Mk ,MlT ) , qkl ∶= µ(Mk ,Ml ) .

The potential is defined as the limiting population loss in Theorem 1: Φ(m, q (1) , q (2) ) ∶= E[L(G)], where (1)

(2)

G ∈ RL×L×K is Gaussian with E[Gijk ] = Cij mk and Cov(Gijk , Gi′ j ′ l ) = Cii′ Cjj ′ qkl + Cij ′ Cji′ qkl . We denote by ∇m Φ and ∇q(r) Φ, r = 1, 2, the vector and matrix gradients of the potential, respectively. (1)

For later use, let Ml

(2)

∶= Ml and Ml K

∶= MlT for 1̄ ∶= 2,

2̄ ∶= 1. Define

2

ψk (m, µ) ∶= 2 ∑ ∑ [∇q(r) Φ(m, q (1) , q (2) )]kl [µ(M (r) ,B ) + µ(A ,M (r) ) ] . l

l=1 r=1

k

k

l

(27)

The crucial new point is that m and µ evolve on different scales. A single SGD step changes m at order αd , whereas the normalized moments change only at order αd /d. This produces the two-timescale limit below. Fast learning phase. On the fast timescale t = nαd , the normalized moments remain asymptotically ̃ frozen, while the mean vector evolves at order one. Denoting the limiting fast trajectory by m(t), the trainable components satisfy ∂Φ dm̃k B (m̃(t), q̄ (1) (0), q̄ (2) (0)) , = − (t̄A k (0) + t̄k (0)) dt ∂mk

k ∈ K.

(28)

The non-trainable components remain fixed. Thus, to leading order, the fast learning phase performs gradient flow in the mean variables while seeing the covariance structure as frozen. Under the local strong-convexity condition introduced below, this flow converges toward the unique critical point m⋆K = m⋆K (q (1) , q (2) ) satisfying ∇mK Φ (m⋆ (q (1) , q (2) ), q (1) , q (2) ) = 0. We denote the corresponding critical manifold by M0 ∶= {(m, µ) ∶ ∇mK Φ(m, q (1) , q (2) ) = 0}. d Slow learning phase. We now introduce the slow timescale τ = dt = nα . After the initial boundary layer, d the trainable means are asymptotically slaved to the critical manifold:

m̄K (τ ) = m⋆K (q̄ (1) (τ ), q̄ (2) (τ )) ,

(m̄(τ ), µ̄(τ )) ∈ M0 .

(29)

The small residual of the fast variables contributes during the slow timescale. This produces the effective bias ψk (m̄, µ̄) β̄k (m̄, µ̄) ∶= − A B . (30) t̄k + t̄k with ψk defined in (27). For a word α = (Z1 , . . . , Zw ), let ki denote the index carried by the letter Zi . For each trainable position i, define the substitution operators ⎧ ⎪ Zi ∈ {Mk , MkT }, ⎪µ(...,Ak +Bk ,...) , Tibias (µ) = ⎨ ⎪ ⎪ ⎩µ(...,Mk +MkT ,...) , Zi ∈ {Ak , Bk },

⎧ µ(...,M (r) B +A M (r) ,...) , Zi = Mk , ⎪ ⎪ k k ⎪ l l ⎪ ⎪ ⎪ ⎪ Zi = MkT , ⎪ ⎪µ(...,Bk Ml(r̄) +Ml(r̄) Ak ,...) , diff Ti,l,r (µ) = ⎨ ⎪ µ(...,M (r) M T +M M (r̄) ,...) , Zi = Ak , ⎪ ⎪ k ⎪ l k l ⎪ ⎪ ⎪ ⎪ µ (r̄) ⎪ T ⎩ (...,M M +M M (r) ,...) , Zi = Bk . l

10

k

k

l

(31)

The notation is understood linearly; for example, µ(...,Ak +Bk ,...) = µ(...,Ak ,...) + µ(...,Bk ,...) . With these operators, the slow evolution of all normalized matrix moments closes as an infinite-dimensional hierarchy w K 2 w ∂Φ dµ̄α diff (µ̄) − 2γ∣α∣K µ̄α = − ∑ β̄ki (m̄, µ̄)Tibias (µ̄) − 2 ∑ ∑ ∑ (r) (m̄, q̄ (1) , q̄ (2) ) Ti,l,r dτ i=1 l=1 r=1 ∂q i=1 ki ∈K

ki ∈K

(32)

ki l

for 0 ≤ τ ≤ T . Together with (29), this determines the limiting trajectory in the slow learning phase. Equations (28)–(32) summarize the fast–slow structure. The mean variables are no longer independent coordinates on the slow scale: they are determined by the current second moments through m⋆ (q). The remaining matrix moments evolve through an infinite hierarchy, as in the tied case, but with the additional effective bias inherited from the fast learning phase. We now state conditions ensuring that this formal separation of timescales is valid and that both phases remain well posed. Assumption 4. The following conditions hold. 1. The loss L ∶ RL×L×K → R is six times continuously differentiable, and L together with its derivatives up to sixth order has at most polynomial growth. √ (0) (0) 2. For every k, ∥Uk ∥F + ∥Vk ∥F = Θ( d). There exists a constant C0 > 0, independent of d, such that (0)

(0)

maxk {∥Uk ∥op , ∥Vk ∥op } ≤ C0 almost surely. (d)

(d)

3. There exist deterministic limits m̃k (0) and µ̄α (0) such that mk (0) → m̃k (0) and µα (0) → µ̄α (0) in probability for every k and every α. B A B 4. For every trainable index k ∈ K, t̄A k (0) > 0, t̄k (0) > 0 and ∣t̄k (0) − t̄k (0)∣ ≥ cbal for some cbal > 0.

5. The regularization is sufficiently large, i.e., γ > γ∗ (C0 , C, L, K, L, T ) for some γ∗ independent of d. B Assumption 4.4 is a technical non-degeneracy condition used to control the denominator tA k + tk of (30) and obtain uniform bounds along the slow trajectory. It is satisfied, for example, by independent Gaussian initializations with slightly different variances for Uk and Vk . Assumption 4.5 plays the same spectral-control role as in the tied case.

Assumption 5 (Local strong convexity). Fix a time horizon T > 0. Let mK denote the trainable mean components, while m∖K is fixed at m∖K (0). Assume that there exist an open convex set M ⊂ R∣K∣ , an open set Q of covariance parameters q = (q (1) , q (2) ), and a constant λ0 > 0 such that: 1. For every (mK , q) ∈ M × Q, ∇2mK Φ(m, q) ⪰ λ0 I. 2. Writing q0 = (q̄ (1) (0), q̄ (2) (0)), we have m̃K (0) ∈ M, q0 ∈ Q, and {mK ∈ M ∶ Φ(mK , m∖K (0), q0 ) ≤ Φ(m̃(0), q0 )} ⋐ M.

(33)

3. For every q ∈ Q, the equation ∇mK Φ(m, q) = 0 has a solution m⋆K (q) ∈ M. The solution (m⋆K (q̄(τ )), q̄(τ )) of (29)-(32) remains in a compact subset of M × Q for 0 ≤ τ ≤ T 3 and sufficiently large γ. Assumption 5 ensures that, for each nearby covariance state, the fast dynamics has a locally unique attracting critical point m⋆ (q). This property allows the fast mean variables to be eliminated on the slow timescale. We now define the empirical interpolations. On the fast timescale, m̃(d) (t) ∶= m(n) ,

t ∈ [nαd , (n + 1)αd ),

3 More precisely, on every interval of local existence contained in [0, T ].

11

(34)

while on the slow timescale, µ̃(d) (τ ) = µ(n) ,

τ ∈[

nαd (n + 1)αd , ). d d

(35)

The following theorem rigorously justifies the separation into a fast boundary layer and a slow infinitedimensional flow. Theorem 3. Suppose that Assumptions 4 and 5 hold for a fixed slow-time horizon T > 0. Assume furc0 thermore that max1≤k≤K dk ≤ dcw for some fixed cw < ∞, and that the learning rate satisfies αd ≤ d log and d −ι αd = Ω(d ) for some constants ι > 0 and sufficiently small c0 > 0 independent of d. Then the online SGD trajectory exhibits the following two-timescale limit. Fix any ϵ > 0. For every finite fast-time horizon T̃ > 0, ⎛ ⎞ lim P sup ∥m̃(d) (t) − m̃(t)∥∞ > ϵ = 0, (36) d→∞ ⎝ ⎠ ̃] t∈[0,T where m̃ is the unique solution of (28). On the slow timescale, let (m̄(τ ), µ̄(τ )) denote the unique admissible solution of (29)–(32). Then for every 0 < T1 < T2 ≤ T , ⎞ ⎛ (37) sup ∥m̃(d) (τ d) − m̄(τ )∥∞ > ϵ = 0 lim P d→∞ ⎝τ ∈[T1 ,T2 ] ⎠ and, for every word α, lim P

d→∞

⎞ ⎛ sup ∣µ̃(d) (τ ) − µ̄α (τ )∣ > ϵ = 0. ⎠ ⎝τ ∈[0,T2 ] α

(38)

Theorem 3 shows that the untied factorization introduces a dynamical structure that has no analogue in the tied model. On the fast scale t = nαd , only the pre-activation means undergo O(1) evolution. On the slow scale τ = nαd /d, these means have already equilibrated and follow the evolving covariance quasi-statically, while the matrix moments undergo their O(1) evolution. This distinction will be essential for weak recovery. The proof is given in Appendix F. Its main additional difficulty relative to the tied case is the singular separation of timescales. We first establish convergence of the fast mean dynamics while the normalized moments remain frozen. We then obtain a refined O(d−1/2 ) control of the residual mean gradient after the boundary layer; this residual accumulates into the effective bias βk on the slow scale. Uniform spectral bounds, tightness of the moment process, and uniqueness then yield the limiting slow trajectory. As in the tied case, the slow infinite-dimensional hierarchy admits finite-order approximations. Appendix F.2 proves an exponentially decaying truncation error.

4.2

Weak recovery and its sample complexity

We now ask what the fast–slow structure implies for feature learning. To make this mechanism explicit, consider one trainable student M1 = U1 V1⊺ and one fixed teacher M2 . Since d1 Tr(Mk ) = learning is naturally measured through the matrix overlaps (1,d)

q12 We write

∶=

1 Tr[M1 M2T ], d

(2,d)

q12

(1,d)

∶=

1 Tr[M1 M2 ]. d (1)

q (d) q12 ∶= ( 12 (2,d) ) , q12

q̄ q̄12 ∶= ( 12 (2) ) . q̄12

(d)

mk √ d

→ 0, feature

(39)

(40)

We impose the following uninformative initialization. Assumption 6. At the initial time, the matrix algebra generated by {M1 , M1T , A1 , B1 } is asymptotically free from the algebra generated by {M2 , M2T , A2 , B2 }.

12

Assumption 6 is the untied analogue of an uninformative initialization. Asymptotic freeness removes not (2) (1) only the pairwise overlaps q12 and q12 , but also the mixed student–teacher moments. (2) (1) q̄22 q̄ Define the teacher covariance as Q2 ∶= ( 22 (2) (1) ). For c > 0, define the weak-recovery sample complexity q̄22 q̄22 by (d,n) (d) (41) Nwr (c) ∶= inf {n ≥ 0 ∶ ∥q12 ∥ ≥ c} . ∞

B ⋆ Let D1 (0) ∶= t̄A 1 (0) + t̄1 (0) > 0, and recall that the fast dynamics drive the trainable mean toward m̄1 ∶= T ∂Φ ∂Φ ⋆ m⋆1 (q̄ (1) (0), q̄ (2) (0)). Define the gradient at the slow-time initial state by g12 ∶= ( ∂q(1) , ∂q(2) ) evaluated at 12 12 ⋆ (1) (2)

(m̄ , q̄ (0), q̄ (0)). The role of the fast phase becomes transparent by evaluating the overlap dynamics at its initial state according to (32): dq̄12 ⋆ ∣ = −2D1 (0)Q2 g12 . (42) dτ τ =0

⋆ Thus the fast dynamics affects feature learning by changing the point at which the gradient g12 is evaluated. ⋆ If Q2 g12 ≠ 0, the slow trajectory leaves the uninformative state with nonzero velocity. The following corollary formalizes this intuition, proven in Appendix F.3, and Appendix F.4 gives a convergence rate for the subsequent strong-recovery phase under an additional Polyak–Łojasiewicz condition.

Corollary 4. Suppose that Assumption 6 and the conditions of Theorem 3 hold. (i) If

⋆ Q2 g12 ≠ 0,

(43)

then there exist constants c > 0 and 0 < T− < T+ < ∞, independent of d, such that lim P (⌊

d→∞

T+ d T− d (d) ⌋ < Nwr (c) ≤ ⌈ ⌉) = 1. αd αd

(44)

1 Consequently, choosing the largest learning-rate scaling αd = Θ ( d log ) covered by Theorem 3 gives d (d) Nwr (c) = ΘP (d2 log d).

(45)

(ii) Let Mfree denote the set of moment sequences for which the algebra generated by {M1 , M1T , A1 , B1 } is free from the algebra generated by {M2 , M2T , A2 , B2 }. Suppose that throughout Mfree , [∇q(r) Φ (m⋆ (q), q (1) , q (2) )]12 = 0,

r = 1, 2.

(46)

Then for every c > 0 and every fixed T > 0, (d) lim P (Nwr (c) ≤ ⌊

d→∞

Td ⌋) = 0. αd

(47)

Thus, for αd = Θ((d log d)−1 ), weak recovery does not occur on the d2 log d sample scale. Corollary 4 identifies a dichotomy. Part (i) states that if the state selected by the fast dynamics exposes a nonzero gradient, weak recovery occurs after Θ(d2 log d) samples. Part (ii) gives the complementary invariant-manifold mechanism: if the relevant cross-gradients vanish throughout the uninformative manifold, then weak recovery does not occur on the same sample scale. The contrast with tied attention is therefore sharp. In the tied case, the PSD geometry itself generates a nonzero escape direction from the original initialization. In the untied case, the fast phase first selects a new mean state, and it is the geometry of the loss at this selected state that determines whether the subsequent slow dynamics can escape.

13

2

The mechanism is particularly transparent for L = 1 and the MSE loss L(z1 , z2 ) = 12 (σ(z1 ) − σ⋆ (z2 )) . (1)

(2)

At an uninformative state q12 = 0, the student and teacher pre-activations Z1 ∼ N (m̄⋆1 , q̄11 + q̄11 ) and (1)

(2)

Z2 ∼ N (m̄2 , q̄22 + q̄22 ) are independent. Price’s theorem gives ∂Φ (1) ∂q12

=

∂Φ (2) ∂q12

= E[

∂2L ] = −E[σ ′ (Z1 )]E[σ∗′ (Z2 )]. ∂z1 ∂z2

(48)

Hence a sufficient condition for weak recovery is E[σ ′ (Z1 )]E[σ∗′ (Z2 )] ≠ 0.

(49)

The two factors play different roles. The teacher-side term E[σ⋆′ (Z2 )] determines whether a first-order teacher signal is available at all. The fast dynamics cannot change this factor because the teacher is fixed. By contrast, the student-side term E[σ ′ (Z1 )] is evaluated at the shifted mean m̄⋆1 selected by the fast phase. The fast learning phase can therefore repair a student-side symmetry: even if E[σ ′ (Z1 )] vanishes initially, it may become nonzero after the mean shifts. When this occurs and the teacher-side first-order signal is nonzero, weak recovery takes Θ(d/αd ) samples. If m̄∗1 fails to break this symmetry, the network stalls. A numerical illustration is presented in Figure 3, which verifies that when Ez2 [σ∗′ (Z2 )] = 0, weak recovery takes substantially more samples. Figure 3 also suggests that the weak recovery time depends on the information exponent of the activation functions, which is left for the future work.

5

Conclusion and future work

This work has two outcomes. First, we develop a high-dimensional framework for attention-indexed models with extensive-rank matrices. In the large-d limit, the population-loss landscape is characterized by a finite collection of trace order parameters. The corresponding training dynamics are substantially richer: online SGD Figure 3: Weak recovery becomes delayed for generates an infinite hierarchy of matrix moments. We higher-order activations. Evolution of the symderive deterministic limits for these dynamics and show metric student–teacher overlap (the overlap bethat the infinite hierarchy can nevertheless be approxitween M1 + M1T and M2 + M2T ) for untied attenmated by finite-order truncations with exponentially detion in a single-token teacher–student model with caying error. MSE loss. The linear activation h1 (x) = x develSecond, this framework reveals that the parameteriops overlap rapidly, whereas the higher-order aczation of the attention matrix can qualitatively change x2 −1 x3 −3x feature learning. For tied attention S = W W ⊺ , the PSD tivations h2 (x) = √2 and h3 (x) = √6 exhibit factorization provides an automatic symmetry-breaking increasingly slow escape from the uninformative mechanism, yielding weak recovery in Θ(d2 log d) sam- state, whose precise sample complexity is left for ples under a mild condition. This behavior contrasts future work. We use d = 800, learning rate 0.01, with direct optimization over S, for which the uninforma- batch size 4096, standard Gaussian initialization, tive manifold can remain invariant. For untied attention and no weight decay. The input data is N (0, Id ). S = U V ⊺ , the mechanism is different. Training separates into a fast and a slower learning phase. The fast phase selects the state from which feature learning begins, and weak recovery on the Θ(d2 log d) scale depends on whether this selected state breaks the relevant symmetry. Thus, tied and untied factorizations do not merely provide different coordinates for the same optimization problem; they induce different routes by which SGD can escape an uninformative state. While the attention-indexed model class is broad enough to represent attention-only architectures, the detailed training dynamics established here, however, only concern tied and untied parameterizations. We do not derive the full dynamics of all query, key, and value matrices in a general deep transformer. Extending the present approach to such parameter dynamics is a natural next step. 14

A second open direction concerns learning on a longer timescale. Corollary 4 characterizes the Θ(d2 log d) sample scale. However, the numerical results suggest that recovery can occur on longer scales controlled by the information exponents. Analyzing these regimes requires tracking finite-d corrections over polynomially long training horizons, since higher-order terms that vanish at Θ(1) time can accumulate over such scales. Developing such a theory would provide a more complete characterization of the parameterization-dependent sample complexity. Several other extensions remain open. The present work studies online SGD, whereas empirical risk minimization, empirical gradient flow, and other finite-sample optimization procedures can exhibit additional correlations and require different dynamical tools; existing results cover only special cases of attentionindexed models [12, 6, 21]. It would also be interesting to seek alternative representations of the infinite moment hierarchy. In particular, a resolvent-based formulation [27] may provide a more direct connection between the training dynamics and the evolving spectral distributions [19, 10] of the attention matrices. More broadly, the results suggest that the parameterization of attention should be regarded as part of the learning mechanism itself. The same loss landscape can support qualitatively different training trajectories depending on the parameterization. Understanding this interaction between architecture, parameterization, and high-dimensional learning dynamics is, in our view, an important step toward a theory of more realistic attention models.

Acknowledgements We acknowledge funding from the Swiss National Science Foundation grants SNSF SMArtNet (grant number 212049), OperaGOST (grant number 200021 200390), and DSGIANGO (grant number 225837). This work was supported by the Simons Collaboration on the Physics of Learning and Neural Computation via the Simons Foundation grant (#1257412 (FK) and #1257413 (LZ)). This work was also supported by the EPFL AI Center PhD Fellowship Program 2026 (YX).

References [1] Luca Arnaboldi, Bruno Loureiro, Ludovic Stephan, Florent Krzakala, and Lenka Zdeborova. Asymptotics of sgd in sequence-single index models and single-layer attention networks. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [2] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on nonconvex losses from high-dimensional inference. Journal of Machine Learning Research, 22(106):1–51, 2021. [3] Jean Barbier, Francesco Camilli, Justin Ko, and Koki Okajima. Phase diagram of extensive-rank symmetric matrix denoising beyond rotational invariance. Physical Review X, 15(2):021085, 2025. [4] Dario Bocchi, Theotime Regimbeau, Carlo Lucibello, Luca Saglietti, and Chiara Cammarota. Escape dynamics and implicit bias of one-pass sgd in overparameterized quadratic networks. arXiv preprint arXiv:2604.03068, 2026. [5] Antoine Bodin and Nicolas Macris. Gradient flow on extensive-rank positive semi-definite matrix denoising. In 2023 IEEE Information Theory Workshop (ITW), pages 365–370. IEEE, 2023. [6] Fabrizio Boncoraglio, Vittorio Erba, Emanuele Troiani, Yizhou Xu, Florent Krzakala, and Lenka Zdeborová. Single-head attention in high dimensions: A theory of generalization, weights spectra, and scaling laws. In Workshop on Scientific Methods for Understanding Deep Learning, 2026. ICML. [7] Fabrizio Boncoraglio, Emanuele Troiani, Vittorio Erba, and Lenka Zdeborova. Bayes optimal learning of attention-indexed models. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [8] Hugo Cui. High-dimensional learning of narrow neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2025(2):023402, 2025. 15

[9] Hugo Cui, Freya Behrens, Florent Krzakala, and Lenka Zdeborová. A phase transition between positional and semantic learning in a solvable model of dot-product attention. Advances in Neural Information Processing Systems, 37:36342–36389, 2024. [10] Leonardo Defilippis, Yizhou Xu, Julius Girardin, Emanuele Troiani, Vittorio Erba, Lenka Zdeborová, Bruno Loureiro, and Florent Krzakala. Scaling laws and spectra of shallow neural networks in the feature learning regime. arXiv preprint arXiv:2509.24882, 2025. [11] Odilon Duranthon, Pierre Marion, Claire Boyer, Bruno Loureiro, and Lenka Zdeborová. Statistical advantage of softmax attention: insights from single-location regression. arXiv preprint arXiv:2509.21936, 2025. [12] Vittorio Erba, Emanuele Troiani, Lenka Zdeborová, and Florent Krzakala. The nuclear route: Sharp asymptotics of erm in overparameterized quadratic networks. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [13] Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019. [14] Ruiquan Huang, Yingbin Liang, and Jing Yang. How transformers learn regular language recognition: A theoretical study on training dynamics and implicit bias. In International Conference on Machine Learning, pages 25516–25549. PMLR, 2025. [15] 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. [16] Daniel Kunin, Giovanni Luca Marchetti, Feng Chen, Dhruva Karkada, James B Simon, Michael R DeWeese, Surya Ganguli, and Nina Miolane. Alternating gradient flows: A theory of feature learning in two-layer neural networks. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. [17] Ashok V Makkuva, Marco Bondaschi, Chanakya Ekbote, Adway Girish, Alliot Nagle, Hyeji Kim, and Michael Gastpar. Local to global: Learning dynamics and effect of initialization for transformers. Advances in Neural Information Processing Systems, 37:86243–86308, 2024. [18] Pierre Marion, Raphaël Berthier, Gérard Biau, and Claire Boyer. Attention layers provably solve singlelocation regression. In ICLR 2025-Thirteenth International Conference on Learning Representations, 2025. [19] Charles H Martin and Michael W Mahoney. Implicit self-regularization in deep neural networks: Evidence from random matrix theory and implications for learning. Journal of Machine Learning Research, 22(165):1–73, 2021. [20] Simon Martin, Francis Bach, and Giulio Biroli. On the impact of overparameterization on the training of a shallow neural network in high dimensions. In International Conference on Artificial Intelligence and Statistics, pages 3655–3663. PMLR, 2024. [21] Simon Martin, Giulio Biroli, and Francis Bach. High-dimensional analysis of gradient flow for extensivewidth quadratic neural networks. arXiv preprint arXiv:2601.10483, 2026. [22] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. In Proceedings of the 41st International Conference on Machine Learning, pages 38018–38070, 2024. [23] Naoki Nishikawa, Yujin Song, Kazusato Oko, Denny Wu, and Taiji Suzuki. Nonlinear transformers can perform inference-time feature learning. In Forty-second International Conference on Machine Learning, 2025. 16

[24] Salim Noreddine and Ivan Nourdin. On the gaussian approximation of vector-valued multiple integrals. Journal of multivariate analysis, 102(6):1008–1017, 2011. [25] Ivan Nourdin, Giovanni Peccati, and Anthony Réveillac. Multivariate normal approximation using stein’s method and malliavin calculus. In Annales de l’IHP Probabilités et statistiques, volume 46, pages 45–58, 2010. [26] David Nualart and Giovanni Peccati. Central limit theorems for sequences of multiple stochastic integrals. The Annals of Probability, 33(1):177–193, 1 2005. [27] Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington. 4+ 3 phases of computeoptimal neural scaling laws. Advances in Neural Information Processing Systems, 37:16459–16537, 2024. [28] Farzad Pourkamali, Jean Barbier, and Nicolas Macris. Matrix inference in growing rank regimes. IEEE Transactions on Information Theory, 70(11):8133–8163, 2024. [29] Farzad Pourkamali and Nicolas Macris. Rectangular rotational invariant estimator for high-rank matrix estimation. Information and Inference: A Journal of the IMA, 14(3):iaaf015, 2025. [30] Chen Siyu, Sheen Heejune, Wang Tianhao, and Yang Zhuoran. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. In The Thirty Seventh Annual Conference on Learning Theory, pages 4573–4573. PMLR, 2024. [31] Bingqing Song, Boran Han, Shuai Zhang, Jie Ding, and Mingyi Hong. Unraveling the gradient descent dynamics of transformers. Advances in Neural Information Processing Systems, 37:92317–92351, 2024. [32] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. Advances in neural information processing systems, 36:71911–71947, 2023. [33] Francois Treves. Ovsyannikov theorem and hyperdifferential operators. Notas de Matematica, 46, 1968. [34] Emanuele Troiani, Hugo Cui, Yatin Dandi, Florent Krzakala, and Lenka Zdeborová. Fundamental limits of learning in sequence multi-index models and deep attention networks: High-dimensional asymptotics and sharp thresholds. In Forty-second International Conference on Machine Learning, 2025. [35] Emanuele Troiani, Vittorio Erba, Florent Krzakala, Antoine Maillard, and Lenka Zdeborová. Optimal denoising of rotationally invariant rectangular matrices. In Mathematical and Scientific Machine Learning, pages 97–112. PMLR, 2022. [36] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. [37] Hongru Yang, Bhavya Kailkhura, Zhangyang Wang, Yingbin Liang, et al. Training dynamics of transformers to recognize word co-occurrence via gradient flow analysis. Advances in Neural Information Processing Systems, 37:46047–46117, 2024. [38] Liu Ziyin, Yizhou Xu, Tomaso Poggio, and Isaac Chuang. Neural quadratic forms: A unified minimal model for sudden learning and scaling laws. arXiv preprint arXiv:2608.13335, 2026.

17

Contents 1 Introduction

1

2 Attention-indexed models and their population loss landscape

2

3 Online SGD of tied attention 3.1 Macroscopic dynamics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 3.2 Weak recovery and its sample complexity . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

5 5 7

4 Online SGD of untied attention 4.1 Macroscopic dynamics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4.2 Weak recovery and its sample complexity . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

9 9 12

5 Conclusion and future work

14

A Notation

19

B Examples of attention-indexed models B.1 Example 1: Multi-layer multi-head attention . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.2 Example 2: Autoregressive sequence model . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.3 Example 3: Multiple-location regression . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.4 Example 4: Multi-step reasoning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5 Example 5: Extensive-rank matrix denoising . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

19 19 20 21 21 22

C Theory in Section 2 C.1 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.2 Tied attention . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.3 Proof of Corollary 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

22 22 24 25

D Theory in Section 3 D.1 Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.1.1 Local well-posedness . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.1.2 Uniform spectral bounds . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.1.3 Moment expansion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.1.4 Compactness of the empirical trajectories . . . . . . . . . . . . . . . . . . . . . . . . . . . D.1.5 Global well-posedness . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.1.6 Final proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.2 Proof of Corollary 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.3 An exactly solvable example . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.4 Proof of Corollary 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.5 Convergence rate of strong recovery . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

26 26 26 27 34 39 40 44 45 47 48 50

E Online SGD over S E.1 Proof of Theorem 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . E.1.1 Uniform spectral bounds . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . E.1.2 Moment expansion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . E.1.3 Final proof of Theorem 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . E.2 Proof of Corollary 7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

52 54 54 57 60 61

F Theory in Section 4 F.1 Proof of Theorem 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . F.1.1 A stronger version of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . F.1.2 Local well-posedness . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . F.1.3 Uniform spectral bounds . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . F.1.4 Fast-dynamics convergence . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

64 64 64 64 65 79

18

F.2 F.3 F.4 F.5

F.1.5 Moment expansion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 81 F.1.6 Limiting slow dynamics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 90 F.1.7 Final proof of Theorem 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 96 Truncation to finite orders . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 97 Proof of Corollary 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 99 Convergence rate of strong recovery . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 100 Proof of Lemma 10 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 102

G Declaration of LLM usage

A

103

Notation

We write [K] ∶= {1, . . . , K} and let K ⊆ [K] denote the set of trainable indices. For a multi-index or word α, ∣α∣ denotes its length, while ∣α∣K denotes the number of positions whose associated index belongs to K. The underlying finite alphabet of α will be clear from context. For matrices, ∥ ⋅ ∥op and ∥ ⋅ ∥F denote the operator and Frobenius norms, respectively, and ⟨A, B⟩F ∶= Tr(A⊺ B). We write SK for the space of real symmetric K × K matrices and SK + for the space of positivesemidefinite matrices. For symmetric matrices, A ⪰ B means that A − B is positive semidefinite. For finite-dimensional vectors, ∥ ⋅ ∥∞ and ∥ ⋅ ∥ denote the ℓ∞ and Euclidean norms. We write 1E for the indicator of an event E, a ∧ b ∶= min{a, b}, and [a]+ ∶= max{a, 0}. For sets A ⊆ B, the notation A ⋐ B means that A is a compact subset of B. We use the standard asymptotic notation O(⋅), o(⋅), Ω(⋅), and Θ(⋅) as d → ∞. For random variables, Xd = OP (ad ) means that Xd /ad is bounded in probability, Xd = oP (ad ) means that Xd /ad Ð → 0, and Xd = P

ΘP (ad ) means that both Xd = OP (ad ) and ad = OP (Xd ). We write Xd Ð → X for convergence in probability. For s > 0, let Xs denote the Banach space of moment sequences µ = (µα )∣α∣≥1 equipped with the norm P

∥µ∥s ∶= supw≥1 sup∣α∣=w ∣µswα ∣ and define Xs0 ∶= {µ ∈ Xs ∶ limw→∞ sup∣α∣=w ∣µswα ∣ = 0}. Equipped with the norm

inherited from Xs , Xs0 is a separable Banach space. C([0, T ]; Xs0 ) denotes the space of continuous functions from [0, T ] to Xs0 . D([0, T ]; Xs0 ) represents the space of all right-continuous functions with left limits mapping from [0, T ] into Xs0 , equipped with the Skorokhod J1 topology. Throughout the paper, a superscript (d) denotes a finite-dimensional quantity, a bar denotes its deterministic high-dimensional limit, and a tilde denotes a piecewise-constant time interpolation, unless specified otherwise.

B

Examples of attention-indexed models

B.1

Example 1: Multi-layer multi-head attention

In this example, we demonstrate that the standard multi-layer multi-head attention network (attention-only transformer) falls within our general theoretical framework. Let the input to the model be Gaussian X0 ∼ N (0, C ⊗ Id ). For layer m = 1, . . . , M with H attention heads, the pre-softmax attention matrix for the h-th head is given by 1 (h) T L×L A(h) m (Xm−1 ) = √ Xm−1 Sm−1 Xm−1 ∈ R d (h)

(h)

(50)

(h)

where Sm−1 ∶= Km−1 (Qm−1 )T ∈ Rd×d is the combined key-query weight matrix. The intermediate representations follow the standard multi-head attention mechanism: H

(h)

L×d Xm = Xm−1 + ∑ softmax (A(h) , m (Xm−1 )) Xm−1 Vm−1 ∈ R

(51)

h=1

(h)

where Vm−1 ∈ Rd×d represents the combined value and output projection matrix for head h. Unrolling the recurrence relation reveals that this network can be expressed as a composite function of the quadratic forms Gk ∶= √1d X0 S̃k X0T . 19

For the first layer (m = 1), the pre-softmax matrix constitutes a bilinear form: 1 (h) (h) (h) A1 (X0 ) = √ X0 S0 X0T ∶= G0 . d

(52) (h)

(h)

The output of the first layer is thus X1 = X0 + ∑H h=1 softmax(G0 )X0 V0 . (h) (h) For the second layer (m = 2), substituting X1 into the pre-softmax attention matrix A2 (X1 ) = √1d X1 S1 X1T yields a sum of four types of cross terms: H 1 1 (h) (h) (i) (i) (h) A2 (X1 ) = √ X0 S1 X0T + ∑ softmax(G0 ) ( √ X0 (V0 S1 ) X0T ) d d i=1 H 1 (j) (h) (j)T ) X0T ) softmax(G0 )T + ∑ ( √ X0 (S1 V0 d j=1

(53)

H H 1 (j) (i) (h) (j)T (i) ) X0T ) σ(G0 )T . + ∑ ∑ softmax(G0 ) ( √ X0 (V0 S1 V0 d i=1 j=1

Every inner term encapsulated by parentheses in the expansion above takes the form √1d X0 S̃X0T . For instance, the attention outputs cross term forms a new bilinear form defined by the effective weight matrix (i) (h) (j)T S̃ ∶= V0 S1 V0 . By induction, any attention matrix Am , and consequently the final output representation, can be written as a finite composite function of a set of base bilinear forms {Gk }K k=1 generated by equivalent effective weight matrices S̃k . These effective matrices S̃k are finite products of the original network weights (S, V ), which falls within the scope of attention-indexed models. Furthermore, because the softmax function is continuous and bounded, the composite function remains continuous and exhibits at most polynomial growth. Therefore, when we assume that both the teacher and the student models are the multi-layer multi-head attention architecture, the final loss satisfies Assumption 1. Note that we do not require the teacher weights to be Gaussian, making the existence of such a teacher model a weak assumption.

B.2

Example 2: Autoregressive sequence model

This example considers the most fundamental pretraining task: next-token prediction. We aim to predict the next token xL using the context x1 , . . . , xL−1 via a simplified single-layer multi-head attention model. The student is a standard attention block, defined as: H L−1

x̂L = ∑ ∑ αj,h Vh xj , h=1 j=1

1 where αj,h = Softmaxj ({ √ xTL−1 W h xk }L−1 k=1 ) . d

(54)

Here, we absorb the query-key matrices into a single weight matrix W h . We assume the data is generated by a sparse contextual dependency process. The next token xL is a composition of specific past tokens that satisfy a relevance criterion H L−1

xL = ∑ ∑ aj,h Vh∗ xj ,

(55)

h=1 j=1

where aj,h ∈ {0, 1} is a gating variable determining if token j is relevant for generating L under head h: L−1

⎛ ⎛ 1 ⎞⎞ aj,h ∼ Bernoulli g { √ xTi M h xj } . ⎝ ⎝ d i,j=1 ⎠⎠

(56)

This models a scenario where the ground truth dependency structure is latent and dynamic, determined by the interaction between tokens.

20

The MSE loss can be expanded as: L ∶= ∥x̂L − xL ∥2 = ∑ xTj Ωjj ′ xj ′ ,

(57)

j,j ′

where Ωjj ′ ∶= ∑h,h′ (aj,h Vh∗ − αj,h Vh )T (aj ′ ,h′ Vh∗′ − αj ′ ,h′ Vh′ ). It is thus an attention-indexed model. We can similarly analyze the standard next-token prediction MSE loss L−1

L ∶= EL,x1 ,⋯,xL ∥x̂L − xL ∥2 = EL EC∣L E(x1 ,⋯,xL−1 )∼N (0,C) ∑ xTj Ωjj ′ xj ′ , j,j ′

(58)

which fits the general framework.

B.3

Example 3: Multiple-location regression

This example considers a classification task (e.g., sentiment analysis), where the goal is to predict the label y ∈ {+1, −1} based on the interaction between tokens (e.g., detecting negation or intensification). We model the student predictor as a one-layer attention: ⎛ ⎞ ŷ = sign ∑ αij σ ({xTi W h xj }ij ) , ⎝ i,j ⎠

(59)

where the model learns αij to weight different token pairs. Moreover, we assume that the true label is generated by a "bag-of-interactions" model: ⎛ ⎞ y = sign ∑ Sij ⋅ ϵij , ⎝ i,j ⎠

(60)

where Sij ∈ {+1, −1} represents the inherent sentiment polarity of the word pair (i, j) (e.g., +1 for "very good", −1 for "not good"), and ϵij indicates whether a meaningful grammatical connection exists between them: ⎧ ⎪ ⎪1, w.p. g({xTi W ∗ xj }ij ) ϵij = ⎨ (61) ⎪ 0, otherwise. ⎪ ⎩ In this setting, ϵij can be interpreted as the ground-truth dependency graph (e.g., i modifies j), which is activated probabilistically by the underlying semantic matching W ∗ .

B.4

Example 4: Multi-step reasoning

This example models complex tasks requiring multi-step reasoning, such as in-context retrieval and composition (e.g., "Context: A:1, B:2, C:3. Query: Calculate A+B"). We model the ground truth y as a hierarchical interaction between two latent tokens. The teacher first identifies two relevant tokens from the context based on the query xL , and then computes their quadratic interaction: y = xTϵ1 U ∗ xϵ2 , (62) where the indices ϵ1 , ϵ2 ∈ {1, . . . , L − 1} are latent pointers. These pointers follow a probabilistic selection mechanism (e.g., a "ground truth" attention): 1 P (ϵh = k) ∝ exp ( √ xTL Wh∗ xk ) , d

for h ∈ {1, 2}.

(63)

This represents the "reasoning logic" where Wh∗ retrieves the necessary operands and U ∗ defines the operation performed on them. To solve this, the student must implement a two-stage process: first filtering/aggregating the context into latent representations, and then combining them. We model the student as a two-head attention 21

structure followed by a bilinear composition layer. The student first generates two latent vectors z1 and z2 by aggregating the context tokens {xk }: L−1

zh = ∑ αj,h xj , j=1

L−1

⎞ ⎛ 1 . where αj,h = Softmaxj { √ xTL W h xk } ⎝ d k=1 ⎠

(64)

Here, zh serves as the "retrieved operand" for head h. The final prediction is computed as the interaction between these aggregated representations: ŷ = z1T W̃ z2 , (65) which takes the form of an attention-indexed model.

B.5

Example 5: Extensive-rank matrix denoising

Consider the observation model:

Y = S ∗ + Z ∈ Rd×d ,

(66)

where S ∗ is an unknown deterministic symmetric matrix with an extensive rank (i.e., limd→∞ ∥S ∗ ∥op /∥S ∗ ∥F = 0), and Z is the noise matrix drawn from the Gaussian Orthogonal Ensemble (GOE), satisfying E[Zab ] = 0 and E[Zab Zcd ] = d1 (δac δbd + δad δbc ). While prior work has extensively studied rotationally invariant priors [35, 28, 29], the general extensive-rank setting remains largely open [3]. A standard approach for inferring S ∗ from Y is via: 1 min L ( Tr[Y S]) + λf (S), S d where L ∶ R → R is the loss function and f (S) represents a regularizer. The population loss can be expressed as an expectation over a standard scalar Gaussian variable g ∼ N (0, 1): √ 1 1 2 ∗ EZ L ( Tr[Y S]) = Eg∼N (0,1) [L ( Tr[S S] + 3/2 ∥S∥F g)] , (67) d d d which can be regarded as a single-token single-index attention-indexed model as a special case of Theorem 1 and Corollary 5.

C

Theory in Section 2

C.1

Proof of Theorem 1

The proof relies on the method of joint cumulants. If all joint cumulants of order m ≥ 3 for a set of random variables vanish in the limit, the variables converge in distribution to a multivariate Gaussian (satisfying Carleman’s condition). k Let Zijk = √1d xTi Sij xj . We concatenate all data vectors into a single vector X = [xT1 , xT2 , . . . , xTL ]T ∈ RdL . Given the covariance structure E[xia xjb ] = Cij δab , we have X ∼ N (0, Σ), where the covariance matrix is the Kronecker product Σ = C ⊗ Id . We can rewrite Zijk as a quadratic form: Zijk = X T A(ijk) X,

(68)

1 1 k T where A(ijk) is a dL×dL symmetric block matrix. Its (i, j)-th block is 2√ S k , its (j, i)-th block is 2√ (Sij ) , d ij d and all other blocks are zero. By the standard expectation of Gaussian quadratic forms, the first moment is:

Cij 1 d k k E[Zijk ] = Tr(A(ijk) Σ) = √ ∑ E[xia xja ](Sij )aa = √ Tr(Sij ). d a=1 d By Assumption 1.2, as d → ∞, this converges to E[Gijk ] = Cij µijk . 22

(69)

To calculate the covariance, we use Wick’s Theorem: ′ ′ ′

Cov(Zijk , Zi′ j ′ k′ ) = 2Tr(A(ijk) ΣA(i j k ) Σ).

(70)

Expanding this into block components yields: Cov(Zijk , Zi′ j ′ k′ ) =

1 k k′ ∑ (S )ab (Si′ j ′ )ce (Cii′ Cjj ′ δac δbe + Cij ′ Cji′ δae δbc ) d a,b,c,e ij ′

′

= Cii′ Cjj ′

k Tr(Sij (Sik′ j ′ )T )

d

+ Cij ′ Cji′

k k Tr(Sij Si′ j ′ )

d

(71)

.

By Assumption 1.2, as d → ∞, this converges to the finite constant Cii′ Cjj ′ Ωijk,i′ j ′ k′ + Cij ′ Cji′ Ψijk,i′ j ′ k′ . Now we calculate the joint cumulant κm ∶= κ(Q1 , ⋯, Qm ) for Qr = X T Ar X, where we denote Ar ∶= (ir jr kr ) A for r ∈ {1, . . . , m}. The joint cumulant generating function reads m

K(t1 , . . . , tm ) = log E [exp ( ∑ tr Qr )] = log E [exp(X T ÃX)] ,

(72)

r=1

where à = ∑m r=1 tr Ar . Using the standard Gaussian integral for quadratic forms in the neighborhood of (t1 , . . . , tm ) = 0, the expectation can be calculated as 1 K(t1 , . . . , tm ) = − log det(I − 2ΣÃ). 2

(73)

1 n Using the identity log det(I − M ) = − ∑∞ n=1 n Tr(M ), the right side can be expanded into a series of traces: n

∞

∞ m 2n−1 2n−1 Tr ((ΣÃ)n ) = ∑ Tr (( ∑ tr ΣAr ) ) . n=1 n n=1 n r=1

K(t1 , . . . , tm ) = ∑

m

K Then the joint cumulant is given by κm ∶= ∂t∂1 ⋯∂t ∣ m

t1 =⋅⋅⋅=tm =0

(74)

. The only term in the series that contributes

to this cross-derivative is n = m. Therefore, after expanding the m-th power, κm is a linear combination of traces over all permutations π in the symmetric group Sm : κm = cm ∑ Tr (Aπ(1) ΣAπ(2) Σ . . . Aπ(m) Σ) ,

(75)

π∈Sm

where cm is a constant depending only on m. Since m is fixed, the sum consists of a finite number of terms. We can bound the absolute value of each individual trace in the sum using the matrix norm inequality ∣Tr(B1 B2 . . . Bm )∣ ≤ (∏m−2 r=1 ∣∣Br ∣∣op ) ∣∣Bm−1 ∣∣F ∣∣Bm ∣∣F . Let Br = A(ir jr kr ) Σ and let us analyze the asymptotic scales of these norms: ∣∣Σ∣∣op = λmax (C ⊗ Id ) = k k λmax (C) = Θ(1). ∣∣A(ijk) ∣∣F = √12d ∣∣Sij ∣∣F = O(1), based on Assumption 1.2. From Assumption 1.1, ∣∣Sij ∣∣op = √ 1 k (ijk) k ∣∣op = 2√d ∣∣Sij ∣∣op = o(1). o(∣∣Sij ∣∣F ) = o( d). Therefore, ∣∣A Thus, for Br , the spectral norm is bounded by ∣∣Br ∣∣op ≤ ∣∣A(ir jr kr ) ∣∣op ∣∣Σ∣∣op = o(1) ⋅ O(1) = o(1), and the Frobenius norm is O(1). Applying this to the trace inequality for any m ≥ 3: κm = O ([o(1)]m−2 ⋅ O(1) ⋅ O(1)) = o(1).

(76)

Therefore, all joint cumulants of order 3 and higher vanish as d → ∞. Because all higher-order cumulants vanish while the first two cumulants converge to finite limits, the vector of quadratic forms {Zijk }L,L,K i,j,k=1 converges weakly in distribution to the multivariate Gaussian vector {Gijk }L,L,K i,j,k=1 . Finally, to pass the limit inside the expectation, we need to show that the sequence L({Zijk }L,L,K i,j,k=1 ) is uniformly integrable. By Assumption 1.3, L is a continuous function with at most polynomial growth, meaning there exist constants c > 0 and p > 0 such that ∣L(Z)∣ ≤ c(1 + ∣∣Z∣∣p ). Therefore, it suffices to show that the random tensor Z = {Zijk }L,L,K i,j,k=1 has bounded joint moments of all orders (up to p + 1) as d → ∞. 23

Recall that any joint moment of a set of random variables can be expressed as a finite polynomial of their joint cumulants. From our previous derivations, the joint cumulants κm converge to finite limits (specifically, bounded constants for m ≤ 2 and 0 for m ≥ 3). Because all joint cumulants of Z are uniformly bounded with respect to d, it immediately follows that all joint moments of Z (up to p + 1) are also uniformly bounded. This uniform boundedness of all moments guarantees that L({Zijk }L,L,K i,j,k=1 ) is uniformly integrable. Thus, we can exchange the limit and the expectation, yielding: L,L,K lim Ex [L({Zijk }L,L,K i,j,k=1 )] = E[L({Gijk }i,j,k=1 )].

d→∞

(77)

This completes the proof.

C.2

Tied attention

k ) = µijk in Theorem 1 might not be satisfied for certain architectures, The assumption limd→∞ √1d Tr(Sij particularly those with tied weights. For instance, if the attention mechanism uses symmetric tied weights (e.g., S = W W T ), its trace scales proportionally to the dimension d, causing the term √1d Tr(S) to diverge as d → ∞. To accommodate such cases, we introduce a centered variant of the attention-indexed model. k ) = µijk is dropped, we Corollary 5. Under the Assumption 1 except that the condition limd→∞ √1d Tr(Sij have ⎡ ⎧ ⎫L,L,K ⎞⎤⎥ k T ⎢ ⎛⎪ ⎪ ⎥ ⎢ ⎪ Tr[Sij (xj xi − Cij Id )] ⎪ ⎟⎥ = E [L({Gijk }L,L,K √ ⎬ (78) lim Ex ⎢L ⎜⎨ i,j,k=1 )] ⎢ ⎝⎪ ⎪ d→∞ d ⎪ ⎪ ⎠⎥⎥ ⎢ ⎩ ⎭ i,j,k=1 ⎦ ⎣

where {Gijk }L,L,K i,j,k=1 is a multivariate Gaussian vector whose mean and covariance are given by: E[Gijk ] = 0, Cov(Gijk , Gi′ j ′ k′ ) = Cii′ Cjj ′ Ωijk,i′ j ′ k′ + Cij ′ Cji′ Ψijk,i′ j ′ k′ .

(79)

k Proof. Let Yijk = √1d Tr[Sij (xj xTi − Cij Id )]. We can rewrite this as a centered quadratic form:

1 k k Yijk = √ (xTi Sij xj − Cij Tr(Sij )) = Zijk − E[Zijk ], d

(80)

k xj as defined in the proof of Theorem 1. where Zijk = √1d xTi Sij By construction, the first moment of Yijk is zero for any d: E[Yijk ] = 0. Because Yijk is a deterministic shift of Zijk , all joint cumulants of order m ≥ 2 are identical to those of Zijk . Specifically, the covariance structure remains:

lim Cov(Yijk , Yi′ j ′ k′ ) = lim Cov(Zijk , Zi′ j ′ k′ ) = Cii′ Cjj ′ Ωijk,i′ j ′ k′ + Cij ′ Cji′ Ψijk,i′ j ′ k′ .

d→∞

d→∞

(81)

Furthermore, for any order m ≥ 3, the joint cumulants of {Yijk } match those of {Zijk } and thus vanish as o(1), following Theorem 1. Since the first two cumulants converge to finite limits and all higher-order cumulants vanish, the sequence L,L,K of random vectors {Yijk }L,L,K i,j,k=1 converges weakly in distribution to {Gijk }i,j,k=1 , which is a multivariate Gaussian vector with mean zero and the asymptotic covariance structure derived above. Because the loss function L is continuous with at most polynomial growth, and {Yijk } has bounded k moments (due to the bounds on the spectral and Frobenius norms of Sij by Assumption 1), the sequence L,L,K L({Yijk }i,j,k=1 ) is uniformly integrable. This allows us to pass the limit inside the expectation: L,L,K

lim Ex [L ({Yijk }i,j,k=1 )] = E [L({Gijk }L,L,K i,j,k=1 )] .

d→∞

This completes the proof.

24

(82)

C.3

Proof of Corollary 1

Corollary 1 is proven under the following regularity assumption. Assumption 7. For a given dimension d, let Ŝ (d) = {Ŝij }L,L i,j=1,k∈K be a sequence of minimizers of the population loss with respect to the learnable weights. Assume that as d → ∞: k,(d)

k,(d)

1. The minimizers satisfy the spectral norm condition uniformly: limd→∞ ∥Ŝij k ∈ K.

k,(d)

∥op /∥Ŝij

∥F = 0 for all

2. The order parameters associated with both the minimizers and the fixed weights, denoted as Q(Ŝ (d) ), converge to a valid point Q̂ ∈ Q. Now we prove Corollary 1. k Let SK denote the collection of learnable weights {Sij }k∈K . Let the population loss in dimension d be denoted by: ⎡ ⎧ T k ⎫L,L,K ⎤ ⎢ ⎛⎪ ⎞⎥⎥ ⎪ ⎢ ⎪ xi Sij xj ⎪ ⎟⎥ . Fd (SK ) ∶= Ex ⎢L ⎜⎨ √ ⎬ (83) ⎢ ⎝⎪ d ⎪ ⎪ ⎠⎥⎥ ⎢ ⎪ ⎭ ⎩ i,j,k=1 ⎦ ⎣ We aim to prove that limd→∞ inf SK Fd (SK ) = inf Q∈Q R(Q) by establishing matching upper and lower bounds. For any arbitrary ϵ > 0, by the definition of the infimum, there exists Qϵ ∈ Q such that: R(Qϵ ) ≤ inf R(Q) + ϵ.

(84)

Q∈Q

(d)

Because Qϵ lies within the feasible domain Q, there exists a sequence of learnable matrices S̃K such (d) that, combined with the fixed matrices, the joint order parameters satisfy limd→∞ Q(S̃K ) = Qϵ , and the k,(d) k,(d) learnable matrices satisfy the norm condition limd→∞ ∥S̃ij ∥op /∥S̃ij ∥F = 0. Since the infimum over all possible learnable weights SK must be less than or equal to the loss evaluated (d) at this specific sequence S̃K , we have: (d)

inf Fd (SK ) ≤ Fd (S̃K ).

(85)

SK

(d)

Taking the limit superior as d → ∞ on both sides, and noting that the complete set of matrices (both S̃K and the fixed matrices) satisfies all conditions of Theorem 1, we can apply Theorem 1 to the right-hand side: (d)

lim sup inf Fd (SK ) ≤ lim Fd (S̃K ) = R(Qϵ ) ≤ inf R(Q) + ϵ. d→∞

SK

Q∈Q

d→∞

(86)

Taking the limit as ϵ → 0 yields the upper bound: lim sup inf Fd (SK ) ≤ inf R(Q). d→∞

SK

(87)

Q∈Q

(d)

(d)

Let ŜK be the sequence of minimizers such that inf SK Fd (SK ) = Fd (ŜK ). According to Assumption 7.1, (d) the sequence ŜK uniformly satisfies the norm ratio condition, and the combined order parameters converge to some state Q̂ ∈ Q. We can then apply Theorem 1 to this joint sequence of minimizers and fixed weights: (d)

lim Fd (ŜK ) = R(Q̂).

(88)

d→∞

Furthermore, since Q̂ is an element of the feasible domain Q, we have R(Q̂) ≥ inf Q∈Q R(Q). Therefore, taking the limit inferior, we obtain: (d)

lim inf inf Fd (SK ) = lim Fd (ŜK ) = R(Q̂) ≥ inf R(Q). d→∞

SK

Q∈Q

d→∞

(89)

Combining the upper bound (87) and the lower bound (89), we get: lim inf Fd (SK ) = inf R(Q).

d→∞ SK

Q∈Q

This concludes the proof. 25

(90)

D

Theory in Section 3

D.1

Proof of Theorem 2

D.1.1

Local well-posedness

We first define a scale of the Banach spaces Xs for s > 0, consisting of infinite-dimensional moment sequences µ = (µα )∣α∣≥1 equipped with the norm ∥µ∥s ∶= supn≥1 sup∣α∣=n ∣µα ∣/sn . By Price’s theorem, the matrix derivative of the potential can be written as ∇Φ(q) = Γkl (q) with Γkl (q) ∶=

L 1 2 ∑ (Cii′ Cjj ′ + Cij ′ Cji′ ) E [∂(ijk),(i′ j ′ l) L(G)] , 2 i,j,i′ ,j ′ =1

1 ≤ k, l ≤ K.

(91)

See Lemma 2 for more detailed derivation. A technical issue is that the potential (10), and hence the gradient Γ (91), are well-defined only for q ⪰ 0, whereas an arbitrary element of Xs need not have a PSD second-moment matrix. To use the Ovsyannikov argument, we extend Γ outside the PSD cone. Define P(q) ∶= ΠSK ( +

q + qT ), 2

(92)

where ΠSK denotes the orthogonal projection onto the closed convex cone of PSD matrices, and set + ̃ Γ(q) ∶= Γ(P(q)),

q ∈ RK×K .

(93)

̃ For q ⪰ 0, this extension agrees with the generalized gradient: Γ(q) = Γ(q). Now we denote the flow as d µ̄(t) = V (µ̄(t)) ∶= F (µ̄(t)) + G(µ̄(t)), dt where

w

K

̃ k l (µ(k ,...,k ,l,k ,...,k ) + µ(k ,...,k ,l,k ,...,k ) ) . [F (µ)](k1 ,...,kw ) ∶= − ∑ ∑ Γ(q) i 1 i−1 i w 1 i i+1 w

(94)

(95)

i=1 l=1 ki ∈K

and

γ [G(µ̄)]α ∶= − ∣α∣K µ̄α 2 are two unbounded operators in the Banach space Xs .

(96)

Lemma 1. Given an initial condition µ̄(0) ∈ Xs0 for some s0 > 0, there exist Tlocal > 0 and a strictly increasing continuous function s(t) with s(0) = s0 such that (94) admits a unique solution µ̄(t) ∈ Xs(t) for t ∈ [0, Tlocal ). K Proof. For every bounded subset of SK is bounded and Lipschitz continuous by definition (91) + , Γ(q) ∈ S and Assumption 2.1 (using Price’s theorem). Therefore, for every R > 0 there exist constants CΓ (R) and LΓ (R) such that sup ∥Γ(q)∥F ≤ CΓ (R), (97) q⪰0,∥q∥F ≤R

and

∥Γ(q) − Γ(q ′ )∥F ≤ LΓ (R)∥q − q ′ ∥F

(98)

whenever q, q ′ ⪰ 0 and ∥q∥F , ∥q ′ ∥F ≤ R. Since the orthogonal projection onto the closed convex cone SK + ̃ satisfies ∥P(q) − P(q ′ )∥F ≤ ∥q − q ′ ∥F , Γ(q) = Γ(P(q)) is locally bounded and locally Lipschitz on RK×K . Now fix 0 < s < s′ and R > 0, and let BR (Xs ) ∶= {µ ∈ Xs ∶ ∥µ∥s ≤ R}. For µ, ν ∈ BR (Xs ), ∣qkl (µ)∣ ≤ Rs2 ,

∣qkl (µ) − qkl (ν)∣ ≤ s2 ∥µ − ν∥s ,

(99)

where we denote qkl (µ) = µ(k,l) Hence, because K is fixed, we have ̃ ∣Γ(q(µ))∣ F ≤ CR,s , 26

(100)

and

̃ ̃ ∥Γ(q(µ)) − Γ(q(ν))∥ F ≤ LR,s ∥µ − ν∥s .

(101)

for constants CR,s and LR,s . Now we prove that the vector field V ∶ Xs → Xs′ is bounded and locally Lipschitz. For the linear term G, we have ∣[G(µ)]α − [G(ν)]α ∣ ≤ γ2 w∣µα − να ∣, and thus s w γw∣µα − να ∣ γ ≤ ∥µ − ν∥s sup [w ( ′ ) ] . ′ w 2(s ) 2 s w≥1 ∣α∣=w≥1

∥G(µ) − G(ν)∥s′ = sup

(102)

1 1 ≤ 1−x for x ∈ (0, 1), we obtain: Using the inequality supw≥1 wxw ≤ e∣ log x∣

s w s′ sup w ( ′ ) ≤ ′ . s s −s w≥1

(103)

′

Thus, ∥G(µ) − G(ν)∥s′ ≤ 2(sγs′ −s) ∥µ − ν∥s . For the non-linear term F (µ), the difference [F (µ)]α −[F (ν)]α for a multi-index α of length w is bounded by ∣[F (µ)]α − [F (ν)]α ∣

(104)

w ⎞ ⎛̃ (1) (2) (1) (2) (1) (2) ̃ ̃ ≤ ∑ ∑ ∣Γ(q(µ)) ki l − Γ(q(ν))ki l ∣ ∣µw+1 + µw+1 ∣ + ∣Γ(q(ν))ki l ∣ ∣(µw+1 + µw+1 ) − (νw+1 + νw+1 )∣ . ⎠ i=1 l ⎝

(105)

ki ∈K

Since the moments of length w + 1 are bounded by ∥µ∥s sw+1 ≤ Rsw+1 and their differences are bounded by ∥µ − ν∥s sw+1 , we can obtain: ∣[F (µ)]α − [F (ν)]α ∣ ≤ wK [LR,s ∥µ − ν∥s (2Rsw+1 ) + CR,s (2∥µ − ν∥s sw+1 )] .

(106)

Computing the norm in Xs′ , we obtain: ∣[F (µ)]α − [F (ν)]α ∣ s w ≤ 2K (RL + C ) s∥µ − ν∥ sup [w ( ) ]. R,s R,s s (s′ )w s′ w≥1 ∣α∣=w≥1

∥F (µ) − F (ν)∥s′ = sup

w

(107)

′

Using the previously established inequality supw≥1 w ( ss′ ) ≤ s′s−s , we consequently arrive at: ∥F (µ) − F (ν)∥s′ ≤

2K (RLR,s + CR,s ) ss′ ∥µ − ν∥s . s′ − s

(108)

Combining F and G, the vector field V ∶ Xs → Xs′ is locally Lipschitz and bounded (using V (0) = 0). Then Lemma 1 is proven by the Ovsyannikov theorem [33]. D.1.2

Uniform spectral bounds

Let Fn = σ(x(0) , x(1) , . . . , x(n−1) ) be the filtration generated by the data samples up to step n − 1. The (n) weights Wk are Fn -measurable. In the online SGD, the update at step n relies on a fresh sample x(n) . Consider the stochastic gradient: (n)

gk (x(n) , W ) ∶= ∇Wk L(G(n) (x(n) , W )), (n)

and the expected gradient ḡk

(109)

(n) = Ex [∑L )Xij Wk ] conditioned on Fn . i,j=1 ∂(i,j,k) L(G

Lemma 2. Fix a constant C∗ > 0. Suppose that, at step n, (n)

max ∥Wl

1≤l≤K (n)

The conditional expected gradient ḡk (n)

(n)

satisfies ḡk

∥Hk ∥op ≤ CH (C∗ ),

∥op ≤ C∗ . (n)

(n)

= d1 Hk Wk (n)

(110) (n)

+ EStein,k , where

∥EStein,k ∥op ≤ CE (C∗ ) d−3/2 ,

for constants CH (C∗ ) and CE (C∗ ) independent of d and n. 27

(111)

Proof. Throughout the proof we condition on Fn and suppress the superscript (n) whenever no confusion can arise. Conditional on Fn , the matrices {Wl , Sl }K l=1 are deterministic, while x = (x1 , . . . , xL ) ∼ N (0, C ⊗ Id ). Recall that 1 Gijk ∶= √ Tr[Sk (xj xTi − Cij Id )] . (112) d Since Sk = SkT , we have ∇Wk Gijk = Xij Wk with Xij ∶= √1d (xj xTi + xi xTj − 2Cij Id ). Consequently, L

gk = ∑ fijk (x)Xij Wk ,

fijk (x) ∶= ∂(ijk) L(G(x)),

(113)

i,j=1

and hence

L

ḡk = ∑ Ex [fijk (x)Xij ]Wk .

(114)

i,j=1

We evaluate Ex [fijk (x)Xij ] using the second-order Stein identity. For any f ∈ C 2 (RdL ) whose derivatives up to second order have at most polynomial growth4 , and for a centered Gaussian vector x with covariance matrix Σ, E[f (x)(xxT − Σ)] = Σ E[∇2x f (x)] Σ. (115) Taking the (j, i) block of this identity for Σ = C ⊗ Id gives L

Ex [fijk (x) (xj xTi − Cji Id )] = ∑ Cju Cvi Ex [∇xu ∇xv fijk (x)] .

(116)

u,v=1

Using the definition of Xij : 1 L Ex [fijk (x)Xij ] = √ ∑ (Cju Cvi + Ciu Cvj ) Ex [∇xu ∇xv fijk (x)] . d u,v=1

(117)

We next compute the Hessian of fijk . Introduce the shorthand a = (i, j, k) and let b = (i′ , j ′ , l) and c = (i′′ , j ′′ , m) as coordinates of the tensor G. By the chain rule, 2 ∇xu ∇xv fa = ∑ ∂a,b L(G) ∇xu ∇xv Gb + Ruv,a ,

(118)

b

where

T

3 Ruv,a ∶= ∑ ∂a,b,c L(G) ∇xu Gb (∇xv Gc ) .

(119)

1 ∇xu Gi′ j ′ l = √ (δui′ Sl xj ′ + δuj ′ Sl xi′ ) , d

(120)

1 ∇xu ∇xv Gi′ j ′ l = √ Sl (δui′ δvj ′ + δuj ′ δvi′ ) . d

(121)

b,c

For b = (i′ , j ′ , l),

and

Substituting the first term in (118) into (117), and then using (121), gives 1 L 2 √ ∑ (Cju Cvi + Ciu Cvj ) ∑ Ex [∂(ijk),(i ′ j ′ l) L(G)] ∇xu ∇xv Gi′ j ′ l ′ ′ d u,v=1 i ,j ,l 2 2 = ∑ (Cii′ Cjj ′ + Cij ′ Cji′ ) Ex [∂(ijk),(i′ j ′ l) L(G)] Sl . d i′ ,j ′ ,l The factor 2 comes from the two Kronecker-delta contributions in (121). 4 This is satisfied under Assumption 2.1.

28

(122)

It remains to control the remainder Ruv,a . Under the spectral bound maxl ∥Wl ∥op ≤ C∗ , we have √ ∥Sl ∥op ≤ C∗2 , ∥Sl ∥F ≤ d C∗2 . (123) Consequently, each Gijk satisfies

sup max Ex ∣Gijk ∣p ≤ Cp (C∗ )

(124)

i,j,k

d

for every fixed p < ∞. Since L and K are fixed, the same is true for any polynomial function of the finitedimensional tensor G. By Assumption 2.1, the third derivatives of L have at most polynomial growth. Hence, for every collection of indices, 2

3 Ex ∣∂a,b,c L(G)∣ ≤ C(C∗ ).

(125)

Now fix arbitrary unit vectors p, q ∈ Rd . Using (119) and (120), every term in pT Ex [Ruv,a ]q is, up to 3 L(G)∣ ∣pT Sl xr ∣ ∣q T Sm xs ∣] for some r, s, l, m. constants depending only on L, K, and C, bounded by d1 Ex [∣∂a,b,c By Hölder’s inequality, 2 1/2

3 3 Ex [∣∂a,b,c L(G)∣ ∣pT Sl xr ∣ ∣q T Sm xs ∣] ≤ (Ex ∣∂a,b,c L(G)∣ )

1/4

(Ex ∣pT Sl xr ∣4 )

1/4

(Ex ∣q T Sm xs ∣4 )

.

(126)

The first factor is bounded by (125). Since xr and xs are Gaussian and ∥Sl ∥op , ∥Sm ∥op ≤ C∗2 , the remaining two factors are also uniformly bounded. Therefore, C(C∗ ) . d

(127)

2 2 ∑ (Cii′ Cjj ′ + Cij ′ Cji′ ) Ex [∂(ijk),(i′ j ′ l) L(G)] Sl + Eijk , d i′ ,j ′ ,l

(128)

∥Ex [Ruv,a ]∥op ≤ Combining (117), (122), and (127), we obtain Ex [fijk (x)Xij ] = where, uniformly over (110),

∥Eijk ∥op ≤ C(C∗ )d−3/2 .

(129)

Substituting (128) into (114), and defining Hk as (n)

Hk

∶= 2

L

∑

K

(n)

∑ (Cii′ Cjj ′ + Cij ′ Cji′ ) Ex [∂(ijk),(i′ j ′ l) L(G 2

i,j,i′ ,j ′ =1 l=1

gives ḡk =

1 Hk Wk + EStein,k , d

(n)

)] Sl

,

(130)

(131)

where EStein,k ∶= ∑L i,j=1 Eijk Wk . By (110) and (129) we have ∥EStein,k ∥op ≤ CE (C∗ )d−3/2 .

(132)

Finally, Assumption 2.1 and the uniform moment bound (124) imply sup

i,j,k,i′ ,j ′ ,l

2 Ex ∣∂(ijk),(i ′ j ′ l) L(G)∣ ≤ C(C∗ ).

(133)

Using (123), together with the fact that K and L are fixed, we therefore obtain ∥Hk ∥op ≤ CH (C∗ ).

(134)

Restoring the superscript (n) completes the proof. To ensure the trajectory does not blow up, we must prove that the spectral norm of the weights remains uniformly bounded across all training steps. 29

Lemma 3. Under the conditions of Theorem 2, there exist constants C∗ > C0 independent of d and T , such that for every fixed T > 0, lim P ( sup

d→∞

(n)

max ∥Wk ∥op ≤ C∗ ) = 1,

(135)

0≤n≤NT 1≤k≤K

where NT ∶= ⌊

T Td ⌋=⌊ ⌋. τd 4αd

(136) (n)

Proof. The weights with k ∉ K remain fixed throughout training and satisfy ∥Wk ∥op ≤ C0 . It therefore suffices to control k ∈ K. Fix a constant C∗ > C0 , to be chosen later, and define the stopping time (n)

J ∶= inf {n ≥ 0 ∶ max ∥Wk ∥op > C∗ } . 1≤k≤K

(137)

(0)

Since maxk ∥Wk ∥op ≤ C0 < C∗ almost surely, J ≥ 1. On the event {n < J}, (n)

(n)

max ∥Wk ∥op ≤ C∗ ,

max ∥Sk ∥op ≤ C∗2 .

k

Step 1: Truncation of the fresh samples.

k

(138)

For every step n, define (n)

An ∶= {∥x(n) ∥2 ≤ cx d} ∩ {max ∣Gijk ∣ ≤ cG log d} , i,j,k

(139)

where cx and cG are sufficiently large constants independent of d and n. Conditional on Fn , the state W (n) is deterministic. On {n < J}, (138) holds. Standard Gaussian concentration gives 1{n<J} P (∥x(n) ∥2 > cx d ∣ Fn ) ≤ e−c1 d . (140) (n)

Moreover, conditional on Fn , each Gijk is a centered quadratic form of x(n) . To make it explicit, let x(n) = (C 1/2 ⊗Id )z (n) , where z (n) ∼ N (0, ILd ) is independent of Fn . For i, j ∈ {1, . . . , L}, define Bij ∶= 21 (ei eTj +ej eTi ), L where {ei }L i=1 is the canonical basis of R . Then from (112) we can write (n)

where

(n)

(n)

Gijk = (z (n) )T Aijk z (n) − Tr(Aijk ),

(141)

1 (n) (n) Aijk ∶= √ (C 1/2 ⊗ Id )(Bij ⊗ Sk )(C 1/2 ⊗ Id ). d

(142)

(n)

Notice that Aijk is an Ld × Ld symmetric matrix. On the event {n < J}, (n)

∥Sk ∥op ≤ C∗2 ,

(n)

∥Sk ∥F ≤

√ 2 d C∗ .

(143)

Since L = Θ(1) and ∥Bij ∥op ≤ 1, ∥Bij ∥F ≤ 1, it follows that ∥C∥op ∥C∥op C 2 (n) (n) ∥Aijk ∥op ≤ √ ∥Bij ∥op ∥Sk ∥op ≤ √ ∗ , d d ∥C∥ ∥C∥ op op (n) (n) (n) ∥Aijk ∥F ≤ √ ∥Bij ⊗ Sk ∥F = √ ∥Bij ∥F ∥Sk ∥F ≤ ∥C∥op C∗2 . d d

(144) (145)

Therefore, uniformly on {n < J}, C (n) ∥Aijk ∥op ≤ √ , d

(n)

∥Aijk ∥F ≤ C, where C depends only on C∗ and C. 30

(146)

Applying the Hanson–Wright inequality to (141), we obtain ⎡ ⎤ ⎫ ⎧ ⎪ ⎪ ⎢ t ⎪⎥⎥ ⎪ t2 (n) 1{n<J} P (∣Gijk ∣ > t ∣ Fn ) ≤ 2 exp ⎢⎢−c min ⎨ (n) , (n) ⎬⎥ . 2 ⎪⎥ ⎪ ⎢ ⎪ ⎭⎦ ⎩ ∥Aijk ∥F ∥Aijk ∥op ⎪ ⎣

(147)

Hence, taking t = cG log d and using (146), (n)

1{n<J} P (∣Gijk ∣ > cG log d ∣ Fn ) ≤ 2 exp [−c min {(log d)2 ,

√

′

d log d}] ≤ 2e−c (log d) . 2

(148)

Since K and L are fixed, a union bound over (i, j, k) gives ′

(n)

1{n<J} P (max ∣Gijk ∣ > cG log d ∣ Fn ) ≤ 2L2 K e−c (log d) . 2

i,j,k

(149)

Therefore there exists pd ≤ C (e−c1 d + e−c2 (log d) ) such that 2

1{n<J} P(Acn ∣ Fn ) ≤ pd .

(150)

Define the event that no truncation failure occurs before the stopping time, NT −1

Adata ∶= ⋂ ({n ≥ J} ∪ An ) .

(151)

n=0

Using (150), NT −1

P(Acdata ) ≤ ∑ E [1{n<J} P(Acn ∣ Fn )] ≤ NT pd .

(152)

n=0

Since αd = Ω(d−ι ), NT ≤ CT d1+ι , and hence P(Acdata ) → 0.

(153)

Step 2: Stopped martingale decomposition and its bounds. truncated stochastic gradient (n) (n) ̂ gk ∶= 1{n<J} gk 1An , its conditional mean and the martingale difference

For k ∈ K, define the stopped and (154)

(n) (n) ̂ ḡk ∶= E[̂ gk ∣ Fn ],

(155)

(n) (n) (n) ξ̂k ∶= ̂ gk − ̂ ḡk .

(156)

(n)

(157)

(n) Then E[ξ̂k ∣ Fn ] = 0.

Let

ḡk Since {n < J} ∈ Fn ,

(n)

∶= E[gk

∣ Fn ].

(n) (n) (n) ̂ ḡk = 1{n<J} (ḡk − Tk ) ,

where

(n)

Tk

(158)

(n)

∶= E [gk 1Acn ∣ Fn ] .

(159)

We next record a uniform second-moment estimate for the stochastic gradient. On {n < J}, (n)

(n)

max {∥E[gk (gk )T ∣ Fn ]∥

(n)

op

(n)

, ∥E[(gk )T gk

31

∣ Fn ]∥ } ≤ Cg (C∗ ). op

(160)

Indeed, for any unit vectors u ∈ Rd and v ∈ Rdk , the two quadratic forms in (160) are respectively bounded by finite sums of expectations of the form (n)

(n)

E [∣∂(ijk) L(G(n) )∣2 ∥(Wk )T (Xij )T u∥2 ∣ Fn ] and

(n)

(n)

E [∣∂(ijk) L(G(n) )∣2 ∥Xij Wk v∥2 ∣ Fn ] .

Similarly to the estimates used in the proof of Lemma 2, Assumption 2.1 together with the spectral bound (138) implies uniform bounds on all fixed moments of G(n) and hence on the corresponding polynomially growing derivatives of L. In particular, by Hölder’s inequality, for every unit vector u ∈ Rd , (n)

(n)

E [∣∂(ijk) L(G(n) )∣2 ∥(Wk )T (Xij )T u∥2 ∣ Fn ] 1/2

≤ (E [∣∂(ijk) L(G(n) )∣4 ∣ Fn ])

(n)

1/2

(n)

(E [∥(Wk )T (Xij )T u∥4 ∣ Fn ])

= O(1),

(161)

uniformly on {n < J}. The same argument applies to the other quadratic variation in (160). Consequently, (n) (n) max {∥E[ξ̂k (ξ̂k )T ∣ Fn ]∥

op

(n) (n) , ∥E[(ξ̂k )T ξ̂k ∣ Fn ]∥ } ≤ Cg (C∗ ). op

(162)

√ (n) (n) On An ∩ {n < J}, ∥Xij ∥op ≤ C d, the first derivatives of L have polynomial growth and maxi,j,k ∣Gijk ∣ ≤ cG log d. Therefore, for some fixed integer r, √ (n) (163) ∥̂ gk ∥op ≤ C d (1 + (log d)r ) =∶ Bd . (n)

The same bound holds for ∥̂ ḡk ∥op , and hence ∥ξ̂k ∥op ≤ 2Bd . (n)

Step 3: Uniform control of the martingale noise. ad ∶= 1 −

(164)

Set

αd γ τd γ =1− . d 4

(165)

For sufficiently large d, ad ∈ (0, 1). Define (n) Dk,n ∶= −αd ξ̂k .

(166)

For every terminal time m ≥ 1, consider the discounted martingale sum m−1

Yk,m ∶= ∑ am−1−n Dk,n . d

(167)

n=0

For each fixed m, the summands in (167) form a rectangular matrix martingale difference sequence. By (164), ∥am−1−n Dk,n ∥op ≤ 2αd Bd =∶ Rd . (168) d Moreover, by (162), the predictable quadratic variation in both rectangular directions is bounded by ∞

Cg αd2 dαd ≤C . 1 − a2d γ

(169)

Rd = o((log d)−1 ).

(170)

σd2 ≤ Cg αd2 ∑ a2r d = r=0

c0 Under αd ≤ d log , we have d

σd2 ≤

Cc0 , γ log d

32

The Matrix Freedman inequality therefore gives for every fixed δ > 0, P(∥Yk,m ∥op ≥ δ) ≤ (d + dk ) exp [−

δ 2 /2 ]. σd2 + Rd δ/3

(171)

For sufficiently large d, (170) implies P(∥Yk,m ∥op ≥ δ) ≤ (d + dk )d−cδ /c0 2

(172)

0 > 0 and set β ∶= max{1, cw }. Since for some constant c > 0 independent of m, T , and d. Choose δ ∶= C∗ −C 4 cw β maxk dk ≤ d , for all sufficiently large d, d + dk ≤ 2d . Moreover, since αd = Ω(d−ι ), we have NT ≤ CT d1+ι . Therefore, by (172) and a union bound over k ∈ K and 1 ≤ m ≤ NT ,

P (max max ∥Yk,m ∥op > δ) ≤ 2KCT d1+ι+β−cδ /c0 . 2

(173)

k∈K 1≤m≤NT

Choose c0 > 0 sufficiently small so that

cδ 2 > 1 + ι + β. c0

(174)

Then the exponent in (173) is negative, and hence P (max max ∥Yk,m ∥op > δ) → 0.

(175)

k∈K 1≤m≤NT

Denote the complementary high-probability event by Anoise . (n) We also control the truncation tail (159). On {n < J}, the stochastic gradient gk has at most polynomial (n) growth in x . Hence, for some finite constant q, (n)

E [∥gk ∥2op ∣ Fn ] ≤ dq .

(176)

Combining Cauchy–Schwarz with (150), (n)

1{n<J} ∥Tk

1/2

(n)

∥op ≤ 1{n<J} (E[∥gk ∥2op ∣ Fn ])

where

1/2

rd ≤ dq/2 pd

(P(Acn ∣ Fn ))

1/2

≤ rd ,

= o(d−M )

(177) (178)

for every fixed M > 0. Step 4: Closing the stopping-time argument. On Adata , for every m ≤ J and every n < m, we have n < J and An occurs. Hence (n) (n) (n) (n) (n) gk = ̂ gk = ḡk − Tk + ξ̂k . (179) The SGD (7) therefore gives, for m ≤ J, (m)

Wk

(0)

= am d Wk

m−1

(n)

− αd ∑ am−1−n ḡk d n=0

m−1

(n)

+ αd ∑ am−1−n Tk d

+ Yk,m .

(180)

n=0

By Lemma 2, uniformly on {n < J}, (n)

ḡk (n)

=

1 (n) (n) (n) H Wk + EStein,k , d k

(181)

(n)

with ∥Hk ∥op ≤ CH (C∗ ) and ∥EStein,k ∥op ≤ CE (C∗ )d−3/2 . Therefore, (n)

∥ḡk ∥op ≤

CH (C∗ )C∗ + CE (C∗ )d−3/2 . d 33

(182)

Since 1 − ad = αdd γ , we have

∞

αd ∑ ard = r=0

d . γ

(183)

Combining (180)–(183), on Adata ∩ Anoise and for every m ≤ J ∧ NT , (m)

∥Wk

(0)

∥op ≤ ∥Wk ∥op +

CH (C∗ )C∗ CE (C∗ ) drd CH (C∗ )C∗ √ + + + δ ≤ C0 + + δ + o(1). γ γ γ γ d

(184)

We now fix C∗ ∶= 2C0 + 1 and choose γ sufficiently large (by Assumption 2.4) so that CH (C∗ )C∗ C∗ − C0 ≤ . γ 4

(185)

For all sufficiently large d, (184) then gives (m)

max ∥Wk k

∥op < C∗

for every m ≤ J ∧ NT

(186)

on Adata ∩ Anoise . Suppose, on this event, that J ≤ NT . Taking m = J in (184) yields (J)

max ∥Wk ∥op < C∗ ,

(187)

k

contradicting the definition (137). Therefore, J > NT on Adata ∩ Anoise for all sufficiently large d. Finally, P(J ≤ NT ) ≤ P(Acdata ) + P(Acnoise ) → 0

(188)

by (153) and (175). Hence (n)

P ( sup max ∥Wk ∥op ≤ C∗ ) Ð→ 1, 0≤n≤NT

(189)

k

which proves the lemma. D.1.3

Moment expansion

Let C∗ be the constant in Lemma 3, and fix T > 0 and ρ > C∗2 . Denote Ed (T ) ∶= { sup

(n)

max ∥Wk ∥op ≤ C∗ } ,

0≤n≤NT 1≤k≤K

NT ∶= ⌊

T ⌋. τd

(190)

Then P(Ed (T )) → 1 by Lemma 3. For the compactness argument later on, it is convenient to introduce the separable closed subspace ⎧ ⎪ ⎪ ⎪ ⎩

⎪ Xρ0 ∶= ⎨µ ∈ Xρ ∶

⎫ ⎪ ∣µα ∣ ⎪ lim sup w = 0⎬ . w→∞ ∣α∣=w ρ ⎪ ⎪ ⎭

(191)

Equipped with the norm inherited from Xρ , Xρ0 is a separable Banach space. Since the empirical moment sequence need not belong to Xρ0 on the exceptional event Ed (T )c , we introduce an auxiliary spectrally bounded process. For the fixed time horizon T , define (n)

W̌k

⎧ (n) ⎪ ⎪W , on Ed (T ), ∶= ⎨ k(0) c ⎪ ⎪ ⎩Wk , on Ed (T ) ,

and let

(n)

Šk

(n)

0 ≤ n ≤ NT ,

(n)

(192)

∶= W̌k (W̌k )T .

(193)

1 (n) (n) Tr [Šk1 ⋯Škw ] , d

(194)

For a multi-index α = (k1 , . . . , kw ), define µ̌(n) α ∶=

34

and let µ̌(d) (t) denote its piecewise-constant interpolation, µ̌(d) (t) = µ̌(n) ,

t ∈ [nτd , (n + 1)τd ).

(195)

Since C∗ > C0 , we have almost surely (n)

max ∥W̌k ∥op ≤ C∗ .

sup

(196)

0≤n≤NT 1≤k≤K

Consequently, for every ∣α∣ = w,

2w ∣µ̌(d) α (t)∣ ≤ C∗ ,

t ∈ [0, T ],

(197)

and hence, almost surely (w.r.t. the initialization), µ̌(d) (t) ∈ Xρ0 Moreover,

for all t ∈ [0, T ].

(198)

P (µ̌(d) (t) = µ̃(d) (t) for all t ∈ [0, T ]) ≥ P(Ed (T )) Ð→ 1.

(199)

All compactness arguments below will be carried out for µ̌(d) , and the resulting convergence will subsequently be transferred to the original empirical trajectory µ̃(d) using (199). We define empirical vector field Vd (µ̃(d) (t)) as 1 w (n) (n) (n) ∑ Tr [Sk1 . . . Akr . . . Skw ] d r=1

[Vd (µ̃(d) )]α =

(200)

kr ∈K

(n)

(n)

(n)

(n)

(n)

(n)

(n)

(n)

where Ak ∶= − 14 (Mk Sk + Sk (Mk )T ) and Mk ∶= Hk + γId with Hk defined in Lemma 2. The following lemma establishes the integral equation for the empirical trajectory (9). Lemma 4. For every t ∈ [0, T ], the auxiliary process satisfies µ̌(d) (t) = µ̌(d) (0) + 1Ed (T ) ∫

t

(d)

Vd (µ̃(d) (s)) ds + Ětotal (t).

0

(201)

For every ρ > C∗2 , all terms in (201) belong to Xρ0 almost surely, and (d)

sup ∥Ětotal (t)∥ Ð → 0. P

(202)

J ∶= inf {n ≥ 0 ∶ max ∥Wk ∥op > C∗ } .

(n)

(203)

P(J > NT ) → 1.

(204)

ρ

t∈[0,T ]

Proof. Define 1≤k≤K

By Lemma 3, All estimates below are uniform on {n < J}.

Step 1: One-step expansion of the attention matrices. (n)

ḡk (n)

=

For k ∈ K, Lemma 2 gives

1 (n) (n) (n) H Wk + EStein,k , d k

(205)

(n)

where ∥Hk ∥op ≤ CH and ∥EStein,k ∥op ≤ CE d−3/2 , uniformly on {n < J}. Define (n)

ξk (n)

Then E[ξk

(n)

∶= gk

(n)

(n)

− ḡk ,

Mk

(n)

∶= Hk

+ γId .

(206)

∣ Fn ] = 0, and the SGD (7) can be written as (n+1)

Wk

(n)

= Wk

−

τd (n) (n) (n) (n) M Wk − αd ξk − αd EStein,k . 4 k 35

(207)

(n+1)

For k ∉ K, Wk

(n)

= Wk . Let

(n)

(n+1)

∶= Wk

δWk

(n+1) (n) (n) (n) (n) T Expanding Sk = (Wk + δWk ) (Wk + δWk ) (n+1)

where, for k ∈ K,

(n)

+ τ d Ak

(n)

(210)

(n)

∶= − (ξk (Wk )T + Wk (ξk )T ) ,

(n)

(n)

+ αd Bk

(n)

1 (n) (n) (n) (n) ∶= − (Mk Sk + Sk (Mk )T ) , 4 (n)

(n)

(n)

(n)

(n)

(n)

(211) (n)

(n)

∶= −αd (EStein,k (Wk )T + Wk (EStein,k )T ) + δWk (δWk )T .

For k ∉ K, we set

(n)

Ak On {n < J},

(n)

(n)

Bk Ck

gives (209)

Ak

and

(n)

(208)

+ Ck ,

Sk

= Sk

(n)

− Wk .

(n)

= Bk

(n)

= Ck

(n)

(212)

= 0.

(213)

(n)

(214)

∥Ak ∥op ≤ CA .

(215)

∥Sk ∥op ≤ C∗2 ,

∥Mk ∥op ≤ CM

for a constant CM independent of d and n. In particular, (n)

We shall also use moment bounds for the stochastic gradient noise. Similarly to the estimates in Lemma 3 we can write (n) 1{n<J} E [∥ξk ∥pF ∣ Fn ] ≤ Cp dp/2 . (216) (n+1)

(n+1)

(n+1)

= d1 Tr [Sk1 ⋯Skw ], we Fix a multi-index α = (k1 , . . . , kw ). Using (209) in every factor of µα separate the expansion into four classes. (n) (n) The zeroth-order term is µα . The terms containing exactly one factor τd Akr define the empirical drift 1 (n) (n) (n) Tr [Sk1 ⋯Akr ⋯Skw ] . d 1≤r≤w

[Vd (µ(n) )]α ∶= ∑

(217)

kr ∈K

(n)

The terms containing exactly one factor αd Bkr define the martingale 1 (n) (n) (n) Tr [Sk1 ⋯Bkr ⋯Skw ] . d 1≤r≤w

ηα(n) ∶= ∑

(218)

kr ∈K

(n)

(n)

(n)

(n)

Since all Sk are Fn -measurable and Bk is linear in ξk , E[ηα ∣ Fn ] = 0. (n) (n) Finally, let Rα denote the sum of all remaining terms in the expansion. Thus Rα contains, in particular, (n)

• every term involving at least one Ckr ; (n)

(n)

• every term involving at least two first-order increments chosen from {τ Akr , αBkr }. With this definition we obtain the expansion: (n) (n) µ(n+1) − µ(n) )]α + αd ηα(n) + Rα . α α = τd [Vd (µ

36

(219)

Step 2: Vanishing of the remainder. We next bound the remainder. For every fixed multi-index α, there exists a constant Cα < ∞ such that, uniformly on {n < J}, (n) E [∣Rα ∣ ∣ Fn ] ≤ Cα (αd2 + αd τd + τd2 + αd d−3/2 ) .

(220)

To see this, first bound terms containing at least two B-factors, let P0 , P1 , ⋯ denote arbitrary products of the remaining S- and A-matrices. On {n < J}, their operator norms are bounded by a constant depending only on α. We obtain, for fixed rB ≥ 2, 1 E [∣Tr(P0 B1 P1 ⋯BrB PrB )∣ ∣ Fn ] ≤ Cα d(rB −2)/2 . d

(221)

Hence the corresponding contribution is bounded by √

Cα αdrB τdrA d(rB −2)/2 ≤ Cα αd2

(222)

for all sufficiently large d, because αd d → 0. A term containing exactly one αd B and at least one τd A is bounded by Cα αd τd , while a term containing at least two τd A factors is bounded by Cα τd2 . (n) It remains to consider terms containing at least one C-factor. By (212), each Ck is a finite sum of terms of the following types: Oop (αd d−3/2 ),

τd2 U1 U2T ,

αd τd U1 ξkT ,

αd2 ξk ξkT ,

(n)

together with terms involving EStein,k that are of smaller order. Here the Ui have uniformly bounded operator √ norm and O( d) Frobenius norm. Applying the same trace inequalities, Hölder’s inequality, and (216), every product containing at least one C-factor is bounded in conditional expectation by Cα (αd2 + αd τd + τd2 + αd d−3/2 ), √ with any additional stochastic increment producing only an additional factor of order αd d = o(1). This proves (220). Define ⌊t/τ ⌋−1

(d) (n) Erem,α (t) ∶= ∑ Rα .

(223)

n=0

(n)

Replace Rα

(n)

temporarily by Rα 1{n<J} . By (220),

NT −1

(n) E [ ∑ 1{n<J} ∣Rα ∣] ≤ NT Cα (αd2 + αd τd + τd2 + αd d−3/2 ) ≤ Cα,T (dαd + αd + n=0

αd + d−1/2 ) . d

(224)

c0 , the right-hand side tends to zero. Hence, by Markov’s inequality, Since αd ≤ d log d (d) sup ∣Erem,α (t)∣ Ð →0 P

(225)

t∈[0,T ]

on {J > NT }. Together with (204), this proves (225) for the original trajectory. Define

Step 3: Vanishing of the first-order martingale noise. (d)

Enoise,α (t) ∶=

⌊t/τd ⌋−1

∑

αd ηα(n) .

(226)

Cα . d2

(227)

n=0

We claim that, for every fixed α, 1{n<J} E [(ηα(n) )2 ∣ Fn ] ≤

(n)

(n)

(n)

(n)

Indeed, every constituent of (218) is a linear functional of ξk . Since ξk = gk − E[gk ∣ Fn ], it suf(n) (n) fices to bound the second moment of the corresponding linear functional of gk . After inserting gk = (n) (n) ∑i,j ∂(ijk) L(G(n) )Xij Wk , each term takes the form ∂(ijk) L(G(n) )Zij (A),

Zij (A) ∶= 37

1 (n) Tr [AXij ] , d

(228)

(n)

where, for fixed α, ∥A∥op ≤ Cα . Since Xij is a centered Gaussian quadratic-form matrix, standard Gaussian moment estimates give Cα E [∣Zij (A)∣4 ∣ Fn ] ≤ 4 . (229) d On the other hand, Assumption 2.1 and the uniform moment bounds for G(n) imply 4

E [∣∂(ijk) L(G(n) )∣ ∣ Fn ] ≤ Cα . Hence, by Cauchy–Schwarz, 2

E [∣∂(ijk) L(G(n) )Zij (A)∣ ∣ Fn ] ≤

(230)

Cα . d2

(231)

Since K, L, and the length of α are fixed, this proves (227). Therefore the stopped process m−1

Mα,m ∶= ∑ αd ηα(n) 1{n<J}

(232)

n=0

is a square-integrable martingale with predictable quadratic variation bounded by E[⟨Mα ⟩NT ] ≤ NT αd2

Cα αd ≤ Cα,T . 2 d d

(233)

Cα,T αd →0 ε2 d

(234)

Doob’s maximal inequality yields, for every ε > 0, P ( max ∣Mα,m ∣ > ε) ≤ m≤NT

Using again P(J > NT ) → 1, we conclude that (d)

sup ∣Enoise,α (t)∣ Ð → 0. P

(235)

t∈[0,T ]

Step 4: From the discrete recursion to the integral equation. from n = 0 to m(t) − 1 gives m(t)−1

For m(t) ∶= ⌊ τt ⌋, summing (219)

(d)

(d) (n) (d) µ̃(d) (t). )]α + Enoise,α (t) + Erem,α α (t) = µ̃α (0) + τd ∑ [Vd (µ

(236)

n=0

Since the interpolation (9) is piecewise constant, m(t)−1

t

(d) (n) (m(t)) )]α . ∫ [Vd (µ̃ (s))]α ds = τd ∑ [Vd (µ )]α + (t − m(t)τd )[Vd (µ 0

(237)

n=0

Define the discretization error (d)

Edisc,α (t) ∶= −(t − m(t)τd )[Vd (µ(m(t)) )]α . (d)

(238)

(d)

Thus we define Ětotal,α ∶= 1Ed (T ) Etotal,α with (d)

(d)

(d)

(d) Etotal,α = Enoise,α + Erem,α + Edisc,α ,

(239)

and (201) follows coordinatewise. On {n < J}, (215) implies that for ∣α∣ = w, ∣[Vd (µ(n) )]α ∣ ≤ BwC∗2w

(240)

for a constant B independent of d, n, and w. Consequently, for every fixed α, (d)

2∣α∣

sup ∣Edisc,α (t)∣ ≤ τ B∣α∣C∗

t∈[0,T ]

→ 0.

(241)

Combining this with (225) and (235), we have, for every fixed multi-index α, (d)

sup ∣Etotal,α (t)∣ Ð → 0.

t∈[0,T ]

38

P

(242)

Step 5: Convergence in Xρ0 . Ed (T ). On Ed (T ),

We now upgrade the coordinatewise convergence to convergence in Xρ0 on 2w sup sup ∣µ̃(d) α (t)∣ ≤ C∗ .

(243)

t∈[0,T ] ∣α∣=w

Together with (240) and

t

(d)

Etotal (t) = µ̃(d) (t) − µ̃(d) (0) − ∫ this yields, on Ed (T ),

0

Vd (µ̃(d) (s)) ds,

(d)

sup sup ∣Etotal,α (t)∣ ≤ (2 + T Bw)C∗2w .

(245)

t∈[0,T ] ∣α∣=w

Since ρ > C∗2 ,

(244)

w

(2 + T Bw) (

C∗2 ) Ð→ 0 ρ

as w → ∞.

(246)

Fix ε > 0. Choose w⋆ sufficiently large that w

C2 ε sup (2 + T Bw) ( ∗ ) < . ρ 2 w>w⋆

(247)

For degrees w ≤ w⋆ , there are only finitely many multi-indices, since K is fixed. Hence (242) and a finite union bound imply (d) ∣Etotal,α (t)∣ P max sup Ð →0 (248) ∣α∣≤w⋆ t∈[0,T ] ρ∣α∣ on Ed (T ). (d) (d) (d) (d) By Ětotal (t) ∶= 1Ed (T ) Etotal (t), Ětotal (t) = 0 on Ed (T )c . On Ed (T ), (245)–(247) show that Ětotal (t) ∈ Xρ0 and that P (d) sup ∥Ětotal (t)∥ Ð → 0. (249) ρ

t∈[0,T ]

Finally, on Ed (T ) we have µ̌(d) = µ̃(d) , and hence (244) gives µ̌(d) (t) = µ̌(d) (0) + 1Ed (T ) ∫

t 0

(d)

Vd (µ̃(d) (s)) ds + Ětotal (t).

(250)

(d)

On Ed (T )c , both the integral term and Ětotal vanish, while by definition µ̌(d) (t) = µ̌(d) (0). Therefore the same identity holds almost surely on the whole probability space. This proves the lemma. D.1.4

Compactness of the empirical trajectories

Lemma 5. Under the conditions of Theorem 2, for any finite T > 0 and any ρ > C∗2 , the sequence of auxiliary moment processes {µ̌(d) (⋅)}d≥1 is C-tight in D([0, T ]; Xρ0 )5 . Proof. Define the continuous approximation µ̂(d) (t) ∶= µ̌(d) (0) + 1Ed (T ) ∫ By Lemma 4, and therefore

t

Vd (µ̃(d) (s)) ds.

(251)

µ̌(d) (t) − µ̂(d) (t) = Ětotal (t),

(d)

(252)

sup ∥µ̌(d) (t) − µ̂(d) (t)∥ρ Ð → 0.

(253)

0

P

t∈[0,T ]

5 We say that a sequence of processes is C-tight in D([0, T ]; X 0 ) if it is tight (relatively compact in distribution) and every ρ subsequential weak limit is supported on C([0, T ]; Xρ0 ).

39

We first establish compactness of the family {µ̂(d) }. By (196) and (240), for every multi-index α of length w, 2w sup ∣µ̂(d) almost surely. (254) α (t)∣ ≤ (1 + T Bw)C∗ =∶ aw t∈[0,T ]

Indeed, on Ed (T ) this follows from Lemma 3, while on Ed (T )c the process µ̂(d) is constant. Since w C∗2 aw = (1 + T Bw) ( ) Ð→ 0, ρw ρ the set

⎧ ⎫ ⎪ ⎪ ⎪ ⎪ KT,ρ ∶= ⎨µ ∈ Xρ0 ∶ sup ∣µα ∣ ≤ aw for every w ≥ 1⎬ ⎪ ⎪ ⎪ ⎪ ∣α∣=w ⎩ ⎭

(255)

(256)

is compact in Xρ0 . This is because for every ε > 0, the coordinates of degree larger than a sufficiently large M contribute at most ε uniformly over KT,ρ , whereas the coordinates of degree at most M form a bounded subset of a finite-dimensional space. Moreover, for 0 ≤ s < t ≤ T and ∣α∣ = w, (d) 2w ∣µ̂(d) α (t) − µ̂α (s)∣ ≤ B∣t − s∣ wC∗ .

(257)

Consequently, w

∥µ̂(d) (t) − µ̂(d) (s)∥ρ ≤ B∣t − s∣ sup w ( w≥1

C∗2 ) =∶ Lρ ∣t − s∣, ρ

(258)

where Lρ < ∞ is independent of d. Thus the trajectories µ̂(d) take values almost surely in the compact set KT,ρ and are uniformly Lipschitz in Xρ0 . By the Arzelà–Ascoli theorem, the family {µ̂(d) } is tight in C([0, T ]; Xρ0 ). Finally, (253) shows that µ̌(d) is asymptotically equivalent to the tight family µ̂(d) in the uniform metric. It follows that {µ̌(d) }d≥1 is tight in D([0, T ]; Xρ0 ), and every subsequential weak limit is supported on C([0, T ]; Xρ0 ). Hence the sequence is C-tight. D.1.5

Global well-posedness

We show below that Vd converges uniformly to V under the spectral bounds in Lemma 3. Lemma 6. Fix C∗ > 0 and ρ > C∗2 . Let Ad (C∗ ) ∶= {(S1 , . . . , SK ) ∶ Sk = Wk WkT , max ∥Wk ∥op ≤ C∗ } . 1≤k≤K

(259)

For S ∈ Ad (C∗ ), let µ(S) denote its associated moment sequence. Then sup S∈Ad (C∗ )

∥Vd (µ(S)) − V (µ(S))∥ρ Ð→ 0

as d → ∞.

(260)

Consequently, on the spectral-bound event Ed (T ), sup ∥Vd (µ̃(d) (t)) − V (µ̃(d) (t))∥ρ Ð → 0. P

t∈[0,T ]

Proof. Define (d)

Γkl ∶=

1 2 (d) ∑ (Cii′ Cjj ′ + Cij ′ Cji′ ) Ex [∂(ijk),(i′ j ′ l) L(G )] . 2 i,j,i′ ,j ′

(261)

(262)

By the definition of Hk , K

(d)

Hk = 4 ∑ Γkl Sl . l=1

40

(263)

Hence the empirical vector field (200) can be written as w

K

(d)

[Vd (µ)]α = − ∑ ∑ Γkr l (µαr,l,− + µαr,l,+ ) − r=1 l=1 kr ∈K

γ ∣α∣K µα , 2

(264)

where αr,l,− and αr,l,+ denote the two multi-indices obtained by inserting l before and after kr , respectively. (d) The limiting vector field has the same form with Γkl replaced by Γkl (q) defined in (91). Therefore, defining (d) δd (C∗ ) ∶= sup max ∣Γkl − Γkl (q(S))∣ , (265) S∈Ad (C∗ ) k,l

we have, for every ∣α∣ = w, Thus,

2(w+1)

∣[Vd (µ) − V (µ)]α ∣ ≤ 2Kw δd (C∗ ) C∗

.

w C∗2 2 ∥Vd (µ) − V (µ)∥ρ ≤ 2KC∗ δd (C∗ ) sup w ( ) . ρ w≥1

(266) (267)

It remains to prove δd (C∗ ) → 0. Suppose otherwise. Then there exist ϵ > 0, a subsequence dm , and S(dm ) ∈ Adm (C∗ ) such that (d )

max ∣Γkl m − Γkl (q (dm ) )∣ ≥ ϵ. k,l

(268)

(d )

Since K is fixed and ∣qkl m ∣ ≤ C∗4 , after taking a further subsequence we may assume q (dm ) → q ∗ ∈ SK + . Define the active index set ∗ I ∶= {k ∈ {1, . . . , K} ∶ qkk > 0}. (269) For every k ∈ I,

1 (d ) (d ) ∗ ∥S m ∥2F = qkkm Ð→ qkk > 0. dm k

(270)

(d )

Since ∥Sk m ∥op ≤ C∗2 , it follows that

(d )

∥Sk m ∥op (d )

∥Sk m ∥F

Ð→ 0.

(271) (d )

Hence Corollary 5 applies to k ∈ I. For k ∉ I, we instead have qkkm → 0, which gives (d )

Gijkm Ð→ 0

in L2 .

(272)

Moreover, all covariances involving an inactive index converge to zero. Combining them, G(dm ) → G∗ ∼ N (0, Σ(q ∗ ))

(273)

in distribution. Therefore, by Assumption 2.1 and uniform integrability, 2 (dm ) 2 ∗ E [∂(ijk),(i )] Ð→ E [∂(ijk),(i ′ j ′ l) L(G ′ j ′ l) L(G )] .

(274)

Consequently, by the definitions of Γ(dm ) and Γ(q ∗ ), (d )

Γkl m Ð→ Γkl (q ∗ ) This contradicts (268), and hence

for every k, l.

δd (C∗ ) → 0.

The claimed convergence in Xρ follows. The next lemma establishes the global well-posedness of (11). 41

(275) (276)

Lemma 7. Under the conditions of Theorem 2, let C∗ be the uniform spectral bound in Lemma 3. For every finite T > 0 and every pair C∗2 < ρ− < ρ+ , there exists a unique admissible trajectory µ̄ ∈ C([0, T ]; Xρ0− ) such that t t ∈ [0, T ], (277) µ̄(t) = µ̄(0) + ∫ V (µ̄(s)) ds, 0

where the identity is understood in Xρ0+ . Moreover, q(µ̄(t)) ⪰ 0, 2∣α∣

and ∣µ̄α (t)∣ ≤ C∗

t ∈ [0, T ],

(278)

for every multi-index α.

Proof. Fix T > 0 and C∗2 < ρ− < ρ+ . By Lemma 5, the sequence {µ̌(d) }d≥1 is C-tight in D([0, T ]; Xρ0− ). Hence, from any subsequence we may extract a further subsequence, still indexed by d, such that µ̌(d) → µ∗

in D([0, T ]; Xρ0− ),

(279)

where µ∗ ∈ C([0, T ]; Xρ0− ) almost surely. We next rewrite the auxiliary equation in terms of the limiting vector field. By Lemma 4, µ̌(d) (t) = µ̌(d) (0) + 1Ed (T ) ∫ with

t 0

(d)

Vd (µ̃(d) (s)) ds + Ětotal (t),

(d)

sup ∥Ětotal (t)∥

t∈[0,T ]

(280)

Ð → 0. P

ρ+

(281)

Define t

(d)

Rd (t) ∶= Ětotal (t) + ∫

0

[1Ed (T ) Vd (µ̃(d) (s)) − V (µ̌(d) (s))] ds.

Then µ̌(d) (t) = µ̌(d) (0) + ∫

0

We claim that

t

(282)

V (µ̌(d) (s)) ds + Rd (t).

(283)

sup ∥Rd (t)∥ρ+ Ð → 0. P

(284)

t∈[0,T ]

Indeed, on Ed (T ) we have µ̌(d) = µ̃(d) , and therefore 1Ed (T ) ∥Vd (µ̃(d) (s)) − V (µ̌(d) (s))∥ρ = 1Ed (T ) ∥Vd (µ̃(d) (s)) − V (µ̃(d) (s))∥ρ . +

+

(285)

Lemma 6 implies 1Ed (T ) sup ∥Vd (µ̃(d) (s)) − V (µ̃(d) (s))∥ρ Ð → 0. P

+

s∈[0,T ]

(286)

On the other hand, by (196), the auxiliary process always satisfies ∣µ̌(d) α (s)∣ ≤ C∗ . 2∣α∣

(287)

̃ is bounded on the corresponding bounded set of second moments, there exists a constant CV = Since Γ CV (C∗ , ρ+ ) < ∞ such that sup sup ∥V (µ̌(d) (s))∥ρ+ ≤ CV d≥1 s∈[0,T ]

almost surely.

(288)

Consequently, → 0, 1Ed (T )c sup ∥V (µ̌(d) (s))∥ρ+ Ð P

s∈[0,T ]

because P(Ed (T )c ) → 0. Together with (281) and (286), this proves (284). 42

(289)

We also have µ̌(d) (0) = µ̃(d) (0). Assumption 2.3 gives convergence in probability of every coordinate, (d) 2∣α∣ while Assumption 2.2 gives ∣µ̌α (0)∣ ≤ C0 . Since ρ− > C∗2 > C02 , the contribution of sufficiently high degrees is uniformly exponentially small, which yields µ̌(d) (0) Ð → µ̄(0) P

in Xρ0− .

(290)

Because Xρ0− is separable, we may apply the Skorokhod representation theorem to the joint laws of (µ̌(d) , µ̌(d) (0), Rd ). After passing to a further subsequence and realizing these random elements on a common probability space, we may assume that, almost surely, µ̌(d) → µ∗

in D([0, T ]; Xρ0− ),

µ̌(d) (0) Ð→ µ̄(0) and

in Xρ0− ,

sup ∥Rd (t)∥ρ+ Ð→ 0.

(291) (292) (293)

t∈[0,T ]

Since µ∗ is continuous, convergence in the Skorokhod topology implies sup ∥µ̌(d) (t) − µ∗ (t)∥ρ Ð→ 0 −

t∈[0,T ]

almost surely.

(294)

The vector field V satisfies the Cauchy estimate established in Lemma 1. Since the convergent trajectories remain in a bounded subset of Xρ0− , there exists L < ∞ such that ∥V (µ) − V (ν)∥ρ+ ≤

L ∥µ − ν∥ρ− ρ+ − ρ−

(295)

LT sup ∥µ̌(d) (s) − µ∗ (s)∥ρ Ð→ 0. − ρ+ − ρ− s∈[0,T ]

(296)

for all states µ, ν under consideration. Hence sup ∥∫

t∈[0,T ]

t 0

[V (µ̌(d) (s)) − V (µ∗ (s))] ds∥

ρ+

≤

Passing to the limit in (283) yields µ∗ (t) = µ̄(0) + ∫

t

V (µ∗ (s)) ds,

0

t ∈ [0, T ],

(297)

as an identity in Xρ0+ . It remains to verify that the limiting trajectory is admissible. For every finite d and every t ∈ [0, T ], 1 (d) (d) Tr [Šk (t)Šl (t)] . d

(298)

almost surely for every d, t.

(299)

qkl (µ̌(d) (t)) = Therefore

q(µ̌(d) (t)) ⪰ 0

Since the PSD cone is closed and convergence in Xρ0− implies convergence of second-order coordinates, q(µ∗ (t)) ⪰ 0, Thus and consequently

t ∈ [0, T ].

(300)

P(q(µ∗ (t))) = q(µ∗ (t)),

(301)

∗ ̃ (t))) = Γ(q(µ∗ (t))). Γ(q(µ

(302)

Therefore µ∗ satisfies the original equation (11). Moreover, by (197), ∣µ̌(d) α (t)∣ ≤ C∗

2∣α∣

almost surely for every d, t. 43

(303)

Passing to the limit gives

∣µ∗α (t)∣ ≤ C∗ , 2∣α∣

α, t ∈ [0, T ].

(304)

We finally establish uniqueness. Suppose that µ1 and µ2 are two admissible solutions of (11). Both are also solutions of (94). By Lemma 1, the solution of (94) is locally unique in the Banach scale. Hence, if µ1 (t0 ) = µ2 (t0 ) at some t0 < T , the two solutions coincide on a finite interval to the right of t0 . Starting from t0 = 0 and applying this local uniqueness with a contradiction argument shows that µ1 (t) = µ2 (t),

t ∈ [0, T ].

(305)

Thus the admissible solution is unique. Since the initial condition µ̄(0) is deterministic, uniqueness implies that the random limit µ∗ is almost surely equal to a deterministic trajectory, which we denote by µ̄ as a unique admissible trajectory. D.1.6

Final proof of Theorem 2

Fix an arbitrary finite horizon T > 0, and choose C∗2 < ρ− < ρ+ .

(306)

By Lemma 5, the auxiliary trajectories {µ̌(d) }d≥1 are C-tight in D([0, T ]; Xρ0− ). By Lemma 7, every subsequential weak limit is the same deterministic path µ̄, and therefore µ̌(d) → µ̄

in D([0, T ]; Xρ0− ).

(307)

Since µ̄ is deterministic, convergence in distribution to µ̄ is equivalent to convergence in probability: dJ1 (µ̌(d) , µ̄) Ð → 0, P

(308)

where dJ1 denotes the Skorokhod J1 metric. Moreover, µ̄ is continuous. Convergence in the Skorokhod topology to a continuous limit implies convergence in the uniform topology. Hence → 0. sup ∥µ̌(d) (t) − µ̄(t)∥ρ Ð P

(309)

−

t∈[0,T ]

For every fixed multi-index α, ∣α∣

sup ∣µ̌(d) α (t) − µ̄α (t)∣ ≤ ρ−

t∈[0,T ]

and therefore

sup ∥µ̌(d) (t) − µ̄(t)∥ρ , −

t∈[0,T ]

sup ∣µ̌(d) → 0. α (t) − µ̄α (t)∣ Ð P

t∈[0,T ]

(310)

(311)

It remains only to transfer this convergence from the auxiliary trajectory to the original trajectory. By (199), for every fixed multi-index α and every ε > 0, P

⎛ ⎞ ⎛ ⎞ sup ∣µ̃(d) (t) − µ̄α (t)∣ > ε ≤ P(Ed (T )c ) + P sup ∣µ̌(d) (t) − µ̄α (t)∣ > ε . ⎝t∈[0,T ] α ⎠ ⎝t∈[0,T ] α ⎠

(312)

The first term tends to zero by Lemma 3, while the second tends to zero by (311). Consequently, lim P

d→∞

⎛ ⎞ sup ∣µ̃(d) α (t) − µ̄α (t)∣ > ε = 0. ⎝t∈[0,T ] ⎠

This proves Theorem 2.

44

(313)

D.2

Proof of Corollary 2

We define a projection operator ΠM ∶ Xs → Xs for a M ≥ 2, such that for any multi-index α: ⎧ ⎪ ⎪µα , if ∣α∣ ≤ M [ΠM µ]α = ⎨ . ⎪ 0, if ∣α∣ > M ⎪ ⎩ Let R ∶= C∗2 . By Lemma 7,

∣µ̄α (t)∣ ≤ R∣α∣ ,

(314)

t ≥ 0.

(315)

In particular, ∥q̄(t)∥F ≤ KR2 . We first show that the finite-dimensional truncation is equivalently the restriction of the extended infinitedimensional ODE (94) to an invariant subspace. Indeed, define YM ∶= {µ ∶ µα = 0 for every ∣α∣ > M } .

(316)

If µ ∈ YM and ∣α∣ > M , then every moment appearing in [V (µ)]α has degree either ∣α∣ or ∣α∣ + 1, and hence vanishes. Therefore, [V (µ)]α = 0 for all ∣α∣ > M . Thus YM is invariant under (94), and µ(M ) is precisely the solution of d (M ) µ (t) = V (µ(M ) (t)), µ(M ) (0) = ΠM µ̄(0). (317) dt We next establish stability with respect to this truncation. Let U ∶= {q ∈ RK×K ∶ ∥q∥F ≤ KR2 + 1} .

(318)

̃ is locally bounded and locally Lipschitz (see Lemma 1), there exist finite constants CΓ , LΓ such that Since Γ ̃kl (q)∣ ≤ CΓ , max ∣Γ

q ∈ U,

k,l

and

(319)

̃kl (q) − Γ ̃kl (q ′ )∣ ≤ LΓ ∥q − q ′ ∥F , max ∣Γ

q, q ′ ∈ U.

k,l

(320)

Choose s0 > R and γ sufficiently large (Assumption 2.4) such that 4KCΓ s0 4K 2 LΓ s20 R + < 1. γ γ

(321)

e(M ) (t) ∶= µ(M ) (t) − µ̄(t)

(322)

κ ∶= Define the error and

EM (t) ∶= ∥e(M ) (t)∥s0 = sup α

(M )

∣eα

(t)∣

∣α∣ s0

.

(323)

Let q (M ) (t) ∶= q(µ(M ) (t)). As long as q (M ) (t) ∈ U, we have ∥q (M ) (t) − q̄(t)∥F ≤ Ks20 EM (t). For a multi-index α = (k1 , . . . , kw ), denote

rα ∶= ∣α∣K .

(324) (325)

Let αi,l,− and αi,l,+ denote the multi-indices obtained by inserting l before and after ki , respectively. Sub̃ tracting the equations for µ(M ) and µ̄, and using Γ(q̄) = Γ(q̄), gives K K d (M ) ) (M ) ̃k l (q (M ) ) (e(M ̃k l (q (M ) ) − Γk l (q̄)] (µ̄αi,l,− + µ̄αi,l,+ ) − γ rα e(M ) . [Γ eα = − ∑ ∑ Γ + e ) − ∑ ∑ i,l,− i,l,+ α i i i α α dt 2 1≤i≤w l=1 1≤i≤w l=1 ki ∈K

ki ∈K

45

(326)

If rα = 0, the right-hand side vanishes and hence ) ) e(M (t) = e(M (0). α α

(327)

Suppose now that rα ≥ 1. Applying the variation-of-constants formula to (326) and using (319) gives (M )

∣eα

(t)∣

sw 0

≤

(M )

∣eα

(0)∣

sw 0 +

+ 2KCΓ rα s0 ∫

t

1 ∫ e sw 0 0

− γ2 rα (t−u)

t 0

γ

e− 2 rα (t−u) EM (u) du

K

̃k l (q ∑ ∑ ∣Γ i

(M )

i∶ki ∈K l=1

(u)) − Γki l (q̄(u))∣ × (∣µ̄αi,l,− (u)∣ + ∣µ̄αi,l,+ (u)∣) du. (328)

By (320), (324), and (315), the second integral is bounded by 2K 2 LΓ s20 R rα (

R w t − γ2 rα (t−u) EM (u) du. ) ∫ e s0 0

Since

t

rα ∫ and R/s0 < 1, we obtain

(M )

∣eα

(t)∣

sw 0

0

≤

2 , γ

(330)

+ κ sup EM (u),

(331)

γ

e− 2 rα (t−u) du ≤ (M )

∣eα

(0)∣

sw 0

(329)

0≤u≤t

with κ defined in (321). The same inequality trivially holds when rα = 0. Taking the supremum over α and over 0 ≤ u ≤ t, we therefore obtain sup EM (u) ≤ EM (0) + κ sup EM (u).

0≤u≤t

0≤u≤t

(332)

Since κ < 1,

1 EM (0). 1 − κ 0≤u≤t It remains to estimate the initial truncation error. By definition, sup EM (u) ≤

⎧ ⎪ ⎪0, ) e(M (0) = ⎨ α ⎪ −µ̄ (0), ⎪ ⎩ α

∣α∣ ≤ M, ∣α∣ > M.

(333)

(334)

Using (315), ∣µ̄α (0)∣ R w R M +1 ≤ sup ( ) = ( ) . sw s0 w>M ∣α∣=w w>M s0 0

EM (0) = sup sup

(335)

Combining (333) and (335) gives sup ∥µ(M ) (t) − µ̄(t)∥s0 ≤

0≤t≤T

1 R M +1 ( ) , 1 − κ s0

(336)

provided q (M ) (t) ∈ U. We finally close this condition by a standard bootstrap argument. Since ∥q (M ) (t) − q̄(t)∥F ≤

Ks20 R M +1 ( ) , 1 − κ s0

(337)

the right-hand side is smaller than 1 for all sufficiently large M . Because ∥q̄(t)∥F ≤ KR2 , it follows that ∥q (M ) (t)∥F < KR2 + 1,

(338)

so q (M ) (t) cannot leave U. Moreover, (336) bounds every coordinate of (13), so its solution cannot undergo finite-time blow-up. Hence, for all sufficiently large M , the solution of (13) exists on [0, T ], and (336) holds throughout the interval. This finishes the proof. 46

D.3

An exactly solvable example

Consider the squared loss L(G1 , G2 ) = (G1 − G2 )2 without weight decay (γ = 0). It is a simplified setting where L = 1 and K = 2. We assume only the first index is learnable (K = {1}). The potential evaluates to: Φ(q) = Ez1 ,z2 ∣∣z1 − z2 ∣∣2 = 2q11 + 2q22 − 4q12 .

(339)

Because the partial derivatives of Φ(q) are constant, the dynamics in (11) form an infinite-dimensional linear ODE hierarchy. Truncating this system at a finite degree N yields a closed linear ODE with polynomial solutions. At truncation degree N = 1, the solution is static: µ̄1 (t) = µ̄1 (0). At N = 2, the dynamics activate: q̄11 (t) = q̄11 (0), q̄12 (t) = q̄12 (0),

(340)

µ̄1 (t) = µ̄1 (0) − 4t[q̄11 (0) − q̄12 (0)]. Extending to N = 3, the solution captures higher-order curvature: q̄11 (t) = q̄11 (0) − 8t[µ̄111 (0) − µ̄112 (0)], q̄12 (t) = q̄12 (0) − 4t[µ̄112 (0) − µ̄122 (0)], µ̄1 (t) = µ̄1 (0) − 4t[q̄11 (0) − q12 (0)] + 8t2 [2µ̄111 (0) − 3µ̄112 (0) + µ̄122 (0)].

(341)

Remarkably, these truncated polynomial solutions match the Taylor expansion of the original matrix flow evaluated around t = 0. To see this, note that the gradient flow for the student forms a matrix Riccati differential equation: d S1 = 2(S1 S2 + S2 S1 ) − 4S12 . (342) dt Assuming that S1 (0) and S2 are invertible, (342) admits a closed-form solution: −1

S1 (t) = [e−2S2 t S1 (0)−1 e−2S2 t + S2−1 (I − e−4S2 t )]

,

(343)

Taking the first-order Taylor expansion of the inner inverse matrix around t = 0 yields: S1 (t)−1 ≈ (I − 2S2 t)S1 (0)−1 (I − 2S2 t) + S2−1 (4S2 t) ≈ S1 (0)−1 + (4I − 2S2 S1 (0)−1 − 2S1 (0)−1 S2 ) t.

(344)

Using the matrix identity (A + tB)−1 ≈ A−1 − A−1 BA−1 t, we recover the first-order expansion for S1 (t): S1 (t) ≈ S1 (0) + (2S1 (0)S2 + 2S2 S1 (0) − 4S1 (0)2 )t.

(345)

Taking the normalized traces of this matrix expansion reproduces the linear terms of the truncated moments: q̄11 (t) = q̄11 (0) − 8t[µ̄111 (0) − µ̄112 (0)], q̄12 (t) = q̄12 (0) − 4t[µ̄112 (0) − µ̄122 (0)], µ̄1 (t) = µ̄1 (0) − 4t[q̄11 (0) − q̄12 (0)].

(346)

Numerical results are shown in Figure 4. The closed-form solution also provides a direct interpretation of why the moment dynamics is intrinsically infinite-dimensional. To see this, consider the special case S2 = cId for some c > 0. Then each eigenvalue λ(t) of S1 (t) evolves independently according to λ̇(t) = 4cλ(t) − 4λ(t)2 , and hence λ(t) =

cλ(0) . λ(0) + (c − λ(0))e−4ct

(347)

Therefore, if ρ0 denotes the limiting empirical spectral distribution of S1 (0), the first moment satisfies µ1 (t) = ∫

cλ dρ0 (λ). λ + (c − λ)e−4ct

(348)

Restricting ρ0 to finitely supported measures ρ0 = ∑M j=0 pj δλj with distinct λj ≠ c, the functions t z→ cλj are linearly independent. Consequently, the trajectory {µ1 (t) ∶ t ≥ 0} uniquely determines the λj +(c−λj )e−4ct weights (p0 , . . . , pM ). 47

Figure 4: Comparison between online SGD and the theory at different truncation orders N for the linear activation, L = 1 and the MSE loss. As N grows larger, the theory is better aligned with the experiments. We choose d = 400, learning rate 0.01, batch size 512 and no weight decay (γ = 0). Suppose, for contradiction, that these trajectories admitted a fixed r-dimensional ODE representation whose initial state depends continuously on the initial spectrum. For each M , the map from the interior of the M -dimensional simplex (p0 , . . . , pM ) to the ODE initial state must then be injective, since identical ODE initial states generate identical trajectories, while the trajectory uniquely determines the weights. For M > r, however, no continuous injective map from an M -dimensional open set into Rr exists (as a consequence of Brouwer’s invariance of domain theorem). Since M is arbitrary, no fixed finite-dimensional ODE with a continuous initialization map can reproduce this family of trajectories. This shows that the infinitedimensionality of (11) is intrinsic and cannot in general be removed by replacing the moment hierarchy with a fixed finite set of order parameters.

D.4

Proof of Corollary 3

Let

p̄12 (t) = q̄12 (t) − t̄1 (t)t̄2 ,

(349)

where t̄2 is constant because the teacher is fixed. We first derive its limiting dynamics from (11). For the order parameter q̄12 = µ̄(1,2) , only the student index 1 is trainable. Therefore, from (11) we have d γ q̄12 = −2Γ11 (q̄)µ̄(1,1,2) − 2Γ12 (q̄)µ̄(1,2,2) − q̄12 . dt 2

(350)

Similarly, for t̄1 = µ̄(1) , d γ t̄1 = −2Γ11 (q̄)q̄11 − 2Γ12 (q̄)q̄12 − t̄1 . dt 2 Subtracting t̄2 times (351) from (350) gives d γ p̄12 = −2Γ11 (q̄) (µ̄(1,1,2) − t̄2 q̄11 ) − 2Γ12 (q̄) (µ̄(1,2,2) − t̄2 q̄12 ) − p̄12 . dt 2

(351)

(352)

Define the centered third-order moments ¯(1,1,2) ∶= lim 1 Tr[S̊ 2 S̊2 ], µ̊ 1 d→∞ d Expanding Sk = S̊k + t̄k I yields and

¯(1,2,2) ∶= lim 1 Tr[S̊1 S̊ 2 ]. µ̊ 2 d→∞ d

(353)

¯(1,1,2) + 2t̄1 p̄12 , µ̄(1,1,2) − t̄2 q̄11 = µ̊

(354)

¯(1,2,2) + t̄2 p̄12 + t̄1 p̄22 . µ̄(1,2,2) − t̄2 q̄12 = µ̊

(355)

48

Consequently, d ¯(1,1,2) + 2t̄1 p̄12 ) − 2Γ12 (q̄) (µ̊ ¯(1,2,2) + t̄2 p̄12 + t̄1 p̄22 ) − γ p̄12 . p̄12 = −2Γ11 (q̄) (µ̊ dt 2 At t = 0, Assumption 3 gives

(356)

¯(1,1,2) (0) = µ̊ ¯(1,2,2) (0) = 0. p̄12 (0) = µ̊

(357)

d p̄12 (t)∣ = −2Γ12 (q̄(0))t̄1 (0)p̄22 . dt t=0 By Assumption 3, this quantity is nonzero. Set

(358)

Hence (356) reduces to

v0 ∶= ∣

d p̄12 (t)∣ ∣ = 2∣Γ12 (q̄(0))∣t̄1 (0)p̄22 > 0. dt t=0

(359)

Since the limiting trajectory satisfies (11) and the relevant finite-order components of V (µ̄(t)) are continuous in t, p̄12 is continuously differentiable. By continuity there exists T0 > 0, independent of d, such that ∣

d v0 p̄12 (t)∣ ≤ , dt 2

0 ≤ t ≤ T0 .

(360)

d d In particular, by continuity dt p̄12 (t) has the same sign as dt p̄12 (0) throughout [0, T0 ]. Since p̄12 (0) = 0,

v0 T0 . 2

(361)

∣p̄12 (T+ )∣ ≥ 2c.

(362)

∣p̄12 (T0 )∣ ≥ Choose c ∶= v04T0 > 0 and T+ ∶= T0 . Then

On the other hand, because p̄12 (0) = 0 and p̄12 is continuous, there exists T− ∈ (0, T+ ) such that c sup ∣p̄12 (t)∣ < . 2 0≤t≤T−

(363)

(d)

Let p̃12 (t) denote the piecewise-constant interpolation of the empirical structural overlap. Since (d)

(d)

(d) (d)

p̃12 = q̃12 − t̃1 t̃2 ,

(364)

Theorem 2 implies that, for every fixed finite T , (d)

sup ∣p̃12 (t) − p̄12 (t)∣ Ð → 0. P

(365)

P (∣p̃12 (T+ )∣ ≥ c) Ð→ 1.

(366)

T+ T+ d ⌉=⌈ ⌉. τd 4αd

(367)

t∈[0,T ]

Applying (365) on [0, T+ ] and using (362), we obtain (d)

Hence, with probability tending to one, (d) Nwr (c) ≤ ⌈

Likewise, by (363) and (365), (d)

P ( sup ∣p̃12 (t)∣ < c) Ð→ 1. 0≤t≤T−

(368)

Therefore, with probability tending to one, (d) Nwr (c) > ⌊

T− T− d ⌋=⌊ ⌋. τd 4αd 49

(369)

Combining (367) and (369) proves (20), and hence d ). αd

(370)

1 ), d log d

(371)

(d) Nwr (c) = ΘP (

Finally, choosing αd = Θ (

which is the largest learning rate covered by Theorem 2, gives (d) Nwr (c) = ΘP (d2 log d).

(372)

This completes the proof.

D.5

Convergence rate of strong recovery

In this section we study the convergence rate after weak recovery. Because the weight decay depends on the first moments of the matrices, it is convenient to introduce tk ∶= µ(k) =

1 Tr[Sk ], d

qkl ∶= µ(k,l) =

1 Tr[Sk Sl ]. d

(373)

The limiting regularized population loss is defined as Ψ(t, q) ∶= Φ(q) +

γ ∑ tk . 2 k∈K

(374)

(1,2)

Let Q+ denote the asymptotically realizable domain of the pair (t, q), namely the closure of all limits generated by sequences of PSD matrices Sk = Wk WkT ⪰ 0, subject to the fixed teacher components when k ∉ K. We define Ψ∗ ∶= inf Ψ(t, q). (375) (1,2)

(t,q)∈Q+

We refer to the convergence to the global minimizers of Ψ as “strong recovery”. Let t̄k (t) ∶= µ̄(k) (t) and q̄kl (t) ∶= µ̄(k,l) (t) denote the corresponding components of the admissible solution of (11). We impose an effective Polyak–Łojasiewicz condition along the trajectory. Assumption 8. There exist constants T0 < ∞ and c > 0 such that, for every t ≥ T0 , −

d Ψ(t̄(t), q̄(t)) ≥ 2c (Ψ(t̄(t), q̄(t)) − Ψ∗ ) . dt

(376)

For a finite d, define the regularized population loss ⎡ L,L,K ⎤ ⎢ ⎛ 1 ⎞⎥⎥ γ 2 Ψd (W ) ∶= Ex ⎢⎢L { √ Tr [Wk WkT (xj xTi − Cij Id )]} ⎥ + 2d ∑ ∥Wk ∥F . ⎠ ⎢ ⎝ ⎥ d k∈K i,j,k=1 ⎦ ⎣

(377)

T0 d Corollary 6. Suppose that the conditions of Theorem 2 and Assumption 8 hold. Let n0 ∶= ⌈ 4α ⌉. Then, for d every fixed target accuracy ϵ > 0, there exists

nϵ = n0 + O ( such that

d 1 log+ ) , αd ϵ

lim P (Ψd (W (nϵ ) ) − Ψ∗ > ϵ) = 0.

d→∞

(378)

(379)

where log+ (x) ∶= max{0, log x}. In particular, if T0 = Θ(1), reaching an ϵ-suboptimal population loss requires at most O ( αdd log 1ϵ ) additional samples. 50

Proof. We first analyze the limiting trajectory in Theorem 2. Define ∆(t) ∶= Ψ(t̄(t), q̄(t)) − Ψ∗ ≥ 0.

(380)

d ∆(t) ≤ −2c∆(t). dt

(381)

By Assumption 8, for every t ≥ T0 , Hence Grönwall’s inequality gives

∆(t) ≤ ∆(T0 )e−2c(t−T0 ) , Let

t ≥ T0 .

∆0 ∶= Ψ(t̄(T0 ), q̄(T0 )) − Ψ∗ .

(382) (383)

If ∆0 = 0, there is nothing to prove. Otherwise, for any fixed ϵ > 0, define Tϵ ∶= T0 +

1 2∆0 log+ ( ). 2c ϵ

(384)

Then

ϵ Ψ(t̄(Tϵ ), q̄(Tϵ )) − Ψ∗ ≤ . (385) 2 We next transfer this bound to the finite-dimensional population loss. Recall that the discrete steps and continuous times are related by t = nτd with τd = 4αd d . For t ≥ 0, define W (d) (t) ∶= W (⌊t/τd ⌋) . For every fixed finite T , Theorem 2, together with the uniform spectral bound of Lemma 3, implies sup ∣Ψd (W (d) (t)) − Ψ(t̄(t), q̄(t))∣ Ð → 0. P

t∈[0,T ]

(386)

This is because on the spectral-bound event, the same argument used in Lemma 6, together with Corollary 5 and uniform integrability, gives uniformly over the spectrally bounded class Ex [L(G(x; W ))] − Φ(q(W )) Ð→ 0.

(387)

γ γ 2 ∑ ∥Wk ∥F = ∑ tk . 2d k∈K 2 k∈K

(388)

The regularization term is exactly

The convergence of the first- and second-order moments supplied by Theorem 2, together with the continuity of Φ, therefore yields (386). For the fixed accuracy ϵ > 0, the time Tϵ in (384) is finite and independent of d. Thus (386), evaluated on [0, Tϵ ], implies ϵ P (∣Ψd (W (d) (Tϵ )) − Ψ(t̄(Tϵ ), q̄(Tϵ ))∣ > ) Ð→ 0. (389) 2 Tϵ d Set nϵ ∶= ⌈ 4α ⌉. Therefore, using (385), d

P (Ψd (W (nϵ ) ) − Ψ∗ > ϵ) Ð→ 0.

(390)

Finally, by (384) and the definition of nϵ , nϵ − n 0 = O (

d 1 log+ ) , αd ϵ

(391)

which proves the claimed sample complexity. Example. Consider the setting γ = 0 and L(G1 , G2 ) = (G1 − G2 )2 with one learnable student S1 = W1 W1T and one fixed teacher S2 . Then Φ(q) = 2q11 + 2q22 − 4q12 = 2 lim

1

d→∞ d

51

Tr[(S1 − S2 )2 ].

(392)

If the teacher is normalized by q22 = 1, the global minimum Φ∗ = 0 is attained at perfect recovery, q11 = q12 = 1. At finite d, up to the normalization induced by the time scaling, the corresponding population loss is proportional to d1 Tr[(S1 − S2 )2 ]. Its gradient with respect to W1 satisfies ∥∇W1 L∥2F ∝

1 Tr[S1 (S1 − S2 )2 ]. d2

(393)

Hence the effective PL condition reduces, up to fixed normalization constants, to a bound of the form Tr[S1 (S1 − S2 )2 ] ≥ c0 Tr[(S1 − S2 )2 ]

(394)

for some c0 > 0. A sufficient condition for (394) is S1 ⪰ c0 Id . More generally, it requires non-degeneracy only along the directions relevant to S1 − S2 .

E

Online SGD over S

In this section, we study online SGD directly over the attention matrices {Sk }K k=1 . At each step n, a fresh sample x(n) ∼ N (0, C ⊗ Id ) is drawn, and for k ∈ {1, . . . , K} we consider the update (n+1)

Sk

(n)

= Sk

− αd χk (∇Sk L(G(n) (x(n) )) +

γ (n) S ), d k

(395)

where χk ∶= 1{k∈K} so that Sk remains fixed whenever k ∉ K. The tensor G in (395) is defined as 1 (n) (n) (n) (n) Gijk ∶= √ Tr [Sk (xj (xi )T − Cij Id )] . d

(396)

We introduce the time as τd ∶= αdd and t ∶= nτd . The order parameters are defined as: (d)

tk ∶=

1 Tr[Sk ], d

S +S T

(d)

Ψkl ∶=

1 Tr[Sk Sl ], d

(d)

Ωkl ∶=

1 Tr[Sk SlT ]. d

(397)

S −S T

kl kl Moreover, writing Sk+ ∶= k 2 k and Sk− ∶= k 2 k , we have Ωkl +Ψ = d1 ⟨Sk+ , Sl+ ⟩F and Ωkl −Ψ = d1 ⟨Sk− , Sl− ⟩F . 2 2 Consequently, the feasible domain of the order parameters is

D ∶= {(Ψ, Ω) ∈ SK × SK ∶

Ω+Ψ ⪰ 0, 2

Ω−Ψ ⪰ 0} . 2

(398)

For (Ψ, Ω) ∈ D, define the potential Φ(Ψ, Ω) ∶= EGΨ,Ω [L(GΨ,Ω )] ,

(399)

where GΨ,Ω ∈ RL×L×K is a centered Gaussian tensor with covariance Cov ((GΨ,Ω )ijk , (GΨ,Ω )i′ j ′ l ) = Cii′ Cjj ′ Ωkl + Cij ′ Cji′ Ψkl . We denote by ∇Ψ Φ and ∇Ω Φ the corresponding matrix derivatives w.r.t. Ψ and Ω. Assumption 9. The following conditions hold. 1. The loss L ∶ RL×L×K → R is four times continuously differentiable, and L together with its derivatives up to fourth order have at most polynomial growth. (0)

2. There exists a constant C0 > 0, independent of d, such that max1≤k≤K ∥Sk ∥op ≤ C0 almost surely. (d)

(d)

3. The initial order parameters converge in probability to deterministic limits: tk (0) → tk (0), Ψkl (0) → (d) Ψ̄kl (0) and Ωkl (0) → Ω̄kl (0) for every k, l. 4. The regularization is sufficiently large, i.e., γ > γ∗ (C0 , C, L, K, L) for some γ∗ independent of d.

52

(d)

Let S̃k (t) denote the piecewise-constant interpolation (d)

(n)

S̃k (t) = Sk ,

t ∈ [nτd , (n + 1)τd ),

(400)

and define t̃(d) , Ψ̃(d) , and Ω̃(d) in the same way according to (397). The following theorem shows that, unlike Theorems 2 and 3, direct optimization over S closes at the level of finitely many first- and second-order parameters. c0 Theorem 4. Under Assumption 9, suppose that the learning rate satisfies αd ≤ d log and αd = Ω(d−ι ) d for some constants ι > 0 and c0 > 0 sufficiently small. Then, for every finite T > 0, the empirical order parameters converge uniformly in probability to a deterministic trajectory (t̄(t), Ψ̄(t), Ω̄(t)). More precisely, for every k, l and every ϵ > 0,

lim P

d→∞

⎛ ⎞ (d) (d) (d) sup ∣t̃ (t) − t̄k (t)∣ + ∣Ψ̃kl (t) − Ψ̄kl (t)∣ + ∣Ω̃kl (t) − Ω̄kl (t)∣ > ϵ = 0. ⎝t∈[0,T ] k ⎠

(401)

The limiting trajectory is the unique solution of K d t̄k = −2χk ∑ ([∇Ω Φ]km + [∇Ψ Φ]km ) t̄m − γχk t̄k , dt m=1

(402)

K K d Ψ̄kl = −2χk ∑ ([∇Ω Φ]km Ψ̄ml + [∇Ψ Φ]km Ω̄ml ) − 2χl ∑ ([∇Ω Φ]lm Ψ̄km + [∇Ψ Φ]lm Ω̄km ) − γ(χk + χl )Ψ̄kl , dt m=1 m=1 (403) K K d Ω̄kl = −2χk ∑ ([∇Ω Φ]km Ω̄ml + [∇Ψ Φ]km Ψ̄ml ) − 2χl ∑ ([∇Ω Φ]lm Ω̄km + [∇Ψ Φ]lm Ψ̄km ) − γ(χk + χl )Ω̄kl , dt m=1 m=1 (404)

where the gradients of Φ are evaluated at (Ψ̄(t), Ω̄(t)) ∈ D for t ∈ [0, T ]. Building upon this uniform convergence, now we consider the weak recovery. First define the traceless parts S̊k ∶= Sk − tk Id , (405) and the structural overlaps

1 (d) (d) (d) Tr[S̊k S̊l ] = Ψkl − tk tl , d 1 (d) (d) (d) (d) pΩ,kl ∶= Tr[S̊k S̊lT ] = Ωkl − tk tl . d It is useful to resolve these overlaps into their symmetric and skew-symmetric channels: (d)

pΨ,kl ∶=

p+,12 ∶=

pΩ,12 + pΨ,12 , 2

p−,12 ∶=

pΩ,12 − pΨ,12 . 2

(406) (407)

(408)

For the teacher, define the corresponding signal strengths 2

pΩ,22 + pΨ,22 1 S̊2 + S̊2T a+ ∶= = lim ∥ ∥ , d→∞ d 2 2 F and

(409)

2

a− ∶=

pΩ,22 − pΨ,22 1 S̊2 − S̊2T = lim ∥ ∥ . d→∞ d 2 2 F

(410)

Thus a+ , a− ≥ 0. The following is the uninformative initialization and informative teacher assumption. Assumption 10. As d → ∞, the initial student S1 (0) and the fixed teacher S2 satisfy t̄1 (0) = 0,

p̄Ψ,12 (0) = 0,

p̄Ω,12 (0) = 0.

Moreover, the teacher is informative in the sense that a+ + a− > 0. 53

(411)

For example, (411) is satisfied by the standard Gaussian initialization with centered i.i.d. entries of variance Θ(d−1 ). Define the effective gradients in the symmetric and skew-symmetric parts by g+ ∶= [∇Ω Φ]12 + [∇Ψ Φ]12 ,

g− ∶= [∇Ω Φ]12 − [∇Ψ Φ]12 .

(412)

For a fixed c > 0, define the weak-recovery sample complexity (d)

(d,n)

(d,n)

Nwr,S (c) ∶= inf {n ≥ 0 ∶ max (∣pΨ,12 ∣, ∣pΩ,12 ∣) ≥ c} .

(413)

The following corollary gives the weak recovery conditions. Corollary 7. Consider (395) with K = {1}, where S1 is the trainable student and S2 is a fixed teacher. Suppose the conditions of Theorem 4 and Assumption 10 hold. 1. If, at initialization, a+ g+ (0) ≠ 0

or

a− g− (0) ≠ 0,

(414)

then there exist constants c > 0 and 0 < T− < T+ < ∞, independent of d, such that lim P (⌊

d→∞

T+ d T− d (d) ⌋ < Nwr,S (c) ≤ ⌈ ⌉) = 1. αd αd

(415)

1 In particular, choosing the largest learning-rate scaling αd = Θ ( d log ) covered by Theorem 4 gives the d weak recovery sample complexity (d) Nwr,S (c) = ΘP (d2 log d). (416)

2. Let the uninformative manifold be U ∶= {(t, Ψ, Ω) ∶ t1 = 0, Ψ12 = 0, Ω12 = 0}. Suppose that [∇Ω Φ(Ψ, Ω)]12 = [∇Ψ Φ(Ψ, Ω)]12 = 0

for every state in U .

(417)

Then for every fixed T < ∞ and every ϵ > 0, lim P

d→∞

⎛ ⎞ (d) (d) sup max (∣pΨ,12 (t)∣, ∣pΩ,12 (t)∣) > ϵ = 0. ⎝t∈[0,T ] ⎠

(418)

In particular, for any learning-rates covered by Theorem 4, S-SGD cannot achieve weak recovery within O(d2 log d) samples.

E.1

Proof of Theorem 4

The proof combines uniform spectral bounds based on the Matrix Freedman inequality with a martingale expansion of the empirical order parameters, and concludes by applying Grönwall’s inequality to the resulting approximate ODE. E.1.1

Uniform spectral bounds

Lemma 8. Under the conditions of Theorem 4, there exists a constant C∗ > C0 , independent of d and T , such that for every fixed T > 0, lim P ( sup

d→∞

(n)

max ∥Sk ∥op ≤ C∗ ) = 1,

0≤n≤NT 1≤k≤K

(419)

where NT ∶= ⌊ τTd ⌋ = ⌊ Tαdd ⌋. The proof follows essentially the same stopping-time and Matrix Freedman argument as Lemma 3. The main difference is that direct optimization over Sk yields an expected gradient of the form d−1 Hk rather than d−1 Hk Wk (see Lemma 2); moreover, since Sk need not be symmetric, Hk contains both Sl and SlT . 54

(n)

Proof. The matrices with k ∉ K remain fixed, and hence satisfy ∥Sk ∥op ≤ C0 . It therefore suffices to control the trainable indices k ∈ K. Fix C∗ > C0 , to be chosen below, and define the stopping time (n)

J ∶= inf {n ≥ 0 ∶ max ∥Sk ∥op > C∗ } .

(420)

(n)

(421)

1≤k≤K

On {n < J},

max ∥Sk ∥op ≤ C∗ . k

As in the proof of Lemma 3, define

Step 1: Truncation of the fresh sample.

(n)

An ∶= {∥x(n) ∥2 ≤ cx d} ∩ {max ∣Gijk ∣ ≤ cG log d} ,

(422)

i,j,k

where cx , cG > 0 are sufficiently large constants. Similarly to Lemma 3, we have 1{n<J} P(Acn ∣ Fn ) ≤ pd ,

(423)

where pd ≤ C(e−c1 d + e−c2 (log d) ). Consequently, defining 2

NT −1

Adata ∶= ⋂ ({n ≥ J} ∪ An ) ,

(424)

P(Acdata ) ≤ NT pd Ð→ 0,

(425)

n=0

we have −ι

because αd = Ω(d ). Step 2: Conditional drift and stopped martingale noise. (n)

gk

For k ∈ K, define the gradient

L

(n)

∶= ∇Sk L(G(n) ) = ∑ ∂(ijk) L(G(n) )Xij ,

(426)

i,j=1

(n)

(n)

(n)

(n)

where Xij ∶= √1d (xi (xj )T − Cij Id ). Let ḡk (n)

ḡk

=

(n)

∶= E[gk

∣ Fn ]. The same calculation as in Lemma 2 gives

1 (n) (n) H + EStein,k , d k

(427)

where (n)

Hk

∶=

L

∑

K

(n)

∑ Ex [∂(ijk),(i′ j ′ l) L(G 2

i,j,i′ ,j ′ =1 l=1

(n)

)] (Cii′ Cjj ′ Sl

(n) T

+ Cij ′ Cji′ (Sl

) ).

(428)

Using Assumption 9.1 and the uniform moment bounds for G(n) under (421), the same estimates as in Lemma 2 yield (n) (n) ∥Hk ∥op ≤ CH (C∗ ), ∥EStein,k ∥op ≤ CE (C∗ )d−3/2 , (429) uniformly on {n < J}. Define the stopped and truncated gradient (n)

̂ gk its conditional mean and the martingale difference

(n)

∶= 1{n<J} gk 1An ,

(430)

(n) (n) ̂ ḡk ∶= E[̂ gk ∣ Fn ],

(431)

(n) (n) (n) ξ̂k ∶= ̂ gk − ̂ ḡk .

(432)

55

Then

(n) (n) (n) ̂ ḡk = 1{n<J} (ḡk − Tk ) ,

(n)

Tk

(n)

∶= E [gk 1Acn ∣ Fn ] .

(433)

As in the proof of Lemma 3, Hölder’s inequality and moment bounds give 1{n<J} max {∥E[ξ̂k (ξ̂k )T ∣ Fn ]∥ (n)

On An ∩ {n < J}, we also have

(n)

, ∥E[(ξ̂k )T ξ̂k (n)

op

(n)

∣ Fn ]∥ } ≤ Cg (C∗ ). op

√ (n) ∥̂ gk ∥op ≤ C d (1 + (log d)r ) =∶ Bd

(434)

(435)

for some fixed integer r, and consequently ∥ξ̂k ∥op ≤ 2Bd . (n)

(436)

Finally, by Cauchy–Schwarz and (423), the truncation error satisfies (n)

1{n<J} ∥Tk

∥op ≤ rd ,

rd = o(d−M )

(437)

for every fixed M > 0. Step 3: Uniform control of the martingale noise. Set ad ∶= 1 − αdd γ = 1 − τd γ. For sufficiently large d, (n) ad ∈ (0, 1). Define Dk,n ∶= −αd ξ̂ and, for each terminal time m ≥ 1, k

m−1

Yk,m ∶= ∑ am−1−n Dk,n . d

(438)

n=0

For each fixed m, the summands form a square matrix martingale difference sequence. By (436), ∥am−1−n Dk,n ∥op ≤ 2αd Bd =∶ Rd , d

(439)

where Rd = o((log d)−1 ) under αd ≤ c0 (d log d)−1 . Moreover, by (434), ∞

σd2 ≤ Cg αd2 ∑ a2r d = r=0

Cg αd2 dαd Cc0 ≤C ≤ . 1 − a2d γ γ log d

(440)

Hence the Matrix Freedman inequality gives, for every fixed δ > 0, P(∥Yk,m ∥op ≥ δ) ≤ 2d exp [−

2 δ 2 /2 ] ≤ 2d1−cδ /c0 2 σd + Rd δ/3

(441)

for all sufficiently large d. Since NT ≤ CT d1+ι and K is fixed, a union bound gives P (max max ∥Yk,m ∥op > δ) ≤ CT d2+ι−cδ /c0 . 2

k∈K 1≤m≤NT

Choosing c0 > 0 sufficiently small yields

(442)

P(Acnoise ) Ð→ 0,

(443)

Anoise ∶= {max max ∥Yk,m ∥op ≤ δ} .

(444)

where k∈K 1≤m≤NT

56

On Adata , for every m ≤ J and every n < m,

Step 4: Closing the stopping-time argument. (n)

gk

(n)

= ḡk

(n)

− Tk

+ ξ̂k . (n)

(445)

For k ∈ K, unrolling (395) therefore gives (m)

Sk

m−1

(0)

(n)

m−1−n = am ḡk d Sk − αd ∑ ad n=0

m−1

(n)

+ αd ∑ am−1−n Tk d

+ Yk,m .

(446)

n=0

d r Since αd ∑∞ r=0 ad = γ , (429) and (437) imply that, on Adata ∩ Anoise and for m ≤ J ∧ NT , (m)

∥Sk

∥op ≤ C0 +

CH (C∗ ) CE (C∗ ) drd CH (C∗ ) √ + + + δ = C0 + + δ + o(1). γ γ γ γ d

(447)

0 . By Assumption 9.4, γ can be chosen sufficiently large so Fix, for instance, C∗ ∶= 2C0 + 1 and δ ∶= C∗ −C 4 that CH (C∗ ) C∗ − C0 ≤ . (448) γ 4 Then, for all sufficiently large d,

(m)

max ∥Sk k

∥op < C∗

for every m ≤ J ∧ NT

(449)

on Adata ∩ Anoise . (J) If J ≤ NT , taking m = J gives maxk ∥Sk ∥op < C∗ , contradicting the definition of J. Hence J > NT on Adata ∩ Anoise for all sufficiently large d. Finally, P(J ≤ NT ) ≤ P(Acdata ) + P(Acnoise ) Ð→ 0 (450) by (425) and (443). This proves the lemma. E.1.2 Let

Moment expansion Q(d) (t) ∶= (t(d) (t), Ψ(d) (t), Ω(d) (t))

(451)

denote the piecewise-constant interpolation of the empirical order parameters, and let F = (Ft , FΨ , FΩ ) denote the finite-dimensional vector field on the right-hand sides of (402)–(404). Lemma 9. Under the conditions of Theorem 4, for every fixed T > 0, Q(d) (t) = Q(d) (0) + ∫

t 0

(d)

F (Q(d) (s)) ds + Etotal (t),

where

(d)

sup ∥Etotal (t)∥

t∈[0,T ]

Ð → 0. P

∞

t ∈ [0, T ],

(452) (453)

The proof follows the same martingale-decomposition and drift-convergence strategy as Lemma 4. The main simplification is that direct optimization over S closes at the finite-dimensional level of (t, Ψ, Ω), so no higher-order moments or weighted moment-space estimates are required; the only additional feature is the simultaneous appearance of Sl and SlT due to the possible asymmetry of Sl . Proof. Let C∗ be the constant in Lemma 8, and define (n)

(454)

P(J > NT ) Ð→ 1.

(455)

J ∶= inf {n ≥ 0 ∶ max ∥Sk ∥op > C∗ } . 1≤k≤K

Then, for NT ∶= ⌊ τTd ⌋ = ⌊ Tαdd ⌋, Lemma 8 gives

All estimates below are uniform on {n < J}. 57

Recall from the proof of Lemma 8 that

Step 1: Conditional drift.

1 (n) (n) H + EStein,k , d k

ḡk

=

(n)

(456)

∑

2 (n) (n) (n) T Ex [∂(ijk),(i )] (Cii′ Cjj ′ Sm + Cij ′ Cji′ (Sm ) ). ′ j ′ m) L(G

(457)

(n)

∥EStein,k ∥op ≤ CE d−3/2 ,

where (n)

Hk

=

i,j,i′ ,j ′ ,m

Define the finite-dimensional coefficients (d)

1 2 (n) ∑ Cii′ Cjj ′ Ex [∂(ijk),(i′ j ′ m) L(G )] , 2 i,j,i′ ,j ′

(458)

(d)

1 2 (n) ∑ Cij ′ Cji′ Ex [∂(ijk),(i′ j ′ m) L(G )] . 2 i,j,i′ ,j ′

(459)

ΓΩ,km ∶= ΓΨ,km ∶= Then (n)

Hk

K

(d)

(d)

(n) (n) T = 2 ∑ (ΓΩ,km Sm + ΓΨ,km (Sm ) ).

(460)

m=1

We claim that, uniformly over the spectrally bounded class, δd (C∗ ) ∶=

(d)

(d)

max ( ∣ΓΩ,km − [∇Ω Φ(Ψ, Ω)]km ∣ + ∣ΓΨ,km − [∇Ψ Φ(Ψ, Ω)]km ∣ ) Ð→ 0.

sup

maxm ∥Sm ∥op ≤C∗ k,m

(461)

The proof is the same subsequence argument as in Lemma 6. Step 2: One-step expansion of the order parameters. (n)

(n)

∶= gk

ξk

Define

(n)

− ḡk .

(462)

Using τd = αd /d, the SGD update can be written as (n)

∆Sk

(n+1)

∶= Sk

(n)

− Sk

(n)

(n)

= −τd χk (Hk

For the normalized trace, (n+1)

tk

(n)

− tk

=

(n)

+ γSk ) − αd χk ξk

(n)

− αd χk EStein,k .

1 (n) Tr[∆Sk ]. d

(463)

(464)

For the two second-order overlaps, 1 1 (n) (n) (n) (n) (n) (n) Tr[∆Sk Sl + Sk ∆Sl ] + Tr[∆Sk ∆Sl ], d d 1 1 (n+1) (n) (n) (n) (n) (n) (n) (n) Ωkl − Ωkl = Tr[∆Sk (Sl )T + Sk (∆Sl )T ] + Tr[∆Sk (∆Sl )T ]. d d (n+1)

Ψkl

(n)

− Ψkl =

(465) (466)

Substituting (460) gives, for example,

and

1 (n) (d) (d) Tr[Hk ] = 2 ∑ (ΓΩ,km + ΓΨ,km ) t(n) m , d m

(467)

1 (n) (n) (d) (n) (d) (n) Tr[Hk Sl ] = 2 ∑ (ΓΩ,km Ψml + ΓΨ,km Ωml ) , d m

(468)

1 (n) (n) (d) (n) (d) (n) Tr[Hk (Sl )T ] = 2 ∑ (ΓΩ,km Ωml + ΓΨ,km Ψml ) . d m

(469)

(n)

Therefore the first-order terms coincide with τd F (Q(n) ), up to an error τd εdrift satisfying (n)

1{n<J} ∥εdrift ∥∞ ≤ C(C∗ )δd (C∗ ). 58

(470)

The first-order martingale terms are χk (n) Tr[ξk ], d 1 (n) (n) (n) (n) (n) ηΨ,kl ∶= − Tr [χk ξk Sl + χl Sk ξl ] , d 1 (n) (n) (n) (n) (n) ηΩ,kl ∶= − Tr [χk ξk (Sl )T + χl Sk (ξl )T ] . d (n)

ηt,k ∶= −

(471) (472) (473)

Collect them into η (n) . The remaining terms in (465)–(466), are denoted by R(n) . Hence (n)

Q(n+1) − Q(n) = τd F (Q(n) ) + αd η (n) + R(n) + τd εdrift .

(474)

Step 3: Vanishing of the martingale noise. As in Step 3 of the proof of Lemma 4, every coordinate (n) of η (n) is a finite sum of linear functionals of ξk . The same moment estimate therefore gives C d2

(475)

Ma(m) ∶= ∑ αd ηa(n) 1{n<J}

(476)

C αd E[⟨Ma ⟩NT ] ≤ NT αd2 2 ≤ CT . d d

(477)

max sup ∣Ma(m) ∣ Ð → 0.

(478)

1{n<J} E [∣ηa(n) ∣2 ∣ Fn ] ≤ for every coordinate a of (t, Ψ, Ω). Thus the stopped process m−1 n=0

is a square-integrable martingale with

By Doob’s maximal inequality, P

a

m≤NT

Here the maximum is over only finitely many coordinates, since K is fixed. Step 4: Vanishing of the higher-order residuals. On {n < J}, the same moment estimates as in Lemma 4 give E [∥R(n) ∥∞ ∣ Fn ] ≤ C (αd2 + αd τd + τd2 + αd d−3/2 ) . (479) The first three terms arise from the products of the drift and stochastic increments in (465)–(466), while the last term comes from the residual in (463). Since NT = O ( αdd ) and τd = αdd , we obtain NT −1

E [ ∑ 1{n<J} ∥R(n) ∥∞ ] ≤ CT (dαd + αd + n=0

αd + d−1/2 ) Ð→ 0, d

(480)

because αd ≤ c0 (d log d)−1 . Hence, by Markov’s inequality, m−1

sup ∥ ∑ R(n) ∥

m≤NT

Ð → 0.

(481)

≤ T C(C∗ )δd (C∗ ) Ð→ 0.

(482)

n=0

P

∞

Likewise, by (470), m−1

(n)

sup ∥τd ∑ εdrift ∥

m≤NT

n=0

∞

59

Step 5: Integral equation.

For m(t) ∶= ⌊ τtd ⌋, summing (474) gives m(t)−1

(d)

(d)

(d) Q(d) (t) = Q(d) (0) + τd ∑ F (Q(n) ) + Enoise (t) + Erem (t) + Edrift (t),

(483)

n=0

(d)

(d)

where Enoise (t) ∶= αd ∑n=0 η (n) , Erem (t) ∶= ∑n=0 Because Q(d) is piecewise constant, m(t)−1

t

∫

0

m(t)−1

m(t)−1 (n) εdrift .

(d)

R(n) and Edrift (t) ∶= τd ∑n=0

m(t)−1

F (Q(d) (s)) ds = τd ∑ F (Q(n) ) + (t − m(t)τd )F (Q(m(t)) ).

(484)

n=0

On {J > NT }, the spectral bound implies that Q(d) (t) remains in a bounded subset, on which F is bounded. (d) Therefore the discretization error Edisc (t) ∶= (t − m(t)τd )F (Q(m(t)) ) satisfies (d)

sup ∥Edisc (t)∥∞ ≤ Cτd Ð→ 0.

(485)

t∈[0,T ]

(d)

Combining (478), (481), (482), and (485), we conclude that supt∈[0,T ] ∥Etotal (t)∥ (d)

(d)

(d)

Ð → 0, where P

∞

(d)

(d) + Edrift + Edisc . Etotal = Enoise + Erem

(486)

This proves the lemma. E.1.3

Final proof of Theorem 4

Fix an arbitrary finite time horizon T > 0, and let C∗ be the constant in Lemma 8. Define Ed (T ) ∶= { sup

(n)

max ∥Sk ∥op ≤ C∗ } ,

0≤n≤NT 1≤k≤K

By Lemma 8,

NT = ⌊

Td ⌋. αd

P(Ed (T )) Ð→ 1.

(487)

(488)

On Ed (T ), the empirical order parameters satisfy (d)

∣tk (t)∣ ≤ C∗ ,

(d)

∣Ψkl (t)∣ ≤ C∗2 ,

(d)

∣Ωkl (t)∣ ≤ C∗2 ,

(489)

uniformly over t ∈ [0, T ]. We first address the fact that the vector field F is naturally defined only on the domain RK × D. For (Ψ, Ω) ∈ SK × SK , write Ω+Ψ Ω−Ψ Q+ ∶= , Q− ∶= , (490) 2 2 and define ̂± ∶= ΠSK (Q± ), ̂ ∶= Q ̂+ + Q ̂− , ̂ ∶= Q ̂+ − Q ̂− . Q Ω Ψ (491) + ̂ Ω) ̂ and ∇Ω Φ(Ψ, ̂ Ω) ̂ in (402)–(404) defines an extension F̃ of F to the finite-dimensional space Using ∇Ψ Φ(Ψ, X ∶= RK × SK × SK . By Assumption 9.1, Price’s theorem, and the non-expansiveness of the PSD projection, F̃ is locally Lipschitz on X . Choose a fixed radius R larger than the bound in (489). By composing F̃ with the projection onto the closed ball BR ⊂ X , we obtain a globally Lipschitz vector field F̂ which agrees with F on all states satisfying (489). Let LR < ∞ denote a global Lipschitz constant of F̂. Consider the globally well-posed ODE d Q̄(t) = F̂(Q̄(t)), dt

Q̄(0) = (t̄(0), Ψ̄(0), Ω̄(0)) , 60

(492)

where Q̄ = (t̄, Ψ̄, Ω̄). By Assumption 9.3 and the fact that K is fixed, ∥Q(d) (0) − Q̄(0)∥∞ Ð → 0. P

(493)

On the other hand, Lemma 9 gives Q(d) (t) = Q(d) (0) + ∫ with

t 0

(d)

F (Q(d) (s)) ds + Etotal (t),

(d)

sup ∥Etotal (t)∥

t∈[0,T ]

Ð → 0. P

∞

(494) (495)

On Ed (T ), every empirical state satisfies (489). Hence F (Q(d) (t)) = F̂(Q(d) (t)),

t ∈ [0, T ].

(496)

Subtracting the integral form of (492) from (494), we obtain on Ed (T ) (d)

∥Q(d) (t) − Q̄(t)∥∞ ≤ ∥Q(d) (0) − Q̄(0)∥∞ + sup ∥Etotal (s)∥ + LR ∫ ∞

s≤T

0

t

∥Q(d) (s) − Q̄(s)∥∞ ds.

(497)

Therefore, by Grönwall’s inequality, ⎛ ⎞ (d) sup ∥Q(d) (t) − Q̄(t)∥∞ ≤ eLR T ∥Q(d) (0) − Q̄(0)∥∞ + sup ∥Etotal (t)∥ ∞ ⎝ ⎠ t∈[0,T ] t∈[0,T ]

(498)

on Ed (T ). Combining (488), (493), and (495) yields sup ∥Q(d) (t) − Q̄(t)∥∞ Ð → 0. P

t∈[0,T ]

(499)

It remains to show that Q̄ is a solution of the original ODE rather than merely a solution of its extension. For every finite d and every t, (Ψ(d) (t), Ω(d) (t)) ∈ D. Since D is closed, (499) implies (Ψ̄(t), Ω̄(t)) ∈ D for t ∈ [0, T ]. Likewise, (489) and (499) imply that Q̄(t) remains inside the region where F̂ = F . Consequently, d Q̄(t) = F (Q̄(t)), dt

t ∈ [0, T ],

(500)

and hence Q̄ satisfies (402)–(404). Finally, the locally Lipschitz extension F̃ implies local uniqueness of the ODE. Thus the solution Q̄ is unique. This proves Theorem 4.

E.2

Proof of Corollary 7

Recall that S1 is trainable and S2 is fixed. For brevity, write Akl (t) ∶= [∇Ω Φ(Ψ̄(t), Ω̄(t))]kl ,

Bkl (t) ∶= [∇Ψ Φ(Ψ̄(t), Ω̄(t))]kl .

(501)

Since S2 is fixed, t̄2 , pΨ,22 , pΩ,22 , a+ , a− are constant in time. We first derive the limiting dynamics of the structural overlaps. From (402), with χ1 = 1 and χ2 = 0, d t̄1 = −2(A11 + B11 )t̄1 − 2(A12 + B12 )t̄2 − γ t̄1 . dt

(502)

d Ψ̄12 = −2A11 Ψ̄12 − 2B11 Ω̄12 − 2A12 Ψ̄22 − 2B12 Ω̄22 − γ Ψ̄12 , dt d Ω̄12 = −2A11 Ω̄12 − 2B11 Ψ̄12 − 2A12 Ω̄22 − 2B12 Ψ̄22 − γ Ω̄12 . dt

(503)

Similarly, (403) and (404) give

61

(504)

Recall

p̄Ψ,12 = Ψ̄12 − t̄1 t̄2 ,

p̄Ω,12 = Ω̄12 − t̄1 t̄2 .

(505)

Subtracting t̄2 times (502) from (503) and (504), respectively, we obtain d p̄Ψ,12 = −2A11 p̄Ψ,12 − 2B11 p̄Ω,12 − 2A12 pΨ,22 − 2B12 pΩ,22 − γ p̄Ψ,12 , dt d p̄Ω,12 = −2A11 p̄Ω,12 − 2B11 p̄Ψ,12 − 2A12 pΩ,22 − 2B12 pΨ,22 − γ p̄Ω,12 . dt It is convenient to use

and recall

(507)

p̄Ω,12 − p̄Ψ,12 . 2

(508)

h+ (t) ∶= A11 (t) + B11 (t),

h− (t) ∶= A11 (t) − B11 (t),

(509)

g+ (t) ∶= A12 (t) + B12 (t),

g− (t) ∶= A12 (t) − B12 (t).

(510)

p̄+ ∶= Define

p̄Ω,12 + p̄Ψ,12 , 2

(506)

p̄− ∶=

Adding and subtracting (506)–(507) yields the decoupled equations d p̄+ = − (2h+ (t) + γ) p̄+ − 2a+ g+ (t), dt d p̄− = − (2h− (t) + γ) p̄− − 2a− g− (t). dt

(511) (512)

By Assumption 10, p̄+ (0) = p̄− (0) = 0. Consequently, p̄˙+ (0) = −2a+ g+ (0),

p̄˙− (0) = −2a− g− (0).

(513)

We now prove the two claims separately. Case 1.

Suppose (414) holds. Then there exists σ ∈ {+, −} such that vσ ∶= p̄˙σ (0) ≠ 0.

(514)

By Theorem 4, the limiting trajectory is continuously differentiable, and hence p̄˙σ (t) is continuous. Therefore there exists T0 > 0, independent of d, such that ∣p̄˙σ (t) − vσ ∣ ≤

∣vσ ∣ , 2

0 ≤ t ≤ T0 .

(515)

In particular, p̄˙σ (t) has the same sign as vσ throughout [0, T0 ] and ∣p̄˙σ (t)∣ ≥

∣vσ ∣ . 2

(516)

∣p̄σ (T0 )∣ ≥

∣vσ ∣T0 . 2

(517)

Since p̄σ (0) = 0, Moreover,

max (∣p̄Ψ,12 ∣, ∣p̄Ω,12 ∣) ≥ max(∣p̄+ ∣, ∣p̄− ∣).

(518)

max (∣p̄Ψ,12 (T+ )∣, ∣p̄Ω,12 (T+ )∣) ≥ 2c.

(519)

Set c ∶= ∣vσ4∣T0 and T+ ∶= T0 . Then

On the other hand, because both structural overlaps vanish at initialization and are continuous, there exists T− ∈ (0, T+ ) such that c sup max (∣p̄Ψ,12 (t)∣, ∣p̄Ω,12 (t)∣) < . (520) 2 0≤t≤T− 62

(d)

(d)

Let p̃Ψ,12 (t) and p̃Ω,12 (t) denote the corresponding empirical piecewise-constant interpolations. Since (d)

(d)

(d) (d)

(d)

Theorem 4 implies

(d)

(d) (d)

p̃Ω,12 = Ω̃12 − t̃1 t̃2 ,

p̃Ψ,12 = Ψ̃12 − t̃1 t̃2 , (d)

(d)

(521)

sup max (∣p̃Ψ,12 (t) − p̄Ψ,12 (t)∣ , ∣p̃Ω,12 (t) − p̄Ω,12 (t)∣) Ð →0 P

t∈[0,T ]

(522)

for every fixed T < ∞. Applying (522) on [0, T+ ] and using (519), we obtain (d)

(d)

(d)

T+ T+ d ⌉=⌈ ⌉. τd αd

P (max (∣p̃Ψ,12 (T+ )∣, ∣p̃Ω,12 (T+ )∣) ≥ c) Ð→ 1.

(523)

Hence, with probability tending to one, Nwr,S (c) ≤ ⌈

(524)

Likewise, (520) and (522) imply (d)

(d)

(d)

T− d ⌋. αd

P ( sup max (∣p̃Ψ,12 (t)∣, ∣p̃Ω,12 (t)∣) < c) Ð→ 1. 0≤t≤T−

(525)

Therefore, with probability tending to one, Nwr,S (c) > ⌊

(526)

(d)

1 Combining (524) and (526) proves (415). In particular, Nwr,S (c) = ΘP ( αdd ). Choosing αd = Θ ( d log ) gives d (d)

Nwr,S (c) = ΘP (d2 log d). Case 2.

(527)

Suppose now that (417) holds. Consider the uninformative manifold U = {(t, Ψ, Ω) ∶ t1 = 0,

Ψ12 = 0,

Ω12 = 0} .

(528)

On U , we have

pΨ,12 = pΩ,12 = p+ = p− = 0. (529) ˙ Moreover, by (417), A12 = B12 = 0. Therefore (502) gives t̄1 = 0 while (511)–(512) give p̄˙+ = p̄˙− = 0. Thus the limiting vector field is tangent to U. By uniqueness of the solution established in Theorem 4, and since Assumption 10 places the initial state on U , the limiting trajectory remains on this manifold: t̄1 (t) = 0, Equivalently,

Ψ̄12 (t) = Ω̄12 (t) = 0,

p̄Ψ,12 (t) = p̄Ω,12 (t) = 0,

t ≥ 0.

t ≥ 0.

(530) (531)

Combining (531) with (522), for every fixed T < ∞ and every ϵ > 0, P

⎛ ⎞ (d) (d) sup max (∣p̃Ψ,12 (t)∣, ∣p̃Ω,12 (t)∣) > ϵ Ð→ 0, ⎝t∈[0,T ] ⎠

which proves (418). Finally, let C > 0 be arbitrary and consider Md ∶= ⌊Cd2 log d⌋ SGD steps, satisfying αd ≤ Cd log d αd ≤ Cc0 Md τd = Md d under the learning-rate condition of Theorem 4. Hence, for every fixed recovery threshold c > 0, (d)

P (Nwr,S (c) ≤ Cd2 log d) Ð→ 0. Since C > 0 is arbitrary, S-SGD cannot achieve weak recovery within O(d2 log d) samples. This completes the proof. 63

(532)

(533)

(534)

F

Theory in Section 4

F.1

Proof of Theorem 3

F.1.1

A stronger version of Theorem 1

We first estimate the error scaling in Theorem 1. It is closely related to existing work on the multivariate Gaussian approximation [25, 24], and we present its independent proof in Appendix F.5 for completeness. Lemma 10. Fix constants C∗ , R∗ < ∞. Consider deterministic matrices Mk ∈ Rd×d satisfying maxk ∥Mk ∥op ≤ (d) (d) C∗ and maxk √1d ∣Tr(Mk )∣ ≤ R∗ . Let Gijk = √1d xTi Mk xj , where x ∼ N (0, C ⊗ Id ), and let GG be a Gaussian

tensor having the same mean and covariance as G(d) . Then for every f ∈ C 3 (RL×L×K ) whose derivatives up to third order have at most polynomial growth, there exists a constant C = C(C∗ , R∗ , C, K, L, f ) independent of d and of the matrices Mk , such that C (d) ∣Ef (G(d) ) − Ef (GG )∣ ≤ √ . d F.1.2

(535)

Local well-posedness

We first introduce an extension of the vector field in (32) outside the set of admissible covariance parameters. Define Q ∶= {(q (1) , q (2) ) ∈ SK × SK ∶ q (1) + q (2) ⪰ 0, q (1) − q (2) ⪰ 0} . (536) Let P ∶ RK×K × RK×K Ð→ Q denote the orthogonal projection onto the closed convex cone Q. In particular, ∥P(q) − P(q ′ )∥F ≤ ∥q − q ′ ∥F .

(537)

We next define the extended ODE for (32). Let q(µ) ∶= (q (1) (µ), q (2) (µ)) and q̂(µ) ∶= P(q(µ)). We set ̃ ⋆K (µ) ∶= m⋆K (̂ m q (µ)),

(538)

with the non-trainable components fixed at m∖K (0), and define (r)

̃ ⋆ (µ), q̂(µ))]kl , akl (µ) ∶= [∇q(r) Φ (m

r = 1, 2.

(539)

ψkext (µ) ∶= 2 ∑ ∑ akl (µ) [µ(M (r) ,B ) + µ(A ,M (r) ) ] ,

(540)

ψ ext (µ) . βkext (µ) ∶= − A k B tk (µ) + tk (µ)

(541)

For k ∈ K, let

K

2

(r)

l

l=1 r=1

and

k

k

l

The extended vector field Vext is defined by w

w

2

(r)

(542)

dµ = Vext (µ) dτ

(543)

i=1 l=1 r=1 ki ∈K

i=1 ki ∈K

We refer to

K

diff [Vext (µ)]α = − ∑ βkext (µ)Tibias (µ) − 2 ∑ ∑ ∑ aki l (µ)Ti,l,r (µ) − 2γ∣α∣K µα . i

as an extension of the original ODE (32). Whenever q(µ) ∈ Q, the projection acts trivially, q̂(µ) = q(µ). Hence, on the admissible covariance set, ̃ ⋆K (µ) = m⋆K (q(µ)), and (543) agrees with the original slow system (29)–(32). m For s > 0, define the open subset Us ∶= {µ ∈ Xs ∶ P(q(µ)) ∈ Q, 64

B min(tA k (µ) + tk (µ)) > 0} . k∈K

(544)

Lemma 11. Let s0 > 0 and µ0 ∈ Us0 . Under Assumptions 4 and 5, there exist Tlocal > 0 and a strictly increasing continuous function s ∶ [0, Tlocal ) → (s0 , ∞) with s(0) = s0 such that (543) admits a unique solution µ(τ ) ∈ Xs(τ ) for 0 ≤ τ < Tlocal satisfying µ(τ ) ∈ Us(τ ) . Proof. Fix s0 > 0 and µ0 ∈ Us0 , and write B d0 ∶= min (tA k (µ0 ) + tk (µ0 )) > 0.

q0 ∶= P(q(µ0 )),

k∈K

(545)

Since q0 ∈ Q, we may choose a neighborhood Vq of q0 in Q whose closure is contained in Q. We first prove the local regularity of the slaved mean m⋆K (q). By Assumption 5, for every q ∈ Vq the map mK ↦ Φ(m, q) is uniformly strongly convex on M and admits a critical point m⋆K (q) ∈ M, which is therefore unique. Moreover, on compact subsets of M × Vq there exists Lq < ∞ such that ∥∇mK Φ(m, q) − ∇mK Φ(m, q ′ )∥ ≤ Lq ∥q − q ′ ∥F .

(546)

⟨∇mK Φ(m, q) − ∇mK Φ(m′ , q), m − m′ ⟩ ≥ λ0 ∥m − m′ ∥2 .

(547)

Uniform strong convexity gives

Applying (546) and (547) to m = m⋆ (q) and m′ = m⋆ (q ′ ) yields ∥m⋆K (q) − m⋆K (q ′ )∥ ≤

Lq ∥q − q ′ ∥F . λ0

(548)

Hence m⋆K is locally Lipschitz on Vq . Since P is non-expansive and the finite-order coordinates of µ are continuous on Xs , there exists a neighborhood Vµ of µ0 in Xs0 such that for every µ ∈ Vµ , P(q(µ)) ∈ Vq ,

B tA k (µ) + tk (µ) ≥

d0 , 2

k ∈ K.

(549)

(r)

It follows from (548), Assumption 4.1, and the definition that akl (µ) is locally bounded and locally Lipschitz on Vµ . The same is true for ψkext (µ), and by the lower bound (549), for βkext (µ). Then following Lemma 1, we can verify that Vext is locally bounded and locally Lipschitz from Xs to Xs′ . The Ovsyannikov theorem therefore yields Tlocal > 0 and a strictly increasing continuous scale s(τ ) for which (543) has a unique solution µ(τ ) ∈ Xs(τ ) on [0, Tlocal ). F.1.3

Uniform spectral bounds

Now we prove that the SGD trajectory does not blow up, i.e. the spectral norm of the weights remains uniformly bounded across the training steps. Step 1: A lower bound for the traces. Lemma 12. Fix T > 0, and let NT ∶= ⌊ Tαdd ⌋. Let J be any stopping time such that, on {n < J}, (n)

(n)

max {∥Uk ∥op , ∥Vk

1≤k≤K

∥op } ≤ C∗ ,

∥m(n) ∥∞ ≤ R∗

(550)

for some constants C∗ , R∗ < ∞ independent of d, T, γ. Under the conditions of Theorem 3, there exists cD > 0, independent of d, such that ⎛ ⎞ A,(n) B,(n) lim P ⎜ inf (tk + tk ) ≥ cD ⎟ = 1. d→∞ T ∧J ⎝0≤n<N ⎠ k∈K

65

(551)

Proof. Fix k ∈ K, and suppress the index k when there is no ambiguity. Let (n)

(n)

gU ∶= ∇U L(G(n) ),

gV

∶= ∇V L(G(n) )

(552)

denote the stochastic gradients. The SGD updates can be written as (n)

(n)

(553)

⟨gU , U (n) ⟩ = ∑ ∂ijk L(G(n) )Gijk = ⟨gV , V (n) ⟩ .

(554)

U (n+1) = ad U (n) − αd gU ,

V (n+1) = ad V (n) − αd gV ,

where ad ∶= 1 − αdd γ . The key observation is that L

(n)

F

(n)

Therefore, defining δk

A,(n)

∶= tk

B,(n)

− tk

(n)

(n)

F

i,j=1

, (553) gives the recursion

δ (n+1) = a2d δ (n) + r(n) ,

r(n) ∶=

αd2 (n) (n) (∥gU ∥2F − ∥gV ∥2F ) . d

(555)

On {n < J}, the bounds in (550) imply uniform bounds on all moments of G(n) . By Assumption 4.1, the same (n) (n) (n) (n) holds for the polynomially growing first derivatives of L. Since gU = √1d ∑L )xi (xj )T V (n) , i,j=1 ∂ijk L(G (n)

and similarly for gV , the estimates in Lemma 2 give (n)

(n)

1{n<J} E [∥gU ∥2F + ∥gV ∥2F ∣ Fn ] ≤ Cd,

(556)

where C depends on C∗ , R∗ , C, L, K, L, but not on d or n. Unrolling (555), for every m < NT ∧ J, m−1

2(m−1−n) (n)

(0) δ (m) − a2m = ∑ ad d δ

r

.

(557)

n=0

Hence NT −1

sup

m<NT ∧J

(0) ∣δ (m) − a2m ∣ ≤ ∑ 1{n<J} ∣r(n) ∣. d δ

(558)

n=0

Using (556), NT −1

E [ ∑ 1{n<J} ∣r(n) ∣] ≤ n=0

αd2 NT −1 ∑ Cd ≤ CT dαd Ð→ 0, d n=0

(559)

because αd ≤ c0 (d log d)−1 . Markov’s inequality proves max

(n)

sup

k∈K 0≤n<NT ∧J

∣δk

(0)

− a2n → 0. d δk ∣ Ð P

(560)

Since K = Θ(1), the convergence holds simultaneously for all k ∈ K. Finally, Assumptions 4.3–4.4 imply (0)

P (min ∣δk ∣ ≥ k∈K

cbal ) Ð→ 1. 2

(561)

Moreover, uniformly for n ≤ NT ,

αd γ 2n ) ≥ e−3γT d for all sufficiently large d. Combining this with (560), with probability tending to one, cbal −3γT (n) inf ∣δ ∣ ≥ e . 0≤n<NT ∧J k 4 a2n d = (1 −

(562)

(563)

k∈K

Since

A,(n)

tk

B,(n)

+ tk

A,(n)

≥ ∣tk

we obtain (551), for example with cD ∶= cbal e−3γT . 4 66

B,(n)

− tk

(n)

∣ = ∣δk ∣,

(564)

Step 2: Finite-dimensional Stein expansion. Lemma 13. Fix a compact set D ⋐ M × Q and a constant C∗ < ∞. Let J be a stopping time such that, on {n < J}, (n) (n) (n) (mK , q (1,n) , q (2,n) ) ∈ D, max {∥Ul ∥op , ∥Vl ∥op } ≤ C∗ . (565) 1≤l≤K

(n)

Write En [⋅] = E[⋅ ∣ Fn ], and define bk (2)

(1)

(n) ∶= ∑L )]. For r = 1, 2, set Cij,i′ j ′ ∶= Cii′ Cjj ′ and i,j=1 Cij En [∂ijk L(G (r)

(r,d,n)

2 (n) )]. Cij,i′ j ′ ∶= Cij ′ Cji′ and define akl ∶= 21 ∑L i,j,i′ ,j ′ =1 Cij,i′ j ′ En [∂(ijk),(i′ j ′ l) L(G Then, uniformly for n < J and k ∈ K, (n)

bk 2 K 2 (r,d,n) (r,n) (n) (n) (n) (n) ḡU,k = √ Vk + ∑ ∑ akl Ml Vk + EU,k , d l=1 r=1 d

(566)

(n)

bk 2 K 2 (r,d,n) (r̄,n) (n) (n) (n) (n) Uk + ∑ ∑ akl ḡV,k = √ Ml Uk + EV,k . d l=1 r=1 d (r,d,n)

(567) (n)

Moreover, there exists C = C(C∗ , D, C, L, K, L) < ∞ such that maxk,l,r ∣akl ∣ + maxk ∣bk ∣ ≤ C, (n) (n) (n) ∂Φ (n) (1,n) (2,n) −3/2 )∣ ≤ √C . ,q maxk∈K {∥EU,k ∥op , ∥EV,k ∥op } ≤ Cd and max1≤k≤K ∣bk − ∂mk (m , q d

Proof. Throughout the proof, we condition on Fn and suppress the superscript (n). Conditional on Fn , all weights are deterministic and x = (x1 , . . . , xL ) ∼ N (0, C ⊗ Id ). Write fijk (x) ∶= ∂ijk L(G(x)). The stochastic gradients are 1 L gU,k = √ ∑ fijk (x)xi xTj Vk , d i,j=1

(568)

1 L gV,k = √ ∑ fijk (x)xj xTi Uk . d i,j=1

(569)

Similarly to the expansion used in Lemma 2, we can obtain 1 L,L,L,L,K 1 L 2 T ḡU,k = √ ∑ Cij E[fijk ]Vk + E [∂(ijk),(i ′ j ′ l) L(G)] (Cii′ Cjj ′ Ml + Cij ′ Cji′ Ml ) Vk + EU,k . ∑ d ′ ′ d i,j=1 i,j,i ,j =1 l=1

and similarly for gV,k . This proves (566) and (567). The remainders are controlled following Lemma 2, which gives ∥EU,k ∥op + ∥EV,k ∥op ≤ Cd−3/2 . We now identify the coefficients. Let GG be the Gaussian tensor having the same mean and covariance as G. Then we have L ∂Φ = ∑ Cij E [∂ijk L(GG )] . (570) ∂mk i,j=1 and [∇q(1) Φ]kl =

L 1 2 ∑ Cii′ Cjj ′ E [∂(ijk),(i′ j ′ l) L(GG )] , 2 i,j,i′ ,j ′ =1

(571)

[∇q(2) Φ]kl =

L 1 2 ∑ Cij ′ Cji′ E [∂(ijk),(i′ j ′ l) L(GG )] . 2 i,j,i′ ,j ′ =1

(572)

by Price’s theorem. Applying Lemma 10 finishes the proof. Step 3: One-step increment.

For any process Y (n) , we define ∆Y (n) ∶= Y (n+1) − Y (n) .

67

Lemma 14. Under the conditions of Theorem 3, let J be a stopping time satisfying (565). Set g (n) ∶= A,(n) B,(n) + tk ) , and ϑ(n) ∶= (tA,(n) , tB,(n) , q (1,n) , q (2,n) ). ∇mK Φ (m(n) , q (1,n) , q (2,n) ) and D(n) ∶= diag (tk k∈K Then, uniformly on {n < J}, e(n) (n) (n) En [∆mK ] = −αd D(n) (g (n) + √ ) + rm , d

(573)

(n)

d where ∥e(n) ∥ ≤ C and ∥rm ∥ ≤ C ( (1+γ)α + αd2 ). Moreover, d

(n)

and

(n)

2

En [∥∆mK − En [∆mK ]∥ ] ≤ Cαd2 ,

(574)

αd (1 + γ)αd ∥En [∆ϑ(n) ]∥ ≤ C ( √ ∥g (n) ∥ + ), d d

(575)

αd2 . d Here C = C(C∗ , D, C, L, K, L) < ∞ is independent of d, n, T , and γ. En [∥∆ϑ(n) ∥2 ] ≤ C(1 + γ)2

(576)

Proof. Throughout the proof, we condition on Fn and suppress the superscript (n). Thus, for any state variable Y (n) , we write Y ∶= Y (n) , Y + ∶= Y (n+1) and ∆Y ∶= Y + − Y . We write ad ∶= 1 − αdd γ . All constants below are uniform over the stopped class (565). The spectral bounds (565) imply max1≤l≤K {∥Ml ∥op , ∥Al ∥op , ∥Bl ∥op } ≤ C∗2 . Together with the compactness assumption on (mK , q), this gives uniform bounds on all fixed moments of G. Assumption 4.1 therefore gives uniform bounds on all fixed moments of the derivatives of L(G). We shall repeatedly use the following elementary conditional Gaussian estimates. If P is Fn -measurable and ∥P ∥op ≤ C, then, for every fixed polynomially growing function F (G) arising below, 2 1 En [∣ F (G)xTi P xj ∣ ] ≤ C, d

and

2 1 C En [∣ 3/2 F (G)xTi P xj ∣ ] ≤ , d d

p

(r)

En ∣xTi Ml xj ∣ ≤ Cp dp/2 , r = 1, 2, √ √ which follows from ∥Ml ∥op = O(1), ∥Ml ∥F = O( d), and TrMl = O( d). For k ∈ K, by definition we have

(577)

(578)

T T ) + αd2 gU,k gV,k . Mk+ = a2d Mk − αd ad (gU,k VkT + Uk gV,k

(579)

α2 αd ad T T En [∆mk ] = (a2d − 1)mk − √ [Tr(ḡU,k VkT ) + Tr(Uk ḡV,k )] , )] + √d En [Tr(gU,k gV,k d d

(580)

Consequently,

(r,d)

where ḡU,k ∶= En [gU,k ], ḡV,k ∶= En [gV,k ]. Let bk and akl (d)

ψk

K

2

(r,d)

∶= 2 ∑ ∑ akl l=1 r=1

be the coefficients from Lemma 13, and define

[µ(M (r) ,B ) + µ(A ,M (r) ) ] . l

k

k

l

(581)

Lemma 13 gives (d)

ψ 1 T √ [Tr(ḡU,k VkT ) + Tr(Uk ḡV,k )] = Dk bk + √k + ρk , d d

68

(582)

C B where Dk ∶= tA k + tk and ∣ρk ∣ ≤ d . Set

ζk ∶=

√

d [bk −

∂Φ (m, q (1) , q (2) )] . ∂mk

(583)

Lemma 13 gives ∣ζk ∣ ≤ C. Moreover, since Ak , Bk ⪰ 0, ∣µ(M (r) ,B ) ∣ ≤ ∥Ml ∥op tB k, (r,d)

The coefficients akl

∣µ(A ,M (r) ) ∣ ≤ ∥Ml ∥op tA k.

k

l

k

l

(584)

(d)

are uniformly bounded (Lemma 13), and hence ∣ψk ∣ ≤ CDk . Define (d)

ek ∶= ζk +

ψk , Dk

(585)

T where the ratio is set equal to zero when Dk = 0. Thus we have ∣ek ∣ ≤ C. Finally, define χk ∶= √1d En [Tr(gU,k gV,k )]. Expanding the two gradients and using (578) gives ∣χk ∣ ≤ C. ζk ∂Φ +√ and (582) into (580), we obtain Substituting bk = ∂m k d

En [∆mk ] = −αd Dk (

∂Φ ek + √ ) + rm,k , ∂mk d

(586)

where rm,k = (a2d − 1)mk + αd (1 − ad )Dk ( By 1 − ad = αdd γ , ∣a2d − 1∣ ≤ 2αdd γ +

α2d γ 2 d2

∂Φ ek + √ ) − αd ad ρk + αd2 χk . ∂mk d

(587)

and (565), ∣rm,k ∣ ≤ C (

(1 + γ)αd + αd2 ) d

(588)

for all sufficiently large d. Summing over the finitely many trainable indices proves (573). Now we define 1 T Lk ∶= √ [Tr(gU,k VkT ) + Tr(Uk gV,k )] , d Thus

1 T Qk ∶= √ Tr(gU,k gV,k ). d

∆mk = (a2d − 1)mk − αd ad Lk + αd2 Qk .

(589)

(590)

By (577), En ∣Lk ∣ ≤ C. The argument used for χ also gives En ∣Qk ∣ ≤ C. Therefore, using ad ≤ 1 for all sufficiently large d, 2 En [∣∆mk − En [∆mk ]∣ ] ≤ Cαd2 + Cαd4 ≤ Cαd2 . (591) 2

2

Since ∣K∣ = O(1), this proves (574). Recall that for a trainable index k, T T ) + αd2 gU,k gV,k Mk+ = a2d Mk − αd ad (gU,k VkT + Uk gV,k ,

(592)

T T ) + αd2 gU,k gU,k A+k = a2d Ak − αd ad (gU,k UkT + Uk gU,k , + 2 T T 2 T Bk = ad Bk − αd ad (gV,k Vk + Vk gV,k ) + αd gV,k gV,k .

(593) (594)

The preceding expansions imply, for every Zk ∈ {Mk , MkT , Ak , Bk }, (1 + γ)αd αd + αd2 ) . ∥En [∆Zk ]∥op ≤ C ( √ ∥g∥ + d d 69

(595)

They also give

En [∥∆Zk ∥2F ] ≤ C(1 + γ)2 αd2 d.

(596)

Using (595), (596), and d1 ∣Tr(P Q)∣ ≤ d1 ∥P ∥F ∥Q∥F , we obtain αd (1 + γ)αd ∥En [∆ϑ]∥ ≤ C ( √ ∥g∥ + + αd2 ) . d d

(597)

The learning-rate condition in Theorem 3 implies αd2 ≤ Cαd /d for all sufficiently large d, and hence proves (575). It remains to estimate the second moments. Expanding (592)–(594) shows that ∆ϑ = αd Λ + Rϑ ,

(598)

1 ∂ijk L(G) xTi P xj for an Fn -measurable where each component of Λ is a finite sum of terms of the form d3/2 matrix P satisfying ∥P ∥op ≤ C. Thus (577) gives

En ∥Λ∥2 ≤

C . d

(599)

The remainder Rϑ consists of the weight-decay terms and terms containing at least two weight increments. Similarly we obtain α2 En [∥Rϑ ∥2 ] ≤ C(1 + γ)2 ( 2d + αd4 ) . (600) d Since αd2 ≤ C/d for all sufficiently large d, (598)–(600) give En [∥∆ϑ∥2 ] ≤ C(1 + γ)2

αd2 . d

(601)

This proves (576). Step 4: A bound for the mean variable. Lemma 15. Under the conditions of Theorem 3, let J be a stopping time satisfying (565). Write q (n) ∶= (q (1,n) , q (2,n) ) and NT ∶= ⌊ Tαdd ⌋. Then there exist constants Cloc = Cloc (C∗ , D, C, L, K, L) < ∞ and cloc = cloc (T, γ, cbal , λ0 ) > 0, independent of d, such that for any ϵ > 0 lim P (

d→∞

sup

0≤n<NT ∧J

(n)

[∥mK − m⋆K (q (n) )∥ − Cloc e−cloc αd n ] > ϵ) = 0. +

(602)

Proof. Define the excess population loss E (mK , q) ∶= Φ(mK , q) − Φ(m⋆K (q), q).

(603)

Let Dq ⋐ Q denote the projection of D onto the q-coordinates. By the implicit function theorem and the positive definiteness of ∇2mK Φ (Assumption 5), the map q z→ m⋆K (q) is continuously differentiable on a neighborhood of Dq . The envelope identity gives ∇mK E (mK , q) = ∇mK Φ(m, q),

(604) ⋆

∇q E (mK , q) = ∇q Φ(m, q) − ∇q Φ(m (q), q).

(605)

Assumption 4.1 implies E has bounded second derivatives on a compact neighborhood of (565). Strong convexity and compactness give constants Λ < ∞ and CE < ∞ such that, uniformly over (565), λ0 Λ 2 2 ∥mK − m⋆K (q)∥ ≤ E (mK , q) ≤ ∥mK − m⋆K (q)∥ , 2 2 2 ∥∇mK Φ(m, q)∥ ≥ 2λ0 E (mK , q), CE ∥∇q E (mK , q)∥ ≤ CE ∥mK − m⋆K (q)∥ ≤ ∥∇mK Φ(m, q)∥ . λ0 70

(606) (607) (608)

The last inequality follows from ∥∇mK Φ(m, q)∥ ≥ λ0 ∥mK − m⋆K (q)∥. (n) A,(n) B,(n) (n) For k ∈ K, let Dk ∶= tk + tk and D(n) ∶= diag (Dk ) . The spectral bound in (565) implies (n)

0 ≤ Dk

k∈K

≤ 2C∗2 . By Lemma 12, there exists cD = cD (T, γ, cbal ) > 0 such that the event ⎧ ⎫ ⎪ ⎪ ⎪ ⎪ ⎪ ⎪ A,(n) B,(n) Bd ∶= ⎨ inf (tk + tk ) ≥ cD ⎬ ⎪ ⎪ 0≤n<N ∧J T ⎪ ⎪ ⎪ ⎪ ⎩ k∈K ⎭

(609)

(n)

satisfies P(Bd ) Ð→ 1. We work on Bd . Set En ∶= E (mK , q (n) ). For n < J, Taylor’s expansion and the boundedness of the second derivatives of E (according to Assumption 4.1) give (n)

(n)

(n)

En [En+1 − En ] = ⟨g (n) , En [∆mK ]⟩ + ⟨∇q E (mK , q (n) ) , En [∆q (n) ]⟩ + RE , where

(n)

(n)

∣RE ∣ ≤ C En [∥∆mK ∥2 + ∥∆q (n) ∥2 ] ≤ C(1 + γ)2 αd2 . (n)

Here and below En [⋅] = E[⋅ ∣ Fn ]. Using Lemma 14, Dk

(610) (611)

≤ 2C∗2 , and (608), we obtain

αd αd + C(1 + γ)2 αd2 . En [En+1 − En ] ≤ −αd (g (n) )T D(n) g (n) + C √ (∥g (n) ∥ + ∥g (n) ∥2 ) + C(1 + γ) d d

(612)

All quantities appearing here are uniformly bounded over D. Since D(n) ⪰ cD I for n < J, the d−1/2 -terms to be absorbed into the leading term for all sufficiently large d. Moreover, αd2 ≤ αd /d under the learning-rate assumption. Consequently, there exist c1 = c1 (T, γ, cbal , λ0 ) > 0 and C1 = C1 (T, γ, C∗ , D) < ∞ such that En [En+1 − En ] ≤ −

cD αd (n) 2 αd αd ∥g ∥ + C1 ≤ −c1 αd En + C1 , 2 d d

(613)

for n < J, where the second inequality uses (607). Define the stopped martingale differences ξ (n+1) ∶= 1{n<J} [En+1 − En − En [En+1 − En ]] .

(614)

Because E has bounded gradient, Lemma 14 implies En [∣ξ (n+1) ∣2 ] ≤ Cαd2 .

(615)

Thus, with ρd ∶= 1 − c1 αd , (613) gives, for n < J, En+1 ≤ ρd En + C1

αd + ξ (n+1) . d

(616)

To control this recursion uniformly over NT = Θ(d/αd ) steps, divide the time interval into blocks of length ℓd ∶= ⌈ α1d ⌉ and define Bd ∶= ⌈ NℓdT ⌉ ≤ CT d. For a block starting at s = bℓd , set Sb,r ∶= ∑s+r−1 ξ (j+1) for 0 ≤ r ≤ ℓd , j=s with the convention that increments after J are zero. Doob’s maximal inequality and (615) give P ( max max ∣Sb,r ∣ > εd ) ≤ 0≤b<Bd 0≤r≤ℓd

CBd ℓd αd2 , ε2d

(617)

where we choose εd ∶= (dαd )1/4 . Since Bd = O(d), ℓd = O(αd−1 ), and dαd → 0, Bd ℓd αd2 ≤ CT (dαd )1/2 Ð→ 0. ε2d

(618)

On the complementary high-probability event, Abel summation gives, for any s < m in the same block, RRRm−1 RR RRR m−1−j (j+1) RRR ξ RRR ≤ 2 max ∣Sb,r ∣ ≤ 2εd . RRR ∑ ρd 1≤r≤m−s RRR RR j=s 71

(619)

1 r Moreover, ρℓdd ≤ e−c1 αd ℓd ≤ e−c1 < 1. Iterating (616) block by block, using (619) and αdd ∑∞ r=0 ρd = c1 d , yields

1 En ≤ C0 ρnd + C (εd + ) , d

0 ≤ n < NT ∧ J,

(620)

where C0 ∶= sup(mK ,q)∈D E (mK , q) < ∞. Using (606) and ρnd ≤ e−c1 αd n , we conclude that, on an event whose probability tends to one, (n)

c1

∥mK − m⋆K (q (n) )∥ ≤ Cloc e− 2 αd n + C [(dαd )1/8 + d−1/2 ]

(621)

uniformly for 0 ≤ n < NT ∧ J. The term in square brackets tends to zero, which proves (602), with cloc = c1 /2. Removing Bd is legitimate because P(Bdc ) → 0. Step 5: A refined bound for the mean gradients. bound for the mean gradients lim P (

d→∞

sup

From Lemma 15, we can obtain the following

(n)

0≤n<NT ∧J

[∥bk ∥ − Cloc e−cloc αd n ] > ϵ) = 0

(622)

+

by the boundedness of the Hessian. However, we would require a refined bound given by the following lemma. Lemma 16. Under the conditions of Theorem 3, fix T > 0 and let J be a stopping time satisfying (565). (n) (n) (n) Set NT ∶= ⌊ Tαdd ⌋ and bK ∶= (bk ) , where bk is defined in Lemma 13. Then there exist constants k∈K Ctr = Ctr (T, γ, C∗ , D) < ∞, ctr = ctr (T, γ, C∗ , D) > 0, Cb = Cb (C∗ , D, C, L, K, L) < ∞, c⋆0 = c⋆0 (T, γ, C∗ , D) > 0 such that, whenever c0 ≤ c⋆0 , lim P (

d→∞

Cb (n) [∥bK ∥ − Ctr e−ctr αd n ] > √ ) = 0. + 0≤n<NT ∧J d

(623)

sup

Proof. Write q (n) ∶= (q (1,n) , q (2,n) ) and D(n) ∶= diag (tk

A,(n)

B,(n)

+ tk

)

k∈K

. All constants below are uniform over A,(n)

(565). By Lemma 12, there exists cD = cD (T, γ, cbal ) > 0 such that the event Bd ∶= {inf 0≤n<NT ∧J (tk k∈K

B,(n)

+ tk

satisfies P(Bd ) Ð→ 1. We work on Bd . Let H⋆ (q) ∶= ∇2mK Φ (m⋆ (q), q) and S(q) ∶= H⋆ (q)1/2 . By Assumption 5, the implicit function theorem, and Assumption 4.1, the maps q z→ m⋆K (q), q z→ H⋆ (q), q z→ S(q) are C 2 on a neighborhood of the projection of D onto the q-coordinates. In particular, on that neighborhood, λ0 I ⪯ H⋆ (q) ⪯ ΛI,

∥S(q)∥op + ∥S(q)−1 ∥op ≤ C,

(624)

and S, S −1 , and m⋆ are Lipschitz. Define (n)

x(n) ∶= mK − m⋆K (q (n) ),

y (n) ∶= S(q (n) )x(n) .

(625)

Taylor expansion gives g (n) ∶= ∇mK Φ (m(n) , q (n) ) = H⋆ (q (n) )x(n) + ρ(n) g ,

(n) 2 ∥ρ(n) ∥ . g ∥ ≤ C∥x

(626)

Fix r0 > 0, to be chosen sufficiently small below. For convenience, write Sn ∶= S(q (n) ), Define

Hn ∶= H⋆ (q (n) ) = Sn2 ,

A(n) ∶= Sn Dn Sn ,

Dn ∶= D(n) .

u(n) ∶= Sn−1 e(n) ,

(627) (628)

(n)

where e is the bounded vector appearing in Lemma 14. On the event Bd , we have cD I ⪯ Dn ⪯ 2C∗2 I. Together with (624), this yields a− I ⪯ A(n) ⪯ a+ I,

a− ∶= λ0 cD > 0, 72

a+ ∶= 2ΛC∗2 .

(629)

) ≥ cD }

Moreover,

sup ∥u(n) ∥ ≤ sup ∥S(q)−1 ∥op sup ∥e(n) ∥ =∶ Cu < ∞. q

n<J

(630)

n<J

We next derive the recursion for y (n) . Define F (mK , q) ∶= S(q) (mK − m⋆K (q)) , so that

(631)

(n)

By Lemma 14,

y (n) = F (mK , q (n) ).

(632)

e(n) (n) (n) , En [∆mK ] = −αd Dn (g (n) + √ ) + rm d

(633)

(n)

d where ∥rm ∥ ≤ C ( (1+γ)α + αd2 ). Furthermore, d

(1 + γ)αd αd ∥En [∆q (n) ]∥ ≤ C ( √ ∥g (n) ∥ + ), d d and En ∥∆q (n) ∥2 ≤ C(1 + γ)2 Also,

(634)

αd2 . d

(635)

(n)

En ∥∆mK ∥2 ≤ Cαd2 .

(636)

Since F is C 2 on a compact neighborhood before the stopping time, a second-order Taylor expansion gives (n)

(n)

(n)

En [y (n+1) − y (n) ] = Sn En [∆mK ] + ∂q F (mK , q (n) ) [En [∆q (n) ]] + ρF ,

(637)

where, uniformly before the stopping time, (n)

(n)

∥ρF ∥ ≤ CEn [∥∆mK ∥2 + ∥∆q (n) ∥2 ] ≤ C(1 + γ)2 αd2 . Using (626),

g (n) = Hn x(n) + ρ(n) g ,

(n) 2 ∥ρ(n) ∥ , g ∥ ≤ C∥x

(638) (639)

and the identities Sn Dn Hn x(n) = Sn Dn Sn y (n) = A(n) y (n) , Sn Dn e(n) = A(n) u(n) , we obtain, whenever n < J and ∥y (n) ∥ ≤ r0 , αd En [y (n+1) − y (n) ] = −αd A(n) y (n) − √ A(n) u(n) + ry(n) , (640) d where (1 + γ)αd αd + (1 + γ)2 αd2 ] . (641) ∥ry(n) ∥ ≤ C [αd ∥y (n) ∥2 + √ ∥y (n) ∥ + d d Here we used (624) and the fact that ∂q F is uniformly bounded before the stopping time. Define ξ (n+1) ∶= y (n+1) − En [y (n+1) ]. Then En [ξ

(n+1)

Consequently,

(642)

] = 0. Since F is Lipschitz before the stopping time, Lemma 14 and (636)– (635) imply En ∥ξ (n+1) ∥2 ≤ Cαd2 .

(643)

αd y (n+1) = (I − αd A(n) ) y (n) − √ A(n) u(n) + ξ (n+1) + ry(n) . d

(644)

We first use Lemma 15 to enter the region in which (644) can be iterated. Let CS ∶= supq ∥S(q)∥op < ∞, where the supremum is taken over a compact neighborhood of the q-projection of D. Choose t0 = t0 (T, γ, r0 ) 73

sufficiently large that CS Cloc e−cloc t0 < r80 and set n0 ∶= ⌈ αt0d ⌉. Applying Lemma 15 with ϵ = r0 /(8CS ), and using αd n0 ≥ t0 , gives r0 P ( sup ∥y (n) ∥ > ) Ð→ 0. (645) 4 n0 ≤n<NT ∧J Here and below, the supremum over an empty index set is understood to be zero. Starting from the entry time n0 , define the local exit time Jloc ∶= inf {n ≥ n0 ∶ ∥y (n) ∥ > r0 } .

(646)

All local estimates below are carried out up to NT ∧ J ∧ Jloc . We divide the time interval after n0 into blocks. Set 1 ⌉, sb ∶= n0 + bℓd , (647) ℓd ∶= ⌈ a− αd and, whenever sb < NT ∧ J ∧ Jloc , let Ib ∶= {sb , . . . , min (sb + ℓd − 1, NT ∧ J ∧ Jloc − 1)} to be the b−th block. We next analyze the dynamics on each block. For a block Ib , set Āb ∶= A(sb ) ,

Pb ∶= I − αd Āb .

(648)

By (629), for all sufficiently large d, 0 ⪯ Pb ⪯ (1 − a− αd )I, For n ∈ Ib , rewrite (644) as where

(n)

ρb For 0 ≤ r ≤ ∣Ib ∣, define

∥Pbj ∥op ≤ 1,

j ≥ 0.

(649)

αd (n) y (n+1) = Pb y (n) − √ Āb u(n) + ξ (n+1) + ρb , d

(650)

αd ∶= −αd (A(n) − Āb )y (n) − √ (A(n) − Āb )u(n) + ry(n) . d

(651)

r−1

(s +j)

Rb,r ∶= ∑ Pbr−1−j ρb b

(652)

j=0

with Rb,0 ∶= 0. Iterating (650) within the block gives r−1 1 r−1 y (sb +r) = Pbr y (sb ) − √ ∑ Pbr−1−j αd Āb u(sb +j) + ∑ Pbr−1−j ξ (sb +j+1) + Rb,r . d j=0 j=0

(653)

We first control the second term of (653). Diagonalize Āb = Qb diag(λ1 , . . . , λp )QTb , where p ∶= ∣K∣. By (629), λi ∈ [a− , a+ ]. For every coordinate i, using (630),

Consequently,

RRRr−1 R r−1 RRR (1 − α λ )r−1−j α λ (QT u(sb +j) ) RRRRR ≤ C r−1−j αd λi ≤ Cu . ∑ (1 − αd λi ) d i d i u b RRR ∑ R i RR j=0 RRj=0 RR

(654)

X X r−1 X X X √ X r−1−j (sb +j) X X X max X X P α Ā u X ≤ p Cu . ∑ d b X X b X X 0≤r≤∣Ib ∣ X X X X Xj=0 X

(655)

We next show that A(n) is approximately frozen on each block. Write ∆ϑ(n) = En [∆ϑ(n) ] + ζ (n+1) ,

En [ζ (n+1) ] = 0.

(656)

By Lemma 14, uniformly before the stopping time, αd (1 + γ)αd ∥En [∆ϑ(n) ]∥ ≤ C ( √ + ), d d 74

(657)

where we used the boundedness of g (n) on (565), while En ∥ζ (n+1) ∥2 ≤ C(1 + γ)2

αd2 . d

(658)

Since ℓd = ⌈ a−1αd ⌉ = O((a− αd )−1 ), the accumulated conditional drift over a single block is o(1). Moreover, Doob’s maximal inequality applied on each block, followed by a union bound over the O(d) blocks, yields max sup ∥ϑ(n) − ϑ(sb ) ∥ Ð → 0. P

b

(659)

n∈Ib

Since A(ϑ) = S(q)D(ϑ)S(q) is Lipschitz on the compact stopped region (565), it follows that δA,d ∶= max sup ∥A(n) − Āb ∥ Ð → 0. P

b

(660)

n∈Ib

We now bound the block remainder Rb,r (the last term of (653)). Recall from (641) that, on the region ∥y (n) ∥ ≤ r0 , αd (1 + γ)αd + (1 + γ)2 αd2 ] . (661) ∥ry(n) ∥ ≤ C [αd ∥y (n) ∥2 + √ ∥y (n) ∥ + d d Using (652), (649), (651), (630), and (661), we obtain, for every 0 ≤ r ≤ ∣Ib ∣, δA,d 1 + γ 1 ∥Rb,r ∥ ≤ Crαd (δA,d + r0 + √ ) max ∥y (sb +j) ∥ + Crαd [ √ + + (1 + γ)2 αd ] . d d 0≤j≤r d

(662)

Since r ≤ ∣Ib ∣ ≤ ℓd and αd ℓd ≤ C/a− , this implies ∥Rb,r ∥ ≤

C C δA,d 1 + γ 1 (r0 + δA,d + √ ) max ∥y (sb +j) ∥ + [√ + + (1 + γ)2 αd ] . a− a− d d 0≤j≤r d

(663)

1 . By (660), δA,d Ð Choose r0 = r0 (T, γ) > 0 sufficiently small so that aC− r0 ≤ 16 → 0, and therefore P

C 1 1 (δA,d + √ ) ≤ a− 16 d

(664)

with probability tending to one. Furthermore, for fixed T, γ, the learning-rate condition αd ≤ c0 (d log d)−1 gives δA,d 1 + γ oP (1) √ + + (1 + γ)2 αd = √ . (665) d d d Consequently, there exists a nonnegative random sequence εd Ð → 0 such that, with probability tending to one, 1 εd 0 ≤ r ≤ ∣Ib ∣, (666) ∥Rb,r ∥ ≤ max ∥y (sb +j) ∥ + √ , 8 0≤j≤r d P

uniformly over all blocks b. Finally, We control the martingale term (the third term of (653)). Truncate the fresh sample at step n (n) by An = {∥x(n) ∥2 ≤ cx d} ∩ {maxi,j,k ∣Gijk ∣ ≤ cG log d}. The probability of any truncation failure before NT ∧ J tends to zero. On the truncated event, ∥ξ (n+1) ∥ ≤ Cαd (1 + (log d)r )

(667)

for some fixed r < ∞. For each block, Freedman’s inequality applied to (653) gives P

RRR ⎛RRRRRr−1 η ⎞ cη 2 a− RRR ∑ (1 − αd λi )r−1−j (QTb ξ (sb +j+1) )i RRRRR > √ ≤ 2 exp [− ] dαd ⎝RRj=0 ⎠ R d R R R 75

(668)

for every fixed η > 0 and all sufficiently large d. Here we used 2

2j T (s +j+1) )i ∣ ] ≤ C ∑ (1 − αd λi ) Esb +j [∣(Qb ξ b

j≥0

αd . a−

(669)

We now make the estimate uniform over all blocks and all terminal times within each block. Let η0 ∶= √1p and p ∶= ∣K∣. The total number of pairs (b, r), with b ranging over the blocks and 0 ≤ r ≤ ∣Ib ∣, is at most NT + 1. Therefore, a union bound gives RRRr−1 RRR ⎛ η0 ⎞ cη 2 a− P max max max RRRRR ∑ (1 − αd λi )r−1−j (QTb ξ (sb +j+1) )i RRRRR > √ ≤ 2p(NT + 1) exp [− 0 ] . dαd ⎝ b 0≤r≤∣Ib ∣ 1≤i≤p RRj=0 RRR d⎠ R

(670)

d By the lower learning-rate bound αd = Ω(d−ι ), NT ≤ CT d1+ι , while dα1 d ≥ log by the upper learning-rate c0 bound. Hence

2p(NT + 1) exp [−

2 cη02 a− ] ≤ CT d 1+ι−cη0 a− /c0 . dαd

(671)

We may therefore choose c⋆0 = c⋆0 (T, γ, C∗ , D) > 0 sufficiently small so that cη02 a− > 1 + ι. c⋆0

(672)

r−1−j (sb +j+1) Whenever c0 ≤ c⋆0 , the right-hand side of (670) converges to zero. Let vb,r ∶= ∑r−1 ξ . Since j=0 Pb T Pb = Qb diag(1 − αd λ1 , . . . , 1 − αd λp )Qb , we have r−1

(QTb vb,r )i = ∑ (1 − αd λi )r−1−j (QTb ξ (sb +j+1) )i .

(673)

j=0

Hence, on the high-probability event obtained by (670), η0 max ∣(QTb vb,r )i ∣ ≤ √ . b,r,i d

(674)

Choosing η0 = p−1/2 and using the orthogonality of Qb , ∥vb,r ∥ = ∥QTb vb,r ∥ ≤ Therefore,

√

1 p max ∣(QTb vb,r )i ∣ ≤ √ . 1≤i≤p d

(675)

X X r−1 X X X 1 X r−1−j (sb +j+1) X X X max max X X P ξ X ≤√ ∑ X X b X b 0≤r≤∣Ib ∣ X X X d X X Xj=0 X

(676)

Cy max ∥y (sb +r) ∥ ≤ 2∥y (sb ) ∥ + √ , 0≤r≤∣Ib ∣ d

(677)

with probability tending to one. Combining (666), (655), and (676), we obtain, on an event whose probability tends to one,

where Cy = Cy (C∗ , D, C, L, K, L) is independent of γ. Moreover, for a full block, i.e. whenever sb +ℓd < NT ∧J ∧Jloc , we additionally have ∥Pbℓd ∥op ≤ (1−a− αd )ℓd ≤ −1 e . Evaluating (653) at the block endpoint and using (677) therefore gives ∥y (sb +ℓd ) ∥ ≤ ∥Pbℓd ∥op ∥y (sb ) ∥ +

C 1 C max ∥y (sb +j) ∥ + √ ≤ θ∥y (sb ) ∥ + √ , 8 0≤j≤ℓd d d

(678)

where we use (666) for the first inequality and choose θ ∶= e−1 + 14 ∈ (0, 1) for the second inequality. Iteration gives Cy √ . ∥y (sb ) ∥ ≤ θb ∥y (s0 ) ∥ + (679) (1 − θ) d 76

We now close the local stopping-time argument. At the entry time, ∥y (n0 ) ∥ ≤ r40 . By (677), for all sufficiently large d, max0≤r≤∣I0 ∣ ∥y (n0 +r) ∥ < r0 . Moreover, by (679), the value at the next block endpoint is again at most r0 /4 for all sufficiently large d. Repeating the same argument block by block shows that Jloc > NT ∧ J with probability tending to one. Hence the local stopping time Jloc may be removed. The number of blocks between n0 and n is bounded below by a− αd (n − n0 ) − 2. Hence (677)– (679) imply Cy′ ′ −c′tr αd n ∥y (n) ∥ ≤ Ctr e +√ , d

n0 ≤ n < NT ∧ J,

(680)

′ where c′tr = −a− log θ > 0. The constants Ctr and c′tr may depend on T, γ, whereas Cy′ is independent of γ. (n)

We finally convert the estimate on y (n) into the desired estimate on bK . By (626) and Lemma 13, C C (n) ∥bK ∥ ≤ ∥g (n) ∥ + √ ≤ C∥y (n) ∥ + C∥y (n) ∥2 + √ . d d

(681)

Combining this with (680), there exist constants C1 < ∞, Cb < ∞, and ctr ∶= c′tr > 0, such that, with probability tending to one, Cb (n) ∥bK ∥ ≤ C1 e−ctr αd n + √ , d

n0 ≤ n < NT ∧ J.

(682)

Here Cb can be chosen independently of γ. It remains to cover the initial interval 0 ≤ n < n0 ∧ J. By Lemma 13, there exists B0 < ∞, independent of d, such that (n) (683) sup ∥bK ∥ ≤ B0 . 0≤n<J

For all sufficiently large d, αd ≤ 1, and hence n < n0

Ô⇒

αd n < αd n0 ≤ t0 + αd ≤ t0 + 1.

(684)

Choose Ctr ≥ max {C1 , B0 ectr (t0 +1) }. Then, for every 0 ≤ n < n0 ∧ J, B0 ≤ Ctr e−ctr αd n . Combining this with (682) yields Cb (n) ∥bK ∥ ≤ Ctr e−ctr αd n + √ , 0 ≤ n < NT ∧ J, (685) d on an event whose probability tends to one. This proves (623). Removing Bd is legitimate because P(Bdc ) → 0. Step 6: Final spectral bounds Lemma 17. Fix T > 0. Let D ⋐ M × Q be a compact set containing the limiting trajectory, and define (n)

Jmac ∶= inf {n ≥ 0 ∶ (mK , q (1,n) , q (2,n) ) ∉ D} .

(686)

Under the conditions of Theorem 3, there exists a constant C∗ > C0 , independent of d, such that lim P (

d→∞

sup

(n)

(n)

max {∥Uk ∥op , ∥Vk

0≤n<NT ∧Jmac 1≤k≤K

∥op } ≤ C∗ ) = 1,

(687)

where NT ∶= ⌊ Tαdd ⌋. Proof. The weights with k ∉ K remain fixed throughout training, so it suffices to consider k ∈ K. Fix C∗ > C0 , to be chosen below, and define the spectral stopping time (n)

(n)

Jsp ∶= inf {n ≥ 0 ∶ max {∥Uk ∥op , ∥Vk 1≤k≤K

∥op } > C∗ } ,

and set J ∶= Jsp ∧ Jmac . All estimates below are uniform on {n < J}. 77

(688)

For k ∈ K, denote the stochastic gradients by (n)

(n)

gU,k ∶= ∇Uk L(G(n) ), (n)

gV,k ∶= ∇Vk L(G(n) ),

(n)

(n)

(689)

(n)

and their conditional means by ḡU,k ∶= E[gU,k ∣ Fn ] and ḡV,k ∶= E[gV,k ∣ Fn ]. Lemma 13 gives (n)

bk 1 (n) (n) (n) (n) (n) ḡU,k = √ Vk + HU,k Vk + EU,k , d d

(690)

(n)

bk 1 (n) (n) (n) (n) (n) ḡV,k = √ Uk + HV,k Uk + EV,k d d (n)

(n)

(n)

(691)

(n)

with maxk∈K {∥HU,k ∥op , ∥HV,k ∥op } ≤ CH and maxk∈K {∥EU,k ∥op , ∥EV,k ∥op } ≤ CE d−3/2 uniformly on {n < J}. Here CH , CE are constants independent of d and n. By Lemma 16: there exist constants Ctr (T, γ), ctr (T, γ) > 0 and Cb (T ) < ∞ such that, with probability tending to one, Cb (T ) (n) max ∣bk ∣ ≤ Ctr (T, γ)e−ctr (T,γ)αd n + √ , 0 ≤ n < NT ∧ J. (692) k∈K d Define (n) An ∶= {∥x(n) ∥2 ≤ cx d} ∩ {max ∣Gijk ∣ ≤ cG log d} , (693) i,j,k

where cx , cG are sufficiently large constants. Thus we have P (∃ n < NT ∧ J ∶ Acn ) Ð→ 0.

(694)

Denote the complementary high-probability event by Adata . (n) (n) (n) (n) Define the stopped and truncated gradients ̂ gU,k ∶= 1{n<J} gU,k 1An , ̂ gV,k ∶= 1{n<J} gV,k 1An , their condi(n) (n) (n) (n) (n) (n) (n) (n) tional means ̂ ḡU,k , ̂ ḡV,k , and the martingale differences ξ̂U,k ∶= ̂ gU,k − ̂ ḡU,k , ξ̂V,k ∶= ̂ gV,k − ̂ ḡV,k . Set ad ∶= 1 − αdd γ . Define the discounted martingale sums m−1

m−1

YU,k,m ∶= −αd ∑ am−1−n ξ̂U,k , d

YV,k,m ∶= −αd ∑ am−1−n ξ̂V,k . d

(n)

n=0

(n)

(695)

n=0

Similarly to Lemma 3 we have ⎞ ⎛ P ⎜ max max {∥YU,k,m ∥op , ∥YV,k,m ∥op } > δ ⎟ Ð→ 0. k∈K ⎠ ⎝1≤m≤N T

(696)

Denote the complementary event by Anoise . Also the truncation-tail corrections are small: (n)

1{n<J} max {∥E [gU,k 1Acn ∣ Fn ∥

op

(n)

, ∥E [gV,k 1Acn ∣ Fn ] ∥op } ≤ rd ,

(697)

where rd = o(d−M ) for every fixed M > 0. On Adata , for every m ≤ J and n < m unrolling the SGD recursion gives (m)

Uk

(0)

= am d Uk

αd m−1 αd m−1 m−1−n (n) (n) (n) (n) − √ ∑ am−1−n bk V k − HU,k Vk + YU,k,m + RU,k,m , ∑a d d n=0 d d n=0

(698)

m−1−n (CE d−3/2 + rd ) = o(1). The analogous where, uniformly for m ≤ J ∧ NT , ∥RU,k,m ∥op ≤ αd ∑m−1 n=0 ad (m)

expansion holds for Vk . We now bound the second term of (698). On the high-probability event where (692) holds, the first term of (692) contributes as αd m−1 C(T, γ) (n) √ ∑ am−1−n Ctr (T, γ)e−ctr (T,γ)αd n ∥Vk ∥op ≤ √ = o(1), d d n=0 d 78

(699)

d r uniformly in m ≤ J ∧ NT , where we use ∑n≥0 e−ctr (T,γ)αd n ≤ C(T, γ)/αd . Using αd ∑∞ r=0 ad = γ , the second term of (692) contributes as

αd m−1 Cb (T ) (n) Cb (T )C∗ √ ∑ am−1−n √ ∥Vk ∥op ≤ . d γ d n=0 d

(700)

The third term of (698) can be bounded as αd m−1 m−1−n (n) (n) CH C∗ ∥HU,k Vk ∥op ≤ . ∑ ad d n=0 γ

(701)

Combining (698)– (701), on Adata ∩ Anoise and the high-probability event (692), (m)

∥Uk

∥op ≤ C0 +

(Cb (T ) + CH )C∗ + δ + o(1), γ

m ≤ J ∧ NT .

(702)

(m)

The same estimate holds for Vk . 0 . By Assumption 4.5, choose γ sufficiently large (which might depend Choose C∗ ∶= 2C0 + 1 and δ ∶= C∗ −C 4 on T ) such that (Cb (T ) + CH )C∗ C∗ − C0 ≤ . (703) γ 4 Therefore, for all sufficiently large d, (m)

max max {∥Uk k∈K

(m)

∥op , ∥Vk

∥op } < C∗

(704)

for every m ≤ J ∧ NT on an event whose probability tends to one. Hence P (Jsp ≤ Jmac ∧ NT ) Ð→ 0.

(705)

This is precisely (687). F.1.4

Fast-dynamics convergence ̃

We now prove the convergence on the fast timescale. Fix a finite fast-time horizon T̃ > 0, and set NTf̃ ∶= ⌊ αTd ⌋. B (1) Let D0 ∶= diag (t̄A (0), q̄ (2) (0)) and let m̃K (t) be the solution of k (0) + t̄k (0))k∈K , q0 ∶= (q̄

dm̃K = −D0 ∇mK Φ (m̃(t), q0 ) , dt

(d)

m̃K (0) = lim mK (0), d→∞

(706)

with the non-trainable components fixed at their initial values. By Assumption 5, d Φ(m̃(t), q0 ) = −∇mK Φ(m̃(t), q0 )T D0 ∇mK Φ(m̃(t), q0 ) ≤ 0. dt

(707)

Hence the trajectory {m̃K (t) ∶ 0 ≤ t ≤ T̃} remains in a compact subset of M. Choose compact neighborhoods Df ⋐ M × Q and C∗ in Lemma 17 such that the deterministic trajectory (m̃K (t), q0 ), 0 ≤ t ≤ T̃, lies in the interior of Df . Define the stopping times (n)

Jf ∶= inf {n ≥ 0 ∶ (mK , q (1,n) , q (2,n) ) ∉ Df } , (n)

(n)

Jsp ∶= inf {n ≥ 0 ∶ max {∥Uk ∥op , ∥Vk k

∥op } > C∗ } ,

(708) (709)

and set J ∶= Jf ∧ Jsp . The spectral estimate of Lemma 17 gives P (Jsp ≤ Jf ∧ NTf̃ ) Ð→ 0. 79

(710)

We first show that all slow coordinates remain frozen on the fast timescale. Recall ϑ(n) ∶= (tA,(n) , tB,(n) , q (1,n) , q (2,n) ). Write (n+1) (n+1) ∆ϑ(n) = En [∆ϑ(n) ] + ζϑ , En [ζϑ ] = 0. (711) On {n < J}, Lemma 14 and the boundedness of ∇mK Φ on Df give αd (1 + γ)αd ∥En [∆ϑ(n) ]∥ ≤ C ( √ + ) d d

(712)

and

αd2 . d Therefore the accumulated conditional drift over NTf̃ = O(αd−1 ) steps satisfies (n+1) 2

En [∥ζϑ

∥ ] ≤ C(1 + γ)2

(713)

1+γ ) = o(1). d

(714)

2⎤ ⎡ m−1 ⎢ α2 αd (n+1) ⎥ Ð→ 0. ∥ ⎥⎥ ≤ CNTf̃ d ≤ CT̃ E ⎢⎢ sup ∥ ∑ 1{n<J} ζϑ d d ⎥ ⎢m≤N f̃ n=0 ⎦ ⎣ T

(715)

m−1

sup ∥ ∑ 1{n<J} En [∆ϑ(n) ]∥ ≤ CT̃ (d−1/2 +

m≤N f̃ T

n=0

On the other hand, Doob’s maximal inequality and (713) imply

Consequently, ∥ϑ(n) − ϑ(0) ∥ Ð → 0.

(716)

∥ϑ(n) − ϑ̄(0)∥ Ð → 0.

(717)

P

sup 0≤n≤N f̃ ∧J T

Together with Assumption 4.3, this yields P

sup 0≤n≤N f̃ ∧J T

In particular, (∥D(n) − D0 ∥ + ∥q (n) − q0 ∥) Ð → 0. P

sup 0≤n≤N f̃ ∧J T

(718)

We next turn to the mean variables. Define F (m) ∶= −D0 ∇mK Φ(m, q0 ). By Lemma 14, on {n < J},

(n)

En [∆mK ] = αd F (m(n) ) + αd ρ(n) m ,

(719) (720)

where, using (718), the local Lipschitz continuity of ∇mK Φ, and the bound ∥e(n) ∥ ≤ C, ∥ρ(n) → 0. m ∥Ð P

sup 0≤n<N f̃ ∧J T

(721)

This is because 1+γ C (n) + αd ) . ∥ρ(n) − D0 ∥ + C∥q (n) − q0 ∥ + √ + C ( m ∥ ≤ C∥D d d Let

(n)

(n)

(n+1) ζm ∶= ∆mK − En [∆mK ].

(722) (723)

(n+1) 2

Again by Lemma 14, En [∥ζm

∥ ] ≤ Cαd2 . Hence Doob’s maximal inequality gives m−1

(n+1) sup ∥ ∑ 1{n<J} ζm ∥Ð → 0,

m≤N f̃ T

P

n=0

80

(724)

because its quadratic variation is bounded by CNTf̃ αd2 ≤ CT̃ αd Ð→ 0. Summing (720), we therefore obtain, uniformly for m ≤ NTf̃ ∧ J, (m)

mK

(0)

m−1

= mK + αd ∑ F (m(n) ) + oP (1),

(725)

n=0

where the error is uniform in m. Since F is Lipschitz on the compact neighborhood Df , the standard discrete error estimate, together with Assumption 4.3 and the Grönwall inequality (see the proof of Theorem 4), gives (n)

sup 0≤n≤N f̃ ∧J T

∥mK − m̃K (nαd )∥ Ð → 0. P

(726)

It remains to remove the stopping time. By construction, the deterministic trajectory {(m̃K (t), q0 ) ∶ 0 ≤ t ≤ T̃} has positive distance from ∂Df . Equations (718) and (726) therefore imply P (Jf ≤ NTf̃ ∧ Jsp ) Ð→ 0.

(727)

Combining (710) and (727), we conclude that P (J ≤ NTf̃ ) Ð→ 0.

(728)

Hence the stopping may be removed from (726). Using the piecewise-constant interpolation (34), we obtain sup ∥m̃(d) (t) − m̃(t)∥∞ Ð → 0, P

(729)

̃] t∈[0,T

which proves (36). F.1.5

Moment expansion (n)

Step 1: One-step expansion (n)

(n)

(n)

(n)

(n)

(n)

(n)

For k ∈ K, define the centered gradient noises ξU,k ∶= gU,k − ḡU,k , ξV,k ∶= (n)

(n)

gV,k − ḡV,k , where ḡU,k ∶= En [gU,k ], ḡV,k = En [gV,k ]. Let ad ∶= 1 − αdd γ . For Zk ∈ {Mk , MkT , Ak , Bk }, define (n) (n) ⎧ −ad (ξU,k VkT + Uk (ξV,k )T ) , Zk = Mk , ⎪ ⎪ ⎪ ⎪ ⎪ (n) T (n) T ⎪ T ⎪ ⎪ ⎪−ad (Vk (ξU,k ) + ξV,k Uk ) , Zk = Mk , (n) Nk (Zk ) ∶= ⎨ (n) (n) ⎪ ⎪ −ad (ξU,k UkT + Uk (ξU,k )T ) , Zk = Ak , ⎪ ⎪ ⎪ ⎪ ⎪ (n) T (n) T ⎪ ⎪ ⎩−ad (ξV,k Vk + Vk (ξV,k ) ) , Zk = Bk . (n)

For k ∉ K, set Nk

(730)

(Zk ) = 0. Define the finite-dimensional raw vector field as w

(d,n) [Vraw ]α ∶= − ∑

i=1 ki ∈K

√

(n)

w

K

2

(r,d,n)

d bki Tibias (µ(n) ) − 2 ∑ ∑ ∑ aki l

diff Ti,l,r (µ(n) ) − 2γ∣α∣K µ(n) α .

(731)

i=1 l=1 r=1 ki ∈K

Lemma 18. Under the conditions of Theorem 3, let J be a stopping time satisfying (565). Fix a finite word α = (Z1 , . . . , Zw ) ∈ Z w , and let ki denote the index carried by the letter Zi . Then, uniformly on {n < J}, αd (d,n) (n) [V ]α + αd ηα(n) + Rα µ(n+1) − µ(n) , (732) α α = d raw where w 1 (n) (n) (n) (n) (n) ηα(n) ∶= ∑ Tr [Z1 ⋯Zi−1 Nki (Zi )Zi+1 ⋯Zw ] (733) d i=1 ki ∈K

(n) is a martingale difference: En [ηα ] = 0. Moreover, for every fixed word α, there exists Cα < ∞, independent

of d and n, such that 1{n<J} En [∣ηα(n) ∣2 ] ≤

Cα , d

and (n) 1{n<J} En [∣Rα ∣] ≤ Cα [(1 + γ)2 αd2 +

81

(734) αd ]. d3/2

(735)

Proof. Throughout the proof we condition on Fn and suppress the superscript (n). Thus Uk , Vk , Mk , Ak , Bk are deterministic conditional on Fn , while the stochastic gradients depend on the fresh sample x(n) . All estimates below are uniform over (565). For k ∈ K, the SGD recursion can be written as Uk+ = ad Uk − αd gU,k and Vk+ = ad Vk − αd gV,k . Therefore T T ) + αd2 gU,k gV,k Mk+ = a2d Mk − αd ad (gU,k VkT + Uk gV,k ,

(736)

T T ) + αd2 gU,k gU,k A+k = a2d Ak − αd ad (gU,k UkT + Uk gU,k , T T ) + αd2 gV,k gV,k . Bk+ = a2d Bk − αd ad (gV,k VkT + Vk gV,k

(737) (738)

Write gU,k = ḡU,k + ξU,k and gV,k = ḡV,k + ξV,k . By Lemma 13,

where

bk 2 K 2 (r,d) (r) ḡU,k = √ Vk + ∑ ∑ akl Ml Vk + EU,k , d l=1 r=1 d

(739)

2 K 2 (r,d) (r̄) bk ḡV,k = √ Uk + ∑ ∑ akl Ml Uk + EV,k , d l=1 r=1 d

(740)

(r,d)

∣bk ∣ + max ∣akl l,r

∣ ≤ C,

∥EU,k ∥op + ∥EV,k ∥op ≤ Cd−3/2 .

(r̄)

(741)

(r)

Substituting (739)–(740) into (736), using (Ml )T = Ml , we obtain αd bk 2αd γ 2αd K 2 (r,d) (r) (r) Mk + αd Nk (Mk ) + Rk (Mk ). (742) ∆Mk = − √ (Ak + Bk ) − ∑ ∑ akl (Ml Bk + Ak Ml ) − d d d l=1 r=1 Taking the transpose gives αd bk 2αd K 2 (r,d) 2αd γ T (r̄) (r̄) ∆MkT = − √ (Ak + Bk ) − Mk + αd Nk (MkT ) + Rk (MkT ). ∑ ∑ akl (Bk Ml + Ml Ak ) − d d d l=1 r=1 (743) Likewise, substituting (739) into (737) gives αd bk 2αd K 2 (r,d) 2αd γ (r) (r̄) (Ml MkT + Mk Ml ) − Ak + αd Nk (Ak ) + Rk (Ak ), (744) ∆Ak = − √ (Mk + MkT ) − ∑∑a d l=1 r=1 kl d d whereas (740) and (738) give 2αd γ αd bk 2αd K 2 (r,d) (r̄) (r) T Bk + αd Nk (Bk ) + Rk (Bk ). (745) ∆Bk = − √ (Mk + MkT ) − ∑ ∑ akl (Ml Mk + Mk Ml ) − d d d l=1 r=1 We next make the remainders Rk (Zk ) explicit. For Zk = Mk , ⎡ ⎤ ⎢ bk 2 K 2 (r,d) (r) (r) ⎥ Rk (Mk ) =αd (1 − ad )⎢⎢ √ (Ak + Bk ) + ∑ ∑ akl (Ml Bk + Ak Ml ) ⎥⎥ d l=1 r=1 ⎢ d ⎥ ⎦ ⎣ αd2 γ 2 T 2 T T − αd ad (EU,k Vk + Uk EV,k ) + 2 Mk + αd gU,k gV,k , d

(746)

α2 γ 2

which follows directly from a2d −1 = − 2αdd γ + dd2 and 1−ad = αdd γ . Moreover, Rk (MkT ) = Rk (Mk )T . Similarly, ⎤ ⎡ ⎢ bk 2 K 2 (r,d) (r) (r̄) ⎥ T T ⎢ Rk (Ak ) = αd (1 − ad )⎢ √ (Mk + Mk ) + ∑ ∑ akl (Ml Mk + Mk Ml ) ⎥⎥ d l=1 r=1 ⎢ d ⎥ ⎣ ⎦ T )+ − αd ad (EU,k UkT + Uk EU,k

αd2 γ 2 T Ak + αd2 gU,k gU,k , d2

82

(747)

and ⎡ ⎤ ⎢ bk 2 K 2 (r,d) (r̄) (r) ⎥ Rk (Bk ) = αd (1 − ad )⎢⎢ √ (Mk + MkT ) + ∑ ∑ akl (Ml Mk + MkT Ml ) ⎥⎥ d l=1 r=1 ⎢ d ⎥ ⎣ ⎦ 2 2 γ α T T ) + d 2 Bk + αd2 gV,k gV,k . − αd ad (EV,k VkT + Vk EV,k d

(748)

We now estimate these remainders in the form needed for the moment expansion. Let P be any Fn measurable d × d matrix satisfying ∥P ∥op ≤ CP . On the stopped (565), all letters satisfy max {∥Ml ∥op , ∥Al ∥op , ∥Bl ∥op } ≤ C∗2 .

1≤l≤K

(749)

Hence every fixed product of the letters appearing in (746)– (748) has uniformly bounded operator norm, i.e., d1 ∣Tr(P Q)∣ ≤ ∥P ∥op ∥Q∥op ≤ C for every such product Q. Since 1 − ad = αdd γ , the terms are therefore bounded by

Likewise,

C C αd (1 − ad ) ( √ + ) ≤ Cαd2 γ (d−3/2 + d−2 ) . d d

(750)

αd2 γ 2 1 αd2 γ 2 ∣Tr(P Z )∣ ≤ C . k d2 d d2

(751)

For the remainders, Lemma 13 gives ∥EU,k ∥op + ∥EV,k ∥op ≤ Cd−3/2 . Since ∥Uk ∥op , ∥Vk ∥op ≤ C∗ , we obtain, for example, 1 ∣Tr(P EU,k VkT )∣ ≤ ∥P EU,k VkT ∥op ≤ CP ∥EU,k ∥op ∥Vk ∥op ≤ Cd−3/2 . d

(752)

The same estimate applies to all the other Stein-remainder terms, and hence their total contribution is bounded by Cαd d−3/2 . Finally, consider the quadratic-gradient terms. For instance, by Cauchy–Schwarz and (556), CP 1 1/2 1/2 T ]∣] ≤ (En ∥gU,k ∥2F ) (En ∥gV,k ∥2F ) ≤ C. En [ ∣Tr [P gU,k gV,k d d The same argument gives

1 1 T T ]∣ + ∣Tr [P gV,k gV,k ]∣] ≤ C. En [ ∣Tr [P gU,k gU,k d d

(753)

(754)

Thus the quadratic-gradient part of the letter remainder contributes at most Cαd2 . Combining (750)– (753), we conclude that, uniformly on {n < J}, 1 αd En [∣ Tr [P Rk (Zk )]∣] ≤ CP [(1 + γ)2 αd2 + 3/2 ] , d d

Zk ∈ {Mk , MkT , Ak , Bk }.

(755)

+ Fix α = (Z1 , . . . , Zw ). Writing Zi+ = Zi + ∆Zi in µ+α = d1 Tr[Z1+ ⋯Zw ] and expanding the product gives w 1 µ+α − µα = ∑ Tr [Z1 ⋯Zi−1 ∆Zi Zi+1 ⋯Zw ] + Rα,mult , i=1 d

(756)

where Rα,mult is the sum of all terms containing increments at two or more positions. Insert (742)– (745) into the first-order sum in (756). By the definition (31), we obtain µ+α − µα =

αd (d) [V ]α + αd ηα + Rα . d raw

(n)

(757)

ηα is defined in (733). By En [Nki (Zi )] = 0 we immediately obtain En [ηα ] = 0. Moreover, on (565) all letters satisfy maxZ∈Z ∥Z∥op ≤ C∗2 , which gives En ∣ηα ∣2 ≤ Cdα and proves (734). Rα is the sum of the first-order letter remainders and Rα,mult . 83

The final step is to control the first-order letter remainders and Rα,mult . By (755), the sum of all first-order letter remainders satisfies RRRR RRRR w R RRR αd 1 (758) En RRR ∑ Tr [Z1 ⋯Zi−1 Rki (Zi )Zi+1 ⋯Zw ]RRRRR ≤ Cα [(1 + γ)2 αd2 + 3/2 ] . d RRR RRR i=1 d RR RRki ∈K We next consider products containing at least two letter increments. The spectral bounds (565), polynomial growth of the derivatives of L, and Gaussian moment estimates imply, for every fixed p < ∞, En [∥∆Zk ∥pF ] ≤ Cp (1 + γ)p αdp dp/2 ,

Zk ∈ {Mk , MkT , Ak , Bk }.

(759)

The same estimate holds with the Frobenius norm replaced by the operator norm, since ∥ ⋅ ∥op ≤ ∥ ⋅ ∥F . Consider a term in Rα,mult containing r ≥ 2 increments. Hölder’s inequality gives 1 En ∣Tr [P0 ∆Zi1 P1 ⋯∆Zir Pr ]∣ ≤ Cα (1 + γ)r αdr d(r−2)/2 , d

(760)

√ where the Pj ’s are products of unchanged letters. Since the length w is fixed and αd d Ð→ 0, under the learning-rate condition of Theorem 3, for every 2 ≤ r ≤ w, √ αdr d(r−2)/2 = αd2 (αd d)r−2 ≤ Cαd2 (761) for all sufficiently large d. Summing over the finitely many multi-increment terms therefore gives En ∣Rα,mult ∣ ≤ Cα (1 + γ)2 αd2 . (n)

Combining (758) and (762), and defining Rα

(762)

to be their sum, yields

(n) En ∣Rα ∣ ≤ Cα [(1 + γ)2 αd2 +

αd ]. d3/2

(763)

This proves (735), and completes the proof. Step 2: Vanishing errors For a fixed word α ∈ Z w , define the stopped interpolation (d)

d (τ )∧J) µ̃α,J (τ ) ∶= µ(m , α

md (τ ) ∶= ⌊

τ ⌋, τd

0 ≤ τ ≤ T.

(764)

Also define the piecewise-constant stopped finite-dimensional vector field by (d,md (τ )) ̃ ]α . [V raw,J (τ )]α ∶= 1{md (τ )<J} [Vraw (d)

(765)

The following lemma proves the integral equation for the interpolated moments, in the spirit of Lemma 4. Lemma 19. Under the conditions of Theorem 3, let J be a stopping time satisfying (565). Fix T > 0, and (d) set τd ∶= αdd , NT ∶= ⌊ τTd ⌋ = ⌊ Tαdd ⌋. Then there exists Eα ∶ [0, T ] → R such that τ

(d)

µ̃α,J (τ ) = µ(0) α +∫ (d)

0

(d) ̃ [V raw,J (s)]α ds + Eα (τ ), (d)

0 ≤ τ ≤ T,

(766)

and supτ ∈[0,T ] ∣Eα (τ )∣ Ð → 0. Proof. Fix a word

P

α = (Z1 , . . . , Zw ).

Throughout the proof all estimates are uniform on (565). Recall from Lemma 18 that, on {n < J}, (d,n) (n) (n) µ(n+1) − µ(n) α α = τd [Vraw ]α + αd ηα + Rα ,

84

(767)

where En [ηα(n) ] = 0, and

Cα , d

(768)

αd ]. d3/2

(769)

En ∣ηα(n) ∣2 ≤

(n) En ∣Rα ∣ ≤ Cα [(1 + γ)2 αd2 +

For m ≤ NT , the stopped process satisfies m−1

(n+1) − µ(n) µ(m∧J) − µ(0) α ). α α = ∑ 1{n<J} (µα

(770)

n=0

Substituting (767) into (770) gives m−1

m−1

m−1

n=0

n=0

n=0

(d,n) (n) (n) µ(m∧J) = µ(0) α α + τd ∑ 1{n<J} [Vraw ]α + ∑ αd 1{n<J} ηα + ∑ 1{n<J} Rα . (d)

m (τ )−1

d Accordingly, define Enoise,α (τ ) ∶= ∑n=0

(n)

αd 1{n<J} ηα

(d)

m (τ )−1

d and Erem,α (τ ) ∶= ∑n=0

(n)

Since J is a stopping time, 1{n<J} is Fn -measurable. Hence αd 1{n<J} ηα sequence. Let m−1

(d) Mα,m ∶= ∑ αd 1{n<J} ηα(n) ,

(771) (n)

1{n<J} Rα .

is a martingale difference

0 ≤ m ≤ NT .

(772)

n=0

Using (768), its predictable quadratic variation satisfies NT −1

αd2 ≤ Cα,T αd . d

(773)

⎞ ⎛ Cα,T αd (d) (d) ∣ > ε) ≤ (τ )∣ > ε ≤ P ( max ∣Mα,m sup ∣E Ð→ 0. m≤NT ε2 ⎠ ⎝τ ∈[0,T ] noise,α

(774)

E [⟨Mα(d) ⟩N ] ≤ ∑ αd2 E [1{n<J} En ∣ηα(n) ∣2 ] ≤ NT Cα T n=0

Since αd → 0, Doob’s maximal inequality gives, for every ε > 0, P Therefore

(d)

→ 0. sup ∣Enoise,α (τ )∣ Ð P

τ ∈[0,T ]

(775)

By (769), NT −1

(n) E [ ∑ 1{n<J} ∣Rα ∣] ≤ NT Cα [(1 + γ)2 αd2 + n=0

αd ] ≤ Cα,T [(1 + γ)2 dαd + d−1/2 ] . d3/2

(776)

c0 For the fixed time horizon T , γ is independent of d. Moreover, dαd ≤ log Ð→ 0 under the learning-rate d assumption. Hence the right-hand side of (776) tends to zero. (d) (n) T −1 Since supτ ∈[0,T ] ∣Erem,α (τ )∣ ≤ ∑N n=0 1{n<J} ∣Rα ∣, Markov’s inequality yields (d) sup ∣Erem,α (τ )∣ Ð → 0. P

τ ∈[0,T ]

(777)

For τ ∈ [0, T ], write m ∶= md (τ ). By the definition (765), τ

∫

0

m−1

(d,n) (d,m) ̃ [V raw,J (s)]α ds = τd ∑ 1{n<J} [Vraw ]α + (τ − mτd )1{m<J} [Vraw ]α . (d)

(778)

n=0

We therefore define the discretization error (d)

(d,m) Edisc,α (τ ) ∶= −(τ − mτd )1{m<J} [Vraw ]α .

85

(779)

By (565), all letters satisfy maxZ∈Z ∥Z∥op ≤ C∗2 . Consequently, for the fixed α, Tibias and Tidiff in (731) are (n) (r,d,n) uniformly bounded. In addition, Lemma 13 gives maxk ∣bk ∣ + maxk,l,r ∣akl ∣ ≤ C uniformly on {n < J}. Hence √ (d,n) ]α ∣ ≤ Cα ( d + 1 + γ) . (780) 1{n<J} ∣[Vraw Since 0 ≤ τ − mτd < τd = αdd , we obtain (d)

sup ∣Edisc,α (τ )∣ ≤

τ ∈[0,T ]

√ αd αd (1 + γ)αd Cα ( d + 1 + γ) ≤ Cα ( √ + ) Ð→ 0. d d d

(781)

Finally, combining (771) and (778) gives τ

(d)

µ̃α,J (τ ) = µ(0) α +∫ (d)

(d)

0

(d)

(d) ̃ [V raw,J (s)]α ds + Enoise,α (τ ) + Erem,α (τ ) + Edisc,α (τ ). (d)

(d)

(d)

(782)

(d)

Defining Eα = Enoise,α + Erem,α + Edisc,α and using (775), (777), (781) finishes the proof.

(d,n)

ψk

K

(n)

∶= tk

+µ

(n)

For k ∈ K, define Dk

Step 3: Effective bias of the fast flow 2

(r,d,n)

∶= 2 ∑ ∑ akl

(n)

[µ

(r)

(Ml

l=1 r=1

,Bk )

A,(n)

(r)

(Ak ,Ml

B,(n)

+ tk

)

and

].

(783)

Define the finite-dimensional effective bias by (d,n)

βk (n)

whenever Dk

(d,n)

> 0, and set βk

(n)

∶= 0 when Dk

(d,n)

∶= −

ψk

(784)

(n)

Dk

= 0. The following lemma proves its convergence.

Lemma 20. Under the conditions of Theorem 3, let J be a stopping time satisfying (565), and fix T > 0. Set τd ∶= αdd and NT ∶= ⌊ Tαdd ⌋. Then, RRR RRR w RR P RRR md (τ )−1 √ (n) (d,n) sup RRRRτd ∑ 1{n<J} ∑ ( d bki − βki ) Tibias (µ(n) )RRRRR Ð → 0, RRR 0≤τ ≤T RR n=0 i=1 R RRR ki ∈K R

(785)

where md (τ ) ∶= ⌊ ατ dd ⌋. (n)

(n)

(n)

Proof. Fix a trainable position i and write k ∶= ki , Hn ∶= Hα,i ∶= Tibias (µ(n) ), Dn ∶= Dk , bn ∶= bk

and

(d,n) βn ∶= βk . All constants below may depend on the fixed word α, but are independent of d and n.

Returning to the proof of Lemma 14, the expansion (582) gives

αd (d,n) (n) (n) En [∆mk ] = −αd Dn bn − √ ψk + rk , d where, uniformly on {n < J}, (n)

∣rk ∣ ≤ C ( (d,n)

Using ψk

(1 + γ)αd + αd2 ) . d

(n+1)

(787)

= −Dn βn , we may rewrite (786) as αd Dn √ (n) (n) En [∆mk ] = − √ ( d bn − βn ) + rk . d

Let ζm,k

(786)

(n)

∶= ∆mk

(788)

(n)

− En [∆mk ]. By Lemma 14, (n+1)

En [∣ζm,k ∣2 ] ≤ Cαd2 . 86

(789)

Therefore (788) gives √ 1 Hn 1 Hn (n+1) 1 Hn (n) (n) ∆mk + √ ζm,k + √ rk . τd Hn ( d bn − βn ) = − √ d Dn d Dn d Dn

(790)

Thus it remains to show that the three terms on the right vanish after summation. By Lemma 12, there exists cD = cD (T, γ, cbal ) > 0 such that ⎛ ⎞ (n) P ⎜ inf Dk ≥ cD ⎟ Ð→ 1. T ∧J ⎝0≤n<N ⎠ k∈K

(791)

Moreover, Lemma 16 gives constants Ctr , ctr , Cb > 0 such that P(

Cb (n) [∥bK ∥ − Ctr e−ctr αd n ] > √ ) Ð→ 0. + 0≤n<NT ∧J d sup

(792)

Now define (n)

JD ∶= inf {n ≥ 0 ∶ min Dk k∈K

< cD } ,

2Cb (n) Jb ∶= inf {n ≥ 0 ∶ ∥bK ∥ > Ctr e−ctr αd n + √ } , d

(793) (794)

and set Ĵ ∶= J ∧ JD ∧ Jb . Equations (791) and (792) imply P (Ĵ < NT ∧ J) Ð→ 0.

(795)

̂ Before J, ̂ Dn ≥ cD , ∣bn ∣ ≤ It therefore suffices to prove the desired convergence with J replaced by J. (n) 2C Ctr e−ctr αd n + √db , and also ∣Hn ∣ ≤ Cα , ∣mk ∣ ≤ C by (565). ̂ To control the second term of (790), consider Define Cn ∶= Hn satisfying ∣Cn ∣ ≤ Cα for n < J. Dn

1 m−1 (n+1) Mm ∶= √ ∑ 1{n<J} ̂ Cn ζm,k . d n=0

(796)

Since Cn 1{n<J} ̂ is Fn -measurable, Mm is a martingale. By (789), E[⟨M ⟩NT ] ≤

Cα NT αd2 ≤ Cα,T αd . d

(797)

Hence Doob’s maximal inequality yields → 0. max ∣Mm ∣ Ð

(798)

1 NT −1 Cα (1 + γ)αd 1+γ √ (n) √ ∑ 1{n<J} + αd2 ) ≤ Cα,T ( √ + d αd ) Ð→ 0, ̂ ∣Cn rk ∣ ≤ √ NT ( d d n=0 d d

(799)

P

m≤NT

For the third term of (790), using (787),

because αd ≤ c0 (d log d)−1 . For the following we control the first term of (790). Since Hn is a finite linear combination of moments A,(n) B,(n) of the same degree, Lemma 18 applies to Hn . The same lemma applies to Dn = tk + tk . Using the H smooth map (H, D) z→ D on the region ∣H∣ ≤ Cα , D ≥ cD , together with a first-order Taylor expansion, we ̂ obtain, for n < J, (n) (n+1) (n) Cn+1 − Cn = AC + ζC + ρC , (800)

87

(n)

(n+1)

where AC is Fn -measurable, ζC

is a martingale difference, and

αd (n) (1 + γ)αd (n) + (1 + γ)2 αd2 ] , ∣AC ∣ ≤ Cα [ √ ∥bK ∥ + d d αd2 (n+1) 2 En [∣ζC ∣ ] ≤ Cα , d αd (n) En [∣ρC ∣] ≤ Cα [(1 + γ)2 αd2 + 3/2 ] . d

(801) (802) (803)

The first term on the right-hand side of (801) comes from (731), while (802) and (803) follow from (734) and (735), together with the quadratic Taylor remainder of H/D. For n < Jb we have NT −1

2Cb αd (n) −ctr αd n + √ ) + Cα,T (1 + γ + dαd ) = O(1), ∑ 1{n<J} ̂ ∣AC ∣ ≤ Cα √ ∑ (Ctr e d n≥0 d n=0

(804)

uniformly in d. Here we use αd ∑n≥0 e−ctr αd n ≤ C and NT αdd = O(1), NT αd2 = O(dαd ) = o(1). Moreover, (803) gives NT −1

(n)

→ 0, ∑ 1{n<J} ̂ ∣ρC ∣ Ð P

(805)

n=0

For every m ≤ NT , m−1

(n)

∑ 1{n<J} ̂ Cn ∆mk

n=0

(m )

(0)

m∗ −1

(n)

= Cm∗ −1 mk ∗ − C0 mk − ∑ mk (Cn − Cn−1 ) ,

(806)

n=1

where m∗ ∶= m ∧ Ĵ with the first term interpreted as zero when m∗ = 0. The boundary terms in (806) are uniformly O(1), and hence become o(1) after multiplication by d−1/2 . (n) For the Cn − Cn−1 term, using the boundedness of mk and (804), m∗ −1 1 (n) (n−1) √ sup ∣ ∑ mk AC ∣ Ð→ 0. d m≤NT n=1

(807)

(n)

The contribution of ρC vanishes by (805). (n) (n−1) (n−1) It remains to treat the martingale part (the second term of (800)). Write mk = mk +∆mk . Then m∗ −1

(n) (n)

m∗ −1

(n−1) (n) ζC +

∑ mk ζC = ∑ mk

n=1

m∗ −1

(n−1) (n) ζC .

∑ ∆mk

(808)

n=1

n=1

The first term of (808) is a martingale with predictable quadratic variation bounded by Cα,T αd , because √ (n−1) mk is uniformly bounded. Hence it is upper bounded by OP ( αd ). For the second term of (808), Lemma (n−1) 2

14 gives En−1 [∣∆mk

∣ ] ≤ Cαd2 before the stopping time. Together with (802) and Cauchy–Schwarz, ⎡NT ∧J−1 ⎤ ̂ √ ⎢ α2 (n−1) (n) ⎥ E ⎢⎢ ∑ ∣∆mk ζC ∣⎥⎥ ≤ CNT √d ≤ CT αd d. ⎢ n=1 ⎥ d ⎣ ⎦

Therefore

(809)

̂

1 NT ∧J−1 (n−1) (n) P √ ζC ∣ Ð → 0, ∑ ∣∆mk d n=1

(810)

since αd → 0. Combining these estimates with (806), we obtain 1 m−1 (n) P sup ∣ √ ∑ 1{n<J} → 0. ̂ Cn ∆mk ∣ Ð m≤NT d n=0 88

(811)

Summing (790) from n = 0 to md (τ ) − 1, and using (798), (799), and (811), gives RRR RRR md (τ )−1 √ P RR → 0. sup RRRRRτd ∑ 1{n<J} ̂ Hn ( d bn − βn )RR Ð R RRR 0≤τ ≤T RR n=0 R

(812)

Finally, (795) allows us to replace Ĵ by J. Summing over the trainable positions proves (785). Step 4: Final results Finally we can write an approximate integral equation for the moments. For a word α = (Z1 , . . . , Zw ) ∈ Z w , define the finite-dimensional effective vector field by (d,n)

[Veff

w

(d,n)

]α ∶= − ∑ βki

K

2

(r,d,n)

diff Ti,l,r (µ(n) ) − 2γ∣α∣K µ(n) α .

(813)

i=1 l=1 r=1 ki ∈K

i=1 ki ∈K

Let

w

Tibias (µ(n) ) − 2 ∑ ∑ ∑ aki l

̃ [V eff,J (τ )]α ∶= 1{md (τ )<J} [Veff (d)

(d,md (τ ))

]α

(814)

and md (τ ) ∶= ⌊ ατ dd ⌋. Lemma 21. Under the conditions of Theorem 3, let J be a stopping time satisfying (565), and fix T > 0. (d) Set τd ∶= αdd and NT ∶= ⌊ Tαdd ⌋. Then, for every α, there exists Eα ∶ [0, T ] → R such that τ

(d)

µ̃α,J (τ ) = µ(0) α +∫ (d)

0

(d) ̃ [V eff,J (s)]α ds + Eα (τ ), (d)

0 ≤ τ ≤ T,

(815)

→ 0. and supτ ∈[0,T ] ∣Eα (τ )∣ Ð P

Proof. By Lemma 19, for every fixed α, τ

(d)

µ̃α,J (τ ) = µ(0) α +∫ (d)

0

(d) ̃ [V raw,J (s)]α ds + Eα (τ ), (d)

(816)

where supτ ∈[0,T ] ∣Eα (τ )∣ Ð → 0. It therefore suffices to prove P

sup ∣∫

τ ∈[0,T ]

τ

0

P ̃(d) (s)]α − [V ̃(d) (s)]α ) ds∣ Ð ([V → 0. raw,J eff,J

(817)

By the definitions (731) and (813), the two vector fields differ only in the bias term. Thus (d,n)

(d,n) [Vraw ]α − [Veff

w √ (n) (d,n) ]α = − ∑ ( d bki − βki ) Tibias (µ(n) ).

(818)

i=1 ki ∈K

Fix τ ∈ [0, T ] and write m ∶= md (τ ) = ⌊ ττd ⌋. Since both vector fields are piecewise constant, τ

∫

0

m−1 w √ ̃(d) (s)]α − [V ̃(d) (s)]α ) ds = −τd ∑ 1{n<J} ∑ ( d b(n) − β (d,n) ) T bias (µ(n) ) + R(d) (τ ), (819) ([V i raw,J eff,J ki ki cell,α n=0

i=1 ki ∈K

where the final term is w √ (m) (d) (d,m) Rcell,α (τ ) ∶= −(τ − mτd )1{m<J} ∑ ( d bki − βki ) Tibias (µ(m) ).

(820)

i=1 ki ∈K

The first term in (819) converges to zero uniformly in τ in probability by Lemma 20. Hence it remains only (d) to control Rcell,α . 89

Before the stopping time (565), all letters satisfy a uniform operator-norm bound. Therefore, for the (n) fixed word α, ∣Tibias (µ(n) )∣ ≤ Cα uniformly for n < J. Moreover, Lemma 13 gives ∣bk ∣ ≤ C uniformly on (d,n)

{n < J}. We also have a uniform bound on βk (n)

∣µ

(r)

(Ml

,Bk )

(n)

∣µ

(r)

(Ak ,Ml

(r,d,n)

The coefficients akl

)

. Indeed, since Ak , Bk ⪰ 0, (r,n)

∥op tk

(r,n)

∥op tk

∣ ≤ ∥Ml ∣ ≤ ∥Ml

B,(n)

,

(821)

A,(n)

.

(822)

(n)

(823)

are uniformly bounded on (565), and hence (d,n)

∣ψk (d,n)

A,(n)

∣ ≤ C (tk

B,(n)

+ tk

) = CDk .

(n)

(d,n)

Consequently, ∣βk ∣ ≤ C whenever Dk > 0, while the same bound is trivial under our convention βk (n) when Dk = 0. Since 0 ≤ τ − mτd < τd = αdd , the bounds above give (d)

sup ∣Rcell,α (τ )∣ ≤ Cα

τ ∈[0,T ]

αd √ αd αd ( d + 1) ≤ Cα ( √ + ) Ð→ 0. d d d

=0

(824)

Combining this with Lemma 20 proves (817). Finally, define Eα(d) (τ ) ∶= Eα(d) (τ ) + ∫

τ

0

̃(d) (s)]α − [V ̃(d) (s)]α ) ds. ([V raw,J eff,J

(825)

Then we finish the proof by (816) and (817). F.1.6

Limiting slow dynamics

We now pass from the finite-dimensional effective dynamics (815) to the limiting dynamics. We first establish compactness of the trajectories. (n)

Step 1: Compactness Fix a compact set D ⋐ M×Q, and recall Jmac ∶= inf {n ≥ 0 ∶ (mK , q (1,n) , q (2,n) ) ∉ D}. (n)

(n)

Let C∗ be the constant in Lemma 17, and define Jsp ∶= inf {n ≥ 0 ∶ max1≤k≤K {∥Uk ∥op , ∥Vk NT ∶= ⌊ Tαdd ⌋, set

∥op } > C∗ }. For

(d) Esp (T ) ∶= {Jsp > Jmac ∧ NT } .

(826)

(d) P (Esp (T )) Ð→ 1.

(827)

By Lemma 17, For 0 ≤ n ≤ NT , define the auxiliary weights (n)

Ǔk (n)

Let M̌k

(n)

⎧ (n∧Jmac ) (d) ⎪ , on Esp (T ), ⎪U ∶= ⎨ k(0) (d) ⎪ on Esp (T )c , ⎪ ⎩Uk ,

(n) T

= Ǔk (V̌k

(n)

) , Ǎk

(n)

(n)

(n)

(n)

V̌k

(n)

= Ǔk (Ǔk )T , B̌k

= V̌k

µ̌(d) (τ ) = µ̌(n) ,

τ ∈[

⎧ (n∧Jmac ) (d) ⎪ , on Esp (T ), ⎪V ∶= ⎨ k(0) (d) ⎪ on Esp (T )c . ⎪ ⎩Vk ,

(828)

(n) T

(V̌k

) , and define the corresponding auxiliary

(n) (n) (n) moments µ̌α ∶= d1 Tr [Ž1 ⋯Žw ], α = (Z1 , . . . , Zw ) ∈ Z w . Their interpolation is denoted by

nαd (n + 1)αd , ). d d

(829)

Lemma 22. Under the conditions of Theorem 3, for every fixed T > 0 and every ρ > C∗2 , the sequence (d) {µ̌(d) (⋅)}d≥1 is C-tight in D([0, T ]; Xρ0 ). Moreover, for every word α of length w, supτ ∈[0,T ] ∣µ̌α (τ )∣ ≤ C∗2w almost surely, and the auxiliary process agrees with the empirical process stopped at Jmac with probability tending to one. 90

Proof. By construction and (826), almost surely (n)

(n)

max {∥Ǔk ∥op , ∥V̌k

sup

0≤n≤NT 1≤k≤K (n)

Hence every letter Žk

(n)

(n)

(n)

∥op } ≤ C∗ .

(n)

(830)

(n)

∈ {M̌k , (M̌k )T , Ǎk , B̌k } satisfies ∥Žk ∥op ≤ C∗2 . Therefore, for every word

(d)

∣α∣ = w, ∣µ̌α (τ )∣ ≤ C∗2w . In particular, µ̌(d) (τ ) ∈ Xρ0 for 0 ≤ τ ≤ T and every ρ > C∗2 . We next record a uniform bound on the finite-d effective vector field. Before the stopping time, Lemma (r,d,n) (d,n) A,(n) B,(n) 13 gives maxk,l,r ∣akl ∣ ≤ C. Moreover, ∣ψk ∣ ≤ C (tk + tk ), and hence, under the convention (d,n)

A,(n)

βk = 0 when tk d, n, w, such that

B,(n)

+ tk

(d,n)

= 0, ∣βk

∣ ≤ C. Consequently, there exists a constant B < ∞, independent of (d,n)

1{n<Jsp ∧Jmac } sup ∣[Veff ∣α∣=w

]α ∣ ≤ BwC∗2w .

(831)

(d)

We now use Lemma 21. Set J ∶= Jsp ∧ Jmac . On Esp (T ), stopping at J agrees with stopping at Jmac throughout the time interval [0, T ]. Therefore, for every fixed word α, the auxiliary process admits the representation τ ̃(d) (s)]α ds + Ě (d) (τ ), µ̌(d) (τ ) = µ̌(d) (0) + 1 (d) (832) ∫ [V α

Esp (T )

α

where

0

α

eff,J

sup ∣Ěα(d) (τ )∣ Ð →0 P

(833)

τ ∈[0,T ] (d)

(d)

for every fixed α. On Esp (T )c , we set Ěα ≡ 0. Define the continuous approximation µ̂(d) (τ ) ∶= µ̌(d) (0) + 1E (d) (T ) ∫ sp

τ 0

̃(d) (s) ds. V eff,J

(834)

The convergence in (833) can be upgraded to convergence in Xρ0 , exactly as in the tied case (Lemma 4): sup ∥µ̌(d) (τ ) − µ̂(d) (τ )∥ρ Ð → 0. P

(835)

τ ∈[0,T ]

It remains to prove compactness of {µ̂(d) }. By (831), for every ∣α∣ = w, 2w sup ∣µ̂(d) α (τ )∣ ≤ (1 + T Bw)C∗ .

(836)

τ ∈[0,T ]

Furthermore, for 0 ≤ s < t ≤ T , w

(d)

∥µ̂

(t) − µ̂

(d)

C2 (s)∥ρ ≤ B∣t − s∣ sup w ( ∗ ) =∶ Lρ ∣t − s∣, ρ w≥1

(837)

where Lρ < ∞ is independent of d. The remainder of the argument is identical to the proof of Lemma 5 in the tied case. Namely, (836) defines a compact subset of Xρ0 , because its tail decays exponentially, while (837) gives uniform equicontinuity. The Arzelà–Ascoli theorem therefore implies that {µ̂(d) }d≥1 is tight in C([0, T ]; Xρ0 ). Together with (835), this proves that {µ̌(d) }d≥1 is C-tight in D([0, T ]; Xρ0 ). (d)

Finally, on Esp (T ), the auxiliary process is precisely the empirical moment process stopped at Jmac . (d) Since P(Esp (T )) → 1, the two processes are asymptotically equivalent on [0, T ]. Step 2: Convergence of the effective vector field Lemma 23. Under the conditions of Theorem 3, fix 0 < T1 < T2 ≤ T , and let J be a stopping time satisfying (d,n) (565). Let C∗ denote the corresponding spectral bound in Lemma 17 and fix ρ > C∗2 . Let Veff be the finite-d effective vector field defined in (813), and let Vext be the extended limiting vector field defined in (542). Then, sup τ ∈[T1 ,T2 ]

(d,md (τ ))

1{md (τ )<J} ∥Veff

− Vext (µ(md (τ )) )∥ Ð → 0. P

ρ

91

(838)

(n)

Proof. Throughout the proof, all estimates are uniform on (565). Write q (n) ∶= (q (1,n) , q (2,n) ) and m⋆ ⋆

(n)

(n) ), where the non-trainable components of m⋆ remain fixed at their initial values.

∶=

m (q For notational convenience, define (r)

Akl (m, q) ∶= [∇q(r) Φ(m, q)]kl ,

r = 1, 2.

(839)

An empirical state q (n) ∈ Q. This is because for arbitrary c = (c1 , . . . , cK ), we can write Mc ∶= ∑K k=1 ck Mk and have 2 1 (840) cT (q (1) + q (2) )c = (Tr[Mc McT ] + Tr[Mc2 ]) = ∥ Sym Mc ∥2F ≥ 0, d d Similarly 2 cT (q (1) − q (2) )c = ∥ Skew Mc ∥2F ≥ 0, (841) d M +M T

where Sym Mc ∶= c 2 c and Skew Mc ∶= trivially along the empirical trajectory: Recall (r,d,n)

akl

∶=

Mc −McT . 2

Hence the projection in Vext (see Appendix F.1.2) acts

q̂(µ(n) ) = q (n) .

(842)

L 1 (r) 2 (n) ∑ Cij,i′ j ′ En [∂(ijk),(i′ j ′ l) L(G )] . 2 i,j,i′ ,j ′ =1

(843)

(r)

We first prove its convergence to Akl (m(n) , q (n) ). (n) Before the stopping time (565), maxk ∥Mk ∥op ≤ C∗2 , while compactness of D gives 1 (n) (n) max ∣mk ∣ = max √ ∣TrMk ∣ ≤ R∗ k k d

(844)

for a constant R∗ < ∞. Therefore Lemma 10 applies uniformly in n < J. 2 For fixed indices (ijk), (i′ j ′ l), apply Lemma 10 to ∂(ijk),(i Assumption 4.1 guarantees the ′ j ′ l) L(G). required regularity and polynomial growth. By Price’s theorem, (r)

Akl (m(n) , q (n) ) =

1 (n) (r) 2 ∑ C ′ ′ E [∂(ijk),(i′ j ′ l) L(GG )] , 2 i,j,i′ ,j ′ ij,i j

(845)

(n)

where GG is Gaussian with the same mean and covariance as G(n) . Consequently, (r,d,n)

sup max ∣akl n<J k,l,r

C (r) − Akl (m(n) , q (n) )∣ ≤ √ . d

(846)

(n)

We next replace m(n) by m⋆ . By Lemma 13, C (n) ∥bK − ∇mK Φ(m(n) , q (n) )∥ ≤ √ . d

(847)

Lemma 16 gives, with probability tending to one, Cb (n) ∥bK ∥ ≤ Ctr e−ctr αd n + √ , d

0 ≤ n < NT ∧ J.

(848)

Hence, uniformly for n ≥ ⌈ Tα1dd ⌉, we have e−ctr αd n ≤ e−ctr T1 d , and therefore sup ⌈T1 d/αd ⌉≤n≤⌊T2 d/αd ⌋ n<J

∥∇mK Φ(m(n) , q (n) )∥ = OP (d−1/2 ).

(849)

(n)

Since ∇mK Φ(m⋆ , q (n) ) = 0, the uniform strong convexity in Assumption 5 implies (n)

λ0 ∥mK − m⋆K (q (n) )∥ ≤ ∥∇mK Φ(m(n) , q (n) )∥ . 92

(850)

Thus

(n)

∥m(n) − m⋆ ∥ = OP (d−1/2 ).

sup ⌈T1 d/αd ⌉≤n≤⌊T2 d/αd ⌋ n<J

(851)

The map (m, q) z→ ∇q(r) Φ(m, q) is Lipschitz on a compact neighborhood of (565) and of {(m⋆ (q), q) ∶ q ∈ projq D}. Therefore, (r)

(r)

(n)

P

(r)

(n)

P

max ∣Akl (m(n) , q (n) ) − Akl (m⋆ , q (n) )∣ Ð → 0.

sup

(852)

⌈T1 d/αd ⌉≤n≤⌊T2 d/αd ⌋ k,l,r n<J

Combining (846) and (852), we conclude that ∆a,d ∶= (d,n)

For k ∈ K, recall ψk

(r,d,n)

max ∣akl

sup

⌈T1 d/αd ⌉≤n≤⌊T2 d/αd ⌋ k,l,r n<J (r,d,n)

2 = 2 ∑K l=1 ∑r=1 akl

(d,n)

B,(n)

(n)

[µ

(r) (Ml ,Bk )

− Akl (m⋆ , q (n) )∣ Ð → 0.

+µ

(n)

(r) (Ak ,Ml )

(d,n)

(853) ψ

] and βk

(d,n)

(n)

= − k(n) , Dk Dk

A,(n)

∶= tk

+

. Now we prove the convergence of βk to βkext Along the empirical trajectory, (842) implies

tk

(d,n)

∣ψk

K

2

(r,n)

(n)

∣µ

(r)

(Ak ,Ml

A,(n)

)

∣ ≤ C∗2 tk

(d,n)

(n)

,Bk )

∣ + ∣µ

(r)

(Ak ,Ml

)

∣) .

∥op ≤ C∗2 . Since Ak , Bk ⪰ 0, we have ∣µ

(854)

(n)

(r)

(Ml

B,(n)

,Bk )

∣ ≤ C∗2 tk

and

. Therefore ∣ψk

Whenever Dk

(r)

(Ml

l=1 r=1

The spectral bound (565) gives ∥Ml

(n)

(n)

− ψkext (µ(n) )∣ ≤ 2∆a,d ∑ ∑ (∣µ

(n)

> 0, division by Dk

(n)

− ψkext (µ(n) )∣ ≤ C∆a,d Dk .

(855)

yields (d,n)

∣βk

− βkext (µ(n) )∣ ≤ C∆a,d .

(856)

(n)

By Lemma 12, Dk is uniformly bounded away from zero before J, for all k ∈ K, with probability tending to one. Hence (853) implies ∆β,d ∶=

sup

(d,n)

max ∣βk

⌈T1 d/αd ⌉≤n≤⌊T2 d/αd ⌋ k∈K n<J

− βkext (µ(n) )∣ Ð → 0. P

(857)

(d,n)

Finally, fix n < J and a word α = (Z1 , . . . , Zw ). The weight-decay terms in Veff and Vext (µ(n) ) coincide exactly. Thus their difference contains only T bias and T diff terms. Under the spectral bounds (565), for 2(w+1) diff every trainable position i, ∣Tibias (µ(n) )∣ ≤ 2C∗2w and ∣Ti,l,r (µ(n) )∣ ≤ 2C∗ . Consequently, (d,n)

∣[Veff

− Vext (µ(n) )]α ∣ ≤ Cw (∆β,d C∗2w + ∆a,d C∗

2(w+1)

),

(858)

where C depends only on K. Dividing by ρw and taking the supremum over all words of degree w ≥ 1 gives w C∗2 (d,n) (n) 2 ∥Veff − Vext (µ )∥ ≤ C (∆β,d + C∗ ∆a,d ) sup w ( ) . ρ ρ w≥1 C2

(859)

w

Because ρ > C∗2 and supw≥1 w ( ρ∗ ) < ∞. Equations (853) and (857) therefore imply sup ⌈T1 d/αd ⌉≤n≤⌊T2 d/αd ⌋ n<J

(d,n)

∥Veff

which proves (838). 93

− Vext (µ(n) )∥ Ð → 0., P

ρ

(860)

Step 3: Global well-posedness We call a pair (m̄, µ̄) an admissible solution on [0, T ] if there exist 0 < s− < s+ < ∞ such that µ̄ ∈ C([0, T ]; Xs0− ), τ

Vext (µ̄(s)) ds

(861)

B min (t̄A k (τ ) + t̄k (τ )) > 0,

(862)

µ̄(τ ) = µ̄(0) + ∫

0

holds as an identity in Xs0+ , and, for every τ ∈ [0, T ], q(µ̄(τ )) ∈ Q, with

k∈K

m̄K (τ ) = m⋆K (q(µ̄(τ ))),

m̄∖K (τ ) = m̄∖K (0).

(863)

Lemma 24. Under the conditions of Theorem 3, the extended slow system (543), initialized at µ̄(0), admits a unique admissible solution µ̄ ∈ C([0, T ]; Xρ0− ) for every pair C∗2 < ρ− < ρ+ , where the integral equation (861) is understood in Xρ0+ . Moreover, µ̄ also satisfies the original ODE (29)–(32). (0)

(0)

(0)

Proof. Fix C∗2 < ρ− < ρ+ . By Assumption 4.2, every initial letter satisfies ∥Mk ∥op , ∥Ak ∥op , ∥Bk ∥op ≤ C02 . (d) Hence, for every word ∣α∣ = w, ∣µα (0)∣ ≤ C02w . Together with Assumption 4.3 and ρ− > C02 , the same tail argument used in the tied case (Lemma 7) gives µ(d) (0) Ð → µ̄(0) P

(864)

in Xρ0− .

By Lemma 22, the auxiliary processes {µ̌(d) } are C-tight in D([0, T ]; Xρ0− ). Hence, from every subsequence, we may extract a further subsequence, still indexed by d, such that µ̌(d) → µ∗

in D([0, T ]; Xρ0− ),

(865)

where µ∗ ∈ C([0, T ]; Xρ0− ) almost surely. As in the proof of Lemma 7 for the tied model, by the Skorokhod representation theorem we may assume sup ∥µ̌(d) (τ ) − µ∗ (τ )∥ρ− Ð→ 0

almost surely.

τ ∈[0,T ]

(866)

Fix 0 < ε < t ≤ T and first work before the stopping time (565). Lemma 21, together with Lemma 22, gives µ̌(d) (t) − µ̌(d) (ε) = ∫ where

t ε

(d)

Veff (s) ds + Rd (ε, t),

sup ∥Rd (ε, s)∥ρ+ Ð → 0. P

(868)

ε≤s≤t

By Lemma 23,

(d)

sup ∥Veff (s) − Vext (µ̌(d) (s))∥

s∈[ε,t]

(867)

Ð → 0. P

ρ+

(869)

The local Lipschitz estimate from Lemma 11 implies, on every bounded subset, L ∥µ − ν∥ρ− . ρ+ − ρ −

(870)

sup ∥Vext (µ̌(d) (s)) − Vext (µ∗ (s))∥ρ Ð→ 0.

(871)

∥Vext (µ) − Vext (ν)∥ρ+ ≤ Hence (866) gives

+

s∈[ε,t]

Passing to the limit in (867) yields µ∗ (t) − µ∗ (ε) = ∫

t ε

94

Vext (µ∗ (s)) ds.

(872)

(d,n)

The effective vector fields Veff are uniformly bounded in Xρ0+ before the stopping time, as established in the proof of Lemma 22. Therefore, letting ε ↓ 0 in (872) and using the continuity of µ∗ gives µ∗ (t) = µ̄(0) + ∫

t

Vext (µ∗ (s)) ds,

0

0 ≤ t ≤ T.

(873)

For every finite d, the empirical covariance belongs to Q (see the proof of Lemma 23). Since Q is closed and convergence in Xρ0− implies convergence of the second-order coordinates, q(µ∗ (τ )) ∈ Q,

0 ≤ τ ≤ T.

(874) A,(n)

Next, Lemma 12 gives a constant cT > 0 such that, with probability tending to one, tk every trainable k before the stopping time. Passing to the limit gives ∗ B ∗ tA k (µ (τ )) + tk (µ (τ )) ≥ cT , (d)

k ∈ K,

B,(n)

+ tk

0 ≤ τ ≤ T.

≥ cT for (875)

2∣α∣

Finally, Lemma 22 gives ∣µ̌α (τ )∣ ≤ C∗ . Passing to the limit yields ∣µ∗α (τ )∣ ≤ C∗ . 2∣α∣

(876)

Thus every subsequential limit is admissible. In particular, P(q(µ∗ (τ ))) = q(µ∗ (τ )), so along the limiting ̃ ⋆K (µ∗ (τ )) = m⋆K (q(µ∗ (τ ))). Hence (873) is precisely the original system (29)–(32). trajectory m We now prove uniqueness. Suppose that µ1 and µ2 are two admissible solutions with the same initial condition. Both solve the extended ODE. By Lemma 11, if µ1 (τ0 ) = µ2 (τ0 ), at some τ0 < T , then the two trajectories agree on a nontrivial interval to the right of τ0 . The standard contradiction argument shows that µ1 (τ ) = µ2 (τ ) for 0 ≤ τ ≤ T . Thus the admissible solution is unique. We denote it by µ̄. It remains to verify that the stopping time Jmac is asymptotically inactive. Recall that D ⋐ M × Q is ̃ chosen to contain both the compact fast-time set Cfast ∶= {(mK , q0 ) ∶ Φ(mK , m∖K (0), q0 ) ≤ Φ(m(0), q0 )} and the compact slow trajectory Cslow ∶= {(m⋆K (q̄(τ )), q̄(τ )) ∶ 0 ≤ τ ≤ T } in its interior. Hence δ ∶= dist(Cfast ∪ Cslow , ∂D) > 0. By C-tightness (Lemma 22) and uniqueness of the subsequential limit, the stopped moment trajectory converges uniformly in probability to µ̄. In particular, sup

∥q (n) − q̄(nαd /d)∥ Ð → 0. P

0≤n≤NT ∧Jmac

(877)

Moreover, for every fixed T1 > 0, Lemma 16 and the strong convexity in Assumption 5 imply (n)

sup T1 d/αd ≤n≤NT ∧Jmac

→ 0. ∥mK − m⋆K (q (n) )∥ Ð P

(878)

Since q ↦ m⋆K (q) is Lipschitz on the compact subset of Q, it follows that (n)

sup T1 d/αd ≤n≤NT ∧Jmac

∥(mK , q (n) ) − (m⋆K (q̄(τn )), q̄(τn ))∥ Ð → 0, P

(879)

d where τn = nα . Therefore, with probability tending to one, the trajectory stays at least δ/2 away from ∂D d for T1 d/αd ≤ n ≤ NT . It remains to control the initial regime and the intermediate regime. Choose a fixed fast-time horizon tf > 0 sufficiently large that Cloc e−cloc tf < δ/8, where Cloc , cloc are the constants in Lemma 15. On 0 ≤ n ≤ tf /αd , the fast-time convergence and the freezing of the slow coordinates proved in Appendix F.1.4 imply that the empirical trajectory remains uniformly close to the fast limiting trajectory. Since the latter stays in Cfast , this part of the trajectory remains at least a δ/4 distance from ∂D with probability tending to one. For tf /αd ≤ n ≤ T1 d/αd , Lemma 15 yields

(n)

∥mK − m⋆K (q (n) )∥ ≤ Cloc e−cloc tf + oP (1)

(880)

uniformly before Jmac . Together with the uniform convergence of q (n) to q̄(nαd /d), this places the empirical state in a δ/4-neighborhood of Cslow throughout the intermediate regime. Combining the three time intervals gives P(Jmac ≤ NT ) Ð→ 0. (881) This completes the proof. 95

F.1.7

Final proof of Theorem 3

The fast-time convergence (36) was established in Appendix F.1.4. It remains to prove the convergence on the slow time scale. Fix 0 < T1 < T2 ≤ T . Choose C∗2 < ρ− < ρ+ , where C∗ is the spectral constant appearing in Lemma 17. By Lemma 22, the auxiliary moment trajectories {µ̌(d) }d≥1 are C-tight in D([0, T2 ]; Xρ0− ). Lemma 24 shows that every subsequential weak limit is almost surely equal to the same deterministic admissible solution µ̄ ∈ C([0, T2 ]; Xρ0− ) of the slow system (29)–(32). Consequently, µ̌(d) → µ̄

in D([0, T2 ]; Xρ0− ).

(882)

Similarly to the tied case, this implies convergence in the uniform topology. Hence sup ∥µ̌(d) (τ ) − µ̄(τ )∥ρ Ð → 0. P

(883)

−

τ ∈[0,T2 ]

In particular, for every fixed word α, ∣α∣

sup ∣µ̌(d) α (τ ) − µ̄α (τ )∣ ≤ ρ−

τ ∈[0,T2 ]

and therefore

sup ∥µ̌(d) (τ ) − µ̄(τ )∥ρ , −

τ ∈[0,T2 ]

sup ∣µ̌(d) → 0. α (τ ) − µ̄α (τ )∣ Ð P

(884)

(885)

τ ∈[0,T2 ]

We now transfer this convergence to the SGD trajectory. Lemma 24 gives P(Jmac ≤ NT2 ) Ð→ 0, where NT2 ∶= ⌊ Tα2dd ⌋, while Lemma 17 gives P (Jsp ≤ Jmac ∧ NT2 ) Ð→ 0. Thus the event Gd ∶= {Jmac > NT2 , Jsp > NT2 }

(886)

satisfies P(Gd ) Ð→ 1. On Gd , the auxiliary trajectory coincides with the original SGD trajectory throughout [0, T2 ]. Hence, for every fixed α and every ε > 0, P

⎞ ⎞ ⎛ ⎛ c sup ∣µ̌(d) sup ∣µ̃(d) α (τ ) − µ̄α (τ )∣ > ε Ð→ 0. α (τ ) − µ̄α (τ )∣ > ε ≤ P(Gd ) + P ⎠ ⎠ ⎝τ ∈[0,T2 ] ⎝τ ∈[0,T2 ]

(887)

This proves (38). It remains to establish the convergence of the mean variables on [T1 , T2 ]. (38) implies sup (∥q (1,d) (τ ) − q̄ (1) (τ )∥F + ∥q (2,d) (τ ) − q̄ (2) (τ )∥F ) Ð → 0. P

τ ∈[0,T2 ]

(888)

On the other hand, recall that (see (878)) (n)

sup ⌈T1 d/αd ⌉≤n≤⌊T2 d/αd ⌋

∥mK − m⋆K (q (n) )∥ Ð → 0. P

(889)

By Assumption 5 and the implicit function theorem, the map q z→ m⋆K (q) is Lipschitz on a compact neighborhood of the limiting covariance trajectory. Combining this with (888) gives sup τ ∈[T1 ,T2 ]

∥m⋆K (q (d) (τ )) − m⋆K (q̄(τ ))∥ Ð → 0. P

(890)

∥m̃(d) (τ d) − m̄(τ )∥∞ Ð → 0.

(891)

Since m̄K (τ ) = m⋆K (q̄(τ )), (889) and (890) imply sup τ ∈[T1 ,T2 ]

P

This proves (37) and completes the proof of Theorem 3. 96

F.2

Truncation to finite orders

We now show that the infinite-dimensional dynamics can be approximated by a finite-order truncation. A truncated moment sequence need not correspond to actual matrices, and in particular its covariance parameters need not belong to the admissible cone Q. For this reason, the truncated system is defined using the extended vector field Vext introduced in (542). For M ≥ 2, define the projection ⎧ ⎪ ⎪µα , ∣α∣ ≤ M, [ΠM µ]α ∶= ⎨ (892) ⎪ 0, ∣α∣ > M. ⎪ ⎩ Let YM ∶= {µ ∶ µα = 0 for every ∣α∣ > M }. Recall that Vext is well defined on the open set Us = {µ ∈ Xs ∶ P(q(µ)) ∈ Q,

B min (tA k (µ) + tk (µ)) > 0} , k∈K

(893)

where P denotes the projection onto the covariance cone Q. We first prove a stability property of the extended slow flow (543). Lemma 25. Let µ̄ be the admissible slow trajectory in Lemma 24, and fix T > 0. Let C∗ be the spectral constant in Lemma 17 and set R ∶= C∗2 . Then there exist constants R < s− < s+ < ∞, εT > 0 and CT < ∞ such that the following holds. If ν0 ∈ Xs0− satisfies ∥ν0 − µ̄(0)∥s− ≤ εT , (894) then the extended ODE

dν = Vext (ν), ν(0) = ν0 , dτ has a unique solution on [0, T ], remains in Us , and satisfies

(895)

sup ∥ν(τ ) − µ̄(τ )∥s+ ≤ CT ∥ν0 − µ̄(0)∥s− .

0≤τ ≤T

Proof. By Lemma 24,

∣µ̄α (τ )∣ ≤ R∣α∣ ,

0 ≤ τ ≤ T.

(896)

(897)

Moreover, admissibility and compactness of the time interval imply B dT ∶= min (t̄A k (τ ) + t̄k (τ )) > 0.

(898)

Kq ∶= {q̄(τ ) ∶ 0 ≤ τ ≤ T }

(899)

0≤τ ≤T k∈K

The covariance trajectory

is a compact subset of Q ∩ Q. Since Q is open, there exists δT > 0 such that {q ′ ∈ Q ∶ dist(q ′ , Kq ) ≤ 2δT } ⋐ Q.

(900)

Because P is non-expansive, whenever ∥q − q̄(τ )∥F < δT for some τ ∈ [0, T ], we also have ∥P(q) − q̄(τ )∥F < δT ,

(901)

and hence P(q) ∈ Q. Choose any s > R. By (897), the trajectory µ̄([0, T ]) is bounded in Xs0 . We may therefore choose an open neighborhood NT of this trajectory such that, for every µ ∈ NT , P(q(µ)) ∈ Q, On NT , the maps

µ z→ m⋆K (P(q(µ))),

B min (tA k (µ) + tk (µ)) ≥ k∈K

(r)

µ z→ akl (µ), 97

dT . 2

µ z→ βkext (µ)

(902)

(903)

are uniformly bounded and locally Lipschitz. In particular, the denominator in βkext is uniformly bounded below by dT /2. The same estimates used in Lemma 11 therefore imply that for every bounded subset of NT and every 0 < s < s′ , there exists LT < ∞ such that ∥Vext (µ) − Vext (ν)∥s′ ≤

LT ∥µ − ν∥s . s′ − s

(904)

The Ovsyannikov theorem, together with (904), gives local existence, uniqueness, and locally Lipschitz dependence on the initial condition in the Banach scale. Since µ̄([0, T ]) is compact and remains a positive distance from the boundary specified in (902), the interval [0, T ] can be covered by finitely many such local existence intervals. Composing the corresponding local stability estimates yields constants εT > 0 and CT < ∞ such that (894) implies existence on the whole interval [0, T ] and (896). (896) guarantees that the perturbed trajectory remains inside NT . Hence it cannot leave the domain of Vext , which closes the continuation argument. We now define the degree-M truncated slow dynamics by d (M ) µ (τ ) = [Vext (µ(M ) (τ ))]α , dτ α with the boundary condition

) µ(M (τ ) ≡ 0, α

∣α∣ ≤ M,

∣α∣ > M,

(905) (906)

and initial condition

⎧ ⎪ ⎪µ̄α (0), ∣α∣ ≤ M, ) µ(M (0) = ⎨ (907) α ⎪ 0, ∣α∣ > M. ⎪ ⎩ Notice that YM is invariant under Vext . Hence (905)–(907) is exactly the restriction of the extended infinitedimensional equation to the invariant finite-dimensional subspace YM . The following corollary proves that the truncated system can approximate the original system arbitrarily well. Corollary 8. Under the conditions of Theorem 3, fix T > 0. Let R = C∗2 , with C∗ the constant in Lemma 17. There exist constants R < s− < s+ < ∞, CT < ∞, and M0 (T ) < ∞, all independent of M , such that for every M ≥ M0 (T ), the truncated system (905)–(907) has a unique solution on [0, T ]. Moreover, for every finite word α, R M +1 ∣α∣ ) sup ∣µ(M (τ ) − µ̄ (τ )∣ ≤ C ( ) s+ . (908) α T α s− 0≤τ ≤T Proof. For M ≥ 2, the projection ΠM leaves all first- and second-order moments unchanged: q(ΠM µ̄(0)) = B A B q(µ̄(0)) and tA k (ΠM µ̄(0)) + tk (ΠM µ̄(0)) = t̄k (0) + t̄k (0) > 0. Hence ΠM µ̄(0) belongs to Us . By (897), ∣µ̄α (0)∣ R w R M +1 ≤ sup ( ) = ( ) . sw s− w>M s− w>M ∣α∣=w −

∥ΠM µ̄(0) − µ̄(0)∥s− = sup sup

(909)

Since R < s− , the right-hand side converges exponentially to zero. Therefore, for all sufficiently large M , ∥ΠM µ̄(0) − µ̄(0)∥s− ≤ εT ,

(910)

where εT is defined in Lemma 25. Apply Lemma 25 with ν0 = ΠM µ̄(0). The resulting solution of the extended ODE exists uniquely on [0, T ], and R M +1 . (911) sup ∥ν(τ ) − µ̄(τ )∥s+ ≤ CT ( ) s− 0≤τ ≤T Since ν0 ∈ YM and YM is invariant, we have ν(τ ) ∈ YM for 0 ≤ τ ≤ T . Consequently, ν is precisely the solution µ(M ) of the truncated system (905)–(907). 98

Finally, for every word α, ∣α∣

) ∣µ(M (τ ) − µ̄α (τ )∣ ≤ s+ ∥µ(M ) (τ ) − µ̄(τ )∥s . α +

(912)

Combining this with (911) proves (908). The slaved mean variables can be reconstructed from the truncated moments by (M )

mK (τ ) ∶= m⋆K (P(q(µ(M ) (τ )))) ,

(M )

m∖K (τ ) = m∖K (0).

(913)

Since m⋆K is Lipschitz on the compact covariance neighborhood in (900), and since P is non-expansive, (911) also gives R M +1 (M ) (914) sup ∥mK (τ ) − m̄K (τ )∥ ≤ CT′ ( ) s− 0≤τ ≤T for some CT′ < ∞.

F.3

Proof of Corollary 4

We first prove part (i). By Theorem 3, the fast dynamics move the mean variable to m⋆1 (q̄ (1) (0), q̄ (2) (0)), while the normalized moments remain asymptotically frozen. The slow trajectory therefore starts from q̄12 (0) = 0 with the mean variables m̄⋆ . B Set D1 (0) ∶= t̄A 1 (0) + t̄1 (0) > 0 and write (r)

a12 ∶= [∇q(r) Φ (m̄⋆ , q̄ (1) (0), q̄ (2) (0))]12 , (1)

r = 1, 2.

(915)

(2)

⋆ Thus g12 = (a12 , a12 )T . (1) Consider first q̄12 = µ̄(M1 ,M2T ) . Since only the first index is trainable, the slow dynamics (32) gives 2 2 d (1) (r) (1) q̄12 = −β̄1 µ̄(A1 +B1 ,M2T ) − 2 ∑ ∑ a1l µ̄(M (r) B +A M (r) ,M T ) − 2γ q̄12 . 1 1 2 l l dτ l=1 r=1

(916)

At τ = 0, Assumption 6 implies freeness of the student and teacher algebras. Moreover, limd→∞ d1 TrM2 = 0. Hence µ̄(A1 +B1 ,M2T ) (0) = 0, (917) and all the l = 1 terms in (916) vanish. For the l = 2 terms, asymptotic freeness gives

Therefore

µ̄(M2 B1 +A1 M2 ,M2T ) (0) = D1 (0)q̄22 ,

(1)

(918)

µ̄(M2T B1 +A1 M2T ,M2T ) (0) = D1 (0)q̄22 .

(2)

(919)

d (1) (1) (1) (2) (2) q̄ ∣ = −2D1 (0) (a12 q̄22 + a12 q̄22 ) . dτ 12 τ =0

(920)

(2)

Similarly, applying (32) to q̄12 = µ̄(M1 ,M2 ) gives d (2) (1) (2) (2) (1) q̄ ∣ = −2D1 (0) (a12 q̄22 + a12 q̄22 ) . dτ 12 τ =0

(921)

d ⋆ q̄12 (τ )∣ = −2D1 (0)Q2 g12 . dτ τ =0

(922)

Combining (920) and (921),

By (43), the right-hand side is nonzero. Define v0 ∶= ∥

d q̄12 (τ )∣ ∥ > 0. dτ τ =0 ∞ 99

(923)

The limiting slow trajectory is continuously differentiable. Hence, there exists an index r0 ∈ {1, 2} and T0 > 0 such that v0 d (r ) 0 ≤ τ ≤ T0 . (924) ∣ q̄120 (τ )∣ ≥ , dτ 2 (r )

(r )

(r )

d d q̄120 (τ ) has the same sign as dτ q̄120 (0) throughout [0, T0 ]. Since q̄120 (0) = 0, In particular, by continuity dτ

v0 T0 . 2

(925)

∥q̄12 (T+ )∥∞ ≥ 2c.

(926)

∥q̄12 (T0 )∥∞ ≥ Set c ∶= v04T0 and T+ ∶= T0 . Then

On the other hand, continuity and q̄12 (0) = 0 imply that there exists T− ∈ (0, T+ ) such that c sup ∥q̄12 (τ )∥∞ < . 2

(927)

0≤τ ≤T− (d)

Let ̃ q12 (τ ) denote the interpolation of the empirical overlap. Theorem 3 gives (d)

sup ∥̃ q12 (τ ) − q̄12 (τ )∥

τ ∈[0,T+ ]

Ð → 0. P

∞

(928)

Using (926) and (927), with probability tending to one, ⌊

T− d T+ d (d) ⌋ < Nwr (c) ≤ ⌈ ⌉. αd αd

(929)

This proves part (i). We next prove part (ii). Let µ̄(0) ∈ Mfree be the limiting initial moment sequence. Consider any state µ ∈ Mfree . In the slow vector field (32), the only substitutions which explicitly insert a teacher letter into diff the evolution of a trainable student letter are Ti,2,r . Their coefficients are [∇q(r) Φ (m⋆ (q), q (1) , q (2) )]12 . By assumption (46), these coefficients vanish on Mfree . The same observation applies to the l = 2 contribution to ψ1 , and hence to the effective bias β1 . Consequently, when the state belongs to Mfree , the limiting slow dynamics do not generate mixed student–teacher moments. Thus, Mfree is invariant under the limiting slow flow. By uniqueness of the admissible solution in Theorem 3, µ̄(τ ) ∈ Mfree for 0 ≤ τ ≤ T . In particular, (1) (2) q̄12 (τ ) = q̄12 (τ ) = 0 for 0 ≤ τ ≤ T . Applying again the uniform convergence in Theorem 3, for every fixed T > 0, (d)

sup ∥̃ q12 (τ )∥

τ ∈[0,T ]

Hence, for every fixed c > 0, (d) P (Nwr (c) ≤ ⌊

Ð → 0. P

∞

(930)

Td ⌋) Ð→ 0. αd

(931)

This finishes the proof.

F.4

Convergence rate of strong recovery

1 B We next study the convergence rate after weak recovery. Recall the first-order moments tA k ∶= d Tr[Ak ], tk ∶= (1) (2) 1 Tr[Bk ] and the covariance parameters qkl ∶= d1 Tr[Mk MlT ], qkl ∶= d1 Tr[Mk Ml ]. d

The limiting regularized

γ A B ∑ (t + tk ) . 2 k∈K k

(932)

population loss is

Ψ(m, tA , tB , q (1) , q (2) ) ∶= Φ(m, q (1) , q (2) ) +

Along the slow trajectory of Theorem 3, the mean variables are slaved to the covariance, m̄K (τ ) = m⋆K (q̄ (1) (τ ), q̄ (2) (τ )) . 100

(933)

Let Quntied denote the asymptotically realizable domain of (m, tA , tB , q (1) , q (2) ), namely the closure of all limiting order parameters generated by sequences of matrices Mk = Uk VkT , Ak = Uk UkT , Bk = Vk VkT subject to the fixed teacher components for k ∉ K. Define Ψ⋆ ∶=

inf

(m,tA ,tB ,q (1) ,q (2) )∈Quntied

Ψ(m, tA , tB , q (1) , q (2) ).

(934)

We refer to convergence to the global minimizers of Ψ as strong recovery. Let ϑ̄(τ ) ∶= (t̄A (τ ), t̄B (τ ), q̄ (1) (τ ), q̄ (2) (τ )) denote the corresponding components of the limiting slow trajectory. As in the tied setting, we impose an effective Polyak–Łojasiewicz condition along this trajectory. Assumption 11. There exist constants T0 > 0 and c > 0 such that, for every τ ≥ T0 , −

d Ψ (m̄(τ ), t̄A (τ ), t̄B (τ ), q̄ (1) (τ ), q̄ (2) (τ )) ≥ 2c [Ψ (m̄(τ ), t̄A (τ ), t̄B (τ ), q̄ (1) (τ ), q̄ (2) (τ )) − Ψ⋆ ] . dτ

(935)

For finite d, define the regularized population loss ⎡ L,L,K ⎤ ⎢ ⎛ 1 ⎞⎥⎥ γ 2 2 Ψd (U, V ) ∶= Ex ⎢⎢L { √ xTi Uk VkT xj } ⎥ + 2d ∑ (∥Uk ∥F + ∥Vk ∥F ) . ⎠ ⎥ ⎢ ⎝ d k∈K i,j,k=1 ⎦ ⎣

(936)

Corollary 9. Suppose that the conditions of Theorem 3 and Assumption 11 hold. Let n0 ∶= ⌈ Tα0dd ⌉. Then, for every fixed target accuracy ϵ > 0, there exists nϵ = n0 + O ( such that

d 1 log+ ) αd ϵ

(937)

lim P (Ψd (U (nϵ ) , V (nϵ ) ) − Ψ⋆ > ϵ) = 0,

(938)

d→∞

where log+ (x) ∶= max{0, log x}. In particular, if T0 = Θ(1), reaching an ϵ-suboptimal population loss requires at most O ( αdd log 1ϵ ) additional samples after entering the strong-recovery regime. For αd = Θ((d log d)−1 ), this becomes O (d2 log d log 1ϵ ).

Proof. Define the limiting excess regularized population loss ∆(τ ) ∶= Ψ (m̄(τ ), t̄A (τ ), t̄B (τ ), q̄ (1) (τ ), q̄ (2) (τ )) − Ψ⋆ ≥ 0. By Assumption 11, for every τ ≥ T0 ,

(939)

d ∆(τ ) ≤ −2c∆(τ ). dτ

(940)

Hence Grönwall’s inequality gives ∆(τ ) ≤ ∆(T0 )e−2c(τ −T0 ) ,

τ ≥ T0 .

(941)

Set ∆0 ∶= ∆(T0 ). If ∆0 = 0, there is nothing to prove. Otherwise, for a fixed ϵ > 0, define Tϵ ∶= T0 + Then

1 2∆0 log+ ( ). 2c ϵ

(942)

ϵ Ψ (m̄(Tϵ ), t̄A (Tϵ ), t̄B (Tϵ ), q̄ (1) (Tϵ ), q̄ (2) (Tϵ )) − Ψ⋆ ≤ . 2

(943)

We next transfer this bound to the finite-dimensional population loss. For τ ≥ 0, define U (d) (τ ) ∶= U (⌊τ d/αd ⌋) and V (d) (τ ) ∶= V (⌊τ d/αd ⌋) . For every fixed 0 < T1 < T2 < ∞, Theorem 3, together with the uniform spectral bounds in Lemma 17, implies sup τ ∈[T1 ,T2 ]

∣Ψd (U (d) (τ ), V (d) (τ )) − Ψ (m̄(τ ), t̄A (τ ), t̄B (τ ), q̄ (1) (τ ), q̄ (2) (τ )) ∣ Ð → 0, P

101

(944)

where we use Theorem 1. Since T0 > 0 and Tϵ < ∞, we may apply (944) on [T0 , Tϵ ]. Consequently, ⎛ ϵ⎞ Ð→ 0. P ∣Ψd (U (d) (Tϵ ), V (d) (Tϵ )) − Ψ (m̄(Tϵ ), t̄A (Tϵ ), t̄B (Tϵ ), q̄ (1) (Tϵ ), q̄ (2) (Tϵ )) ∣ > 2⎠ ⎝

(945)

Set nϵ ∶= ⌈ Tαϵdd ⌉. Using (943) and (945), we obtain P (Ψd (U (nϵ ) , V (nϵ ) ) − Ψ⋆ > ϵ) Ð→ 0.

(946)

Finally, by (942), nϵ − n0 = O (

d 1 log+ ) , αd ϵ

(947)

which proves the claimed sample complexity.

F.5

Proof of Lemma 10

We concatenate the indices (i, j, k) into a single index a ∈ {1, . . . , P }, where P = L2 K = Θ(1), and write (d) G(d) = (Ga )P a=1 . As in the proof of Theorem 1, concatenate the input vectors into X = [xT1 , . . . , xTL ]T ∼ N (0, Σ), where (d) Σ = C ⊗ Id . For each a = (i, j, k), there exists a symmetric Ld × Ld matrix Aa such that Ga = X T Aa X. 1/2 1/2 1/2 ̃a ∶= Σ Aa Σ , we obtain a Gaussian quadratic form Writing X = Σ z, with z ∼ N (0, ILd ), and setting A T ̃ G(d) a = z Aa z.

(948)

By the same norm estimates as in the proof of Theorem 1, the assumptions ∥Mk ∥op ≤ C∗ imply, uniformly over a, C ̃a ∥op ≤ √ ̃a ∥F ≤ C, ∥A , ∥A (949) d (d)

where C depends only on C∗ , C, L, K. Moreover, maxk √1d ∣Tr(Mk )∣ ≤ R∗ implies that the means µa ∶= (d) ̃a ) are uniformly bounded. EGa = Tr(A Let (d) T ̃ ̃ Wa(d) ∶= G(d) (950) a − µa = z Aa z − Tr(Aa ), (d)

and let Z (d) be a centered Gaussian vector having the same covariance as W (d) . Thus GG = µ(d) + Z (d) . As in Theorem 1, Wick’s formula gives (d) (d) ̃a A ̃b ). Ξab ∶= Cov(Wa(d) , Wb ) = 2Tr(A

(951)

In particular, Ξ(d) is uniformly bounded. We now compare W (d) with Z (d) through Gaussian interpolation. Let Z (d) be independent of z, and define √ √ Yt ∶= µ(d) + t W (d) + 1 − t Z (d) , t ∈ [0, 1], (952) (d)

and F (t) ∶= E[f (Yt )]. Then F (1) = E[f (G(d) )] and F (0) = E[f (GG )]. (d) ̃a A ̃b z. The Gaussian integration by parts gives For the centered quadratic forms Wa , define Γab ∶= 2z T A E [Wa(d) ∂a f (Yt )] =

√

P

2 t ∑ E [Γab ∂ab f (Yt )] .

(953)

b=1

On the other hand, Gaussian integration by parts for Z (d) gives E [Za(d) ∂a f (Yt )] =

√

P

(d)

2 1 − t ∑ Ξab E [∂ab f (Yt )] . b=1

102

(954)

Differentiating F (t) and combining (953)–(954) yields F ′ (t) =

1 P (d) 2 ∑ E [(Γab − Ξab ) ∂ab f (Yt )] . 2 a,b=1

(955)

(d) ̃a and A ̃b are symmetric, It remains to control the fluctuation of Γab . By (951), E[Γab ] = Ξab . Since A we have (d) ̃a A ̃b ∥2 ≤ C∥A ̃a ∥2 ∥A ̃b ∥2 ≤ C , E ∣Γab − Ξab ∣ ≤ C∥A F op F d 2

(956)

where we used (949) in the last step. As in the final part of the proof of Theorem 1, for every fixed p < ∞, sup E∥W (d) ∥p < ∞.

(957)

d

The Gaussian vector Z (d) has uniformly bounded covariance, so its moments are uniformly bounded as well. Since the means µ(d) are uniformly bounded, it follows that sup sup E∥Yt ∥p < ∞

(958)

d t∈[0,1]

for every fixed p. Since the second derivatives of f grow polynomially, (958) therefore implies 2

2 sup sup E ∣∂ab f (Yt )∣ ≤ Cf .

(959)

d t∈[0,1]

Applying Cauchy–Schwarz to (955), and using (956) and (959), gives C ∣F ′ (t)∣ ≤ √ , d

t ∈ [0, 1],

(960)

uniformly over all admissible matrix collections. Integrating from 0 to 1, we conclude that C (d) ∣E[f (G(d) )] − E[f (GG )]∣ = ∣F (1) − F (0)∣ ≤ √ . d

(961)

This proves the lemma.

G

Declaration of LLM usage

We used LLMs in the following stages of the preparation of this work: editing (e.g., grammar, spelling, word choice), drafting sections of the paper, facilitating or running experiments, visualizing results for submission. All LLM-generated material used in this manuscript was manually verified by the authors to ensure that it performs as intended, in particular with regard to the proof and the numerical experiments.

103

Related documents

• On State Reduction in Linear Attention#1141482
• Infinite-Width Limit of a Single Attention Layer: Analysis via Tensor Programs#1129081
• Why shared attention vectors fail: a case for outcome-indexed tuning#985950
• Problems in High-Dimensional Estimation and Large Language Models#643872
• Modeling attention and binding in the brain through bidirectional recurrent gating.#159659