PaperlayerView exact revision on arXiv

Source-linked AI summary

High-Dimensional Learning Dynamics of Attention-Indexed Models

Yizhou Xu, Margarita Sagitova, Lenka Zdeborová, Florent Krzakala

arXiv:2609.03858v1cs.LGstat.ML

TL;DR

Attention training dynamics remain poorly understood in the high-dimensional extensive-rank regime. The paper analyzes attention-indexed models, deriving a finite-dimensional population-loss description and richer SGD dynamics. It finds that tied and untied parameterizations have different symmetry-breaking behavior and that untied recovery depends on the state selected by fast dynamics.

  • Problem

    Training dynamics for attention, especially with extensive-rank matrices in high-dimensional feature learning, remain poorly understood.

    Source: §1 · p. 1

  • Method

    The paper studies attention-indexed models and derives high-dimensional population-loss and online-SGD descriptions for tied and untied attention parameterizations.

    Source: §1 · p. 2

  • Results

    Tied attention yields weak recovery in Θ(d^2 log d) samples under a mild condition, while untied recovery on that scale occurs when the fast-selected state breaks the relevant symmetry.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Attention parameterization acts as an architectural implicit bias, with tied and untied factorizations overcoming or preserving uninformative symmetries through different mechanisms.

    Source: §5 · p. 14

    The detailed training dynamics cover tied and untied parameterizations, not the full query, key, and value dynamics of a general deep transformer.

    Source: §5 · p. 14

Abstract

from arXiv · show

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\in\mathbb{R}^{d\times d}$ can remain trapped in an uninformative state. Tied attention ($S=WW^\top$) induces an automatic symmetry-breaking mechanism and yields weak recovery in $\Theta(d^2\log d)$ samples. For untied attention, $S=UV^\top$, 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 $\Theta(d^2\log d)$ scale occurs when the state selected by the fast dynamics breaks the initial symmetry.

1 Introduction

The paper develops a high-dimensional account of extensive-rank attention, where population loss has a finite-dimensional description but online SGD has richer dynamics. It shows that attention parameterization changes weak-recovery behavior through distinct tied and untied mechanisms.

1 Introduction

  • Existing feature-learning analyses often omit a high-dimensional limit or focus on finite-rank attention, leaving extensive-rank dynamics poorly understood.The paper studies this regime because attention-matrix ranks can diverge with embedding dimension.Source: §1 · p. 1
  • Attention-indexed models describe losses through quadratic token interactions and can represent multi-layer, multi-head attention-only architectures.Earlier analyses of these models focused on static Bayes-optimal or empirical-risk minimization, while tractable extensive-rank dynamics were largely limited to simpler quadratic models.Source: §1 · p. 1
  • Direct optimization of S can remain trapped at an uninformative state, whereas tied and untied factorizations can overcome such symmetries through different mechanisms.The tied factorization provides automatic symmetry breaking; the untied factorization selects a state through its fast dynamics.Source: §1 · p. 2
  • The high-dimensional population-loss landscape is finite-dimensional, but online SGD generates an infinite hierarchy of matrix moments that admits exponentially accurate finite-order truncations.The resulting empirical trajectory converges to a deterministic infinite-dimensional ODE.Source: §1 · p. 2
  • Tied attention achieves weak recovery in Θ(d^2 log d) samples under a mild condition, while untied recovery on this scale depends on whether fast dynamics break the relevant symmetry.For untied attention, pre-activation means evolve quickly and matrix moments more slowly; when symmetry survives, dynamics remain uninformative on this sample scale.Source: §1 · p. 2

2 Attention-indexed models and their population loss landscape

This section establishes a finite-dimensional description of the population-loss landscape for attention-indexed models in the high-dimensional limit. The framework covers attention-only architectures, while the parameterization determines whether training dynamics close on those loss order parameters.

2 Attention-indexed models and their population loss landscape

  • Attention-indexed models cover standard multi-layer, multi-head attention-only networks and losses built from quadratic token interactions.The setup includes a teacher-student example with Gaussian input data and does not require Gaussian teacher weights.Source: §2 · p. 3
  • The loss reduction uses extensive-rank matrices and a finite collection of first- and second-order trace statistics to characterize the limiting distribution of the quadratic forms.The section assumes fixed-scale sequence length and index count, with regularity conditions on the loss.Source: §2 · pp. 3–4
  • In the high-dimensional limit, the population loss depends only on a finite collection of order parameters, despite the growing matrix parameter space.The limiting joint distribution of the quadratic forms is Gaussian under the stated assumptions.Source: §2 · p. 4
  • Under an appropriate realizability condition, optimizing over the high-dimensional matrices is asymptotically equivalent to optimizing over their finite set of limiting order parameters.The optimal asymptotic population loss is an optimization over the finite-dimensional feasible set Q.Source: §2 · p. 4
  • A finite-dimensional loss landscape does not imply finite-dimensional training dynamics: direct optimization over S closes on finitely many quantities, but tied and untied factorizations generate additional matrix moments.This distinction affects weak recovery.Source: §2 · p. 4

3 Online SGD of tied attention

The section derives tied-attention online-SGD dynamics and uses them to characterize weak recovery. Although the limiting loss is finite-dimensional, the training dynamics require an infinite moment hierarchy that can be approximated by truncation.

3 Online SGD of tied attention

  • The section first derives the tied-attention macroscopic dynamics, then uses them to compare weak recovery with direct optimization over S.Weak recovery is measured through the overlap of traceless components, which removes the potentially nonzero isotropic contribution.Source: §3 · p. 5, §3.2 · p. 7

3.1 Macroscopic dynamics

  • The limiting potential depends on finitely many order parameters, but tied-factor online SGD evolves through an infinite hierarchy of joint matrix moments.The hierarchy is not generally closed by any fixed finite collection of moments.Source: §3.1 · p. 6
  • Under Theorem 2's conditions, online SGD converges uniformly in probability to a unique admissible solution of the infinite-dimensional ODE.The theorem establishes a deterministic high-dimensional limit for the moment trajectory.Source: §3.1 · p. 6
  • The infinite-dimensional description is not merely a coordinate artifact: an exactly solvable quadratic example admits no finite-dimensional ODE with a continuous initialization map for all trajectories.This example rules out a general finite-dimensional closure of the dynamics.Source: §3.1 · p. 6
  • Truncating the moment hierarchy at degree M approximates the true dynamics with exponentially small error, and increasing M improves empirical agreement.In the reported L = 2 softmax, MSE setting, the degree-M truncations are compared with online SGD.Source: §3.1 · pp. 6–7, Figure 1 · p. 7
  • The truncation result is uniform in time under the stated regularity condition; without it, the error-bound constants generally depend on T.The condition is used to obtain a uniform-in-time truncation error bound.Source: §3.1 · p. 7

3.2 Weak recovery and its sample complexity

  • Tied attention yields weak recovery in Θ(d^2 log d) samples when its symmetry-breaking mechanism gives the initially uninformative state a nonzero informative-direction velocity.The nonzero velocity is tied to a nonzero cross-gradient of the population loss at initialization.Source: §3.2 · p. 8
  • For higher-order activations, tied attention escapes the uninformative state while direct optimization over S can remain uninformative on the displayed sample scale.For linear activation, both parameterizations develop overlap.Source: §3.2 · p. 8, Figure 2 · p. 9

4 Online SGD of untied attention

This section characterizes online SGD for untied attention, focusing on fast mean relaxation and slower matrix-moment evolution. It describes how weak recovery depends on the fast phase producing a nonzero signal that allows slower dynamics to escape the uninformative state.

4 Online SGD of untied attention

  • The fast phase does not create O(1) teacher overlap but can shift the pre-activation mean, breaking student-side symmetry so slow dynamics escape the uninformative manifold.If the relevant symmetry survives fast relaxation, this recovery mechanism is absent.Source: §4 · p. 9

4.1 Macroscopic dynamics

  • On the fast timescale, only pre-activation means undergo O(1) evolution; on the slow timescale, those means follow evolving covariance while matrix moments undergo O(1) evolution.The slow moment dynamics form an infinite hierarchy and include an effective bias inherited from the fast phase.Source: §4.1 · pp. 10–12

4.2 Weak recovery and its sample complexity

  • Weak recovery occurs in Θ(d^2 log d) samples when fast relaxation leaves a nonzero cross-gradient; if cross-gradients vanish throughout the uninformative manifold, recovery does not occur on that scale.This dichotomy is characterized under the uninformative-initialization and macroscopic-dynamics conditions.Source: §4.2 · p. 13
  • Θ(d/α_d) samples yield weak recovery when the fast phase makes E[σ′(Z₁)] nonzero and the teacher-side first-order signal is nonzero.The fast phase can shift the mean so that E[σ′(Z₁)] becomes nonzero even if it vanishes initially.Source: §4.2 · p. 14
  • If the state selected by the fast phase fails to break the student-side symmetry, the network stalls.Source: §4.2 · p. 14
  • When E_{z₂²}[σ′*(Z₂)] = 0, weak recovery takes substantially more samples, and its time may depend on the activation's information exponent.The dependence on the information exponent is suggested by the numerical illustration and left for future work.Source: §4.2 · p. 14

5 Conclusion and future work

The conclusion presents a finite description of the loss landscape alongside richer but finitely approximable SGD dynamics, and shows that attention parameterization changes recovery routes. It also identifies the limits of the analyzed dynamics and several open extensions.

5 Conclusion and future work

  • The high-dimensional framework gives a finite trace-order-parameter loss landscape, while online SGD has an infinite moment hierarchy approximable by finite truncations with exponentially decaying error.The conclusion also presents deterministic limits for these dynamics.Source: §5 · p. 14
  • Tied attention yields weak recovery in Θ(d^2 log d) samples; untied recovery on that scale depends on fast symmetry breaking, unlike direct optimization's possibly uninformative state.The fast phase selects the state from which feature learning begins.Source: §5 · p. 14
  • Scope remains tied and untied dynamics; general deep-transformer query, key, and value dynamics and other optimizers remain unanalysed, while longer-timescale and higher-order recovery scales remain open.Longer horizons require tracking finite-d corrections over polynomial training times; the figure reports delayed recovery for higher-order activations.Source: §5 · pp. 14–15, Figure 3 · p. 14

B.1 Example 1: Multi-layer multi-head attention

  • Standard multi-layer, multi-head attention-only transformers fit the framework as finite composites of quadratic forms with effective weight matrices.The construction assumes continuous, bounded softmax and does not require Gaussian teacher weights.Source: §B.1 · pp. 19–20

B.2 Example 2: Autoregressive sequence model

  • The simplified autoregressive example casts next-token prediction with a single-layer, multi-head attention block as an attention-indexed model.The task predicts the next token from the preceding context.Source: §B.2 · pp. 20–21

B.3 Example 3: Multiple-location regression

  • The classification example represents labels through token-pair interactions, with polarity and grammatical relevance encoding sentiment-style dependencies.The dependency graph is activated probabilistically by semantic matching.Source: §B.3 · p. 21

B.4 Example 4: Multi-step reasoning

  • The multi-step reasoning example models latent token retrieval followed by quadratic interaction, with a two-head student aggregating operands before composition.The student forms two latent vectors from context tokens and predicts through their interaction.Source: §B.4 · pp. 21–22

B.5 Example 5: Extensive-rank matrix denoising

  • Extensive-rank matrix denoising is a single-token, single-index special case of the framework.Source: §B.5 · p. 22

C.1 Proof of Theorem 1

The proof establishes Gaussian limits for the quadratic-form statistics by showing that higher-order joint cumulants vanish. For tied attention, a centered variant accommodates cases where the uncentered mean diverges.

C.1 Proof of Theorem 1

  • The limiting first moment is E[G_ijk] = C_ijµ_ijk, and the covariance converges to a finite expression in C, Ω, and Ψ.The covariance limit is C_ii′C_jj′Ω_ijk,i′j′k′ + C_ij′C_ji′Ψ_ijk,i′j′k′.Source: §C.1 · pp. 22–23
  • As d →∞, all joint cumulants of order 3 and higher vanish, yielding a multivariate Gaussian limit for the quadratic-form vector.The proof uses joint cumulants and convergence of the first two cumulants; it then establishes uniform integrability to exchange the limit and population-loss expectation.Source: §C.1 · pp. 22–24

C.2 Tied attention

  • Tied weights such as S = WW^T can make the uncentered mean term diverge, so the proof introduces a centered attention-indexed variant.The centered variables have zero mean for every d, retain the covariance structure of the original variables, and have vanishing higher-order joint cumulants.Source: §C.2 · p. 24
  • The centered variables converge to a multivariate Gaussian with mean zero and the derived asymptotic covariance structure.Their higher-order joint cumulants match those of the original variables and vanish as d grows.Source: §C.2 · p. 24

C.3 Proof of Corollary 1

The proof establishes the limiting optimum by matching upper and lower bounds under conditions on minimizers. It also obtains local well-posedness for the infinite moment-sequence flow.

C.3 Proof of Corollary 1

  • Under the stated conditions, matching upper and lower bounds establish lim_d→∞ inf_S_K F_d(S_K) = inf_Q∈Q R(Q).The minimizers must satisfy a uniform spectral-to-Frobenius norm ratio condition, and their order parameters with the fixed weights must converge to a feasible point.Source: §C.3 · p. 25

D.1.1 Local well-posedness

  • The local well-posedness argument extends the gradient outside the positive-semidefinite cone by projection, while retaining agreement with the original gradient on the cone.This extension addresses that the truncated moment sequence need not have a positive-semidefinite second-moment matrix.Source: §D.1.1 · p. 26
  • For an initial moment sequence in X_s0, the flow has a unique local solution in X_s(t), with s(t) strictly increasing.The local existence and uniqueness result follows by establishing boundedness and local Lipschitz continuity of the vector field, then applying the Ovsyannikov theorem.Source: §D.1.1 · pp. 26–27

D.1.2 Uniform spectral bounds

The section proves a uniform spectral bound for the weights over any fixed training horizon. Its stopping-time argument controls sample truncation and martingale noise.

D.1.2 Uniform spectral bounds

  • For every fixed T, Lemma 3 asserts a uniform spectral bound with constants C* > C0 independent of d and T.The proof closes by showing the stopping time exceeds NT on the joint data and noise control events.Source: §D.1.2 · pp. 30, 34
  • The stopping-time proof truncates fresh samples and controls the stochastic-gradient martingale noise before ruling out a bound violation.On the intersection of the data and noise events, the stopping time J exceeds NT for sufficiently large d.Source: §D.1.2 · pp. 30, 32, 33, 34

D.1.3 Moment expansion

The moment expansion establishes an empirical-trajectory integral equation with vanishing errors, C-tight auxiliary moment processes, and uniform convergence of the empirical vector field under spectral bounds. Under the stated conditions, it also yields a unique admissible trajectory and convergence in probability to it, uniformly on each finite time horizon.

D.1.3 Moment expansion

  • The auxiliary moment process satisfies an integral equation with the empirical drift, and vanishing errors transfer the relation to the original trajectory.The proof bounds the expansion remainder and uses Markov’s inequality to establish the result for the original trajectory.Source: §D.1.3 · pp. 35, 37, 39

D.1.4 Compactness of the empirical trajectories

  • The auxiliary empirical moment processes are C-tight in D([0,T];X0_ρ) under the conditions of Theorem 2.The proof uses compactness and uniform Lipschitz bounds for an approximation asymptotically equivalent to the auxiliary processes.Source: §D.1.4 · pp. 39–40

D.1.5 Global well-posedness

  • Under the spectral bounds, the empirical vector field Vd converges uniformly to V, and the claimed convergence in Xρ follows.Source: §D.1.5 · pp. 40–41
  • The original equation has a unique admissible trajectory, which is deterministic when the initial condition is deterministic.The proof obtains a solution of the original equation, establishes uniqueness, and uses the deterministic initial condition to identify the limit as a deterministic trajectory.Source: §D.1.5 · pp. 43–44

D.1.6 Final proof of Theorem 2

  • Theorem 2 establishes convergence in probability, uniformly over each finite time horizon, from the original trajectory to the deterministic limiting trajectory.The proof first obtains convergence to the deterministic limit and then transfers it from the auxiliary trajectory to the original one.Source: §D.1.6 · p. 44

D.2 Proof of Corollary 2

The proof views finite-degree truncation as an invariant restriction of the extended moment ODE, then establishes stability and existence over the interval.

D.2 Proof of Corollary 2

  • The finite-dimensional truncation is the restriction of the extended ODE to an invariant subspace, initialized by projecting the initial moments.Moments above degree M have zero vector-field coordinates on that subspace.Source: §D.2 · p. 45
  • The proof establishes stability of the truncation by comparing its trajectory with the full moment solution under the stated bounds and regularity conditions.It uses local boundedness and local Lipschitz continuity of the vector field and a bootstrap argument.Source: §D.2 · pp. 45–46
  • For sufficiently large truncation degree M, the truncated solution exists throughout [0,T] and remains bounded by ∥q̄(t)∥F ≤ KR^2.The proof also rules out leaving the designated set U and finite-time blow-up.Source: §D.2 · pp. 45–46

D.3 An exactly solvable example

In a simplified exactly solvable setting, increasing truncation order improves agreement with online SGD, while the moment dynamics can be intrinsically infinite-dimensional. The accompanying Corollary 3 proof establishes movement away from zero overlap under its stated nonzero-condition argument.

D.3 An exactly solvable example

  • Truncations at degrees N = 2 and N = 3 activate dynamics and capture higher-order curvature, matching the matrix-flow Taylor expansion near t = 0.At N = 1, the first-moment solution is static; normalized traces of the matrix expansion reproduce the truncated moments’ linear terms.Source: §D.3 · p. 47
  • For S2 = cI, each eigenvalue of S1 follows a scalar flow, and the first-moment trajectory can uniquely determine a finite discrete initial spectrum.The resulting family rules out a fixed finite-dimensional ODE representation with continuous initialization for all such spectra.Source: §D.3 · pp. 47–48
  • As truncation order N grows, theory aligns better with online SGD for the linear-activation, L = 1, MSE example.Figure 4 uses d = 400, learning rate 0.01, batch size 512, and no weight decay.Source: §D.3 · p. 47, Figure 4 · p. 48

D.4 Proof of Corollary 3

  • Under the proof’s nonzero-condition argument, the cross-overlap derivative keeps its initial sign over a fixed interval, moving the overlap away from zero.The proof then combines its bounds to establish the stated Corollary 3 conclusion.Source: §D.4 · pp. 49–50

D.5 Convergence rate of strong recovery

The section bounds the additional samples needed for strong recovery under an effective PL condition along the limiting trajectory. In the normalized-teacher example, perfect recovery attains zero loss, with non-degeneracy needed only in relevant directions.

D.5 Convergence rate of strong recovery

  • When T0 = Θ(1), reaching an ϵ-suboptimal population loss requires at most O(d/ϵ) additional samples.The bound is stated under the conditions of Theorem 2 and Assumption 8.Source: §D.5 · p. 50
  • The convergence-rate argument transfers a finite-time bound on the limiting trajectory to the finite-dimensional regularized population loss.It uses convergence of first- and second-order moments and continuity of Φ.Source: §D.5 · p. 51
  • With normalized teacher q22 = 1, the global minimum Φ* = 0 is attained at perfect recovery, q11 = q12 = 1.The PL condition requires non-degeneracy only along directions relevant to S1 − S2, not necessarily across the full matrix space.Source: §D.5 · p. 52

E Online SGD over S

Direct online SGD over attention matrices has a finite-dimensional limiting description under the theorem’s assumptions and learning-rate conditions. Its weak-recovery behavior depends on whether the uninformative manifold is invariant: in that case, recovery does not occur on the O(d^2 log d) sample scale.

E Online SGD over S

  • Under Assumption 9 and the stated learning-rate range, direct S-SGD order parameters converge uniformly in probability to a deterministic trajectory.The trajectory closes at finitely many first- and second-order parameters.Source: E Online SGD over S · p. 53
  • If the relevant cross-gradients vanish throughout the uninformative manifold, that manifold is invariant and S-SGD cannot achieve weak recovery within O(d^2 log d) samples.This contrasts with tied factorization, which can escape an initially uninformative state on that sample scale.Source: E Online SGD over S · p. 54

E.1 Proof of Theorem 4

  • The proof combines uniform spectral bounds, martingale expansion of empirical order parameters, and Grönwall’s inequality to obtain the limiting ODE.Direct optimization closes at (t, Ψ, Ω), without higher-order moments or weighted moment-space estimates.Source: §E.1 · p. 54, §E.1.2 · p. 57

E.1.1 Uniform spectral bounds

  • The uniform spectral-bound argument controls trainable attention matrices over each fixed finite time interval.The theorem’s convergence result is proved using these bounds alongside the moment expansion.Source: §E.1.1 · pp. 54, 57

E.1.2 Moment expansion

  • The moment expansion shows that first-order terms match the finite-dimensional vector field, while martingale noise and higher-order residuals vanish.These estimates yield uniform convergence of the empirical order-parameter trajectory.Source: §E.1.2 · pp. 57–60

E.1.3 Final proof of Theorem 4

The proof establishes the empirical order-parameter result uniformly over a finite time interval, then verifies that the limiting trajectory solves the original ODE uniquely.

E.1.3 Final proof of Theorem 4

  • The empirical order-parameter result holds uniformly over t in [0,T].The proof derives this conclusion on the event E_d(T).Source: §E.1.3 · pp. 60–61
  • A globally Lipschitz extension of the vector field agrees with the original field on states satisfying the bound in (489).The extension is formed by projecting onto a closed ball whose radius exceeds that bound.Source: §E.1.3 · p. 60
  • Gronwall’s inequality controls the difference between the empirical and ODE quantities on E_d(T).The proof obtains the comparison by subtracting the integral form of (492) from (494).Source: §E.1.3 · p. 61
  • Closedness of D and the bounds in (489) ensure the limiting trajectory remains in the region where the extended field equals the original field.This places the limiting trajectory in D and lets the proof conclude it satisfies (402)–(404).Source: §E.1.3 · p. 61
  • Local Lipschitz continuity of the extension implies uniqueness of the ODE solution, completing Theorem 4.The extension is locally Lipschitz by Assumption 9.1, Price’s theorem, and non-expansiveness of the PSD projection.Source: §E.1.3 · pp. 60–61

E.2 Proof of Corollary 7

This section proves Corollary 7 by developing high-probability controls and finite-dimensional approximations for the mean variables and matrix moments, then establishing convergence to the unique admissible limiting solution. Together, these results show that S-SGD cannot achieve weak recovery within $O(d^2\log d)$ samples.

E.2 Proof of Corollary 7

  • Since C > 0 is arbitrary, the cited result establishes that S-SGD cannot achieve weak recovery within O(d2 log d) samples.Source: §E.2 · p. 63

F.1.1 A stronger version of Theorem 1

  • An independent proof in Appendix F.5 addresses the error scaling in Theorem 1 using a multivariate Gaussian approximation.The passage relates the estimate to existing work on multivariate Gaussian approximation.Source: §F.1.1 · p. 64

F.1.2 Local well-posedness

  • Under Assumptions 4 and 5, the extended ODE has a unique local solution that remains in the specified evolving neighborhoods.The result supplies a local time interval and a strictly increasing continuous scale function.Source: §F.1.2 · p. 65

F.1.3 Uniform spectral bounds

  • The proof establishes that SGD trajectories do not blow up: the weights’ spectral norms remain uniformly bounded across training steps.The section develops trace lower bounds, a finite-dimensional Stein expansion, and a one-step increment estimate in support of this bound.Source: §F.1.3 · pp. 65, 67, 68
  • Lemma 15 proves a uniform bound on the excess population loss up to the stopped horizon, with a positive decay rate.The proof identifies the rate as cloc = c1/2 and removes a high-probability event whose probability tends to one.Source: §F.1.3 · pp. 70, 72
  • Lemma 16 establishes a refined bound for the mean gradients under Theorem 3’s conditions and a stopping-time constraint.Its constants are specified in terms of the stated parameters, with some bounds chosen independently of γ.Source: §F.1.3 · pp. 72, 77
  • Blockwise control keeps the local dynamics within the required region with probability tending to one, so the local stopping time can be removed.The proof controls changes in A and martingale fluctuations across blocks before concluding the region is maintained through the horizon.Source: §F.1.3 · pp. 75–77
  • The uniform spectral estimate is extended to the initial training interval, completing (623) on an event whose probability tends to one.The proof bounds the initial interval using Lemma 13 and combines that bound with (682).Source: §F.1.3 · p. 77

Step 6: Final spectral bounds

  • Uniform spectral bounds hold for the trainable weights through the finite training horizon, with probability tending to one.The horizon is NT = ⌊T d^αd⌋; the proof combines the stopped estimates to establish the stated bound.Source: Step 6: Final spectral bounds · pp. 77, 79

F.1.4 Fast-dynamics convergence

  • On the fast timescale, slow coordinates remain frozen and the mean variables converge uniformly to their deterministic trajectory.The limiting trajectory dissipates Φ; the proof removes the stopping time before applying the piecewise-constant interpolation.Source: §F.1.4 · pp. 79–81

F.1.5 Moment expansion

  • For every fixed finite word, the one-step moment expansion has centered noise and a remainder controlled uniformly under the stopped spectral bounds.The proof bounds first-order remainders and shows contributions containing multiple letter increments vanish under Theorem 3’s learning-rate condition.Source: §F.1.5 · pp. 81, 83, 84
  • The interpolated moments satisfy an integral equation with errors that vanish over a fixed time horizon under the learning-rate assumption.The martingale error is controlled using its predictable quadratic variation and Doob’s maximal inequality.Source: §F.1.5 · pp. 84–85
  • Under Theorem 3’s conditions, the moments satisfy an approximate integral equation driven by a finite-dimensional effective vector field.The proof first controls the effective bias and uses that convergence to establish the moment approximation.Source: §F.1.5 · pp. 86, 89

F.1.6 Limiting slow dynamics

  • On fixed time intervals, the stopped empirical moment trajectories are tight, and the auxiliary and empirical processes agree with probability tending to one.The argument obtains compactness through uniform bounds and equicontinuity, then establishes asymptotic equivalence of the processes.Source: §F.1.6 · pp. 90–91

Step 2: Convergence of the effective vector field

  • Before the stopping time, the finite-dimensional effective vector field converges to the extended limiting vector field along empirical trajectories.The comparison uses uniform control of the component terms and a lower bound on the denominators.Source: Step 2: Convergence of the effective… · pp. 91, 93
  • Under Theorem 3’s conditions, the extended slow system has a unique admissible solution, which also satisfies the original ODE.The solution is constructed on the stated time interval in the relevant function space.Source: Step 2: Convergence of the effective… · p. 94
  • Every subsequential limit is admissible, solves the original system, and is the unique admissible solution with the given initial condition.Uniqueness follows because admissible solutions solve the extended ODE and trajectories agreeing at a time must agree thereafter.Source: Step 2: Convergence of the effective… · p. 95
  • The stopped moment trajectory converges uniformly in probability to the unique admissible solution.This follows from C-tightness and uniqueness of the subsequential limit.Source: Step 2: Convergence of the effective… · p. 95
  • The stopping time is asymptotically inactive because the empirical trajectory stays separated from the boundary across the late, initial fast-time, and intermediate regimes.The proof places the trajectory at positive distance from the boundary in each regime with probability tending to one.Source: Step 2: Convergence of the effective… · p. 95

F.1.7 Final proof of Theorem 3

The proof establishes convergence of the slow dynamics and transfers it to the SGD trajectory, then treats finite-order approximation and an invariant-state case.

F.1.7 Final proof of Theorem 3

  • Auxiliary moment trajectories converge uniformly to a unique deterministic admissible limit, and coincide with the original SGD trajectory through [0,T2] with probability tending to one.Source: §F.1.7 · p. 96

F.2 Truncation to finite orders

  • For sufficiently large M, the degree-M truncated system has a unique solution on [0,T] and can approximate the original system arbitrarily well.The truncation error bound decays exponentially with M.Source: §F.2 · p. 98

F.3 Proof of Corollary 4

  • When the limiting state lies in Mfree, the slow flow preserves it and generates no mixed student–teacher moments, so the specified overlaps remain zero.Source: §F.3 · pp. 99–100

F.4 Convergence rate of strong recovery

The section derives a conditional post-entry convergence rate for strong recovery and closes with a Gaussian-interpolation comparison used in Lemma 10.

F.4 Convergence rate of strong recovery

  • O(d/ϵ) additional samples suffice to reach an ϵ-suboptimal population loss after entering strong recovery when T0=Θ(1).Source: §F.4 · p. 101

F.5 Proof of Lemma 10

  • The Lemma 10 proof compares W(d) with a centered Gaussian Z(d) of matching covariance via Gaussian interpolation, with Ξ(d) uniformly bounded.Source: §F.5 · p. 102
Loading 2609.03858v1…