The Multiscale Single-Index Model: An MSE Analysis of Online SGD Dynamics

arXiv:2611.09682v1 [cs.LG] 6 Nov 2026

The Multiscale Single-Index Model: An MSE Analysis of Online SGD Dynamics

Gordon Dai New York University td2568@nyu.edu

Joan Bruna New York University jb4496@nyu.edu

Abstract We consider the Multiscale Single-Index Model (MSIM), first introduced in [27], as a stylized model for hierarchical feature learning with scale separation. Each layer extracts a shared single-index feature at one physical scale and passes it to the next, giving a tractable setting in which to study how deep architectures learn local multiscale representations. Under non-degeneracy and delocalization assumptions on the link function and planted features respectively, for the two hidden-layer case (K=2), we analyze joint online spherical SGD under MSE loss, which is a canonical regression objective, in discrete time and show that it achieves weak recovery from random initialization in O(d log2 d) samples, followed by strong recovery to accuracy ε in O(d log(1/ε)) additional samples. The main technical challenge arises from the non-Gaussian learned features observed by the second layer. We address this challenge through a quantitative Gaussian comparison using delocalization, combined with sharp trajectory-level control of the coupled SGD dynamics. This work is a companion to [5], which studies recovery under a correlation objective for general K.

1

Introduction

Feature learning, informally defined as the ability to extract informative low-dimensional representations out of high-dimensional data, is the defining characteristic of Neural Networks trained with gradient descent —and the basis for the modern AI paradigm of pretraining then fine-tuning. Over the recent years, a feature learning theory has emerged, focusing on a stylized model for high-dimensional learning based on multi-index models [2, 16, 12, 13, 14, 3, 11]; see [6] for a recent survey, demonstrating the ability of shallow neural networks to perform efficient feature learning, and even identifying sharp phase transitions [25]. A natural shortcoming for such theory is that shallow NNs are limited idealizations of modern NNs, where depth is a critical component. Towards addressing this gap, several works [1, 26, 28, 15, 31] have studied feature learning in deeper architectures, raising novel questions such as (i) elucidating the hierarchy between features at different layers, or (ii) how to adapt gradient methods to facilitate learning across all layers. In this work, motivated by applications in data domains such as images or text, we consider a stylized model for deep feature learning based on the principle of scale separation, where each layer of the network can be associated with a different physical scale. In the context of computer vision, the model becomes a Convolutional Neural Network (CNN) with disjoint receptive fields, as originally introduced in [27]; the first hidden-layer extracts local features at each patch, and subsequent layers capture local information at the next scale, progressively coarse-graining the domain [7, 22]. 40th Conference on Neural Information Processing Systems (NeurIPS 2026).

We instantiate each feature learning stage using a single-index model. Specifically, we consider a setting with Gaussian inputs X ∈ (Rd )⊗K , capturing a multiscale structure with K scales and ‘patches’ of size d, and a target obtained by a K-layer network, whereby each layer k ∈ [K] coarsegrains the feature map Z (k−1) ∈ (Rd )⊗K−k+1 as Z (k) = ϕ(Z (k−1) θk∗ ) ∈ (Rd )⊗(K−k) ; see Figure 1. The model X 7→ fθ∗ (X) := Z (K) ∈ R is thus parametrized by hidden directions or features ∗ θ1∗ , . . . , θK ∈ Sd−1 , and a non-linear activation function ϕ : R → R. By considering the Wiener chaos expansion of the target f (X), this model can be viewed as a ∗ non-linear analogue of Tensor PCA [24], where the planted directions θ1∗ , . . . , θK may be estimated from a spike in the associated chaos decomposition. This naturally leads to the study of two different algorithmic frameworks: (i) a spectral method, motivated by its efficiency in related problems, eg single- and multi-index models [23, 13, 20], spiked tensor PCA, and tensor-decomposition methods for non-overlapping CNNs [27]; and (ii) stochastic gradient-descent, inspired by the practical connection with deep learning. The latter setting is not addressed by these tensor-decomposition analyses and therefore is the main focus of our work. Our analysis makes two key structural assumptions: we consider the setting where ϕ is non-degenerate, in the sense that EG∼N (0,1) [ϕ′ (G)] := κ ̸= 0 1 , and further assume that the planted directions are delocalized, unaligned with the basis on which the non-linear activations act. Exploiting these two assumptions, we analyze online SGD in the case (K=2), the immediate multiscale extension of the single-index setting, we establish a sample complexity of n = O(d(log(d)2 + log ε−1 )) to achieve error ε, nearly matching the spectral estimator up to logarithmic factor. It is important to emphasize that the machinery of Ben Arous et al. [2] does not directly apply in this setting, due to the non-Gaussian features seen by the second hidden layer activations.

Related Works. Understanding the benefits of depth in Neural Networks is a longstanding question in the field, with initial progress centered on approximation questions [30, 17]. The efforts then moved towards incorporating optimization aspects, with early notable works such as [29, 21], and later developed in [32, 26, 1]. One aspect shared by all these works is the fact that the input highdimensional distribution is rotation-invariant, ie with no particular physical structure. Early works that also exploited tensor representations of data to capture scale separation are [10], which focused on expressivity. The closest works to ours are [15] and [28], which we discuss briefly. Dandi√et al. [15] study Single/Multi-Index Gaussian Hierarchical Targets f ⋆ (x) = g ⋆ (a⋆⊤ Pk (W ⋆ x)/ dϵ1 ) with W ⋆ ∈ ϵ1 Rd ×d , where the hierarchy is one of dimensions d → dϵ1 → 1. They train a three-layer MLP layer-wise and propagate Hermite decompositions across layers via the asymptotic Gaussianity of the inner statistic h⋆ , which holds because IE(Pk ) = 2. Our setup differs in three essentials: (1) the hierarchy is one of physical scales (within- vs. across-patch), with both layers acting in the same dimension d on single-index teachers β ⋆ , θ⋆ ∈ Sd−1 ; (2) both layers are trained jointly by online SGD on the same timescale, so the second-layer student sees a moving, non-Gaussian feature map throughout; and (3) at IE(ϕ) = 1 the inner statistic is not asymptotically Gaussian, so we replace the CLT by a quantitative Stein/Lindeberg comparison driven by the delocalization. Ren et al. [28] prove that L-layer convolutional networks trained layer-wise learn L-level Random Hierarchy Models with O(m(1+o(1))L ) samples, via a problem-independent “shallow-to-deep chaining” reduction: layer-wise training succeeds whenever the target correlates with lower-level latents, the signal to lower layers is clean, and the lower-level features are weakly identifiable. Their setting is complementary to ours: The data are discrete tokens drawn from a probabilistic context-free grammar rather than continuous Gaussian patches, and the targets are categorical class labels rather than scalar regressions, so the relevant geometry is a tree of synonym classes rather than the product sphere. Our companion work  [5] shows that under correlation loss between teacher model and student model E fθ (X)fθ⋆ (X) when initialized in a favorable basin of constant-in-d probability, online SGD achieves achieves weak e K−1 ) samples, and continues to 1 − od (1) overlap after a lower-order continuation recovery in O(d phase, for any fixed depth K. Yet it remains unclear how the online SGD behaves under standard loss such as Mean Square Error (MSE) loss, which is the main focus of this work.

1 This corresponds to information exponent 1 in the language of [2].

2

Figure 1: Schema of the Multiscale SIM model. Each patch within the same layer k shares the same local feature, parametrized by a SIM of direction θk . Notation: All expectations EX are with respect to X, with the directions held fixed, unless explicitly stated otherwise. For unit vectors u, v, define the sign-invariant distance dist± (u, v) := min{∥u − v∥2 , ∥u + v∥2 }.

2

The Multiscale Single-Index Model

Let K ≥ 1 be fixed. We model the input as a K-th order Gaussian tensor of dimension d, ie X ∈ (Rd )⊗K have independent N (0, 1) entries. Given unit vectors θ1 , . . . , θK ∈ S d−1 , define Z (0) = X, and recursively (r)

Yir+1 ,...,iK =

d X

(r−1)

θr (j)Zj,ir+1 ,...,iK ,

1 ≤ r ≤ K,

(1)

j=1

Z (r) = ϕ(Y (r) )

entrywise.

At the final layer Z (K) ∈ R, and we write fθ (X) := Z (K) . As illustrated in Figure 1, this is a stylized model of a NN with physical structure: each layer computes local and translation-invariant features at a fixed spatial scale. In the case where K = 2, for instance, the first layer extracts a single shared ‘local’ feature θ1∗ capturing internal correlations within each patch, while the second layer extracts a coarse feature θ2∗ capturing correlations across patches. Our algorithmic approaches will be based on two main assumptions: Assumption 2.1 (Non-degeneracy and regularity, informal). Let G ∼ N (0, 1). We assume Eϕ(G) = 0,

Eϕ(G)2 = 1,

κ := Eϕ′ (G) ̸= 0.

(A1)

Further, ϕ satisfies certain standard regularity assumptions, which differ slightly on each setting (discussed in each section). In the language of Ben Arous et al. [2], non-degeneracy corresponds to the case where the information exponent of ϕ is 1. In the shallow setting K = 1, this assumption greatly simplifies the geometry of the model, ensuring a descent direction from random initializations. In our setting K > 1, the non-degeneracy is instrumental for gradient methods to escape the mediocrity, as will be shown in later sections. Assumption 2.2 (Incoherence, informal). The planted directions are delocalized, in the sense that ∥θk ∥∞ ≪ 1. Incoherence corresponds to the notion that no patch exerts an dominant influence on the output, and enables a mean-field approximation, or equivalently a Gaussian universality in our SGD analysis, similarly as in [18]. 3

3

Main Result

We now turn to verify that, indeed, K = 2 is a tractable setting for SGD dynamics to achieve strong recovery under the natural MSE loss rather than linear correlation in our companion paper [5]. For ease of presentation, we name θ1∗ = θ∗ and θ2∗ = β ∗ in the remainder of this section. 3.1

Setup

Given the teacher direction (θ∗ , β ∗ ) ∈ Sd−1 × Sd−1 we want to study the ability of online stochastic gradient descent (online SGD) to minimize the population loss L(βt , θt ) := ∥fβt ,θt − fβ ∗ ,θ∗ ∥2 = Ex∼γD [fβt ,θt (x) − fβ ∗ ,θ∗ (x)]2 , (2) ∗ ∗ where the label is generated by zi := ϕ(Xi · β ), fβ ∗ ,θ∗ (X) := ψ(z · θ ), i ∈ [d]. To recover the teacher, we run online SGD on Sd−1 × Sd−1 : βt − δ∇Sβ ℓ(βt , θt ; xt ) θt − δ∇Sθ ℓ(βt , θt ; xt ) , θ = , (3) βt+1 = t+1 βt − δ∇Sβ ℓ(βt , θt ; xt ) θt − δ∇Sθ ℓ(βt , θt ; xt ) for t = 0, 1, 2, . . . , where (xt ∈ Rd×d ) are fresh samples independent of the past, δ is the step-size, and the spherical gradient is defined as: ∇Sβ ℓ(β, θ; x) := (I − ββ ⊤ )∇β ℓ(β, θ; x),

∇Sθ ℓ(β, θ; x) := (I − θθ⊤ )∇θ ℓ(β, θ; x).

Following [2, 8], throughout the learning process, we track the evolution of summary statistics, defined as overlaps ρt := βt · β ∗ and χt := θt · θ∗ for the first and second layer’s parameter respectively. Theorem 3.1 (Informal). With our assumptions (A1)-(A4) (see Section B), for any ε > 0, under online SGD training (3), we have: • Weak Recovery: starting from initialization, with probability at least 1 − Kε, choosing step size δ ≍ 1/(d log d), SGD achieves weak recovery: min{ρt , χt } = Θ(1), using O(d log2 d) samples. • Strong Recovery: conditioned on the weak recovery event, with probability at least 1 − O (log(1/ε)/d) , choosing step size δ ≍ 1/d, SGD achieves strong recovery: min{ρt , χt } ≥ 1 − ε, using O (d log(1/ε)) additional samples. We build our analysis on top of the analytical framework introduced in Ben Arous et al. [2] by modeling the learning process as drifts with fluctuations. The main technical point is to leverage delocalization (which, as remarked in B.2, will be true w.h.p. under (A4)) to control the gaussian approximation LG of the true loss L to derive the continuous-population level on the gradient flow dynamics: Based on Stein’s Coupling (see Theorem E.1), for each layer, in Lemma C.2, we bound the distance between the spherical gradient of the true loss ∇S,β L and its Gaussian proxy ∇S,β LG by some constant factor of the incoherence ηt := ∥θt ∥∞ + ∥θ∗ ∥∞ that holds true uniformly in time: (1)

(1)

|⟨∇S,β (L − LG )(θt , βt ), β ∗ ⟩| ≤ Cβ ηt ,

|⟨∇S,θ (L − LG )(θt , βt ), θ∗ ⟩| ≤ Cθ ηt .

(4)

Continuous-time Analysis In continuous time, we consider the population-level spherical gradient flow defined as β̇t = −∇S,β L(θt , βt ), θ̇t = −∇S,θ L(θt , βt ). (5) Based on (4), in Lemma C.3 we derive a system of coupled ODE inequalities that captures the evolution of overlaps as ρ̇t ≥ aρ χt (1 − ρ2t ) − bρ ηt , χ̇t ≥ aχ ρt (1 − χ2t ) − bχ ηt , (6) where aρ , aχ , bρ , bχ > 0 are independent of d and (θt , βt ). From here, we can see that for the RHS to be non-vacuous and provide enough signal for the growth of (ρt , θt ), terms bρ ηt and bχ ηt must not exceed aρ χt (1 − ρ2t ) and aχ ρt (1 − χ2t ), respectively. This requires us to provide a quantitative control of ηt , derived in Lemma C.4: η̇t ≤ C1 ρt ∥θ∗ ∥∞ + C2 χt ρt ηt + O(ηt2 ) . (7) We make the following remarks on our setting: 4

• On the Delocalization hypothesis: our proof crucially relies on the Gaussian comparison in order to obtain an overlap growth differential inequality. Delocalization is not necessary for efficient recovery, e.g. if ∥θ∗ ∥0 = Od (1) the target is a multi-index model, for which gradient-based guarantees exist [12, 4]. A natural follow-up question is to understand whether our SGD can obtain recovery guarantees uniform on θ∗ . • On non-degeneracy The key mechanism for learning is described in (6). Whenever the information exponent of ϕ is s > 1, one would instead obtain ρ̇t ≥ aρ χt ρs−1 (1 − ρ2t ) − bρ ηt , t

χ̇t ≥ aχ ρt χs−1 (1 − χ2t ) − bχ ηt , t

(8)

in which the perturbative ηt term dominates the positive drift at initialization. This difficulty is not merely technical; our numerical simulations in Section A show that online SGD does not manage to learn in this setting. • On Gaussian inputs For simplicity we state our result in the Gaussian setting, but using delocalization in the first layer feature, our arguments naturally extend to the class of well-behaved product measures for which Stein’s CLT coupling applies. 3.2

Main Elements of the Proof

Below, we present the high-level roadmap of proofs in getting Theorem 3.1. 3.2.1

Weak Recovery

The proof has two dynamical regimes, Phase I and Phase II, separated by τA := inf{t : Zt2 ≥ Aηt },

Zt := ρt + χt .

In Phase I, Zt is still below the incoherence scale. In this phase, the crude O(ηt ) comparison error is too large under spherical initialization; in Lemma C.5 we bound the deviation of true population drifts from its Gaussian proxy by a refined third-moment quantity ωt :=

d X

3 |θt,i | + |θi∗ | ,

(9)

i=1

which itself is controlled in Lemma C.6 and Lemma C.7. Under the small-overlap regime of weak recovery, the dynamics is well approximated by the coupled drift inequalities ρ̇t ≥ aχt − bωt ,

χ̇t ≥ aρt − bωt .

(10) −1/2

In Lemma C.8, we showed that (i) starting from initialization, (10) pulls Zt from d to the signal-dominant scale Zt2 ∼ ηt in O(log d) time; (ii) after τA , the incoherence error is of a lower order compared with the positive Z-drift. The sum Zt then grows from order d−1/4 to a fixed small constant, while the imbalance Dt = ρt − χt contracts. The entry estimate DτA = o(ZτA ) ensures that both overlaps become positive constants simultaneously under another O(log d) time, obtaining weak recovery. Extension to Discrete-Time SGD The discrete-time weak-recovery proof mirrors the same tworegime mechanism, with additional martingale and discretization controls: We start from the one-step expansion in Lemma D.2, (ρ)

ρt+1 = ρt + δFt

(ρ)

+ δMt

(ρ)

(χ)

+ Rt

χt+1 = χt + δFt

(χ)

+ δMt

(χ)

+ Rt ,

(11)

where we take the step size δ ≍ 1/d log d. During Phase I, plugging the ω-controlled population drift in Lemma C.6 into (11), we obtained the aggregate one-step update expression (Z)

Zt+1 ≥ (1 + βδ)Zt − Cδωt + εt , (Z)

(12)

where εt as the one-step increment of martingale and discretization errors, with their time(Z) accumulation is controlled by Abel summation from the event ET . Besides, for the third-moment term wt in (12), under Lemma D.12 as the discretized √ version of third-moment control made in Lemma C.7 we showed that for Phase I, ωt ≲ (1 + tδ)/ d, making the weighted forcing remains 5

of order d−1/2 . Therefore Zt grows from its initialization scale√ d−1/2 to the signal-dominant scale −1/4 −1 Z τA ≳ d in O(δ log d) iterations, while ηt ≲ (log d)/ d. Same as continuous case, we arrived at an entry-balance estimate DτA = o(ZτA ) using the small-overlap D-drift derived in D.15. Entering Phase II, as shown in Lemma D.16, once Zt2 ≳ ηt , the incoherence error is negligible compared with the positive Z-drift, so Zt grows from order d−1/4 to a fixed constant cZ within another O(δ −1 log d) iterations. At the same time, the D-recursion is contractive up to lowerorder forcing, which keeps |Dt | ≤ Zt /2. Hence, as shown in Theorem D.17, at the hitting time tesc := inf{t : Zt ≥ cZ }, both overlaps are bounded below with high probability: min{ρtesc , χtesc } = 3.2.2

Ztesc − |Dtesc | cZ ≥ . 2 4

Strong Recovery

Continuous-time Strong Recovery Results After weak recovery, we show in Lemma C.9 that under additional O(1) time, the population dynamics drives (θt , βt ) into a fixed local cap of arbitrary level around y ∗ := (θ∗ , β ∗ ), the global minimizer of the true population loss. While the small-overlap estimate from weak recovery no longer holds, we employ the uniform bounds in Lemma C.3 and Lemma C.4 to show ηt = od (1), which guarantees both overlaps enter the claimed local cap. Once inside this local cap, the proof switches from overlap drift to local geometry. We analyze the Hessian of the loss function by decomposing it as:

HL (y) = HL (y) − HL (y ∗ ) + HL (y ∗ ) − HLG (y ∗ ) + HLG (y ∗ ) . (13) | {z } | {z } | {z } (A)

(B)

(C)

From (13), we derive Lemma C.12 to lower bound term (B) by some constant factor of −∥θ∗ ∥∞ , ∗ and p Lemma C.13 to lower bound term (A) by some constant factor of − distM (y, y ) := − arccos(ρ)2 + arccos(χ)2 , taken arbitrarily small. Finally, as (C) is lower bounded by a Hermitecoefficient quantity that dominates the other constants, in Lemma C.14, we obtain the positivity of the Hessian for the true loss HL (y) throughout the local cap. Hence, under this implied local strong convexity, in Lemma C.15, by exploiting the the log-drift inequality 1 1 arccos(χt )2 + arccos(ρt )2 , (14) 2 2 we make a standard Riemannian gradient-flow argument to obtain the exponential contraction of distM (y, y ∗ ), in the sense that ∇M L(yt ), Logyt (y ∗ ) ≤ −2µloc Qt ,

Qt :=

distM (yt , y ∗ )2 ≤ e−2µloc (t−tesc ) distM (ytesc , y ∗ )2 .

(15)

Converting the product distance back to overlap gives the strong recovery of min{ρt , χt } ≥ 1 − ε

(16)

after an additional O(log(1/ε)) time from weak recovery. We formalize this result in Theorem C.16. Extension to Discrete-Time SGD Same as the extension to discrete-time SGD in Section 3.2.1, here we must control stochastic fluctuations of SGD inside the strongly-convex basin to ensure the local contraction. Finishing the discrete-time weak recovery gave by Theorem D.17, we switch the SGD step size to δloc = cloc d , that is optimal to strong recovery as remarked in RemarkD.18. Mirroring Lemma C.9, Proposition D.21 shows that after an additional O(d) iterations beyond weak recovery, the trajectory enters the local strongly-convex cap in discrete-time. Inside this cap, we track the local distance 1 1 (17) Qt := arccos(χt )2 + arccos(ρt )2 . 2 2 From its one-step recursion in Lemma D.22: 2 Qt+1 ≤ Qt + δloc ⟨ḡt , Lt ⟩ + δloc Ξt + Cret δloc Gt ,

6

(18)

we bound the population drift ⟨ḡt , Lt ⟩ by the log-drift inequality (14) in continuous time, and we bound both the martingale fluctuation term Ξt := ⟨gt − ḡt , Lt ⟩ and the quadratic gradient term Gt := ∥gt ∥2 in Lemma D.23 by the sharp local self-quenching estimate: E[Ξ2t | Ft ] ≲ Q2t , and E[Gt | Ft ] ≲ dQt ,

(19)

which show that both the gradient energy and the martingale variance vanish proportionally to the current distance to the teacher. Thus the SGD noise becomes weaker as the iterates approach the −1 teacher: Proposition D.24 proves that one local SGD epoch of length O(δloc ) contracts the distance scale by a factor of 1/2 with probability 1 − O(δloc ), while the trajectory remains inside the local
cap. Iterating dyadic epochs in Proposition D.25 yields Qt ≤ ε after O d log 1ε additional samples. We conclude our discrete-time result in Theorem D.26, where we obtained (16), namely the strong
recovery under discrete-time SGD, with a total sample complexity of O(d log2 d) + O d log 1ε .

4

Conclusions and Future Directions

In this work, we studied the Multiscale Single-Index Model, in which a dK -dimensional input with multiscale structure is iteratively coarse-grained via local, translation-invariant single-index models, and where the number of scales K corresponds directly to the depth of the associated neural network. Provided the non-linearities are non-degenerate and the planted directions are delocalized, we proved weak and strong recovery of online SGD in the three-layer setting K = 2 using O(d(log(d)2 ) samples. Our work leaves several open questions ahead of us: the first, most pressing one is to extend the SGD recoverability to general depth K > 2. While [5] provides d-asymptotic strong recovery for general K under the correlation loss, it remains unclear how a general K model behaves under canonical regression objectives such as MSE loss. In our current SGD analysis, the restriction to K = 2 is tied to the weak-recovery phase near random initialization. Indeed, the MSE objective decomposes as       LMSE (Θ) = E fΘ (X)2 − 2E fΘ (X)fΘ⋆ (X) + E fΘ⋆ (X)2 . (20) Q When mj = ⟨θj , θj⋆ ⟩ ≍ d−1/2 , the teacher-dependent correlation drift in mode i has size ℓ̸=i mℓ ≍ d−(K−1)/2 . In contrast, our current first-order bounds allow teacher-independent finite-d self-energy corrections at scale d−1/2 . For K = 2, these terms occur at the same scale, and our refined thirdmoment and incoherence estimates allow us to prove that the resulting two-overlap dynamics still has a positive cooperative drift to escape mediocrity. For K > 2, however, the planted drift is of strictly smaller order, and a first-order Gaussian or norm comparison no longer determines the sign of the overlap evolution. Moreover, unlike the perturbations in the correlation objective, the self-energy corrections are not naturally accompanied by planted-overlap factors. Figure 3, Section A shows the empirical evidence of hardness of MSE loss for K > 2. The next natural question is to consider the regime of fixed d (the so-called ‘receptive field’) and diverging depth K; this will also require substantial technical innovation, since our CLT/universality arguments do not directly apply. Finally, an important limitation of our work is that it imposes a form of scale separation: each layer lives on its dedicated scale. While this is a valid assumption for certain physical data (e.g images), most popular deep architectures do not have scale separation; for instance in ResNets [19]/ NeuralODEs [9] the features remain at the same scale throughout several layers. Building toy models that explain feature learning in these settings remains an outstanding challenge.

5

Acknowledgments

GD was partially supported by the Dean’s Undergradaute Research Fund from NYU’s College of Arts and Science. This work was completed during the JB sabbatical visit at the Flatiron Institute, which also provided the computing resources used in the project. The authors would also like to thank Hengzhi Li for helping them running experiments, and Dr. Denny Wu, Yunwei Ren, Hanzhe Wu for reviewing our draft and providing helpful feedback. 7

References [1] Z. Allen-Zhu and Y. Li. What Can ResNet Learn Efficiently, Going Beyond Kernels? In Advances in Neural Information Processing Systems, 2019. [2] G. B. Arous, R. Gheissari, and A. Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference, 2021. [3] J. Barbier, F. Krzakala, N. Macris, L. Miolane, and L. Zdeborová. Optimal errors and phase transitions in high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, 116(12):5451–5460, 2019. [4] A. Bietti, J. Bruna, and L. Pillaud-Vivien. On learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310.19793, 2023. [5] J. Bruna. The multiscale single-index model: A stylized model for hierarchical feature learning, 2026. [6] J. Bruna and D. Hsu. Survey on algorithms for multi-index models, 2025. [7] J. Bruna and S. Mallat. Invariant scattering convolution networks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 35(8):1872–1886, 2013. [8] J. Bruna, L. Pillaud-Vivien, and A. Zweig. On single index models beyond gaussian data, 2023. [9] R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. Neural ordinary differential equations, 2019. [10] N. Cohen, O. Sharir, and A. Shashua. On the expressive power of deep learning: A tensor analysis. In Conference on learning theory, pages 698–728. PMLR, 2016. [11] H. Cui, L. Pesce, Y. Dandi, F. Krzakala, Y. M. Lu, L. Zdeborová, and B. Loureiro. Asymptotics of feature learning in two-layer networks after one gradient-step. arXiv preprint arXiv:2402.04980, 2024. [12] A. Damian, J. Lee, and M. Soltanolkotabi. Neural Networks can Learn Representations with Gradient Descent. In Proceedings of the Thirty Fifth Conference on Learning Theory, 2022. [13] A. Damian, L. Pillaud-Vivien, J. D. Lee, and J. Bruna. Computational-statistical gaps in gaussian single-index models. arXiv preprint arXiv:2403.05529, 2024. [14] Y. Dandi, F. Krzakala, B. Loureiro, L. Pesce, and L. Stephan. Learning two-layer neural networks, one (giant) step at a time. arXiv preprint arXiv:2305.18270, 2023. [15] Y. Dandi, L. Pesce, L. Zdeborová, and F. Krzakala. The computational advantage of depth: Learning high-dimensional hierarchical functions with gradient descent, 2025. [16] R. Dudeja and D. Hsu. Learning single-index models in gaussian space. In Conference On Learning Theory, pages 1887–1930. PMLR, 2018. [17] R. Eldan and O. Shamir. The power of depth for feedforward neural networks. In Conference on learning theory, pages 907–940. PMLR, 2016. [18] R. Gheissari and A. Jagannath. Universality of high-dimensional scaling limits of stochastic gradient descent. arXiv preprint arXiv:2512.13634, 2025. [19] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition, 2015. [20] N. Joshi, T. Misiakiewicz, and N. Srebro. On the complexity of learning sparse functions with statistical and gradient queries. arXiv preprint arXiv:2407.05622, 2024. [21] E. Malach, G. Yehudai, S. Shalev-Schwartz, and O. Shamir. The connection between approximation, depth separation and learnability in neural networks. In Conference on Learning Theory, pages 3265–3295. PMLR, 2021. 8

[22] S. Mallat. Understanding deep convolutional networks. Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences, 374(2065):20150203, 2016. [23] M. Mondelli and A. Montanari. Fundamental limits of weak recovery with applications to phase retrieval. In Conference On Learning Theory, pages 1445–1450. PMLR, 2018. [24] A. Montanari and E. Richard. A statistical model for tensor pca. Advances in neural information processing systems, 27, 2014. [25] A. Montanari and Z. Wang. Phase transitions for feature learning in neural networks, 2026. [26] E. Nichani, A. Damian, and J. D. Lee. Provable guarantees for nonlinear feature learning in three-layer neural networks. Advances in Neural Information Processing Systems, 36, 2023. [27] S. Oymak and M. Soltanolkotabi. Learning a deep convolutional neural network via tensor decomposition. Information and Inference: A Journal of the IMA, 10(3):1031–1071, 2021. [28] Y. Ren, Y. Dandi, F. Krzakala, and J. D. Lee. Provable learning of random hierarchy models and hierarchical shallow-to-deep chaining, 2026. [29] I. Safran and J. Lee. Optimization-based separations for neural networks. In P.-L. Loh and M. Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 3–64. PMLR, 02–05 Jul 2022. [30] M. Telgarsky. Benefits of depth in neural networks. In Conference on learning theory, pages 1517–1539. PMLR, 2016. [31] Z. Wang, E. Nichani, and J. D. Lee. Learning hierarchical polynomials with three-layer neural networks. arXiv preprint arXiv:2311.13774, 2023. [32] Z. Wang, E. Nichani, and J. D. Lee. Learning hierarchical polynomials with three-layer neural networks, 2023.

9

Appendix Table of Contents A Numerical Experiments

10

B Formalized Assumptions

11

C Proof of Continuous-Time Population Dynamics C.1 Initialization Event . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.2 General Results on Dynamics . . . . . . . . . . . . . . . . . . . . . . . . . . . C.3 Weak Recovery . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . C.4 Strong Recovery . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

13 13 14 20 32

D Proof of Discrete-Time Online SGD Dynamics D.1 Weak Recovery Setups . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.2 Weak Recovery Dynamics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.3 Strong Recovery Setups . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . D.4 Strong Recovery Dynamics . . . . . . . . . . . . . . . . . . . . . . . . . . . .

46 46 60 75 89

E Auxiliary Lemmas 97 E.1 Stein Bounds . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 97 E.2 Gradient Bounds . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 101 E.3 Summation by Parts . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 104

A

Numerical Experiments

We illustrate the performance of online SGD on several representative experiments. We use minibatch online SGD. The figure captions and experiment descriptions specify the batch sizes and learning rates. We initialize the parameters (θ∗ , β ∗ ) drawn i.i.d. from the uniform measure, and we initialize (θ0 , β0 ) uniformly at random from the non-negative-overlap hemispheres {θ; θ · θ∗ ≥ 0}, {β; β · β ∗ ≥ 0}. Sample complexity In Figure 2 we plot the recovery performance of online SGD for d ∈ {20, 50, 100, 200}, using normalized GELU. We track min(|χt |, |ρt |) over 10 independent seeds and show variability across seeds. As predicted by our theory, the algorithm successfully recovers the planted parameters, at a rate that is roughly n ≃ d log d. Comparing different K, d pair choices To empirically see the potential of online SGD under MSE loss extends beyond K > 2 and compare its performance with correlation loss as studied in our companion work [5], we ran a set of experiments with varying K. Specifically, based on our model set-up, we consider a fixed dimension of model input, dK = 16777216, and tested pair choices (K, d) ∈ {(2, 4096), (3, 256), (4, 64), (6, 16)} for both correlation and MSE loss. All experiments use normalized ϕ = GeLU activations, a learning rate of η = 0.000391, minibatches of size 16, and 227 training examples across 10 random seeds. As the figure shows, only for K = 2 do we see reliable recovery for both losses. Their performance diverges substantially as depth increases, with correlation loss exhibits meaningful recovery among all remaining trials. Testing non-degeneracy assumption Finally, we investigate the effect of removing the nondegeneracy assumption, by using instead ϕ = He2 , which has information exponent 2. Figure 4 shows that, in this setting, online SGD does not make significant progress, even at timescales n ≫ d2 .

10

Figure 2: Minimum absolute overlap for K = 2 with normalized GELU at d ∈ {20, 50, 100, 200}. We use η = 18.4207/(d log d), batch size 256, and 10 seeds.

B

Formalized Assumptions

Let G ∼ N (0, 1), and write the Hermite expansions ϕ = normalized Hermite basis. We assume:

P

k≥0 ak hk and ψ =

P

k≥0 αk hk in the

(A1) Normalization. Eϕ(G) = 0 and Eϕ(G)2 = 1. (A2) Non-degeneracy. The first Hermite coefficients are nonzero: a1 = E[Gϕ(G)] = E[ϕ′ (G)] ̸= 0,

α1 = E[Gψ(G)] = E[ψ ′ (G)] ̸= 0.

(A3) Regularity of activation function. The inner activation ϕ ∈ C 3 (R) and outer activation ψ ∈ Cb4 (R) have finite regularity constants Lϕ := ∥ϕ′ ∥∞ + ∥ϕ′′ ∥∞ + ∥ϕ′′′ ∥∞ ,

Lψ := max ∥ψ (j) ∥∞ . 0≤j≤4

Here Cb4 (R) requires ψ and its derivatives through order four to be bounded; ϕ itself need not be bounded. (A4) Initialization. The initial second layer student parameter and the teacher parameter are uniformly sampled from the unit sphere: θ0 , θ∗ ∼ Unif(Sd−1 ) and set the initial overlap to be of the order d−1/2 The initialization distribution of the overlap satisfies n o Eov := ρ0 , χ0 ∈ [cinit d−1/2 , Cinit d−1/2 ] . Below, we made two remarks, for (A3) and for (A4) respectively: Remark B.1. Assumption (A3) implies all moment conditions used below. Indeed, since ∥ϕ′ ∥∞ < ∞, |ϕ(x)| ≤ |ϕ(0)| + ∥ϕ′ ∥∞ |x|. Hence ϕ(G) has finite Gaussian moments of every order. Moreover, for 1 ≤ j ≤ 3, the derivatives ϕ(j) are bounded. Therefore, for every finite p ≥ 1, h i max E (1 + |G|)p |ϕ(j) (G)|p < ∞, G ∼ N (0, 1), 0≤j≤3

where ϕ(0) = ϕ. In particular, E|ϕ′ (G)|8 < ∞,

E|ϕ(G)|8 < ∞, and

  E (1 + |ϕ(G)|2 )|ϕ′ (G)|2 < ∞, 11

  E |ϕ(G)| |ϕ′ (G)|2 < ∞.

Figure 3: Comparison between MSE and Correlation Loss, for various pairs of K, d with dK = 16777216, using ϕ = GeLU.

12

Figure 4: Minimum absolute overlap for K = 2 with normalized GELU and with He2 . We use d = 64, η = 0.02, batch size 64, and 10 seeds. Finally, if Y is sub-Gaussian, then ϕ(Y ) is sub-Gaussian, with
∥ϕ(Y )∥ψ2 ≤ C |ϕ(0)| + Lϕ ∥Y ∥ψ2 , for a universal constant C. This follows from the same Lipschitz bound. Remark B.2. By standard concentration of the uniform measure (e.g. [4]), the event Eov in (A4) will happen with probability at least p0 > 0 independent of d for the uniform initialization. Fix K > 0. We work on the event E0 (K) := Eov ∩ Edeloc (K), where ( Edeloc (K) :=

r ∥θ0 ∥∞ + ∥θ∗ ∥∞ ≤ CK

log d , d

) ∥θ0 ∥33 + ∥θ∗ ∥33 ≤ CK d−1/2

.

By Lemma C.1, given θ0 , θ∗ ∼ Unif(Sd−1 ), Edeloc (K) holds with probability at least 1 − d−K .

C

Proof of Continuous-Time Population Dynamics

C.1

Initialization Event

Lemma C.1 (Delocalization of random spherical vectors). Let u ∼ Unif(Sd−1 ). Then, for every fixed K > 0, there exists a constant CK < ∞, independent of d, such that ! r log d P ∥u∥∞ ≤ CK and ∥u∥33 ≤ CK d−1/2 ≥ 1 − d−K . (21) d Consequently, if θ0 , θ∗ are independent uniform random vectors on Sd−1 , then, after increasing CK , with probability at least 1 − d−K , r log d ∗ ∥θ0 ∥∞ + ∥θ ∥∞ ≤ CK , (22) d and ∥θ0 ∥33 + ∥θ∗ ∥33 ≤ CK d−1/2 . 13

(23)

Proof. Write

g iid , g = (g1 , . . . , gd ), gi ∼ N (0, 1). ∥g∥2 Standard Gaussian concentration gives, for every fixed K > 0, after choosing CK large enough, p d ∥g∥22 ≥ max |gi | ≤ CK log d, 1≤i≤d 2 u=

with probability at least 1 − d−K . Hence maxi |gi | ≤ CK ∥u∥∞ = ∥g∥2

r

log d . d

For the third moment, d X

∥u∥33 =

3

Pd

|ui | =

i=1

3 i=1 |gi | . 3 ∥g∥2

By a standard concentration bound for sums of iid variables with finite moments, for every fixed K > 0 there exists CK < ∞ such that d X

|gi |3 ≤ CK d

i=1

with probability at least 1 − d

−K

. Combining this with ∥g∥22 ≥ d/2, we get ∥u∥33 ≤

CK d ≤ CK d−1/2 . (d/2)3/2

This proves (21). By a union bound applied to θ0 and θ∗ , we get (22) and (23), respectively. C.2

General Results on Dynamics

Lemma C.2 (Uniform first-order gradient comparison). Assume (A1) and (A3), then there exist (1) (1) constants Cβ , Cθ > 0, independent of d, such that uniformly for all (θ, β) with ρ, χ ≥ 0, (1)

|⟨∇S,β (L − LG )(θ, β), β ∗ ⟩| ≤ Cβ η.

(24)

(1) |⟨∇S,θ (L − LG )(θ, β), θ∗ ⟩| ≤ Cθ η.

(25)

Proof. Throughout the proof, constants may change from line to line but depend only on the model and moment constants, not on d, θ, θ∗ , β, β ∗ . Set

Ui∗ := x(i) · β ∗ ,

Ui := x(i) · β,

so that (Ui , Ui∗ )di=1 are iid centered Gaussian pairs with correlation ρ. Define Xi∗ := ϕ(Ui∗ ),

Xi := ϕ(Ui ), and S :=

d X

S ∗ :=

θi Xi ,

i=1

Then

E[Xi ] = E[Xi∗ ] = 0,

d X

θi∗ Xi∗ .

i=1

E[Xi2 ] = E[(Xi∗ )2 ] = 1,

E[Xi Xi∗ ] = κ(ρ).

Let (Γi , Γ∗i )di=1 be iid Gaussian pairs with the same covariance matrix E[Γ2i ] = E[(Γ∗i )2 ] = 1, and write SG :=

d X

E[Γi Γ∗i ] = κ(ρ), ∗ SG :=

θi Γi ,

i=1

d X i=1

14

θi∗ Γ∗i .

Step 1: the θ-gradient comparison. Let w := Pθ⊥ θ∗ = θ∗ − χθ. Then

" ⟨∇S,θ L(θ, β), θ∗ ⟩ = 2E (ψ(S) − ψ(S ∗ ))ψ ′ (S)

d X

# wi Xi .

i=1

The same formula holds for LG , with (X, X ∗ ) replaced by (Γ, Γ∗ ). Hence it is enough to compare the expectations of
h(x, x∗ ) := 2 ψ(x · θ) − ψ(x∗ · θ∗ ) ψ ′ (x · θ)(x · w). We apply Corollary E.2. Fix a block j, condition on all other blocks, and write t := x∗j ,

s := xj ,

ℓ3 = θj∗ t + S3 ,

ℓ1 = θj s + S1 , ℓ2 = wj s + S2 , where S1 , S2 , S3 are the leave-out sums. Then

h(s, t) = 2(ψ(ℓ1 ) − ψ(ℓ3 ))ψ ′ (ℓ1 )ℓ2 . Since ψ ∈ Cb4 , every third partial derivative ∂sa ∂tb h, a + b = 3, is a finite sum of terms of the two following types: C |θj |a |θj∗ |b |ℓ2 |, a + b = 3, and ′ ′ C |wj | |θj |a |θj∗ |b , a′ + b′ = 2. At every hybrid vector in the Lindeberg interpolation, E|ℓ2 | ≤ C∥w∥2 ≤ C. Therefore

! E



∂sa ∂tb h(W )



≤C

X

|θj | |θj∗ |b + |wj | a

a′

|θj |

′ |θj∗ |b

.

a′ +b′ =2

Summing over j and over a + b = 3, we get   X X X X ′ ′ |θj |a |θj∗ |b + |wj | |θj |a |θj∗ |b  . |Eh(X, X ∗ ) − Eh(Γ, Γ∗ )| ≤ C  j

The first sum is bounded by

j

a+b=3

a′ +b′ =2

C(∥θ∥∞ + ∥θ∗ ∥∞ ).

Indeed, X

|θj |3 ≤ ∥θ∥∞ ,

X

j

|θj |2 |θj∗ | ≤ ∥θ∗ ∥∞ ,

j

and similarly for the remaining degree-three monomials. For the second sum, use ∥w∥∞ = ∥θ∗ − χθ∥∞ ≤ ∥θ∗ ∥∞ + ∥θ∥∞ = η, and ∥w∥2 , ∥θ∥2 , ∥θ∗ ∥2 ≤ C, to obtain X X ′ ′ |wj | |θj |a |θj∗ |b ≤ Cη. j

a′ +b′ =2

Thus

|⟨∇S,θ (L − LG )(θ, β), θ∗ ⟩| ≤ Cη. This proves the θ-gradient comparison. Step 2: exact formula for the β-gradient. Let b := Pβ ⊥ β ∗ = β ∗ − ρβ. 15

(26)

Then

⟨∇S,β L(θ, β), β ∗ ⟩ = ⟨∇β L(θ, β), b⟩. Differentiating S in the direction b gives Db S =

d X

θi ϕ′ (Ui )Hi ,

Hi := x(i) · b.

i=1

Hence

" ⟨∇S,β L(θ, β), β ∗ ⟩ = 2E (ψ(S) − ψ(S ∗ ))ψ ′ (S)

d X

# θi ϕ′ (Ui )Hi .

i=1

For each i, Hi = Ui∗ − ρUi , so Hi is centered Gaussian with variance 1 − ρ2 , independent of Ui . Moreover Ui∗ = ρUi + Hi . Applying Gaussian integration by parts in the Hi -variable gives E [(ψ(S) − ψ(S ∗ ))ψ ′ (S)θi ϕ′ (Ui )Hi ] = −(1 − ρ2 )θi θi∗ E [ψ ′ (S)ψ ′ (S ∗ )ϕ′ (Ui )ϕ′ (Ui∗ )] . Therefore − ⟨∇S,β L(θ, β), β ∗ ⟩ = 2(1 − ρ2 )

d X

θi θi∗ E [ψ ′ (S)ψ ′ (S ∗ )ϕ′ (Ui )ϕ′ (Ui∗ )] .

(27)

i=1

Step 3: leave-one-out reduction for the true β-gradient. Set F (r, s) := ψ ′ (r)ψ ′ (s),

qi := ϕ′ (Ui )ϕ′ (Ui∗ ),

Also define

ci := θi θi∗ .

∗ S−i := S ∗ − θi∗ Xi∗ .

S−i := S − θi Xi , ∗ Since S−i , S−i are independent of (Ui , Ui∗ ),

∗ ∗ E[F (S−i , S−i )qi ] = E[F (S−i , S−i )]E[qi ].

By the Hermite identity,

E[qi ] = E[ϕ′ (Ui )ϕ′ (Ui∗ )] = κ′ (ρ). Moreover, since F is Lipschitz and the stated moment assumptions imply sup E [(|Xi | + |Xi∗ |)|qi |] < ∞, ρ∈[0,1]

we have

∗ E[F (S, S ∗ )qi ] − E[F (S−i , S−i )qi ] ≤ C(|θi | + |θi∗ |).

Multiplying by |ci | and summing, d X

|ci |(|θi | + |θi∗ |) ≤

i=1

Thus

X

|θi |2 |θi∗ | +

X

i d X

i

ci E[F (S, S ∗ )qi ] = κ′ (ρ)

i=1

|θi ||θi∗ |2 ≤ ∥θ∗ ∥∞ + ∥θ∥∞ = η.

d X

∗ ci E[F (S−i , S−i )] + O(η).

i=1

Similarly,

∗ E[F (S, S ∗ )] − E[F (S−i , S−i )] ≤ C(|θi | + |θi∗ |), and the same summation estimate gives ! d d X X ∗ ci E[F (S−i , S−i )] = ci E[F (S, S ∗ )] + O(η). i=1

Since

P

i ci = θ · θ

∗

i=1

= χ, we conclude d X

ci E[F (S, S ∗ )qi ] = κ′ (ρ)χ E[F (S, S ∗ )] + O(η).

i=1

16

(28)

Step 4: Stein comparison for the two-dimensional statistic. Apply Corollary E.2 to h(x, x∗ ) := F (x · θ, x∗ · θ∗ ). For a fixed block j, the third partials have the form ∂sa ∂tb h = θja (θj∗ )b ∂1a ∂2b F (ℓ1 , ℓ2 ), Since F ∈ Cb3 (R2 ), Therefore

a + b = 3.

  E |∂sa ∂tb h(W )| ≤ C|θj |a |θj∗ |b .

∗ |E[F (S, S ∗ )] − E[F (SG , SG )]| ≤ C

X X j

|θj |a |θj∗ |b ≤ Cη.

a+b=3

Combining this with (28), and using that κ′ (ρ) and χ are bounded on the entry corridor, gives d X

∗ ci E[F (S, S ∗ )qi ] = κ′ (ρ)χ E[F (SG , SG )] + O(η).

(29)

i=1

Step 5: identify the Gaussian β-gradient. In the Gaussian equivalent model, X   LG (θ, β) = Φ(χ, ρ), Φ(χ, ρ) = 2 αk2 1 − (κ(ρ)χ)k . k≥1 ∗ ), Equivalently, by Price’s identity applied to the Gaussian pair (SG , SG ∗ −Φρ (χ, ρ) = 2κ′ (ρ)χE[ψ ′ (SG )ψ ′ (SG )].

Since

∇S,β LG = Φρ (χ, ρ)Pβ ⊥ β ∗ ,

we have

∗ − ⟨∇S,β LG (θ, β), β ∗ ⟩ = 2(1 − ρ2 )κ′ (ρ)χE[ψ ′ (SG )ψ ′ (SG )].

(30)

Comparing (27), (29), and (30), we obtain |⟨∇S,β (L − LG )(θ, β), β ∗ ⟩| ≤ C(1 − ρ2 )η ≤ Cη. This proves the β-gradient comparison. Combining the two block estimates gives the lemma. Lemma C.3 (Uniform positive-overlap drift). Assume Lemma C.2. On the positive overlap region ρ, χ ≥ 0, the population gradient flow satisfies ρ̇t ≥ aρ χt (1 − ρ2t ) − bρ ηt ,

χ̇t ≥ aχ ρt (1 − χ2t ) − bχ ηt ,

where aρ , aχ , bρ , bχ > 0 are independent of d. Proof. With the population flow defined in (5), we have ρ̇t = − ⟨∇S,β L(θt , βt ), β ∗ ⟩ ,

χ̇t = − ⟨∇S,θ L(θt , βt ), θ∗ ⟩ .

We first compute the Gaussian drifts exactly. Recall that LG (θ, β) = Φ(χ, ρ),

Φ(χ, ρ) = 2

X

  αk2 1 − (κ(ρ)χ)k ,

k≥1

where κ(ρ) =

2 ℓ ℓ≥1 aℓ ρ .

P

For the first layer, Hence

∇S,β LG = Φρ (χ, ρ)Pβ ⊥ β ∗ = Φρ (χ, ρ)(β ∗ − ρβ). − ⟨∇S,β LG , β ∗ ⟩ = −Φρ (χ, ρ)(1 − ρ2 ). 17

(31)

Moreover, X

−Φρ (χ, ρ) = 2χκ′ (ρ)

kαk2 (κ(ρ)χ)k−1 .

k≥1

On the positive-overlap region ρ, χ ≥ 0, all terms are nonnegative. Keeping only the k = 1 term gives −Φρ (χ, ρ) ≥ 2α12 χκ′ (ρ). Since κ′ (ρ) =

X

ℓa2ℓ ρℓ−1 ≥ a21

(ρ ≥ 0),

ℓ≥1

we obtain − ⟨∇S,β LG , β ∗ ⟩ ≥ 2α12 a21 χ(1 − ρ2 ). Similarly, ∇S,θ LG = Φχ (χ, ρ)Pθ⊥ θ∗ = Φχ (χ, ρ)(θ∗ − χθ), so − ⟨∇S,θ LG , θ∗ ⟩ = −Φχ (χ, ρ)(1 − χ2 ). Now −Φχ (χ, ρ) = 2κ(ρ)

X

kαk2 (κ(ρ)χ)k−1 .

k≥1

Again all terms are nonnegative on ρ, χ ≥ 0, so −Φχ (χ, ρ) ≥ 2α12 κ(ρ). Since κ(ρ) =

X

a2ℓ ρℓ ≥ a21 ρ

(ρ ≥ 0),

ℓ≥1

we get − ⟨∇S,θ LG , θ∗ ⟩ ≥ 2α12 a21 ρ(1 − χ2 ). Now transfer the Gaussian drift to the true population loss. By (24) from Lemma C.2, ρ̇t = − ⟨∇S,β LG (θt , βt ), β ∗ ⟩ − ⟨∇S,β (L − LG )(θt , βt ), β ∗ ⟩ (1)

≥ 2α12 a21 χt (1 − ρ2t ) − Cβ ηt . Likewise, by (25) from Lemma C.2, χ̇t = − ⟨∇S,θ LG (θt , βt ), θ∗ ⟩ − ⟨∇S,θ (L − LG )(θt , βt ), θ∗ ⟩ (1)

≥ 2α12 a21 ρt (1 − χ2t ) − Cθ ηt . Thus the claimed inequalities hold with aρ = aχ := 2α12 a21 ,

(1)

bρ := Cβ ,

(1)

bχ := Cθ .

Lemma C.4 (ℓ∞ control). Under our assumptions, η̇t ≤ C1 ρt ∥θ∗ ∥∞ + C2 χt ρt ηt + O(ηt2 ) .

(32)

Proof. By definition, we have L = E(z,z∗ )∼ϕ# γβ·β∗ [|ψ(z · θ) − ψ(z ∗ · θ∗ )|2 ] so ∇θ L = E(z,z∗ )∼ϕ# γβ·β∗ [(ψ(z · θ) − ψ(z ∗ · θ∗ ))ψ ′ (z · θ)z] . Call ν the law of ϕ(w), where w ∼ N (0, 1). The random vector z is then drawn from ν ⊗d . Let us now analyze the two terms in the RHS. 18

The first term reads E[F (z · θ)z] for F of the form F = ψψ ′ . Assume that F ′ is Lipschitz. A second-order Taylor expansion of F yields F (z · θ) = F (z1 θ1 + z−1 · θ−1 )

(33)

1 = F (z−1 · θ−1 ) + F ′ (z−1 · θ−1 )θ1 z1 + θ12 z12 F ′′ (M ) , 2 and therefore, for all i ∈ [d],

(34)

E[F (z · θ)zi ] = θi E[F ′ (z−i · θ−i )] + O(θi2 ) . (35) P Now, for each θ and i ∈ [d], consider µi = Law( j̸=i θj zj ) ∈ P(R), where zj ∼ ν iid. We claim that sup W1 (µi , µi′ ) = O(∥θ∥∞ ) .

(36)

i,i′

Indeed, fix i, i′ and consider P the coupling (Yi , Yi′ ) ∼ Γµi ,µi′ given by Yi = W +θi Z, Yi′ = W +θi′ Z, where Z ∼ ν and W = j̸=i,i′ θj zj , with zj ∼ ν independently of Z. We then have p W1 (µi , µi′ ) ≤ E|Yi − Yi′ | = |θi − θi′ |E|Z| ≤ 2 E[|ϕ(G)|2 ]∥θ∥∞ ≤ 2∥θ∥∞ . Now, since F ′ is Lipschitz, we have sup |E[F ′ (z−i · θ−i )] − E[F ′ (z−j · θ−j )]| ≤ Lip(F ′ ) sup W1 (µi , µi′ ) = O(∥θ∥∞ ) ,

(37)

i,i′

i,j

and therefore, denoting by υ(θ) = d1

′ i E[F (z−i · θ−i )], we have

P

E[F (z · θ)z] = θυ(θ) + ε , with |εi | ≤ C|θi |∥θ∥∞ .

(38)

Pθ⊥ E[F (z · θ)z] = Pθ⊥ ε = ε − (ε · θ)θ

(39)

∥Pθ⊥ E[F (z · θ)z]∥∞ ≤ C̃∥θ∥2∞ .

(40)

It follows that satisfies ∗

Now we repeat the argument for a term of the form E[G(z · θ, z̃ · θ )z] for G of the form G(t1 , t2 ) = ψ ′ (t1 )ψ(t2 ). Doing a second-order Taylor expansion leads to ∗ G(z · θ, z̃ · θ∗ ) = G(z1 θ1 + z−1 · θ−1 , z̃1 θ1∗ + z̃−1 · θ−1 ) ∗ ∗ = ψ (z−1 · θ−1 )ψ(z̃−1 · θ−1 ) + ψ ′′ (z−1 · θ−1 )ψ(z̃−1 · θ−1 )θ1 z1 ′ ′ ∗ ∗ + ψ (z−1 · θ−1 )ψ (z̃−1 · θ−1 )z̃1 θ1

(41)

′

1 + [z1 θ1 z̃1 θ1∗ ]⊤ ∇2 G(M )[z1 θ1 z̃1 θ1∗ ] , 2 and thus, for all i ∈ [d], denoting α = E[zi z̃i ], and assuming that ∇G is Lipschitz, we have

(42) (43)

∗ E[G(z · θ, z̃ · θ∗ )zi ] = θi E[ψ ′′ (z−i · θ−i )ψ(z̃−i · θ−i )] ∗ + αθi∗ E[ψ ′ (z−i · θ−i )ψ ′ (z̃−i · θ−i )] + O(θi2 + (θi∗ )2 ) .

(44)

∗

Following our previous argument, we P Pnow consider for each θ, θ and each i ∈ [d] the bivariate distributions µ̄i = Law( j̸=i θj zj , j̸=i θj∗ z̃j ), where (zi , z̃i ) are drawn i.i.d. from ν̄. We now have sup W1 (µ̄i , µ̄i′ ) = O(∥θ∥∞ + ∥θ∗ ∥∞ ) ,

(45)

E[G(z · θ, z̃ · θ∗ )z] = θυ1 (θ, θ∗ ) + αθ∗ υ2 (θ, θ∗ ) + ε ,

(46)

i,i′

leading to with |εi | ≤ (C1 |θi | + C2 |α||θi∗ |)(∥θ∥∞ + ∥θ∗ ∥∞ ) and υ2 = O(1). We then arrive at Pθ⊥ E[G(z · θ, z̃ · θ∗ )z] = αθ∗ υ2 (θ, θ∗ ) − χαυ2 θ + ε − (ε · θ)θ .

(47)

∗

Again, since |ε · θ| = O(∥θ∥∞ + ∥θ ∥∞ ), we have ∥∇Sθ L∥∞ ≤ C1 ρ∥θ∗ ∥∞ + C2 ρχ∥θ∥∞ + O(∥θ∥2∞ + ∥θ∗ ∥2∞ ) , proving the lemma. 19

(48)

C.3

Weak Recovery

At the weak recovery, the total overlap Zt is still below the incoherence scale. In this regime the crude O(ηt ) comparison error from Lemma C.3 is too large under spherical initialization. We instead
3 Pd track the third moment quantity ωt := i=1 |θt,i | + |θi∗ | and show it is controlled, at population √ . level, under the order 1+t d Lemma C.5 (First-order comparison in the small-overlap regime). Assume the hypotheses of Lemma C.2 and the regularity needed for Corollary E.2. Define ω(θ, θ∗ ) :=

d X

3 |θi | + |θi∗ | .

i=1

Let

ρ̇ := − ⟨∇S,β L(θ, β), β ∗ ⟩ , χ̇ := − ⟨∇S,θ L(θ, β), θ∗ ⟩ , and define ρ̇G , χ̇G analogously with LG . Then, uniformly on the small-overlap region 0 ≤ ρ, χ and Z := ρ + χ ≤ csm , |ρ̇ − ρ̇G | + |χ̇ − χ̇G | ≤ Cω(θ, θ∗ ),

(49)

where C is independent of d. Proof. Throughout, constants may change from line to line and depend only on the model and moment constants. Set Ui := x(i) · β, Ui∗ := x(i) · β ∗ , so that (Ui , Ui∗ )di=1 are iid centered Gaussian pairs with correlation ρ. Let Xi∗ := ϕ(Ui∗ ),

Xi := ϕ(Ui ), and S :=

d X

S ∗ :=

θi Xi ,

i=1

d X

θi∗ Xi∗ .

i=1

Let (Γi , Γ∗i )di=1 be iid centered Gaussian pairs with the same block covariance as (Xi , Xi∗ ): E[Γ2i ] = E[(Γ∗i )2 ] = 1, and write SG :=

d X

E[Γi Γ∗i ] = κ(ρ), ∗ SG :=

θi Γi ,

i=1

d X

θi∗ Γ∗i .

i=1

We prove the two block comparisons separately. The θ-block. Let d−1

Since w ∈ Tθ S

w := Pθ⊥ θ∗ = θ∗ − χθ. , ⟨∇S,θ (L − LG )(θ, β), θ∗ ⟩ = ⟨∇θ (L − LG )(θ, β), w⟩ .

Thus the θ-comparison is the difference between the true and Gaussian expectations of a function of the form
h(x, x∗ ) = 2 ψ(x · θ) − ψ(x∗ · θ∗ ) ψ ′ (x · θ) (x · w). Apply Corollary E.2 to h. For a fixed block j, every third partial derivative ∂sa ∂tb h, a + b = 3, is a finite sum of terms of the following two types: C|θj |a |θj∗ |b |ℓ2 |, and

′

a + b = 3,

′

C|wj | |θj |a |θj∗ |b , a′ + b′ = 2, where ℓ2 denotes the linear form w · W evaluated at the corresponding Lindeberg hybrid vector. Since ∥w∥2 ≤ 1, the moment assumptions imply E|ℓ2 | ≤ C. 20

Therefore |⟨∇S,θ (L − LG )(θ, β), θ∗ ⟩| ≤ C

d X X

|θj |a |θj∗ |b + C

j=1 a+b=3

Now

d X

|wj |

X

′

′

|θj |a |θj∗ |b .

a′ +b′ =2

j=1

|wj | = |θj∗ − χθj | ≤ |θj∗ | + |θj |,

because |χ| ≤ 1. Hence, for every j, X

3 |θj |a |θj∗ |b ≤ C |θj | + |θj∗ | ,

a+b=3

and |wj |

3 ′ ′ |θj |a |θj∗ |b ≤ C |θj | + |θj∗ | .

X a′ +b′ =2

Summing over j, we obtain |⟨∇S,θ (L − LG )(θ, β), θ∗ ⟩| ≤ Cω(θ, θ∗ ). Equivalently,

|χ̇ − χ̇G | ≤ Cω(θ, θ∗ ).

The β-block. Let

b := Pβ ⊥ β ∗ = β ∗ − ρβ.

Then

⟨∇S,β L(θ, β), β ∗ ⟩ = ⟨∇β L(θ, β), b⟩ .

As in the proof of Lemma C.2, Gaussian integration by parts in the b-direction gives the exact formula − ⟨∇S,β L(θ, β), β ∗ ⟩ = 2(1 − ρ2 )

d X

θi θi∗ E[ψ ′ (S)ψ ′ (S ∗ )ϕ′ (Ui )ϕ′ (Ui∗ )] .

i=1

Set

F (s, t) := ψ ′ (s)ψ ′ (t),

qi := ϕ′ (Ui )ϕ′ (Ui∗ ),

ci := θi θi∗ .

We compare d X

ci E[F (S, S ∗ )qi ]

i=1

with its Gaussian analogue. Let

∗ S−i := S ∗ − θi∗ Xi∗ .

S−i := S − θi Xi , ∗ Since (S−i , S−i ) is independent of (Ui , Ui∗ ),

∗ ∗ E[F (S−i , S−i )qi ] = E[F (S−i , S−i )]E[qi ].

By the Hermite identity,

E[qi ] = E[ϕ′ (Ui )ϕ′ (Ui∗ )] = κ′ (ρ).

Moreover, by the Lipschitz bound on F and the moment assumptions, ∗ E[F (S, S ∗ )qi ] − E[F (S−i , S−i )qi ] ≤ C(|θi | + |θi∗ |).

Multiplying by |ci | = |θi θi∗ | and summing yields X X |ci |(|θi | + |θi∗ |) = |θi θi∗ |(|θi | + |θi∗ |) ≤ Cω(θ, θ∗ ). i

i

Hence d X i=1

∗

′

ci E[F (S, S )qi ] = κ (ρ)

d X i=1

21

∗ ci E[F (S−i , S−i )] + O(ω).

The same leave-one-out argument gives d X

∗ ci E[F (S−i , S−i )] =

i=1

d X

! ci

E[F (S, S ∗ )] + O(ω).

i=1

Since

d X

ci = θ · θ∗ = χ,

i=1

we obtain

d X

ci E[F (S, S ∗ )qi ] = κ′ (ρ)χ E[F (S, S ∗ )] + O(ω).

i=1 ∗ It remains to replace (S, S ∗ ) by (SG , SG ). Applying Corollary E.2 to

h(x, x∗ ) = F (x · θ, x∗ · θ∗ ) gives ∗ |E[F (S, S ∗ )] − E[F (SG , SG )]| ≤ C

d X X

|θj |a |θj∗ |b ≤ Cω(θ, θ∗ ).

j=1 a+b=3 ′

Since κ (ρ) and χ are bounded, we conclude d X

∗ ci E[F (S, S ∗ )qi ] = κ′ (ρ)χ E[F (SG , SG )] + O(ω).

i=1

In the Gaussian equivalent model, Price’s identity gives ∗ − ⟨∇S,β LG (θ, β), β ∗ ⟩ = 2(1 − ρ2 )κ′ (ρ)χE[F (SG , SG )].

Comparing the last two displays and using 0 ≤ 1 − ρ2 ≤ 1, we get |⟨∇S,β (L − LG )(θ, β), β ∗ ⟩| ≤ Cω(θ, θ∗ ). Equivalently,

|ρ̇ − ρ̇G | ≤ Cω(θ, θ∗ ).

Combining the θ- and β-block estimates proves |ρ̇ − ρ̇G | + |χ̇ − χ̇G | ≤ Cω(θ, θ∗ ). Lemma C.6 (Small-overlap drift estimates). Under the assumptions of Lemma C.2, with the same definition of ω(θ, θ∗ ) in C.5. There exist constants csm , a, b, C > 0, independent of d, such that whenever ρ, χ ≥ 0, Z := ρ + χ ≤ csm , the population overlap drifts satisfy ρ̇ ≥ aχ − bω(θ, θ∗ ),

χ̇ ≥ aρ − bω(θ, θ∗ ).

(50)

e := 2α2 a2 > 0, one has the refined two-sided expansion Moreover, with C 1 1 e χ(1 − ρ2 ) + Rρ , e ρ(1 − χ2 ) + Rχ , ρ̇ = C χ̇ = C
where |Rρ | + |Rχ | ≤ C Z 2 + ω(θ, θ∗ ) . Consequently, for D := ρ − χ,
|RD | ≤ C Z 2 + ω(θ, θ∗ ) .

e + ρχ)D + RD , Ḋ = −C(1

(51)

(52)

Proof. Let LG denote the Gaussian-equivalent population loss. As in the proof of Lemma C.3, the Gaussian overlap drifts are ρ̇G = 2χ κ′ (ρ)A(q)(1 − ρ2 ), where q := κ(ρ)χ,

A(q) :=

X

χ̇G = 2κ(ρ)A(q)(1 − χ2 ), kαk2 q k−1 ,

k≥1

κ(ρ) :=

X ℓ≥1

22

a2ℓ ρℓ .

Since a1 ̸= 0 and α1 ̸= 0, κ′ (0) = a21 > 0,

A(0) = α12 > 0.

First we prove the cooperative lower bounds. On the positive-overlap region ρ, χ ≥ 0, all terms in the Gaussian drift formulas are nonnegative. Keeping only the leading terms gives ρ̇G ≥ 2α12 a21 χ(1 − ρ2 ),

χ̇G ≥ 2α12 a21 ρ(1 − χ2 ).

Choose csm > 0 sufficiently small so that Z ≤ csm implies ρ, χ ≤ csm and hence 1 1 1 − ρ2 ≥ , 1 − χ2 ≥ . 2 2 Therefore ρ̇G ≥ c0 χ, χ̇G ≥ c0 ρ for c0 := α12 a21 > 0. It remains to transfer these estimates from LG to the true population loss. By Lemma C.5, |ρ̇ − ρ̇G | + |χ̇ − χ̇G | ≤ Cω(θ, θ∗ ). Hence, after setting a := c0 /2 and increasing b, ρ̇ ≥ aχ − bω(θ, θ∗ ),

χ̇ ≥ aρ − bω(θ, θ∗ ),

which proves (50). We now prove the two-sided expansion. Since Z ≤ csm , we have |ρ| + |χ| ≤ csm . ′

Taylor expanding κ, κ , and A near the origin gives κ′ (ρ) = a21 + O(|ρ|),

κ(ρ) = a21 ρ + O(ρ2 ), and since

q = κ(ρ)χ = O(Z 2 ),

also

A(q) = α12 + O(Z 2 ). Substituting these expansions into the exact Gaussian drift formulas yields ρ̇G = 2α12 a21 χ(1 − ρ2 ) + O(Z 2 ), and

χ̇G = 2α12 a21 ρ(1 − χ2 ) + O(Z 2 ). Indeed, in the first formula the error from κ′ (ρ) − a21 is O(|ρ||χ|) = O(Z 2 ), while the error from A(q) − α12 is O(|χ|Z 2 ) = O(Z 3 ). In the second formula, the error from κ(ρ) − a21 ρ is O(ρ2 ) = O(Z 2 ), and the error from A(q) − α12 is O(|ρ|Z 2 ) = O(Z 3 ).

e := 2α2 a2 , we have Thus, with C 1 1 e χ(1 − ρ2 ) + O(Z 2 ), ρ̇G = C

e ρ(1 − χ2 ) + O(Z 2 ). χ̇G = C

Using again the refined comparison estimate |ρ̇ − ρ̇G | + |χ̇ − χ̇G | ≤ Cω(θ, θ∗ ), we obtain

e χ(1 − ρ2 ) + Rρ , e ρ(1 − χ2 ) + Rχ , ρ̇ = C χ̇ = C
with |Rρ | + |Rχ | ≤ C Z 2 + ω(θ, θ∗ ) . This proves (51). Finally, subtract the two expansions. Since χ(1 − ρ2 ) − ρ(1 − χ2 ) = χ − ρ − ρ2 χ + ρχ2 = −(ρ − χ) − ρχ(ρ − χ) = −(1 + ρχ)D, we get e + ρχ)D + RD , Ḋ = −C(1 where

|RD | ≤ |Rρ | + |Rχ | ≤ C Z 2 + ω(θ, θ∗ ) .

This proves (52). 23

Pd Lemma C.7 (Third-moment control at population level). Let Ωt := ∥θt ∥33 = i=1 |θt,i |3 , ωt :=
Pd ∗ 3 i=1 |θt,i | + |θi | . Work on the initialization event of Lemma C.1. Then, on the positive-overlap region ρt , χt ≥ 0, with Zt := ρt + χt , the population flow satisfies
Ω̇t ≤ C Zt ∥θ∗ ∥∞ + Zt2 ηt + ηt2 . (53) Consequently, suppose that on an interval 0 ≤ t ≤ T , with T ≤ CT log d, one has √ log d + 1 + t 2 √ Zt ≤ Aηt , ηt ≤ Cη . d Then, for all sufficiently large d, 1+t Ωt ≤ CΩ √ , d

1+t ωt ≤ Cω √ , d

0 ≤ t ≤ T.

The constants CΩ , Cω may depend on K, A, Cη , CT , but not on d. Proof. The population θ-flow is θ̇t = −∇S,θ L(θt , βt ). Let

vt = (|θt,i |θt,i )di=1 .

Since x 7→ |x|3 is C 1 , we have X d Ωt = 3 |θt,i |θt,i θ̇t,i = −3⟨vt , ∇S,θ L(θt , βt )⟩. dt i Moreover, ∥vt ∥1 =

X

|θt,i |2 = 1.

i

Hence Ω̇t ≤ 3∥∇S,θ L(θt , βt )∥∞ . By Lemma C.4, more precisely by (48), we have ∥∇S,θ L(θt , βt )∥∞ ≤ C1 |ρt |∥θ∗ ∥∞ + C2 |ρt χt |∥θt ∥∞ + C3 ηt2 . On the positive-overlap region, ρt , χt ≥ 0, and hence |ρt χt | ≤ Zt2 ,

|ρt | ≤ Zt , Therefore

∥θt ∥∞ ≤ ηt .

Ω̇t ≤ C Zt ∥θ∗ ∥∞ + Zt2 ηt + ηt2 ,

which proves (53). Now with the conditions Zt2 ≤ Aηt ,

√ ηt ≤ Cη

log d + 1 + t √ , d

Then

0 ≤ t ≤ T ≤ CT log d. r

∗

Zt ∥θ ∥∞ ≤

p

Aηt CK

log d . d

Using the envelope for ηt , s√ ∗

Zt ∥θ ∥∞ ≤ C

log d + 1 + t √ d

Since t ≤ CT log d, the right-hand side is bounded by C

log d C ≤√ d3/4 d 24

r

log d . d

(54)

for all sufficiently large d. Also,

Zt2 ηt ≤ Aηt2 .

Therefore, throughout the Phase-I interval, C Ω̇t ≤ √ + Cηt2 . d Integrating from 0 to t ≤ T , we get Ct Ωt ≤ Ω0 + √ + C d

Z t

ηs2 ds.

0

By Lemma C.1, Ω0 = ∥θ0 ∥33 ≤ CK d−1/2 . Moreover,

Z t

C d

ηs2 ds ≤

0

For t ≤ CT log d,

Z t p ( log d + 1 + s)2 ds. 0

Z t p ( log d + 1 + s)2 ds ≤ C(log d)3 . 0

Thus

Z t

ηs2 ds ≤

0

Hence

C(log d)3 = o(d−1/2 ). d

1+t Ωt ≤ CΩ √ , d

0 ≤ t ≤ T.

Finally, by the elementary inequality (a + b)3 ≤ 4a3 + 4b3 we have ωt =

X

(a, b ≥ 0),

3 |θt,i | + |θi∗ | ≤ 4∥θt ∥33 + 4∥θ∗ ∥33 .

i

Using the bound on Ωt and the initialization event ∥θ∗ ∥33 ≤ CK d−1/2 , we obtain

1+t ωt ≤ Cω √ . d

This proves (54). Lemma C.8 (Escape from mediocrity at population level). Assume Lemmas C.4, C.6, and C.7. Work on the random-spherical delocalization event of Lemma C.1, namely r log d ∗ ∥θ0 ∥∞ + ∥θ ∥∞ ≤ CK , ∥θ0 ∥33 + ∥θ∗ ∥33 ≤ CK d−1/2 , d and the favorable initialization event defined in (A4). Then there exist constants cesc , Cesc > 0, independent of d, such that for all sufficiently large d, there is a time tesc ≤ Cesc log d

(55)

min{ρtesc , χtesc } ≥ cesc .

(56)

with

Moreover, throughout the escape window, log d ηt ≤ C √ = od (1). d 25

(57)

Proof. Fix a large constant CT > 0, to be chosen below, and set T := CT log d. Let A > 0 be a fixed large constant and define the stopping times τA := inf{t : Zt2 ≥ Aηt }, τB := inf{t : ρt ≥ 1/2 or χt ≥ 1/2}, τ+ := inf{t : ρt < 0 or χt < 0}. We first work on Phase I, defined as when t ∈ [0, σ] where σ := τA ∧ τB ∧ τ+ ∧ T. Step 1: Incoherence Control in Phase I We begin with the Phase-I control of ηt . By Lemma C.4, on the positive-overlap region, η̇t ≤ CZt ∥θ∗ ∥∞ + CZt2 ηt + Cηt2 . 2 For t < τA , we have Zt < Aηt . Therefore p Zt ∥θ∗ ∥∞ ≤ Aηt ∥θ∗ ∥∞ . Define √ log d + 1 + t √ HI (t) := KI , d where KI is a large constant. We claim that ηt ≤ HI (t) (0 ≤ t ≤ σ). Indeed, at initialization, r log d 1 ∗ η0 = ∥θ0 ∥∞ + ∥θ ∥∞ ≤ CK ≤ HI (0) d 2 after choosing KI large. Suppose the estimate holds up to a first violation time. Since t ≤ T = CT log d, the bootstrap bound gives log d ηt ≤ HI (t) ≤ C √ . d Thus r p log d C log d ∗ ≤ C 3/4 ≤ √ Zt ∥θ ∥∞ ≤ AHI (t) CK d d d for all sufficiently large d. Also, Zt2 ηt ≤ Aηt2 . Consequently, up to the first violation time, C η̇t ≤ √ + Cηt2 . d Integrating, Z t Ct √ ηt ≤ η0 + +C ηs2 ds. d 0 Under the bootstrap ηs ≤ HI (s),  √ Z Z t C(log d)3 log d + 1 + t C CT log d p √ ( log d + 1 + s)2 ds ≤ =o ηs2 ds ≤ . d 0 d d 0 Choosing KI large gives 1 ηt ≤ HI (t) 2 at the putative first violation, a contradiction. Hence for 0 ≤ t ≤ σ, ηt ≤ HI (t). So by Lemma C.7, we have for 0 ≤ t ≤ σ, 1+t Ωt := ∥θt ∥33 ≤ CΩ √ , d and 1+t ωt ≤ Cω √ . d 26

Step 2: Phase I Positivity, Exponential Pull-Up, and Balancing. We next prove positivity and exponential pull-up in Phase I. On t < σ, we have 0 ≤ ρt , χt ≤ 1/2. By Lemma C.6, since 0 ≤ ρt , χt ≤ 1/2 on the stopped Phase-I interval and d is large, after decreasing a > 0 and increasing b > 0, ρ̇t ≥ aχt − bωt ,

χ̇t ≥ aρt − bωt .

Let



⊤

ut := (ρt , χt ) ,

J :=

 1 , 0

0 1

1 := (1, 1)⊤ .

Then, componentwise, u̇t ≥ aJut − bωt 1. aJr

Since e

has nonnegative entries and eaJr 1 = ear 1,

variation of constants gives, for 0 ≤ t ≤ σ, aJt

ut ≥ e

Z t u0 − b

eaJ(t−s) ωs 1 ds.

0

Using ρ0 , χ0 ≥ zinit d−1/2 , from (A4), we get componentwise zinit eaJt u0 ≥ eat √ 1. d Also, Z t b

eaJ(t−s) ωs 1 ds = beat

Z t

0

 e−as ωs ds 1.

0

Since

1+s ωs ≤ Cω √ , d

we have

Z ∞ 0

C e−as ωs ds ≤ √ . d

Let

Z ∞ Mω := bCω

e−as (1 + s) ds.

0

By choosing zinit > 2Mω , we obtain ρt , χt ≥ eat

c− zinit − Mω √ ≥ √ eat d d

(0 ≤ t ≤ σ),

for some c− > 0. In particular, 2c− Zt ≥ √ eat d

(0 ≤ t ≤ σ),

and τ+ cannot be the first stopping time. Choose t0 := C0 log d with C0 > 0 sufficiently large, and then choose CT > C0 . Here we show that σ ≤ t0 . Suppose instead that σ > t0 . Then t0 < τA , so Zt20 < Aηt0 . 27

(58)

But the lower bound above gives c 2at0 = c d2aC0 −1 , e d

Zt20 ≥ whereas the Phase-I η-bound gives

√

log d + 1 + t0 log d √ ≤C √ . d d For C0 sufficiently large, this is impossible for large d. Hence (58) holds. Aηt0 ≤ C

Since t0 < T and τ+ cannot occur first, the first stopping time is either τA or τB . We claim it must be τA . If σ = τB ≤ τA , then by continuity and the definition of τB , 1 Z τB ≥ . 2 On the other hand, since τB ≤ τA , log d Zτ2B ≤ AητB ≤ C √ , d and hence p ZτB ≤ C log d d−1/4 = o(1), a contradiction. Therefore τA ≤ C0 log d,

σ = τA ,

τA < τB .

−1/2

Since Z0 ≤ Cinit d by assumption and given the lower bound η0 ≥ ∥θ∗ ∥∞ ≥ d−1/2 , we have 2 Z0 < Aη0 for all large d. Thus, by continuity, Zτ2A = AητA . Moreover,

ητA ≥ ∥θ∗ ∥∞ ≥ d−1/2 ,

and from the Phase-I bound,

Therefore

ZτA ≥ cA d−1/4 ,

and

log d η τA ≤ C √ . d p ZτA ≤ CA log d d−1/4 = o(1), log d η τA ≤ C √ . d

We also need the two overlaps to be balanced at τA . On [0, τA ], by the D-drift consequence (52) of Lemma C.6, we have e + ρt χt )Dt + O(Zt2 + ωt ). Ḋt = −C(1 (59) Integrating (59) on [0, τA ], and using Zt2 ≤ Aηt , gives Z τA |DτA | ≤ |D0 | + C (Zs2 + ωs ) ds. 0

Using

√ |D0 | ≤ Cinit d

−1/2

and τA ≤ C0 log d, we obtain C |DτA | ≤ √ + C d Since ZτA ≥ cA d

−1/4

,

ηs ≤ C

Z C0 log d  √ 0

log d + 1 + s √ , d

log d + 1 + s 1 + s √ + √ d d

, we have |DτA | = o(ZτA ).

In particular, for all sufficiently large d, |DτA | ≤ 28

1 Zτ . 8 A

1+s ωs ≤ C √ , d  ds ≤ C

(log d)2 √ . d

Step 3: Phase II Incoherence Control and Drift Pull-Up Exiting Phase I, we next show that Zt will be pulled from the d−1/4 -scale to a fixed constant. We define Phase II as when t ∈ [τA , t∗ ], where t∗ := inf{t ≥ τA : Zt ≥ cZ }. Choose a small constant cZ > 0, to be fixed below, and define log d H := K2 √ , d

T2 := τA + C2 log d,

where K2 , C2 are large constants. Define the bootstrap stopping time   1 σ2 := inf t ≥ τA : ηt > H or Zt > 2cZ or |Dt | > Zt or t = T2 . 2 At t = τA , none of the first three clauses holds for large d. Hence σ2 > τA . For t < σ2 , we have |Dt | ≤ and therefore

1 Zt , 2

1 3 Zt ≤ ρt , χt ≤ Zt . 4 4

Also Zt ≤ 2cZ . We first obtain a crude bound on ωt on the Phase-II bootstrap interval. From the third-moment differential inequality,
Ω̇t ≤ C Zt ∥θ∗ ∥∞ + Zt2 ηt + ηt2 . For t < σ2 , we have Zt ≤ 2cZ , ηt ≤ H, and t − τA ≤ C2 log d. Thus Z t Z t Z t Ωt ≤ ΩτA + C∥θ∗ ∥∞ Zs ds + C Zs2 ηs ds + C ηs2 ds. τA

τA

τA

Using 1 + τA log d ≤C √ , Ωτ A ≤ C √ d d r log d ∥θ∗ ∥∞ ≤ CK , d and the crude bounds Zs ≤ 2cZ , ηs ≤ H, we get log d (log d)3 (log d)3/2 (log d)2 √ Ωt ≤ C √ + CcZ . + Cc2Z √ +C d d d d After increasing the constant, this yields Ωt ≤ C

(log d)2 √ d

(t < σ2 ).

Since ωt ≤ 4Ωt + 4∥θ∗ ∥33 and ∥θ∗ ∥33 ≤ CK d−1/2 by Lemma C.1, we obtain ωt ≤ Wd := C

(log d)2 √ d

(t < σ2 ).

We next prove that Zt grows exponentially on the Phase-II bootstrap interval. Let τ− := inf{t ≥ τA : Zt < ZτA /2}. On [τA , σ2 ∧ τ− ), we have Zt ≥ ZτA /2, the balance bounds imply ρt , χt ≥ Zt /4, and, by choosing cZ small, ρt , χt ≤ 1/2. Hence the refined drift gives Żt = ρ̇t + χ̇t ≥ aZ Zt − bZ ωt ≥ aZ Zt − bZ Wd . Since

√ Wd (log d)2 / d (log d)2 ≤C = C 1/4 = o(1), −1/4 Z τA d d 29

we have, for all sufficiently large d, bZ Wd ≤

1 aZ Zt 2

on [τA , σ2 ∧ τ− ).

Thus Żt ≥ µZt (τA ≤ t < σ2 ∧ τ− ) for some µ > 0. This implies τ− > σ2 , and hence Zt ≥ ZτA eµ(t−τA )

(τA ≤ t < σ2 ).

In particular, for s ≤ t < σ2 , Zs ≤ e−µ(t−s) Zt . We now close the η-bootstrap before t∗ . For τA ≤ t < σ2 ∧ t∗ , Lemma C.4 gives Z t Z t Z t ∗ 2 ηt ≤ ητA + C∥θ ∥∞ Zs ds + C Zs ηs ds + C ηs2 ds. τA

τA

τA

∗

Since t < t , we have Zt < cZ . Using the exponential comparison Zs ≤ e−µ(t−s) Zt , we obtain Z t Z t Zs ds ≤ CcZ , Zs2 ds ≤ Cc2Z . τA

τA

Therefore log d ηt ≤ C √ + CcZ d Choose cZ > 0 small enough that

r

log d + Cc2Z H + CH 2 log d. d

Cc2Z H ≤

1 H. 8

Then choose K2 large enough so that log d 1 C √ ≤ H. 4 d Moreover, r

log d = o(H), d Thus, for all sufficiently large d, CcZ

H 2 log d = K22

(log d)3 = o(H). d

1 H (τA ≤ t < σ2 ∧ t∗ ). 2 By a first-violation argument, the clause ηt > H cannot trigger before t∗ . ηt ≤

Next we close the D/Z-balance bootstrap. On the balanced small-overlap band, the same subtraction estimate as above gives e + ρt χt )Dt + O(Z 2 + ωt ). Ḋt = −C(1 t Since cZ is small, 1 + ρt χt ≥ 1/2. Hence, for τA ≤ t < σ2 ∧ t∗ , Z t |Dt | ≤ e−c(t−τA ) |DτA | + C e−c(t−s) (Zs2 + ωs ) ds. τA

The entry term is bounded by e−c(t−τA ) |DτA | ≤ |DτA | ≤

1 1 Zτ ≤ Zt . 8 A 8

The ω-term is bounded by CWd = o(ZτA ) ≤ o(Zt ). For the Zs2 -term, using

Zs ≤ e−µ(t−s) Zt , 30

we get Z t

e−c(t−s) Zs2 ds ≤ CZt2 ≤ CcZ Zt .

τA

Choosing cZ smaller if necessary, and then d large, gives 3 1 Zt < Zt (τA ≤ t < σ2 ∧ t∗ ). 8 2 Thus the clause |Dt | > Zt /2 cannot trigger before t∗ . |Dt | ≤

We now show that t∗ ≤ T2 . Suppose, for contradiction, that t∗ > T2 . Then Zt < cZ for all t ≤ T2 . The clauses ηt > H and |Dt | > Zt /2 cannot trigger before t∗ , as just shown, and the clause Zt > 2cZ cannot trigger while Zt < cZ . Hence T2 < σ2 . The exponential lower bound then gives ZT2 ≥ ZτA eµC2 log d ≥ cA d−1/4 dµC2 . Choosing C2 > 1/(4µ) large enough gives ZT2 ≥ cZ for all sufficiently large d, contradicting t∗ > T2 . Therefore t∗ ≤ T2 . By continuity, Zt∗ = cZ . The bootstrap balance gives |Dt∗ | ≤ Therefore min{ρt∗ , χt∗ } =

1 Z t∗ . 2

Zt∗ − |Dt∗ | cZ ≥ . 2 4

Set tesc := t∗ ,

cesc :=

cZ . 4

Since τA ≤ C0 log d and

t∗ ≤ τA + C2 log d,

we have tesc ≤ (C0 + C2 ) log d. Thus the theorem holds with Cesc := C0 + C2 . Finally, the escape-window ηt -bound follows by combining the Phase-I estimate √ log d + 1 + t log d √ ηt ≤ KI ≤C √ (0 ≤ t ≤ τA ) d d with the Phase-II bootstrap bound log d ηt ≤ H = K 2 √ d

(τA ≤ t ≤ t∗ ).

Hence, throughout weak recovery window 0 ≤ t ≤ tesc , log d ηt ≤ C √ = od (1). d This completes the proof. 31

C.4

Strong Recovery

Building on Lemma C.8, we show that after weak recovery the population dynamics enters the fixed local strong-convexity cap around the teacher in an additional O(1) time. Once inside this cap, local strong convexity gives exponential contraction to the teacher, achiving strong recovery. Lemma C.9 (Entry into the local cap after weak recovery). Let rloc > 0 be the radius from Lemma C.14. After possibly shrinking rloc , assume π rloc < . 2 Set sloc := cos(rloc /4). Assume C.8 with tesc , cesc , and Cesc . Under the same assumptions as Lemmas C.3 and C.4, for all sufficiently large d, there exists a time tloc ≤ tesc + Cin where Cin < ∞ is independent of d, such that ρtloc ≥ sloc ,

χtloc ≥ sloc .

Consequently, distM (ytloc , y ∗ )2 ≤

2 rloc . 8

In particular, since tesc ≤ Cesc log d, tloc ≤ Cloc log d for some constant Cloc < ∞ independent of d. Proof. The proof is deterministic on the favorable event supplied by Lemma C.8. By that lemma, there are constants cesc ∈ (0, 1) and Cesc < ∞, independent of d, and a time tesc such that tesc ≤ Cesc log d,

ρtesc ≥ cesc ,

Set c− :=

χtesc ≥ cesc ,

ηtesc = od (1).

cesc . 2

Since rloc > 0, we have sloc = cos(rloc /4) < 1. Hence

mloc := 1 − s2loc > 0.

By Lemma C.3, on the positive-overlap region ρt , χt ≥ 0, the flow satisfies ρ̇t ≥ aρ χt (1 − ρ2t ) − bρ ηt , χ̇t ≥ aχ ρt (1 − χ2t ) − bχ ηt , where aρ , aχ , bρ , bχ > 0 are independent of d. Choose ε∗ > 0, independent of d, so small that 1 aρ c− mloc , 2

bχ ε∗ ≤

1 aχ c− mloc , 2

1 aρ c− (1 − c2− ), 2

bχ ε∗ ≤

1 aχ c− (1 − c2− ). 2

1 aρ c− mloc , 2

vχ :=

1 aχ c− mloc . 2

(sloc − cesc )+ , vρ

Tχ :=

(sloc − cesc )+ , vχ

bρ ε∗ ≤ and also bρ ε∗ ≤ Define vρ := Let Tρ :=

32

and set Cin := 1 + max{Tρ , Tχ }. This constant depends only on the fixed population constants and on rloc , and therefore is independent of d. We first show that ηt remains smaller than ε∗ on the fixed window [tesc , tesc + Cin ]. By Lemma C.4,

η̇t ≤ C1 |ρt |∥θ∗ ∥∞ + C2 |χt ρt |ηt + C3 ηt2 . Since |ρt |, |χt | ≤ 1, on the present fixed window this gives η̇t ≤ Ad + Bηt + Cηt2 ,

where

r ∗

Ad := C1 ∥θ ∥∞ ≤ C

log d = od (1), d

and B, C < ∞ are independent of d. Define τη := inf{t ≥ tesc : ηt ≥ ε∗ } ∧ (tesc + Cin ). Since ηtesc = od (1), for all sufficiently large d, ηtesc < ε∗ . Thus τη > tesc . On [tesc , τη ], we have ηt ≤ ε∗ , and therefore for almost every t ∈ [tesc , τη ], η̇t ≤ Ad + (B + Cε∗ )ηt . Set B∗ := B + Cε∗ . By Gronwall’s inequality, for every t ∈ [tesc , τη ], Z t ηt ≤ eB∗ (t−tesc ) ηtesc + Ad eB∗ (t−u) du. tesc

Hence

ηt ≤ eB∗ Cin ηtesc + Ad Cin eB∗ Cin . The right-hand side is od (1). Therefore, for all sufficiently large d, ηt < ε∗

for all t ∈ [tesc , τη ].

By the definition of τη , this forces τη = tesc + Cin . Consequently, ηt ≤ ε∗

for all t ∈ [tesc , tesc + Cin ].

Next we prove that the rectangle R− := {(ρ, χ) : ρ ≥ c− , χ ≥ c− } is forward invariant on the same window. At time tesc , ρtesc ≥ cesc > c− ,

χtesc ≥ cesc > c− .

Suppose, for contradiction, that the trajectory exits R− before tesc + Cin . Let τ− := inf {t ∈ [tesc , tesc + Cin ] : min{ρt , χt } ≤ c− } . Then τ− < ∞, and by continuity, ρτ− ≥ c− ,

χτ− ≥ c− ,

with at least one equality. 33

If ρτ− = c− , then ρτ− , χτ− ≥ 0, so Lemma C.3 applies. Using also ητ− ≤ ε∗ , we get ρ̇τ− ≥ aρ c− (1 − c2− ) − bρ ε∗ ≥

1 aρ c− (1 − c2− ) > 0. 2

Thus ρt cannot cross downward through c− at τ− . Similarly, if χτ− = c− , then χ̇τ− ≥ aχ c− (1 − c2− ) − bχ ε∗ ≥

1 aχ c− (1 − c2− ) > 0, 2

so χt cannot cross downward through c− . This contradicts the definition of τ− . Therefore, ρt , χt ≥ c−

for all t ∈ [tesc , tesc + Cin ].

We now show that ρt reaches sloc within time Tρ and then cannot fall below sloc during the window. Define τρ := inf{t ≥ tesc : ρt ≥ sloc }. If ρtesc ≥ sloc , then τρ = tesc . Otherwise, before ρt reaches sloc , we have ρt < sloc . On this interval, the rectangle invariance gives χt ≥ c− , and the previous step gives ηt ≤ ε∗ . Hence Lemma C.3 implies ρ̇t ≥ aρ c− (1 − ρ2t ) − bρ ε∗ . Since ρt < sloc ,

1 − ρ2t ≥ 1 − s2loc = mloc .

Therefore ρ̇t ≥ aρ c− mloc − bρ ε∗ ≥

1 aρ c− mloc = vρ . 2

Thus ρt reaches sloc no later than tesc + Tρ . Moreover, after ρt reaches sloc , it cannot cross downward through sloc inside the same window. Indeed, at a hypothetical first downward crossing through sloc , we would have ρt = sloc ,

χt ≥ c− ,

ηt ≤ ε∗ .

Therefore ρ̇t ≥ aρ c− (1 − s2loc ) − bρ ε∗ ≥ vρ > 0, which is incompatible with a first downward crossing. The same argument applies to χt . Define τχ := inf{t ≥ tesc : χt ≥ sloc }. If χtesc ≥ sloc , then τχ = tesc . Otherwise, before χt reaches sloc , we have χt < sloc , while ρt ≥ c− and ηt ≤ ε∗ . Hence χ̇t ≥ aχ c− (1 − χ2t ) − bχ ε∗ ≥ aχ c− mloc − bχ ε∗ ≥ vχ . Thus χt reaches sloc no later than tesc + Tχ . Once it reaches sloc , it cannot cross downward through sloc , because at such a crossing the same estimate gives χ̇t ≥ vχ > 0. Now set tloc := max{τρ , τχ }. By the previous estimates, tloc ≤ tesc + max{Tρ , Tχ } ≤ tesc + Cin , and by the no-downcrossing arguments, ρtloc ≥ sloc ,

χtloc ≥ sloc . 34

It remains to convert the overlap bounds into a product-manifold distance bound. Since ρtloc = ⟨βtloc , β ∗ ⟩ ≥ sloc = cos(rloc /4), and since arccos is decreasing on [−1, 1], distSd−1 (βtloc , β ∗ ) = arccos(ρtloc ) ≤

rloc . 4

distSd−1 (θtloc , θ∗ ) = arccos(χtloc ) ≤

rloc . 4

Similarly,

Using the product metric on M = Sd−1 × Sd−1 , we get distM (ytloc , y ∗ )2 = distSd−1 (θtloc , θ∗ )2 + distSd−1 (βtloc , β ∗ )2 . Therefore, distM (ytloc , y ∗ )2 ≤

r

loc

4

2

+

r

loc

4

2

=

2 rloc . 8

Finally, since Lemma C.8 gives tesc ≤ Cesc log d, we have tloc ≤ Cesc log d + Cin . For d ≥ e, Cin ≤ Cin log d, and hence tloc ≤ (Cesc + Cin ) log d. Thus the last claim holds with Cloc := Cesc + Cin . The proof is complete. Curvature near the teacher. Here, we prove the local curvature estimate used for the contraction phase. The comparison with the Gaussian model is needed only at the teacher; the extension to a small cap uses Hessian continuity of the true loss. Lemma C.10 (Gaussian local strong convexity). Let M := Sd−1 × Sd−1 , y = (θ, β), y ∗ = (θ∗ , β ∗ ). Let LG be the Gaussian population loss h i 2 LG (θ, β) = E (ψ(g · θ) − ψ(g ∗ · θ∗ )) , where (g, g ∗ ) ∈ Rd ×Rd is centered Gaussian with E[gi gj ] = δij , E[gi∗ gj∗ ] = δij , E[gi gj∗ ] = κ(ρ)δij , with X ρ := β · β ∗ , κ(ρ) := a2ℓ ρℓ . ℓ≥1

Assume ∥ϕ∥L2 (γ) = 1, so κ(1) = 1, and assume the information-exponent-one conditions α1 ̸= 0, a1 ̸= 0, where X X ψ= αk hk , ϕ= aℓ hℓ . k≥0

ℓ≥1

Assume also enough regularity so that the function below is C 2 in a neighborhood of (1, 1). Then there exist constants rG > 0 and µG > 0, independent of d, such that for every  2 (θ, β) ∈ UrG := (θ, β) : distS (θ, θ∗ )2 + distS (β, β ∗ )2 ≤ rG , and every product tangent vector ξ = (u, v) ∈ Tθ Sd−1 × Tβ Sd−1 , one has
HessM LG (θ, β)[ξ, ξ] ≥ µG ∥u∥2 + ∥v∥2 . 35

(60)

Proof. In the Gaussian model, the scalar pair (g · θ, g ∗ · θ∗ ) is centered Gaussian with unit variances and correlation m = κ(ρ)χ. Therefore, using the Hermite expansion of ψ, X E[ψ(g · θ)ψ(g ∗ · θ∗ )] = αk2 (κ(ρ)χ)k . k≥0

Since both marginals have variance one, E[ψ(g · θ)2 ] = E[ψ(g ∗ · θ∗ )2 ] =

X

αk2 .

k≥0

Hence LG (θ, β) = 2

X

  αk2 1 − (κ(ρ)χ)k ,

(61)

k≥0

with k = 0 term is zero. We define Φ(χ, ρ) := 2

X

  αk2 1 − (κ(ρ)χ)k = LG (θ, β).

k≥1

Now we first compute the Hessian point (θ, β) = (θ∗ ,P β ∗ ) : we have χ = 1, ρ = P at the planted 2 ′ 1, κ(1) = 1. Also Φχ (1, 1) = −2 k≥1 kαk , and Φρ (1, 1) = −2κ (1) k≥1 kαk2 . Define X λθ := −Φχ (1, 1) = 2 kαk2 , k≥1

and

λβ := −Φρ (1, 1) = 2κ′ (1)

X

kαk2 .

k≥1

The information-exponent-one condition α1 ̸= 0 implies λθ > 0. Moreover, X κ′ (1) = ℓa2ℓ > 0, ℓ≥1

because a1 ̸= 0. Hence λβ > 0. Now let u ∈ Tθ∗ Sd−1 , v ∈ Tβ ∗ Sd−1 and let (θt , βt ) be the product geodesic with (θ0 , β0 ) = (θ∗ , β ∗ ), Then and Thus

(θ̇0 , β̇0 ) = (u, v).

χ(t) := θt · θ∗ = cos(t∥u∥), ρ(t) := βt · β ∗ = cos(t∥v∥). χ′ (0) = 0,

χ′′ (0) = −∥u∥2 ,

and

ρ′ (0) = 0, ρ′′ (0) = −∥v∥2 . Since the curve is a product geodesic, the Riemannian Hessian is computed by the second derivative along this curve: d2 HessM LG (y ∗ )[(u, v), (u, v)] = 2 Φ(χ(t), ρ(t)). dt t=0 By the chain rule and the fact that χ′ (0) = ρ′ (0) = 0, d2 Φ(χ(t), ρ(t)) = Φχ (1, 1)χ′′ (0) + Φρ (1, 1)ρ′′ (0). dt2 t=0 Therefore

HessM LG (y ∗ )[(u, v), (u, v)] = λθ ∥u∥2 + λβ ∥v∥2 .

Consequently, with λ0 := min{λθ , λβ } > 0, 36

we have HessM LG (y ∗ )[(u, v), (u, v)] ≥ λ0 (∥u∥2 + ∥v∥2 ). It remains to extend this lower bound to a small product cap around y ∗ . We do this explicitly. Let u ∈ Tθ Sd−1 ,

(θ, β) ∈ M,

v ∈ Tβ Sd−1 .

Let (θt , βt ) be the product geodesic starting from (θ, β) with initial velocity (u, v). Then χ′ (0) = u · θ∗ ,

χ′′ (0) = −χ∥u∥2 ,

ρ′ (0) = v · β ∗ ,

ρ′′ (0) = −ρ∥v∥2 .

and Therefore HessM LG (θ, β)[(u, v), (u, v)]

(62)

∗ 2

∗

∗

∗ 2

= Φχχ (χ, ρ)(u · θ ) + 2Φχρ (χ, ρ)(u · θ )(v · β ) + Φρρ (χ, ρ)(v · β ) − χΦχ (χ, ρ)∥u∥2 − ρΦρ (χ, ρ)∥v∥2 .

(63)

We now estimate the terms in this expression. Choose a small neighborhood N := {(χ, ρ) : |χ − 1| ≤ r0 , |ρ − 1| ≤ r0 } on which all second derivatives of Φ are bounded: |Φχχ | + |Φχρ | + |Φρρ | ≤ M. By continuity of Φχ , Φρ , after shrinking r0 if necessary, we also have −χΦχ (χ, ρ) ≥

λθ , 2

−ρΦρ (χ, ρ) ≥

λβ 2

throughout N . Now suppose distS (θ, θ∗ ) ≤ r,

distS (β, β ∗ ) ≤ r,

with r > 0 small enough so that (χ, ρ) ∈ N . Since u ⊥ θ, |u · θ∗ | = |u · Pθ⊥ θ∗ | ≤ ∥u∥∥Pθ⊥ θ∗ ∥ = ∥u∥

p

1 − χ2 ≤ r∥u∥.

Similarly, |v · β ∗ | ≤ r∥v∥. Hence the second-derivative terms are bounded in absolute value by M r2 ∥u∥2 + 2M r2 ∥u∥∥v∥ + M r2 ∥v∥2 ≤ 2M r2 (∥u∥2 + ∥v∥2 ). Therefore, HessM LG (θ, β)[(u, v), (u, v)] ≥

λθ λβ ∥u∥2 + ∥v∥2 − 2M r2 (∥u∥2 + ∥v∥2 ). 2 2

2 Choose rG > 0 small enough that 2M rG ≤ 41 min{λθ , λβ }. Then, for every r ≤ rG ,

HessM LG (θ, β)[(u, v), (u, v)] ≥

1 min{λθ , λβ }(∥u∥2 + ∥v∥2 ). 4

Thus the lemma holds with µG :=

1 min{λθ , λβ } > 0. 4

37

Lemma C.11 (Teacher Hessian block formula). Let y ∗ := (θ∗ , β ∗ ). Define Gi := x(i) · β ∗ ,

S :=

d X

θi∗ ϕ(Gi ).

i=1

For u ∈ Tθ∗ Sd−1 , set Au :=

d X

ui ϕ(Gi ).

i=1

Then, for u1 , u2 ∈ Tθ∗ Sd−1 and v1 , v2 ∈ Tβ ∗ Sd−1 , the Hessian blocks of L at the teacher satisfy   HL,θθ (y ∗ )[u1 , u2 ] = 2E ψ ′ (S)2 Au1 Au2 , (64) HL,θβ (y ∗ )[u, v] = 0,

(65)

HL,ββ (y ∗ )[v1 , v2 ] = 2⟨v1 , v2 ⟩

d X

  (θi∗ )2 E ψ ′ (S)2 ϕ′ (Gi )2 .

(66)

i=1

Consequently, for every (u, v) ∈ Tθ∗ Sd−1 × Tβ ∗ Sd−1 , d X     HessM L(y ∗ )[(u, v), (u, v)] = 2E ψ ′ (S)2 A2u + 2∥v∥2 (θi∗ )2 E ψ ′ (S)2 ϕ′ (Gi )2 .

(67)

i=1

Proof. We first compute the full quadratic form at the teacher and then polarize. Let yt = (θt , βt ) be any product geodesic on M = Sd−1 × Sd−1 with y0 = y ∗ ,

ẏ0 = (u, v).

Write ft (x) := fβt ,θt (x), Since f0 = f∗ pointwise,

f∗ (x) := fβ ∗ ,θ∗ (x).

L(yt ) = E[(ft − f∗ )2 ],

and therefore

d2 L(yt ) = 2E[f˙02 ] + 2E[(f0 − f∗ )f¨0 ] = 2E[f˙02 ]. dt2 t=0 Because yt is a geodesic, this second derivative is exactly the Riemannian Hessian: HessM L(y ∗ )[(u, v), (u, v)] = 2E[f˙02 ]. Now compute f˙0 . At y ∗ , ft (x) = ψ

d X

! θt,i ϕ(x(i) · βt ) .

i=1

Thus f˙0 = ψ ′ (S)

" d X

ui ϕ(Gi ) +

i=1

d X

# θi∗ ϕ′ (Gi )(x(i) · v) .

i=1

Define, for the proof, Bv :=

d X

θi∗ ϕ′ (Gi )Hi (v),

Hi (v) := x(i) · v.

i=1

Then

f˙0 = ψ ′ (S)(Au + Bv ),

and hence   HessM L(y ∗ )[(u, v), (u, v)] = 2E ψ ′ (S)2 (Au + Bv )2 . 38

(68)

Since the right-hand side of (68) is a quadratic form in (u, v), its associated symmetric bilinear form is obtained by polarization:   HessM L(y ∗ )[(u1 , v1 ), (u2 , v2 )] = 2E ψ ′ (S)2 (Au1 + Bv1 )(Au2 + Bv2 ) . (69) Now, we derive the three Hessian blocks: First, take v1 = v2 = 0 in (69). This gives   HL,θθ (y ∗ )[u1 , u2 ] = 2E ψ ′ (S)2 Au1 Au2 , which proves (64), the θθ-block. Second, take (u1 , v1 ) = (u, 0) and (u2 , v2 ) = (0, v). Then   HL,θβ (y ∗ )[u, v] = 2E ψ ′ (S)2 Au Bv . We show this expectation is zero. Since v ∈ Tβ ∗ Sd−1 , we have v ⊥ β ∗ . For each i, the pair Gi = x(i) · β ∗ ,

Hi (v) = x(i) · v

is jointly Gaussian with covariance E[Gi Hi (v)] = ⟨β ∗ , v⟩ = 0. Hence Gi and Hi (v) are independent. Moreover, conditional on (G1 , . . . , Gd ), the variables Hi (v) are centered Gaussian and independent across i. Therefore E[Bv | G1 , . . . , Gd ] =

d X

θi∗ ϕ′ (Gi )E[Hi (v) | G1 , . . . , Gd ] = 0.

i=1 ′

2

Since ψ (S) Au is measurable with respect to σ(G1 , . . . , Gd ), we obtain   E[ψ ′ (S)2 Au Bv ] = E ψ ′ (S)2 Au E[Bv | G1 , . . . , Gd ] = 0. Thus we have (65), the θβ-block:

HL,θβ (y ∗ )[u, v] = 0.

Finally, take u1 = u2 = 0 in (69). Then   HL,ββ (y ∗ )[v1 , v2 ] = 2E ψ ′ (S)2 Bv1 Bv2 . Conditioning again on G := (G1 , . . . , Gd ), B v1 B v2 =

d X

θi∗ θj∗ ϕ′ (Gi )ϕ′ (Gj )Hi (v1 )Hj (v2 ).

i,j=1

For i ̸= j, conditional independence and centering give E[Hi (v1 )Hj (v2 ) | G] = 0. For i = j, since v1 , v2 ⊥ β ∗ , E[Hi (v1 )Hi (v2 ) | G] = ⟨v1 , v2 ⟩. Hence E[Bv1 Bv2 | G] = ⟨v1 , v2 ⟩

d X

(θi∗ )2 ϕ′ (Gi )2 .

i=1

Because ψ ′ (S)2 is G-measurable, E[ψ ′ (S)2 Bv1 Bv2 ] = ⟨v1 , v2 ⟩

d X

  (θi∗ )2 E ψ ′ (S)2 ϕ′ (Gi )2 .

i=1

Therefore we have the final ββ-block in (66): HL,ββ (y ∗ )[v1 , v2 ] = 2⟨v1 , v2 ⟩

d X   (θi∗ )2 E ψ ′ (S)2 ϕ′ (Gi )2 . i=1

Combining the θθ-block, the vanishing mixed block, and the ββ-block with u1 = u2 = u, v1 = v2 = v, gives the final quadratic formula. This completes the proof. 39

Stein inputs for the teacher Hessian. We next state the two blockwise teacher-level comparisons used below. Their proofs are deferred to Subsection E.1. Lemma C.12 (Teacher Hessian comparison). Let Aψ := E[ψ ′ (G)2 ],

Bϕ := E[ϕ′ (G)2 ] = κ′ (1),

λθ := 2Aψ ,

λβ := 2Aψ Bϕ .

Under the regularity assumptions of Lemma E.4 and Lemma E.6, there exists C∗ > 0, independent of d, such that sup |HessM L(y ∗ )[ξ, ξ] − HessM LG (y ∗ )[ξ, ξ]| ≤ C∗ ∥θ∗ ∥∞ .

(70)

∥ξ∥=1

Moreover, HessM LG (y ∗ )[(u, v), (u, v)] = λθ ∥u∥2 + λβ ∥v∥2 . Proof. Let

(71)

∥ξ∥2 = ∥u∥2 + ∥v∥2 .

ξ = (u, v) ∈ Ty∗ M,

First, the Gaussian teacher Hessian formula follows from the explicit Gaussian loss: X
LG (θ, β) = 2 αk2 1 − (κ(ρ)χ)k . k≥1 ∗

At y , tangent geodesics satisfy χ′′ (0) = −∥u∥2 , Therefore

ρ′′ (0) = −∥v∥2 ,

χ′ (0) = ρ′ (0) = 0.

HessM LG (y ∗ )[(u, v), (u, v)] = λθ ∥u∥2 + λβ ∥v∥2 ,

with λθ = 2

X

kαk2 = 2Aψ ,

λβ = 2κ′ (1)

k≥1

X

kαk2 = 2Aψ Bϕ .

k≥1

For the true loss, decompose the Hessian into blocks. By Lemma C.11, HL,θβ (y ∗ ) = 0. The same geodesic computation above gives HLG ,θβ (y ∗ ) = 0. For the θθ-block, Lemma E.4 at (θ, β) = (θ∗ , β ∗ ) gives |HL,θθ (y ∗ )[u, u] − HLG ,θθ (y ∗ )[u, u]| ≤ Cθ ∥θ∗ ∥∞ ∥u∥2 . For the ββ-block, Lemma E.6 gives |HL,ββ (y ∗ )[v, v] − HLG ,ββ (y ∗ )[v, v]| ≤ Cβ ∥θ∗ ∥∞ ∥v∥2 . Combining the three block estimates, |HL (y ∗ )[ξ, ξ] − HLG (y ∗ )[ξ, ξ]| ≤ (Cθ + Cβ )∥θ∗ ∥∞ (∥u∥2 + ∥v∥2 ). Taking ∥ξ∥ = 1 and setting C∗ := Cθ + Cβ proves the claim. Lemma C.13 (Uniform Hessian continuity near the teacher). Assume ψ ∈ Cb3 (R), and assume that ϕ ∈ C 3 (R) satisfies h i max E (1 + |G|)12 |ϕ(j) (G)|12 < ∞, G ∼ N (0, 1). 0≤j≤3

Assume also E[ϕ(G)] = 0, E[ϕ(G)2 ] = 1. Then there exist constants r0 > 0 and C3 > 0, independent of d, such that for every y ∈ Ur0 , ∥Py→y∗ HessM L(y)Py∗ →y − HessM L(y ∗ )∥op ≤ C3 distM (y, y ∗ ). Here Py→y∗ denotes parallel transport along the minimizing product geodesic from y to y ∗ . 40

(72)

Proof. It is enough to prove that the covariant derivative of the Hessian is uniformly bounded on a fixed product cap. Indeed, let γ : [0, ℓ] → M be the minimizing product geodesic from y ∗ to y, parametrized by arclength, so that ℓ = distM (y, y ∗ ). Let ξ0 ∈ Ty∗ M be a unit vector, and let ξt be its parallel transport along γ. Then d HessM L(γ(t))[ξt , ξt ] = (∇ HessM L)(γ(t))[γ̇(t), ξt , ξt ]. dt Thus, if sup

|(∇ HessM L)(γ(t))[ζ, ξ, ξ]| ≤ C3 ,

sup

γ(t)∈Ur0 ∥ζ∥=∥ξ∥=1

then integration along γ gives the desired transported-Hessian Lipschitz bound. We prove this uniform bound. Write sy (x) :=

d X

θi ϕ(x(i) · β),

s∗ (x) :=

i=1

d X

θi∗ ϕ(x(i) · β ∗ ),

i=1

so that L(y) = E

h

ψ(sy (x)) − ψ(s∗ (x))

2 i

.

Let yt = (θt , βt ) be any unit-speed product geodesic contained in Ur0 , and set st := syt (x). We first record the basic moment estimate   (3) sup ∥ṡt ∥L6 + ∥s̈t ∥L3 + ∥st ∥L1 ≤ C,

(73)

t

where C is independent of d and of the chosen unit-speed product geodesic. To prove (73), observe that st =

d X

Ut,i := x(i) · βt .

θt,i ϕ(Ut,i ),

i=1

Along a unit-speed spherical geodesic, (3)

(3)

∥θ̇t ∥ + ∥θ̈t ∥ + ∥θt ∥ ≤ C,

∥β̇t ∥ + ∥β̈t ∥ + ∥βt ∥ ≤ C.

Moreover, U̇t,i = x(i) · β̇t ,

Üt,i = x(i) · β̈t ,

(3)

(3)

Ut,i = x(i) · βt .

The variables coming from distinct blocks i are independent, and each block has moments bounded uniformly in t. Differentiating θt,i ϕ(Ut,i ) up to order three gives a finite sum of terms whose coefficient vectors have uniformly bounded ℓ2 -norm and whose block variables have uniformly bounded moments by the assumptions on ϕ. The only point that needs care is centering. For each i, E[ϕ(Ut,i )] = E[ϕ(G)] = 0 for every t, because Ut,i ∼ N (0, 1). Differentiating this identity in t shows that the differentiated (3) block variables appearing in ṡt , s̈t , st are centered. Therefore Rosenthal’s inequality applies to each differentiated sum and gives the bounds in (73), uniformly in d. Now define Rt := ψ(st ) − ψ(s∗ ). Since ψ and its derivatives up to order three are bounded, Ṙt = ψ ′ (st )ṡt , 41

R̈t = ψ ′′ (st )ṡ2t + ψ ′ (st )s̈t , and

(3)

Rt

(3)

= ψ (3) (st )ṡ3t + 3ψ ′′ (st )ṡt s̈t + ψ ′ (st )st .

Since L(yt ) = E[Rt2 ],

h i d3 (3) L(y ) = 2E 3 Ṙ R̈ + R R . t t t t t dt3 Using the boundedness of ψ, ψ ′ , ψ ′′ , ψ (3) , together with (73) and Hölder’s inequality, we obtain sup t

d3 L(yt ) ≤ C. dt3

The same argument, applied to geodesic variations and then polarized, gives sup

|(∇ HessM L)(y)[ζ, ξ, ξ]| ≤ C.

sup

y∈Ur0 ∥ζ∥=∥ξ∥=1

Equivalently, the transported Hessian is C-Lipschitz along every minimizing product geodesic in Ur0 . Integrating this bound along the geodesic from y ∗ to y gives ∥Py→y∗ HessM L(y)Py∗ →y − HessM L(y ∗ )∥op ≤ C distM (y, y ∗ ). This proves the lemma, with C3 = C. Lemma C.14 (True local product strong convexity). Assume the teacher Hessian comparison of Lemma C.12 and the uniform Hessian continuity of Lemma C.13. On the random-spherical delocalization event of Lemma C.1, so that r log d ∗ ∥θ ∥∞ ≤ CK , d there exist constants rloc > 0 and µloc > 0, independent of d, such that for all sufficiently large d and every y ∈ Urloc , HessM L(y) ⪰ µloc I.

(74)

Proof. Let λ0 := min{λθ , λβ } = 2Aψ min{1, Bϕ }. By the information-exponent-one assumptions, Aψ > 0 and Bϕ > 0, so λ0 > 0. By Lemma C.12, for every ξ ∈ Ty∗ M, HessM L(y ∗ )[ξ, ξ] ≥ HessM LG (y ∗ )[ξ, ξ] − C∗ ∥θ∗ ∥∞ ∥ξ∥2 . Since by Lemma C.10 we obtain

HessM LG (y ∗ )[ξ, ξ] ≥ λ0 ∥ξ∥2 , HessM L(y ∗ )[ξ, ξ] ≥ (λ0 − C∗ ∥θ∗ ∥∞ ) ∥ξ∥2 .

By Lemma C.1, r ∗

∥θ ∥∞ ≤ CK

log d . d

Hence, for all sufficiently large d, r ∗

C∗ ∥θ ∥∞ ≤ C∗ CK Hence HessM L(y ∗ ) ⪰

log d λ0 ≤ . d 4 3λ0 I. 4

Now choose rloc ≤ r0

and 42

C3 rloc ≤

λ0 . 4

Let y ∈ Urloc , and let ξ ∈ Ty M be a unit tangent vector. Parallel transport ξ to Ty∗ M, writing ξe := Py→y∗ ξ. e = 1. By Lemma C.13, Then ∥ξ∥ e ξ] e − C3 distM (y, y ∗ ). HessM L(y)[ξ, ξ] ≥ HessM L(y ∗ )[ξ, Since y ∈ Urloc ,

distM (y, y ∗ ) ≤ rloc .

Therefore HessM L(y)[ξ, ξ] ≥

λ0 3λ0 − C3 rloc ≥ . 4 2

Thus the lemma holds with

λ0 . 2

µloc :=

Strong Recovery Dynamics. Lemma C.15 (Local contraction under product strong convexity). Let M = Sd−1 × Sd−1 equipped with the product Riemannian metric, and let y ∗ = (θ∗ , β ∗ ) ∈ M. Fix r ∈ (0, π/2) and suppose L : M → R is C 2 and satisfies for all y ∈ Ur := {y ∈ M : distM (y, y ∗ ) ≤ r},

HessM L(y) ⪰ µI

for some µ > 0 (Lemma C.14, with r = rloc and µ = µloc ), and that y ∗ is a critical point of L, i.e., ∇M L(y ∗ ) = 0. Let yt be the population gradient flow ẏt = −∇M L(yt ),

yt0 ∈ M,

distM (yt0 , y ∗ ) ≤ r/2.

Then yt ∈ Ur for all t ≥ t0 , and distM (yt , y ∗ )2 ≤ e−2µ(t−t0 ) distM (yt0 , y ∗ )2 .

(75)

Proof. Write Dt := distM (yt , y ∗ ),

Et :=

1 2 D . 2 t

Since r < π/2, every point y = (θ, β) ∈ Ur satisfies distS (θ, θ∗ ) < π/2,

distS (β, β ∗ ) < π/2.

Thus neither component is antipodal to its target. Hence the minimizing geodesic from y to y ∗ is unique and is obtained by taking the componentwise minimizing geodesics on the two spheres. In particular, Logy (y ∗ ) is well-defined, and ∥ Logy (y ∗ )∥ = distM (y, y ∗ ). Moreover, if γ : [0, 1] → M is this minimizing product geodesic from y to y ∗ , then distM (γ(s), y ∗ ) = (1 − s) distM (y, y ∗ ), so γ(s) ∈ Ur whenever y ∈ Ur . Define the exit time τ := inf{t ≥ t0 : Dt > r}. For t ∈ [t0 , τ ), we have yt ∈ Ur . We will prove d Et ≤ −2µEt dt Once this is shown, Gronwall gives

for all t ∈ [t0 , τ ).

Et ≤ e−2µ(t−t0 ) Et0

for all t ∈ [t0 , τ ). 43

Since Dt0 ≤ r/2, this implies Dt ≤ r/2 on [t0 , τ ). If τ < ∞, then by continuity Dτ ≤ r/2 < r, contradicting the definition of τ . Hence τ = ∞, and the trajectory stays in Ur for all t ≥ t0 . It remains to prove the differential inequality. Fix t < τ , and let γ : [0, 1] → M be the minimizing geodesic from yt to y ∗ , parametrized with constant speed: γ(0) = yt , Then

γ(1) = y ∗ ,

∥γ ′ (s)∥ ≡ Dt .

γ ′ (0) = Logyt (y ∗ ).

Set φ(s) := L(γ(s)). Since γ is a geodesic,

φ′′ (s) = HessM L(γ(s))[γ ′ (s), γ ′ (s)]. The geodesic remains in Ur , so the Hessian lower bound gives φ′′ (s) ≥ µ∥γ ′ (s)∥2 = µDt2 . Integrating from 0 to 1,

φ′ (1) − φ′ (0) ≥ µDt2 .

∗

Because y is critical,

φ′ (1) = ⟨∇M L(y ∗ ), γ ′ (1)⟩ = 0.

Therefore

φ′ (0) ≤ −µDt2 .

But

φ′ (0) = ∇M L(yt ), Logyt (y ∗ ) .

Thus

∇M L(yt ), Logyt (y ∗ ) ≤ −µDt2 .

Finally, the first-variation formula for the squared distance gives d Et = − ẏt , Logyt (y ∗ ) . dt Since ẏt = −∇M L(yt ), d Et = ∇M L(yt ), Logyt (y ∗ ) ≤ −µDt2 = −2µEt . dt This proves the desired differential inequality, and hence the contraction estimate. Theorem C.16 (Continuous-time strong recovery). Assume the weak-recovery conclusion of Lemma C.8, together with the assumptions of Lemmas C.3 and C.4. Assume also the regularity and incoherence assumptions under which Lemmas C.12, C.13, C.14, and C.15 hold. Let yt := (θt , βt ) ∈ M := Sd−1 × Sd−1 , y ∗ := (θ∗ , β ∗ ), and write Rt := distM (yt , y ∗ ). Then there exist constants rloc > 0, µloc > 0, Cloc < ∞, independent of d, such that for all sufficiently large d, on the favorable initialization event of (A4), there exists a time tloc ≤ Cloc log d for which Rt2loc ≤

2 rloc . 8

(76)

Moreover, for every t ≥ tloc , Rt2 ≤ e−2µloc (t−tloc ) Rt2loc . Consequently, 2 rloc e−2µloc (t−tloc ) . 16 In particular, there exists ε0 > 0, independent of d, such that for every ε ∈ (0, ε0 ),

(1 − χt ) + (1 − ρt ) ≤

min{χt , ρt } ≥ 1 − ε whenever t ≥ C log d + C log(1/ε), where C < ∞ is independent of d and ε. 44

(77)

Proof. By Lemma C.14, there exist constants rloc > 0 and µloc > 0, independent of d, such that for all sufficiently large d, HessM L(y) ⪰ µloc I

for every y ∈ Urloc := {y ∈ M : distM (y, y ∗ ) ≤ rloc }.

After shrinking rloc , if necessary, we may assume rloc <

π . 2

The same Hessian lower bound remains valid on the smaller neighborhood. Apply Lemma C.9 with this value of rloc . On the favorable event of Lemma C.8, it gives a time tloc satisfying tloc ≤ tesc + Cin , where Cin < ∞ is independent of d, and ρtloc ≥ cos(rloc /4),

χtloc ≥ cos(rloc /4).

Equivalently Rt2loc = distM (ytloc , y ∗ )2 ≤

2 rloc , 8

which proves (76). Since Lemma C.8 gives tesc ≤ Cesc log d, we have, for d ≥ e, tloc ≤ Cesc log d + Cin ≤ (Cesc + Cin ) log d. Thus the claimed bound holds with Cloc := Cesc + Cin . Next, observe that y ∗ is a critical point of L. Indeed, at the planted pair the residual vanishes pointwise, so L(y ∗ ) = 0. Since L ≥ 0, y ∗ is a global minimizer of L, and therefore ∇M L(y ∗ ) = 0. The entry estimate gives

rloc rloc . Rtloc ≤ √ < 2 2 2 Hence we may apply Lemma C.15 with r = rloc ,

µ = µloc ,

t0 = tloc .

It follows that the trajectory remains in Urloc for all t ≥ tloc , and that Rt2 ≤ e−2µloc (t−tloc ) Rt2loc

for every t ≥ tloc .

We now convert the product-manifold distance decay into overlap decay. Since χt = ⟨θt , θ∗ ⟩ = cos (distSd−1 (θt , θ∗ )) , and ρt = ⟨βt , β ∗ ⟩ = cos (distSd−1 (βt , β ∗ )) , the elementary bound 1 − 2 cos x ≤ x2 (x ≥ 0) gives 1 1 − χt ≤ distSd−1 (θt , θ∗ )2 , 2 and 1 1 − ρt ≤ distSd−1 (βt , β ∗ )2 . 2 Summing these two estimates and using the product metric, we have (1 − χt ) + (1 − ρt ) ≤

1 1 distM (yt , y ∗ )2 = Rt2 . 2 2 45

Therefore, 1 −2µloc (t−tloc ) 2 e Rtloc . 2

(1 − χt ) + (1 − ρt ) ≤ Using 76 that we proved, we obtain

(1 − χt ) + (1 − ρt ) ≤

2 rloc e−2µloc (t−tloc ) . 16

Finally, set 2 rloc . 16

ε0 := Let ε ∈ (0, ε0 ). If

2 rloc e−2µloc (t−tloc ) ≤ ε, 16

then (1 − χt ) + (1 − ρt ) ≤ ε. Since both summands are nonnegative, this implies 1 − χt ≤ ε,

1 − ρt ≤ ε,

and hence we arrive at (77): min{χt , ρt } ≥ 1 − ε. The displayed exponential condition is equivalent to t − tloc ≥

r2 1 log loc . 2µloc 16ε

Using tloc ≤ Cloc log d, we conclude that min{χt , ρt } ≥ 1 − ε whenever t ≥ Cloc log d +

1 r2 log loc . 2µloc 16ε

Since rloc and µloc are fixed constants independent of d and ε, this is of the form t ≥ C log d + C log(1/ε), for a constant C < ∞ independent of d and ε. This proves the theorem.

D

Proof of Discrete-Time Online SGD Dynamics

We now move to analyze the learning dynamics under SGD discretization. D.1

Weak Recovery Setups

Assumption D.1 (SGD regime and initialization for weak recovery). Fix 0 < ε ≤ ε∗ . Consider online blockwise spherical SGD on Sd−1 × Sd−1 follows (3). The horizon and step-size follows T δ ≤ CT log d,

δ≤

c0 ε , d log d

0 < ε ≤ ε∗ .

for an absolute constant c0 > 0, so we have the weak recovery horizon T := ⌈CT δ −1 log d⌉ ∈ N. Initialization condition (A4) also follows. We start the discretizations analysis by deriving the one-step evolution of the overlaps under SGD: 46

Lemma D.2 (One-step evolution of the overlaps under SGD). Under Assumption D.1, on the event (T )

Et

(β)

(θ)

:= {δ∥gt ∥ ≤ 1/2, δ∥gt ∥ ≤ 1/2}, (ρ)

(χ)

(ρ)

+ δMt

there exist Ft -measurable predictable drifts Ft , Ft (ρ) (χ) remainders Rt , Rt such that for all t,

(χ)

(79)

(χ) (χ) (χ) χt+1 = χt + δFt + δMt + Rt

(80)

(·)

(ρ)

where E[Mt | Ft ] = 0 for each component, Rt (β)

(ρ)

, martingale increments Mt , Mt , and (ρ)

ρt+1 = ρt + δFt

Proof. Let gt

(78)

(θ)

:= ∇Sβ l(βt , θt ; xt ) and gt

(ρ)

+ Rt

(β)

(χ)

= O(δ 3 ∥gt ∥3 ), and Rt

(θ)

= O(δ 3 ∥gt ∥3 ).

:= ∇Sθ l(βt , θt ; xt ) be the spherical gradients taken with (β)

respect β and θ respectively, which implies gt the normalized update as

(θ)

⊥ βt and gt

⊥ θt .From (3) and under Et , we have

(β)

δ 2 (β) (β) (β) = βt − δgt − ∥gt ∥2 βt + O(δ 3 ∥gt ∥3 ) (β) 2 ∥βt − δgt ∥ (θ) δ 2 (θ) 2 θt − δgt (θ) (θ) θt+1 = − = θ − δg ∥gt ∥ θt + O(δ 3 ∥gt ∥3 ) t t (θ) 2 ∥θt − δgt ∥

βt+1 =

βt − δgt

(81) (82)

Hence taking overlap with the teacher direction β ∗ and θ∗ , respectively, we have δ 2 (β) 2 (β) ∥g ∥ ⟨βt , β ∗ ⟩ + O(δ 3 ∥gt ∥3 ) 2 t δ 2 (θ) (θ) (θ) ⟨θt+1 , θ∗ ⟩ = ⟨θt , θ∗ ⟩ − δ⟨gt , θ∗ ⟩ − ∥gt ∥2 ⟨θt , θ∗ ⟩ + O(δ 3 ∥gt ∥3 ), 2 (β)

⟨βt+1 , β ∗ ⟩ = ⟨βt , β ∗ ⟩ − δ⟨gt , β ∗ ⟩ −

(83) (84)

where ρt := ⟨βt , β ∗ ⟩ and χt := ⟨θt , θ∗ ⟩ Now, we decompose the middle two terms involving spherical gradient as the population drift plus the martingale increment, namely,   δ (β) 2 δ (β) 2 (β) (β) ∗ ∗ ⟨gt , β ⟩ + ∥gt ∥ ρt = E ⟨gt , β ⟩ + ∥gt ∥ ρt |Ft 2 2 {z } | (ρ)

−Ft

 δ (β) 2 (β) ∗ ⟨g , β ⟩ + ∥g ∥ ρ t  t  2 t +  δ (β) (β) 2 ∗ − E[⟨gt , β ⟩ + ∥gt ∥ ρt |Ft ] | {z 2 } 

(85)

(ρ)

−Mt

  δ (θ) 2 δ (θ) 2 (θ) ∗ (θ) ∗ ⟨gt , θ ⟩ + ∥gt ∥ χt = E ⟨gt , θ ⟩ + ∥gt ∥ χt |Ft 2 2 | {z } (χ)

−Ft

  δ (θ) 2 (θ) ∗ ⟨g , θ ⟩ + ∥g ∥ χ t  t  2 t + . δ (θ) (θ) − E[⟨gt , θ∗ ⟩ + ∥gt ∥2 χt |Ft ] | {z 2 }

(86)

(χ)

−Mt

Therefore, we can write (83) and (84) respectively as (ρ)

(ρ)

Defining Rt

(ρ)

(β)

+ O(δ 3 ∥gt ∥3 )

(87)

(χ) (χ) (θ) χt+1 = χt + δFt + δMt + O(δ 3 ∥gt ∥3 )

(88)

ρt+1 = ρt + δFt

(β)

(χ)

:= O(δ 3 ∥gt ∥3 ) and Rt

+ δMt

(θ)

:= O(δ 3 ∥gt ∥3 ) gives the claimed expression. 47

Next we provide a class of concentration inequalities to better control the iteration of (ρt , χt ) that has been used in studying the aggregate SGD dynamics, on top of (79) and (80). In particular, Lemma D.5 controls the aggregate martingale term, and Lemma D.7 controls the aggregate discretization remainder term. Besides, As the one-step expansion in Lemma D.2 relies on the event Et , we provide Lemma D.3 to show that Et is true with high probability. All of these inequalities are obtained via using the technical bounds specified in Lemma E.7 in E.2. D.1.1

Taylor Expansion Error Control

Lemma D.3 (Uniform Taylor event on the Weak-Recovery Horizon). Assume the moment bounds of Lemma E.7 and following Assumption D.1 suppose c0 ε T δ ≤ CT log d, δ≤ , 0 < ε ≤ ε∗ . d log d Then there is a constant KTay , depending only on model constants, CT , c0 , ε∗ , such that for all (T ) sufficiently large d, the Taylor Event ET satisfies: (T )

P((ET )c ) ≤ 32KT δ 4 d2 ≤ KTay

ε3 ≤ KTay ε. d log2 d

(89)

(T )

Moreover, on ET , the one-step expansions (79) and (80) in Lemma D.2 are valid simultaneously for every t < T , with remainders (ρ)

(β)

(χ)

|Rt | ≤ Cδ 3 ∥gt ∥3 ,

(θ)

|Rt | ≤ Cδ 3 ∥gt ∥3 .

Proof. By Lemma E.7, conditionally on Ft , we have h i h i (β) (θ) E ∥gt ∥4 | Ft ≤ Kd2 , E ∥gt ∥4 | Ft ≤ Kd2 .   h i (β) (β) Hence, by Markov’s inequality, P δ∥gt ∥ > 21 | Ft ≤ 16δ 4 E ∥gt ∥4 | Ft ≤ 16Kδ 4 d2 . The same bound holds for the θ-block. Therefore (T ) c

P(Et

) ≤ 32Kδ 4 d2 .

Union bounding over t = 0, . . . , T − 1, (T ) c

P(ET Since T δ ≤ CT log d,

) ≤ 32KT δ 4 d2 .

T δ 4 d2 = (T δ)δ 3 d2 ≤ CT δ 3 d2 log d.

0ε Using δ ≤ d clog d , we get

δ 3 d2 log d ≤

c30 ε3 . d log2 d

Thus

ε3 ≤ KTay ε d log2 d for all sufficiently large d. This proves the lemma. (T ) c

P(ET

D.1.2

) ≤ KTay

Martingale and Discretization Error Control for Zt

The Martingale Error for Zt depends on this intermediate Corollary, derived from Lemma E.7. (ρ)

(χ)

Corollary D.4 (Negligible remainder for drift). With the martingale increments Mt and Mt we defined in Lemma D.2 and adopting the technical gradient bounds in Lemma E.7, we have h i h i (ρ) (χ) ρ χ E (Mt )2 | Ft ≤ KM , E (Mt )2 | Ft ≤ KM , (90) ρ χ where KM , KM absolute for large d under step size δ ≤ c0 ε/d log d.

48

(ρ)

(β)

(β)

(ρ)

Proof. Let At = ⟨gt , β ∗ ⟩ + 2δ ∥gt ∥2 ρt , then the martingale term writes as Mt (ρ) (ρ) E[At | Ft ]), so the conditional second moment of Mt can be written as  (ρ)
2  (ρ) (ρ) E[(Mt )2 | Ft ] = E At − E[At | Ft ] | Ft

(ρ)

= −(At

−

(91)

Since we have: (ρ)

At

(ρ)

− E[At

(β)

where we defined ∆gt be written as:

δ (β) (β) | Ft ] = ⟨gt , β ∗ ⟩ + ∥gt ∥2 ρt 2  (β) ∗  δ (β) − E ⟨gt , β ⟩ + ∥gt ∥2 ρt | Ft 2
(β) (β) ∗ = ⟨gt , β ⟩ − E[⟨gt , β ∗ ⟩ | Ft ]
δ (β) (β) + ρt ∥gt ∥2 − E[∥gt ∥2 | Ft ] by regrouping 2 δ (β) (β) = ⟨∆gt , β ∗ ⟩ + ρt ∆qt , 2 (β)

:= gt

(β)

− E[gt

(β)

| Ft ], and ∆qt

(β)

(β)

:= ∥gt ∥2 − E[∥gt ∥2 | Ft ], (91) can

δ (ρ) (β) (β)
2 | Ft ] (92) E[(Mt )2 | Ft ] = E[ ⟨∆gt , β ∗ ⟩ + ρt ∆qt 2 δ2 (β) (β) = 2E[⟨∆gt , β ∗ ⟩2 | Ft ] + ρ2t E[(∆qt )2 ], given (a + b)2 ≤ 2a2 + 2b2 (93) 2 By Lemma E.7, we can bound the first term of (93) by K. For the second term of (93), note that
2 (β) (β) (β) (∆q)2 = ∥gt ∥2 −E[∥gt ∥2 | Ft ] ≤ ∥gt ∥4 , we can use the forth moment bound in Lemma E.7 2 to obtain the Kd bound, so (δρt )2 Kd2 2 δ2 ≤ 2K + Kd2 . 2

(ρ)

E[(Mt )2 | Ft ] ≤ 2K +

(94) (95)

c2 ε2

2

δ 2 0ε 0 Given our step size δ ≤ d clog d by Assumption D.1, 2 Kd ≤ 2 log d K, enabling us to write the bound (ρ)

ρ E[(Mt )2 | Ft ] ≤ KM ,

(96)

ρ for an absolute KM that depends only on model constants and c0 , ε∗ . The exact argument for χ gives (χ)

χ E[(Mt )2 | Ft ] ≤ KM ,

(97)

χ for an absolute KM .

Lemma D.5 (Martingale inequality for Zt ). Let (Z) Mt :=

t−1 X

(ρ)

(χ)

Mk + M k

,

t ≥ 0.

k=0

Under Assumption D.1 and the bounds of Corollary D.4, for all T ∈ N and all λ > 0,   KT δ 2 (Z) P sup δMt ≥ λ ≤ , λ2 0≤t≤T

(98)

for some constant K > 0 depending only on the model. Pt−1 (ρ) (χ) (Z) Proof. Let Dk := Mk + Mk , so that Mt = k=0 Dk is a martingale with respect to (Ft ). By Doob’s L2 maximal inequality, h i  (Z) 2 (Z) 2  E sup Mt ≤ 4 E MT . (99) 0≤t≤T

49

Moreover, using (a + b)2 ≤ 2a2 + 2b2 and Corollary D.4, (ρ)

(χ)

E[Dk2 | Fk ] ≤ 2E[(Mk )2 | Fk ] + 2E[(Mk )2 | Fk ] ≤ 4K. Since martingale increments are orthogonal in L2 , T −1 T −1 X X  (Z) 2  = E MT E[Dk2 ] ≤ 4K = 4KT. k=0

k=0

Combining with (99) and applying Markov’s inequality yields i   δ2 h (Z) 2 (Z) P sup δMt ≥ λ ≤ 2 E sup Mt λ 0≤t≤T 0≤t≤T 2  δ 16KT δ 2 (Z) 2  ≤ 2 · 4 E MT ≤ . λ λ2 Absorbing the numerical factor into the constant completes the proof. √ −1 0ε Remark D.6. With δ ≤ d clog log d⌉ and picking λ = C0 / d, we have d and T := ⌈CT δ  P

sup 0≤t≤T

(Z)

δMt

C0  KT δ 2 KT δ 2 d KCT δd log d KCT c0 ≥ √ ≤ = ≤ ≤ ε. 2 2 2 λ C0 C0 C02 d

Lemma D.7 (Discretization terms for Zt ). Under Assumption D.1 and Lemma E.7, there exists K > 0 such that for all T ∈ N and all λ > 0,    KT δ 3 d3/2 (T ) (Z) P ET ∩ sup |Rt | ≥ λ ≤ . λ t≤T Consequently, under the assumptions of Lemma D.3,     KT δ 3 d3/2 (Z) (T ) . P sup |Rt | ≥ λ ≤ P (ET )c + λ t≤T

(100)

Proof. Let (Z) (ρ) (χ) rk := Rk + Rk ,

(Z) Rt :=

t−1 X (Z) rk . k=0

(T )

On the Taylor event ET , Lemma D.3 gives, simultaneously for all k < T , (ρ)

(β)

(χ)

|Rk | ≤ Cδ 3 ∥gk ∥3 ,

(θ)

|Rk | ≤ Cδ 3 ∥gk ∥3 .

(T )

Therefore, on ET ,

  (Z) (β) (θ) |rk | ≤ Cδ 3 ∥gk ∥3 + ∥gk ∥3 .

Hence (T ) ET ∩



(



(Z) sup |Rt | ≥ λ t≤T

⊆

Cδ

3

X

(β) (θ) ∥gk ∥3 + ∥gk ∥3



) ≥λ .

k<T

By Markov’s inequality, h i (β) 3 (θ) 3    Cδ 3 P E ∥g ∥ + ∥g ∥ k<T k k (T ) (Z) P ET ∩ sup |Rt | ≥ λ ≤ . λ t≤T By Hölder’s inequality and Lemma E.7, i  h i3/4 h (β) (β) ≤ (Kd2 )3/4 ≤ Kd3/2 . E ∥gk ∥3 | Fk ≤ E ∥gk ∥4 | Fk (θ)

The same bound holds for gk . Therefore i X h (β) (θ) E ∥gk ∥3 + ∥gk ∥3 ≤ KT d3/2 . k<T

50

Thus

   KT δ 3 d3/2 (T ) (Z) P ET ∩ sup |Rt | ≥ λ ≤ . λ t≤T

Finally, 

(Z) sup |Rt | ≥ λ t≤T



(T ) ⊆ (ET )c ∪

   (T ) (Z) ET ∩ sup |Rt | ≥ λ , t≤T

so

  KT δ 3 d3/2  (Z) (T ) . P sup |Rt | ≥ λ ≤ P (ET )c + λ t≤T The second displayed bound follows from Lemma D.3. √ Remark D.8. Taking T δ ≤ c2 log d, δ ≤ c0 ε/(d log d), and λ = C0 / d, Lemma D.7 gives 

KT δ 3 d3/2 KT δ 3 d2 Kc2 2 2 Kc2 c20 ε2 = ≤ δ d log d ≤ . λ C0 C0 C0 log d Together with Lemma D.3,   ε3 (T ) . P (ET )c ≤ KTay d log2 d Hence the discretization contribution is absorbed into the existing O(ε) good-event failure probability. Good Events for Zt Now, we formalize the definition and the high-likelihood of good events for Zt iterate: From the two-layer aggregate martingale and discretization remainder for each time step (Z) (ρ) (χ) (Z) (ρ) (χ) defined as Mt := Mt + Mt , rt := Rt + Rt , we write their time accumulation as X X (Z) (Z) (Z) (Z) (Z) (Z) (Z) Mt := Rt := Nt := δMt + Rt . Mk , rk , k<t

k<t

√ For a threshold λZ = C0 / d, define   (Z,M ) (Z) ET := sup |δMt | ≤ λZ ,

(Z,R) ET :=

t≤T



(Z) sup |Rt | ≤ λZ t≤T

 ,

and (Z)

ET

(Z,M )

:= ET

(Z,R)

∩ ET

.

(101)

(Z) Then on ET , (Z)

sup |Nt

| ≤ 2λZ .

(102)

t≤T

By Lemma D.5 and Lemma D.7, together with Lemma D.3, we have (Z)
(Z,M ) c
(Z,R) c
P (ET )c ≤ P (ET ) + P (ET ) ≤ Kε, for some constant K.

(103)

Good Events for Dt Besides, we keep track of the good events for Dt := ρt − χt : for a ∈ {ρ, χ}, we have the individual-layer time accumulation as: X (a) X (a) (a) (a) (a) (a) (a) Mt := Mk , Rt := Rk , Nt := δMt + Rt . k<t

k<t

(D) (ρ) (χ) , We then track their difference Nt := Nt − Nt : With parameter λρ = λχ := C√ind d

λρ + λχ , we define the good events here as     (ρ,χ) (ρ) (χ) ET := sup |Nt | ≤ 2λρ ∩ sup |Nt | ≤ 2λχ , t≤T

(ρ,χ)

So on ET

λD :=

(104)

t≤T

, (D)

sup |Nt

t≤T

(ρ)

(χ)

| ≤ sup |Nt | + sup |Nt | ≤ 2(λρ + λχ ) = 2λD . t≤T

t≤T

By the individual versions of Lemmas D.5 and D.7, choosing Cind large enough gives (ρ,χ)
P (ET )c ≤ Kε. 51

(105)

D.1.3

Martingale and Discretization Error Control for ηt

Deriving the discrete-time evolution ηt requires a different stragegy compares to that of Zt . To start with, we derive the one-step evolution ηt : (θ)

Lemma D.9 (Evolution of ηt ). Work under Assumption D.1. Let gt := ∇Sθ ℓ(βt , θt ; xt ) so that (θ) (θ) (θ) (θ) (θ) (θ) gt ⊥ θt , and define the mean-field spherical gradient ḡt := E[gt | Ft ] and ∆gt := gt − ḡt . Then for all t, we have ∥θt ∥∞ ≤ ∥θ0 ∥∞ + δ

X

∥ḡk ∥∞ + max δ i

k<t

X

ξk,i +

k<t

δ2 X δ3 X ∥gk ∥22 ∥θk ∥∞ + ∥gk ∥32 . 2 2 k<t

k<t

(106) (θ)

Proof. Since gt

q (θ) (θ) ⊥ θt and ∥θt ∥2 = 1, we have at := 1/∥θt − δgt ∥2 = 1/ 1 + δ 2 ∥gt ∥22 and (θ)

θt+1 = at (θt − δgt ). Define the "telescoping correlation" (θ)

(θ)

rt = θt+1 − (θt − δgt ) = (at − 1)(θt − δgt ), (θ)

so as θt+1 = θt − δgt

(107)

+ rt , under time-cumulation we have X (θ) X θt = θ0 − δ gk + rk . k<t

k<t

(θ)

(θ)

(θ)

(θ)

Now we decompose the gradient gk as the mean-field and martingale noise gk = ḡk + ξk , (θ) where ḡk := E[gk | Fk ] and ξk := gk − ḡk . Thus, taking the l∞ -norm X (θ) X (θ) X (θ) ∥θt ∥∞ ≤ ∥θ0 ∥∞ + δ ∥ḡk ∥∞ + max δ ξk,i + ∥rk ∥∞ (108) i

k<t

k<t

k<t

(θ)

To bound ∥rk ∥∞ , taking l∞ -norm of (107): (θ)

(θ)

∥rt ∥∞ ≤ |at − 1| · ∥θt − δgt ∥∞

(109)

(θ) ≤ |1 − at | · (∥θt ∥∞ + δ∥gt ∥∞ ) √ δ 2 (θ) (θ) ≤ ∥gt ∥22 · (∥θt ∥∞ + δ∥gt ∥∞ ), since for x ≥ 0, 1 + x ≥ 1 + x/2.

2 δ 2 (θ) 2 δ 3 (θ) (θ) (θ) = ∥gt ∥2 ∥θt ∥∞ + ∥gt ∥32 , given ∥gt ∥∞ ≤ ∥gt ∥. 2 2

(110) (111) (112)

So we have the bound ∥θt ∥∞ ≤ ∥θ0 ∥∞ + δ

X k<t

(θ)

∥ḡk ∥∞ + max δ i

X (θ) δ 2 X (θ) 2 δ 3 X (θ) 3 ξk,i + ∥gk ∥2 ∥θk ∥∞ + ∥gk ∥2 . 2 2 k<t

k<t

k<t

(113)

Now, we analyze the aggregation through time: let Θt := maxs≤t ∥θs ∥∞ to be the maximum over the trajectory, we then have for any t ≤ T, X (θ) X (θ)  δ 2 X (θ)  δ 3 X (θ) 3 ξk,i + Θt ≤ ∥θ0 ∥∞ + δ ∥ḡk ∥∞ + max sup δ ∥gk ∥22 Θt + ∥gk ∥2 . i 2 2 s≤t k<t k<s k<t k<t | {z } | {z } | {z } | {z } 1

2

3

4

(114) 52

To bound term 3 in (114), we have by Markov Inequality, (θ) 2   δ2 X  E δ2 P δ 2 T Kd (θ) 2 k<T ∥gk ∥2 2 P ∥gk ∥2 ≥ λ ≤ ≤ 2 λ λ k<T

Taking λ = 1/4, we have P

 δ2 X 2

(θ)

∥gk ∥22 ≥

k<T

1 = O(ε), 4

(115)

giving that w.p. 1 − O(ε), Θt ≤ ∥θ0 ∥∞ + δ

X

(θ)

X (θ) 1 δ 3 X (θ) 3 ξk,i + Θt + ∥gk ∥2 . 4 2 s≤t k<s k<t {z } | {z }

∥ḡk ∥∞ + max sup δ i

k<t

|

{z

|

}

1

2

4

3 X

X (θ) X (θ) δ 1 (θ) ξk,i + (1 − )Θt ≤ ∥θ0 ∥∞ + δ ∥ḡk ∥∞ + max sup δ ∥gk ∥32 . i 4 2 s≤t k<s k<t k<t {z } | {z } | {z } | 1

Θt ≤

4 3

∥θ0 ∥∞ + δ

X

2

 X (θ) δ (θ) ξk,i + ∥gk ∥32 . 2 s≤t k<s k<t {z } | {z }

(θ)

∥ḡk ∥∞ + max sup δ i

k<t

|

4

3 X

{z 1

}

|

2

(116)

4

We start by bounding term 2 using a Freedman Inequality. Yet to show it, we need this corollary below: Corollary D.10 (Coordinate variance and 8th-moment bounds for ξ (θ) ). Work under Assumption D.1 (θ) (θ) (θ) (θ) and Lemma E.7. Let ξt := gt − E[gt | Ft ] and write ξt,i for its i-th coordinate. Then there exist constants σ 2 > 0 and K8 > 0 (depending only on model constants) such that for all t and all i ≤ d, h i h i (θ) (θ) E (ξt,i )2 | Ft ≤ σ 2 , E |ξt,i |8 | Ft ≤ K8 . (117) Proof. Conditional variance. Lemma E.7 gives the directional variance bound   (θ) sup E ((∆gt ) · u)2 | Ft ≤ K. ∥u∥=1

(θ)

Taking u = ei yields E[(ξt,i )2 | Ft ] ≤ K. Set σ 2 := K. 8-th moment. Recall that in Lemma E.7, from the per-sample gradient ∇θ ℓ(βt , θt ; xt ) = At zt with At := 2(ft − ft∗ )ψ ′ (zt · θt ) defined in (215), using the boundedness of ψ, ψ ′ we showed |At | ≤ CA in (219). From there, the spherical gradient is
(θ) gt = (I − θt θt⊤ )∇θ ℓ = At zt − (zt · θt )θt , so for each coordinate

(θ) |gt,i | ≤ |At | |zt,i | + |zt · θt | ≤ CA |zt,i | + |zt · θt | .

(118)

Condition on Ft . Then θt is deterministic with ∥θt ∥2 = 1 and zt = (zt,1 , . . . , zt,d ) has i.i.d. coordinates zt,j = ϕ(γt,j ) with γt,j ∼ N (0, 1). In particular, 2 µ2 := E[zt,1 ] = E[ϕ(γ)2 ] < ∞,

µ8 := E[|zt,1 |8 ] = E[|ϕ(γ)|8 ] < ∞.

Let Xj := θt,j zt,j . Then (Xj )j≤d are independent and mean-zero, and d X j=1

E[Xj2 | Ft ] =

d X

2 2 θt,j E[zt,1 ] = µ2

j=1

d X j=1

53

2 θt,j = µ2 .

By Rosenthal’s inequality for p = 8, there exists a constant C8 > 0 such that d d d i  X h X 4 X  8 Ft ≤ C8 E Xj E[Xj2 | Ft ] + E[|Xj |8 | Ft ] . j=1

j=1

j=1

Moreover, d X j=1

E[|Xj |8 | Ft ] =

d X

|θt,j |8 E[|zt,1 |8 ] = µ8

j=1

d X

|θt,j |8 ≤ µ8

j=1

d X

2 θt,j = µ8 ,

j=1

since |θt,j | ≤ 1 implies |θt,j |8 ≤ |θt,j |2 . Therefore, d i h X 8   8 Ft ≤ C8 (µ42 + µ8 ) =: C, E |zt · θt | | Ft = E Xj j=1 (θ)

which is a constant independent of t and d. Hence, plugging into (118), we have E[|gt,i |8 | Ft ] ≤ K8′ uniformly. Finally, by Jensen and (a − b)8 ≤ 27 (a8 + b8 ), (θ)

(θ)

E[|ξt,i |8 | Ft ] ≤ 28 E[|gt,i |8 | Ft ] ≤ 28 K8′ . Set K8 := 28 K8′ gives the claimed result. Now we move to prove the Freedman bound: Lemma D.11 (Coordinate martingale control for the η-bootstrap). Assume Corollary D.10. Let T δ ≤ CT log d and δ ≤ c0 ε/(d log d). Then there is a constant CM such that ! r X (θ) log(40d/ε) P max sup δ ξk,i > CM ≤ ε/5. (119) i≤d t≤T d k<t

Proof. Let 40d . ε The case T = 0 is trivial, so assume T ≥ 1. Let σ 2 , K8 be the constants from Corollary D.10. Define the truncation level  1/8 40dT K8 B := . ε For k < T and i ≤ d, write (θ) ξk,i := ξk,i . L := log

Define the centered truncated martingale difference   Yk,i := ξk,i 1{|ξk,i |≤B} − E ξk,i 1{|ξk,i |≤B} | Fk . Then E[Yk,i | Fk ] = 0,

|Yk,i | ≤ 2B,

and

2 2 E[Yk,i | Fk ] ≤ E[ξk,i | Fk ] ≤ σ 2 . Hence, for each fixed i, the process X (i) Mt := Yk,i k<t

is a martingale with increments bounded by 2B and predictable quadratic variation at most σ 2 t ≤ σ2 T . By the maximal form of Freedman’s inequality, for every x > 0,     x2 (i) P sup Mt ≥ x ≤ exp − , 2(σ 2 T + 2Bx/3) t≤T 54

(i)

and the same estimate holds for −Mt . Set √ 4 xT := 2σ 2 T L + BL. 3 Then x2T ≥ L. 2 2(σ T + 2BxT /3) √ Indeed, writing a = σ 2 T , u = 2aL, and v = 34 BL, we have xT = u + v, while   2BxT 2L a + = 2aL + v(u + v). 3 Since

x2T = (u + v)2 = 2aL + 2uv + v 2 ,

we get   2BxT x2T − 2L a + = uv ≥ 0. 3 Therefore

 P

(i) sup |Mt | ≥ xT t≤T



≤ 2e−L .

Taking a union bound over i ≤ d, and using L = log(40d/ε), gives   ε (i) . P max sup |Mt | ≥ xT ≤ 2de−L = i≤d t≤T 20 It remains to compare the original martingale to the truncated one. Let   J := max max |ξk,i | ≤ B . k<T i≤d

By Corollary D.10 and Markov’s inequality, XX K8 ε P(J c ) ≤ P(|ξk,i | > B) ≤ dT 8 = . B 40 k<T i≤d

Now define   bk,i := E ξk,i 1{|ξk,i |>B} | Fk . Since E[ξk,i | Fk ] = 0,   E ξk,i 1{|ξk,i |≤B} | Fk = −bk,i . Hence Yk,i = ξk,i 1{|ξk,i |≤B} + bk,i . On J , the indicator 1{|ξk,i |≤B} is identically one for all k < T and i ≤ d. Therefore, on J , ξk,i = Yk,i − bk,i . Moreover, by the conditional eighth-moment bound,   E[|ξk,i |8 | Fk ] K8 |bk,i | ≤ E |ξk,i |1{|ξk,i |>B} | Fk ≤ ≤ 7. B7 B Thus, on J , sup t≤T

X

(i)

ξk,i ≤ sup |Mt | + t≤T

k<t

T K8 . B7

Combining the previous estimates, with probability at least 1 − ε/20 − ε/40, we have max sup i≤d t≤T

X

ξk,i ≤ xT +

k<t

55

T K8 . B7

Multiplying by δ, it remains to show that T K8 δxT + δ 7 ≤ CM B

r

L . d

First, √ δ

r 2σ 2 T L =

p p 2σ 2 L δ(δT ) ≤ 2σ 2 CT c0 ε

L ≤C d

r

L , d

because T δ ≤ CT log d, δ ≤ c0 ε/(d log d), and ε ≤ 1. Next, using T ≤ CT (log d)/δ,  δBL = δL

40dT K8 ε

1/8 ≤ CL δ

7/8



d log d ε

1/8 .

Using again δ ≤ c0 ε/(d log d), this gives 7/8

δBL ≤ Cc0 ε3/4 L d−3/4 (log d)−3/4 . Since L = log(40d/ε) and 0 < ε ≤ 1, ε3/4 L1/2 ≤ C(log d)1/2 for d ≥ 3. Hence 7/8 δBL ≤ Cc0 L1/2 d−3/4 (log d)−1/4 ≤ C

r

L . d

Finally,  7/8 T K8 ε = δT K ≤ CδT 1/8 ε7/8 d−7/8 . 8 B7 40dT K8 Using T ≤ CT (log d)/δ and δ ≤ c0 ε/(d log d), we get r T K8 L 7/8 7/4 −7/4 −3/4 δ 7 ≤ Cc0 ε d (log d) ≤C . B d δ

Therefore, after increasing CM to absorb constants depending only on σ 2 , K8 , CT , c0 , we have r r T K8 L log(40d/ε) δxT + δ 7 ≤ CM = CM . B d d Since ε ε ε + < , 20 40 5 the desired bound follows. We plug this proved Freedman bound into (116) to obtain X (θ) 4 δ 3 X (θ) 3  ∥θ0 ∥∞ + δ ∥ḡk ∥∞ +Λη + ∥gk ∥2 . Θt ≤ 3 2 k<t k<t | | {z } {z } 1

∗

(120)

4

∗

Since ηt = ∥θt ∥∞ + ∥θ ∥∞ ≤ Θt + ∥θ ∥∞ , we re-write it as: X (θ) 4 δ 3 X (θ) 3  ηt ≤ η0 + δ ∥ḡk ∥∞ +Λη + ∥gk ∥2 . 3 2 k<t k<t | | {z } {z } 1

3

Next we bound 4: notice that since δ2

4

(θ) 3 k<t ∥gk ∥2 is non decreasing in t, we have

P

δ 3 X (θ) 3 δ 3 X (θ) 3 ∥gk ∥2 = ∥gk ∥2 . 2 t≤T 2 sup

k<t

k<T

56

(121)

On the other hand,  (θ)   (θ) 3/4 E ∥gk ∥32 ≤ E ∥gk ∥42 by Holder inequality 2 3/4

≤ (Kd )

(122)

by the technical bounds in Lemma E.7

(123)

= K 3/4 d3/2 .

(124)

Therefore, for a given λcub ,   δ3 X   δ 3 X (θ) 3 (θ) ∥gk ∥2 ≥ λcub = P ∥gk ∥32 ≥ λcub by monotonicity P sup 2 t≤T 2 k<t k<T  δ3 P (θ) 3  E 2 k<t ∥gk ∥2 ≤ , by Markov Inequality λcub

(125)

(126)

3

K 3/4 d3/2 T δ2 , following (124) ≤ λcub

(127)

√ Taking λcub = Ccub / d, and plugging in T ≤ c2 (log d)/δ and δ ≤ c0 ε/(d log d), we have ! δ 3 X (θ) 3 c2 , c20 , K 3/4 Ccub ε2 P sup ∥gk ∥2 ≥ √ . ≤ · 2Ccub log d d t≤T 2

(128)

k<t

2

ε Choose Ccub large enough (absorbing constants) and take d sufficiently large so that log d ≤ ε/10, we obtain ! δ 3 X (θ) 3 ε Ccub P sup ≤ ∥gk ∥2 ≥ √ (129) 2 10 d t≤T k<t

Plugging it in the η bound in (116) it becomes  X (θ) Ccub  ηt ≤ Cη η0 + δ ∥ḡk ∥∞ +Λη + √ , d k<t {z } |

(130)

1

with high probability. The bound of 1 is contingent to dynamics’s phase that we will derive it later. We later bound them in Step 1 of the Phase I proof (Lemma D.13) and Step 1(a) of the Phase II proof (Lemma D.16). Good event for η iteration

Define (η)

(η,M )

ET := ET

(η,2)

∩ ET

(η,3)

∩ ET

,

(131)

where )

( (η,M ) ET :=

X (θ) max sup δ ξk,i ≤ Λη i≤d t≤T

( (η,2) ET :=

,

(132)

k<t

1 δ 2 X (θ) 2 ∥gk ∥2 ≤ 2 4

) ,

(133)

k<T

( (η,3) ET :=

δ 3 X (θ) 3 Ccub sup ∥gk ∥2 ≤ √ d t≤T 2

) .

(134)

k<t

Therefore, with Bound (119) from Lemma E.7 to ensure (132), Bound (115) to ensure (133), and Bound (129) to ensure (134), we obtain a union bound, (η)
(η,M ) c
(η,2)
(η,3)
P (ET )c ≤ P (ET ) + P (ET )c + P (ET )c ≤ Kε, (135) after increasing K. 57

D.1.4

Martingale and Discretization Error Control for Ωt and ωt

As a discrete analogy of Lemma C.7, we next provide the control of third moment quantities
3 Pd Pd Ωt := ∥θt ∥33 = i=1 |θt,i |3 , and ωt := i=1 |θt,i | + |θi∗ | , upon their martingale and discretizatio noises. Lemma D.12 (Third-moment control for SGD through Phase I). Let Ωt and ωt follow the definition made in Lemma C.7. Fix T ∈ N with T δ ≤ CT log d. Under Assumption D.1, Lemma E.7, and (Ω) Corollary D.10, there is an event ET satisfying (Ω)
P (ET )c ≤ KΩ ε (Ω)

such that the following holds. Work on E0 ∩ ET . Let τA := inf{t : Zt2 ≥ Aηt }, τB := inf{t : ρt ≥ 1/2 or χt ≥ 1/2}, τ+ := inf{t : ρt < 0 or χt < 0}, and set σ := τA ∧ τB ∧ τ+ ∧ T. Suppose that, for some constant Cη > 0, √ log d + 1 + tδ √ ηt ≤ Cη d

(0 ≤ t ≤ σ).

(136)

Then, for all sufficiently large d, 1 + tδ Ωt ≤ CΩ √ , d

1 + tδ ωt ≤ Cω √ , d

0 ≤ t ≤ σ.

(137)

The constants KΩ , CΩ , Cω depend only on model constants, A, Cη , CT , and the step-size regime. (θ)

Proof. Write gt := gt , ḡt := E[gt | Ft ], and ξt := gt − ḡt . Since gt ⊥ θt , the normalized spherical update is θt+1 = at (θt − δgt ), at = (1 + δ 2 ∥gt ∥22 )−1/2 ≤ 1. Hence

Ωt+1 = ∥θt+1 ∥33 ≤ ∥θt − δgt ∥33 .

We use the elementary inequality
|x − y|3 ≤ |x|3 − 3|x|x y + C |x|y 2 + |y|3 ,

x, y ∈ R.

d Applying it coordinatewise with x = θt,i and y = δgt,i , and letting vt := |θt,i |θt,i i=1 , gives Ωt+1 ≤ Ωt − 3δ⟨vt , ḡt ⟩ − 3δ⟨vt , ξt ⟩ + Cδ 2

d X

2 |θt,i |gt,i + Cδ 3 ∥gt ∥33 .

(138)

(Ω)

(139)

i=1

Since ∥vt ∥1 =

2 i |θt,i | = 1,

P

−3δ⟨vt , ḡt ⟩ ≤ 3δ∥ḡt ∥∞ .

Summing (138) gives Ωt ≤ Ω0 + 3δ

X

(Ω)

∥ḡk ∥∞ + Mt

+ CQt

(Ω)

+ CCt

,

k<t

where

(Ω)

Mt

:= −3δ

X

⟨vk , ξk ⟩,

k<t (Ω)

Qt

:= δ 2

d XX

2 |θk,i |gk,i ,

k<t i=1

(Ω)

Ct

:= δ 3

X k<t

58

∥gk ∥33 .

(Ω)

We first define the good event ET . Let CΩ,M λΩ,M := √ , d

CΩ,Q λΩ,Q := √ , d

CΩ,C λΩ,C := √ , d

where the constants will be chosen large. Set       (Ω) (Ω) (Ω) (Ω) ET := sup |Mt | ≤ λΩ,M ∩ sup Qt ≤ λΩ,Q ∩ sup Ct ≤ λΩ,C . t≤T

t≤T

(140)

t≤T

(Ω)

We verify that this event has high probability. The process Mt is a martingale. Moreover, X X ∥vk ∥22 = |θk,i |4 ≤ |θk,i |2 = 1. i

i

By Lemma E.7,   E ⟨vk , ξk ⟩2 | Fk ≤ K. Thus Doob’s L2 inequality gives   KT δ 2 d KCT c0 KT δ 2 (Ω) = ≤ ε. P sup |Mt | ≥ λΩ,M ≤ 2 2 2 λΩ,M CΩ,M CΩ,M t≤T Choosing CΩ,M large makes this contribution O(ε). Next, by the coordinate moment bound proved in Corollary D.10 and (118), we have 2 E[gk,i | Fk ] ≤ K

Therefore

" E

uniformly in i, k.

# X

2 |θk,i |gk,i

Fk ≤ K

i

Hence

X

√ |θk,i | ≤ K d.

i

√ (Ω) E[QT ] ≤ KT δ 2 d.

By Markov’s inequality, √   KT δ 2 d KT δ 2 d KCT c0 (Ω) √ = P sup Qt ≥ λΩ,Q ≤ ≤ ε. C CΩ,Q CΩ,Q / d Ω,Q t≤T Choosing CΩ,Q large makes this contribution O(ε). Similarly, the same coordinate moment bounds imply E[|gk,i |3 | Fk ] ≤ K, Thus

E[∥gk ∥33 | Fk ] ≤ Kd.

(Ω)

E[CT ] ≤ KT δ 3 d. Again by Markov’s inequality,   KT δ 3 d3/2 KT δ 3 d (Ω) √ = . P sup Ct ≥ λΩ,C ≤ CΩ,C CΩ,C / d t≤T Since T δ ≤ CT log d and δ ≤ c0 ε/(d log d), the last display is o(ε) for d large. Therefore, after increasing the constants if necessary, (Ω)
P (ET )c ≤ KΩ ε. (141) (Ω)

It remains to prove the pathwise bound on Ωt . Work on E0 ∩ ET . From (139) and (140), Ωt ≤ Ω0 + 3δ

C ∥ḡk ∥∞ + √ , d k<t

X

59

t ≤ T.

(142)

We now bound the mean-field term for t ≤ σ. Since k < σ implies k < τ+ , we have ρk , χk ≥ 0. Since k < σ also implies k < τA , we have Zk2 < Aηk . By (48) in Lemma C.4, ∥ḡk ∥∞ ≤ CZk ∥θ∗ ∥∞ + CZk2 ηk + Cηk2 . On E0 , ∥θ∗ ∥∞ ≤ C

p

log d/d. Using Zk2 < Aηk and the assumed envelope for ηk , we obtain s√ r p log d + 1 + kδ log d ∗ ∗ √ . Zk ∥θ ∥∞ ≤ Aηk ∥θ ∥∞ ≤ C d d

Since kδ ≤ T δ ≤ CT log d, the right-hand side is bounded by C(log d)d−3/4 ≤ Cd−1/2 for all sufficiently large d. Also, Zk2 ηk ≤ Aηk2 . Hence, for k < σ, C (143) ∥ḡk ∥∞ ≤ √ + Cηk2 . d Therefore, for t ≤ σ, X X tδ δ ∥ḡk ∥∞ ≤ C √ + Cδ ηk2 . d k<t k<t Using the assumed η-envelope, X
2 C X p C(log d)3 ηk2 ≤ δ log d + 1 + kδ ≤ = o(d−1/2 ). δ d d k<t k<t √ Plugging this into (142), and using Ω0 ≤ C/ d on E0 , gives 1 + tδ Ωt ≤ C √ , 0 ≤ t ≤ σ. d This proves the desired Ωt -bound. Finally, ωt =

X

|θt,i | + |θi∗ |

3

≤4

√i On E0 , ∥θ∗ ∥33 ≤ C/ d. Therefore

X

|θt,i |3 + 4

i

X

|θi∗ |3 = 4Ωt + 4∥θ∗ ∥33 .

i

1 + tδ ωt ≤ Cω √ , d

0 ≤ t ≤ σ.

This completes the proof. D.2

Weak Recovery Dynamics

Good Events for Weak Recovery Now, we define the good event for weak recovery by taking (T ) (Z) (ρ,χ) (η) the intersection of Edeloc from B.2 ET from (78), ET from (101), ET from (104), ET from (Ω) (131), and ET from (140): (T )

(Z)

(ρ,χ)

(η)

(Ω)

GTweak := E0 ∩ ET ∩ ET ∩ ET ∩ ET ∩ ET . (144) With their respective probabilistic controls made in (21), (89), (103), (105), (135), (141), they are true with probability (1 − Kalg ε)(1 − d−Kini .) D.2.1

Phase I Dynamics, Weak Recovery

Lemma D.13 (Phase I: pull-up to the signal-dominant scale). Work on GTweak . Then, up to τA ∧ τB , √ log d + 1 + tδ 1 + tδ √ ηt ≤ Cη , ωt ≤ Cω √ . (145) d d Moreover, τA < τB , τ+ > τA , τA δ ≤ C log d, ZτA ≥ cA d−1/4 . (146) Equivalently, ρt , χt > 0 for all 0 ≤ t ≤ τA . (147) Finally, p ZτA ≤ CA log d d−1/4 + o(1). (148) 60

Proof. Work throughout on GTweak . We use the favorable initialization event in the form r z− log d √ , η0 ≤ C ρ0 , χ0 ≥ , d d where z− > 0 is fixed large enough below. Recall that τ+ := inf{t : ρt < 0 or χt < 0}. Set σ := τA ∧ τB ∧ τ+ ∧ T. It is enough to prove the desired estimates up to σ, and then show that σ = τA . Step 1: Phase-I η-envelope.

Let √

Θt := max ηs , 0≤s≤t

HI (t) := KI

log d + 1 + tδ √ , d (η)

where KI will be chosen large. By the discrete η-estimate (130), on ET , ! X (θ) Ccub ηt ≤ C η0 + δ ∥ḡk ∥∞ + Λη + √ , t ≤ T. d k<t (η)

The definition of ET gives r Λη ≤ C

log(40d/ε) ≤C d

√

log d + 1 + tδ √ , d

t ≤ T,

For k < σ, we have k < τ+ , hence ρk , χk ≥ 0, and also k < τA , hence Zk2 < Aηk . The population ℓ∞ -gradient bound (48) gives (θ)

∥ḡk ∥∞ ≤ CZk ∥θ∗ ∥∞ + CZk2 ηk + Cηk2 . Using Zk2 < Aηk , this becomes √ (θ) ∥ḡk ∥∞ ≤ C ηk ∥θ∗ ∥∞ + Cηk2 . p Since ∥θ∗ ∥∞ ≤ C log d/d, and since throughout the bootstrap kδ ≤ T δ ≤ CT log d, the implication ηk ≤ HI (k) gives log d C √ ηk ∥θ∗ ∥∞ ≤ C 3/4 ≤ √ d d for all sufficiently large d. Thus, under the bootstrap, C (θ) ∥ḡk ∥∞ ≤ √ + Cηk2 . d Let τH := inf{t ≤ σ : Θt > HI (t)}. Suppose for contradiction that τH ≤ σ. Then for every k < τH , ηk ≤ Θk ≤ HI (k), and therefore X X (θ) τH δ δ ∥ḡk ∥∞ ≤ C √ + Cδ HI (k)2 . d k<τH k<τH Moreover, √  X
2 CKI2 X p C(log d)3 log d + 1 + τH δ √ δ HI (k)2 ≤ δ log d + 1 + kδ ≤ =o . d d d k<τH

k<τH

Combining the previous estimates with the bound on ηt , we get √ √  log d + 1 + τH δ log d + 1 + τH δ √ √ Θτ H ≤ C +o . d d 61

Choosing KI large enough, and then d large enough, gives 1 HI (τH ), 2 contradicting the definition of τH . Hence τH > σ, and therefore √ log d + 1 + tδ √ ηt ≤ Cη , 0 ≤ t ≤ σ, d ΘτH ≤

(149)

which is exactly the assumption (136) made in Lemma D.12. Alongside the condition Zt2 < Aηt for t < τA , we can apply Lemma D.12 to obtain 1 + tδ ωt ≤ Cω √ , d

0 ≤ t ≤ σ.

(150)

Step 2: Positivity and scalar Z-pull-up. For t < σ, we have t < τA , hence Zt2 < Aηt . By Step 1, √ log d log d + 1 + tδ √ ηt ≤ Cη ≤C √ . d d Thus

p Zt ≤ C log d d−1/4 = o(1). For d large, Zt ≤ csm , so we can apply Lemma C.6 ρ̇t ≥ aχt − bωt ,

χ̇t ≥ aρt − bωt . (T )

Combining this with Lemma D.2, the Taylor event ET , and Lemma E.7, we obtain   (u) ut+1 ≥ (1 − Cδ 2 d)I + aδJ ut − Cδωt 1 + εt , ut := (ρt , χt )⊤ , where  J=

0 1

 1 , 0

(u)

1 = (1, 1)⊤ ,

Let

εt

(ρ)

(ρ)

(χ)

(χ)
⊤

= Nt+1 − Nt , Nt+1 − Nt

.

B := (1 − Cδ 2 d)I + aδJ.

For d large, B has nonnegative entries. Its eigenvalues in the 1- and (1, −1)⊤ -directions are γ+ := 1 + aδ − Cδ 2 d,

γ− := 1 − aδ − Cδ 2 d.

After decreasing a, we may assume γ+ ≥ 1 + βδ,

0 < γ− < 1,

β > 0.

Unrolling the previous cooperative recursion and applying the summation-by-parts bounds to the Z-noise and D-noise gives, componentwise, " # X −(k+1) z− t ρt , χt ≥ γ+ √ − C γ+ δωk − C(λZ + λρ + λχ ) , 0 ≤ t ≤ σ. (151) d k<t Indeed, the Z-noise is controlled by Corollary E.9, while the contracting D-noise is controlled by Corollary E.10. By (150),   X X C C −(k+1) −(k+1) −(k+1)  γ+ δωk ≤ √ δ γ+ + δ2 kγ+ ≤√ . d d k<t k≥0 k≥0

X Since

λZ + λρ + λχ = O(d−1/2 ), choosing z− larger than the constants in the bracket yields 0 ≤ t ≤ σ.

ρt , χt > 0, 62

(152)

In particular, σ ̸= τ+ . Now sum the two one-step inequalities. For (Z)

εt

(Z)

(Z)

:= Nt+1 − Nt

,

we get (Z)

Zt+1 ≥ (1 + βδ)Zt − Cδωt + εt , Set γ := 1 + βδ. Multiplying by γ

t < σ.

−(t+1)

and summing, X X (Z) γ −t Zt ≥ Z0 − C γ −(k+1) δωk + γ −(k+1) εk . k<t

k<t (Z)

The weighted ω-sum is O(d−1/2 ) by (150). Also, on ET , sup |Nu(Z) | ≤ 2λZ ,

u≤T

so Corollary E.9 gives X

(Z)

γ −(k+1) εk

≤ 4λZ .

k<t

Since Z0 ≥ 2z− d−1/2 , increasing z− if necessary gives a constant c− > 0 such that c− Zt ≥ √ γ t , 0 ≤ t ≤ σ. d

(153)

Step 3: Some stopping event occurs by O(δ −1 log d). Let   t0 := c6 δ −1 log d , where c6 will be chosen large. Choose the weak-recovery horizon constant CT so that t0 < T . Suppose, for contradiction, that σ > t0 . Then t0 < τA , hence Zt20 < Aηt0 . Using (153) and (149), c2− 2t0 γ ≤ Zt20 < Aηt0 ≤ C d

√

log d + 1 + t0 δ log d √ ≤C √ . d d

Since log γ ≥ cδ for some c > 0, we have γ 2t0 ≥ d2cc6 +o(1) . Choosing c6 so large that 2cc6 > 1/2 gives a contradiction for all large d. Therefore σ ≤ t0 ,

σδ ≤ c6 log d + O(δ).

(154)

Step 4: The first stopping event is τA . Since t0 < T , (154) implies σ ̸= T . By (152), σ ̸= τ+ . Thus σ is either τA or τB . Suppose for contradiction that σ = τB ≤ τA . Since τB ≤ t0 , at time τB − 1 we have Zτ2B −1 < AητB −1 . Thus, by (149),

√ ZτB −1 ≤ C

log d = o(1). d1/4

The one-step expansion gives (Z)

(Z)

ZτB = ZτB −1 + δFτB −1 + ετB −1 . By Lemma E.7, (Z)

|Ft

√ | ≤ C d + Cδd, 63

and hence (Z)

δ|Ft

√ | ≤ Cδ d + Cδ 2 d = o(1).

Also, (Z)

(Z)

|ετB −1 | ≤ |Nτ(Z) | + |NτB −1 | ≤ 4λZ = o(1). B Therefore ZτB = o(1). But (152) gives ρτB , χτB > 0, and by definition of τB , one of ρτB , χτB is at least 1/2. Hence ZτB ≥ 1/2, a contradiction. Therefore σ = τA , The bounds

τA δ ≤ C log d.

τA < τB ,

√

1 + tδ log d + 1 + tδ √ , ωt ≤ Cω √ d d now hold for all t ≤ τA , and since τA < τB , also up to τA ∧ τB . ηt ≤ Cη

Step 5: Size of ZτA .

By definition of τA , Zτ2A ≥ AητA .

Since θτA ∈ Sd−1 , ∥θτA ∥∞ ≥ d−1/2 , and therefore ητA ≥ d−1/2 . Hence Z τA ≥ √ This proves the lower bound with cA := A.

√

A d−1/4 .

For the upper bound, first note that τA ≥ 1 for large d, since Z02 = O(d−1 ) while η0 ≥ d−1/2 . Hence τA − 1 < τA , and so Zτ2A −1 < AητA −1 . Using the η-bound and τA δ ≤ C log d, √ ZτA −1 ≤ C

log d . d1/4

The one-step expansion gives (Z)

(Z)

ZτA ≤ ZτA −1 + δ|FτA −1 | + |ετA −1 |. As above,

√ (Z) δ|FτA −1 | ≤ Cδ d + Cδ 2 d = o(1),

(Z)

|ετA −1 | ≤ 4λZ = o(1).

Therefore ZτA ≤ CA

p log d d−1/4 + o(1).

This completes the proof. Lemma D.14 (Entry balance at τA ). Work on GTweak , and assume the conclusions of Lemma D.13. Let Dt := ρt − χt , λD := λρ + λχ . Then |DτA | ≤ C

√ (log d)2 log d √ + CλD + Cδd 1/4 . d d

In particular, |DτA | = o(ZτA ). Hence, for all sufficiently large d, |DτA | ≤ 64

1 Zτ . 8 A

(155)

Proof. By Lemma D.13, on 0 ≤ t ≤ τA we have τA < τB ,

τA δ ≤ C log d,

τ+ > τA ,

and for t ≤ τA ,

√ ηt ≤ Cη

Moreover, for t < τA ,

log d + 1 + tδ √ . d

Zt2 < Aηt .

Therefore, for t < τA ,

√

log d = o(1). d1/4 Thus, throughout the interval 0 ≤ t < τA , the overlaps are nonnegative and lie in the small-overlap regime. Hence, by (52) from Lemma C.6, Zt ≤ C

|Et | ≤ C(Zt2 + ωt ).

e + ρt χt )Dt + Et , Ḋt = −C(1 Now use the one-step SGD expansion. Let (β)

(θ)

at := E[∥gt ∥2 | Ft ],

bt := E[∥gt ∥2 | Ft ].

Subtracting the one-step overlap expansions gives Dt+1 = Dt + δ(ρ̇t − χ̇t ) − where

(D)

εt

(D)

(D)

δ2 (D) (at ρt − bt χt ) + εt , 2 (D)

(ρ)

(χ)

:= Nt+1 − Nt

,

at ρt − bt χt =

at + bt at − bt Dt + Zt . 2 2

Nt

:= Nt

− Nt .

Decompose

Therefore

(D)

Dt+1 = (1 − αt δ)Dt + δet + εt where

,

(1)

e + ρt χt ) + δ (at + bt ), αt := C(1 4

and

δ et := Et − (at − bt )Zt . 4

By Lemma E.7, at , bt ≤ Kd. Hence |et | ≤ C(Zt2 + ηt ) + Cδd Zt .

(2)

Also, since ρt , χt ≥ 0, e αt ≥ C. After reducing the step-size constant, we may assume 0 ≤ 1 − αt δ ≤ 1 − α0 δ < 1 for a fixed α0 > 0. Unroll (1) up to t = τA . Define t−1 Y

Pm,t :=

(1 − αℓ δ),

Pt,t := 1.

ℓ=m

Then DτA = P0,τA D0 +

τX A −1

Pk+1,τA δek +

k=0

τX A −1 k=0

Since 0 ≤ Pk+1,τA ≤ 1, |P0,τA D0 | ≤ |D0 | ≤ Cd−1/2 , 65

(D)

Pk+1,τA εk .

by the favorable initialization event. For the first forcing term, using Zk2 ≤ Aηk for k < τA , X X ηk . (Zk2 + ηk ) ≤ Cδ δ k<τA

k<τA

By the Phase-I η-bound and τA δ ≤ C log d, X
C X p (log d)2 ηk ≤ √ δ δ log d + 1 + kδ ≤ C √ . d k<τA d k<τA

(3)

For the discretization forcing term, since Pk+1,τA ≤ (1 − α0 δ)τA −1−k , we have δ

X



 X (1 − α0 δ)j ≤ C sup Zk . sup Zk δ

Pk+1,τA Zk ≤

k<τA

k<τA

k<τA

j≥0

As observed above,

√ sup Zk ≤ C k<τA

log d . d1/4

Therefore

√ δ

X

Pk+1,τA δd Zk ≤ Cδd

k<τA

log d . d1/4

(4)

It remains to control the noise term. For fixed τA , the weights 0 ≤ k ≤ τA − 1,

Pk+1,τA ,

(ρ,χ)

lie in [0, 1] and are nondecreasing in k. On ET

,

sup |Nu(D) | ≤ 2(λρ + λχ ) = 2λD .

u≤T

Thus Lemma E.8 gives τX A −1

(D)

Pk+1,τA εk

≤ 4λD .

k=0

Combining (2)–(5), and absorbing |D0 | ≤ Cd−1/2 into the first term, gives √ log d (log d)2 |DτA | ≤ C √ + CλD + Cδd 1/4 . d d Finally, Lemma D.13 gives ZτA ≥ cA d−1/4 . Hence

p |DτA | (log d)2 λD ≤ C 1/4 + C −1/4 + Cδd log d. Z τA d d

Since λD = O(d−1/2 ) and δd ≤

c0 ε , log d

the right-hand side tends to zero. Therefore |DτA | = o(ZτA ). In particular, for all sufficiently large d, |DτA | ≤

66

1 Zτ . 8 A

(5)

Before Phase II, we provide another lemma below that controls the difference of drift Lemma D.15 (Balanced D-drift on the Phase-II band). Assume Lemma C.2. Let D := ρ − χ, Z := ρ + χ, and define e := 2α12 a21 > 0. C There exist constants cbd ∈ (0, 1/8) and CE > 0, depending only on the model, such that the following holds. Suppose 1 0 < Z ≤ 2cbd , |D| ≤ Z. (156) 2 Then ρ, χ ≥ 0, |ρ|, |χ| ≤ 3cbd /2, and the population drifts satisfy e + ρχ)D + E, ρ̇ − χ̇ = −C(1

(157)

where the error obeys the sharper bound |E| ≤ CE (Z 2 + η).

(158)

In particular, since Z ≤ 2cbd , one also has the weaker bound |E| ≤ CE (c2bd + η).

(159)

Proof. From Z = ρ + χ and D = ρ − χ, Z +D , 2 The balance condition |D| ≤ Z/2 gives

χ=

ρ=

Z −D . 2

1 3 Z ≤ ρ, χ ≤ Z. 4 4 Thus ρ, χ ≥ 0, and since Z ≤ 2cbd , we also have |ρ|, |χ| ≤

3 3 1 Z ≤ cbd < . 4 2 2

(160)

Let LG denote the Gaussian-equivalent population loss. As in the proof of Lemma C.3, X   LG (θ, β) = Φ(χ, ρ), Φ(χ, ρ) = 2 αk2 1 − (κ(ρ)χ)k , k≥1

where κ(ρ) =

2 ℓ ℓ≥1 aℓ ρ . Set

P

q := κ(ρ)χ,

A(q) :=

X

kαk2 q k−1 .

k≥1

Then A(0) = α12 . By the Gaussian β-gradient identity (30), equivalently by differentiating Φ in ρ, and by the corresponding χ-derivative displayed in the proof of Lemma C.3, the Gaussian-model overlap drifts are ρ̇G = 2χ κ′ (ρ)A(q)(1 − ρ2 ),

χ̇G = 2κ(ρ)A(q)(1 − χ2 ).

(161)

We now expand these identities on the small band (156). Since |ρ| ≤ 3cbd /2 < 1/2, the Hermite covariance function satisfies κ′ (ρ) = a21 + O(|ρ|).

κ(ρ) = a21 ρ + O(ρ2 ),

(162)

Moreover, by (160), q = κ(ρ)χ = O(Z 2 ), so A(q) = α12 + O(Z 2 ). Using (161)–(163),

(163)

ρ̇G = 2α12 a21 χ(1 − ρ2 ) + O(Z 2 ), because the error from κ′ (ρ)−a21 is O(|χ||ρ|) = O(Z 2 ), and the error from A(q)−α12 is O(|χ|Z 2 ) = O(Z 3 ). Similarly, χ̇G = 2α12 a21 ρ(1 − χ2 ) + O(Z 2 ), 67

because κ(ρ) − a21 ρ = O(ρ2 ) = O(Z 2 ). e := 2α2 a2 , Thus, with C 1 1 e χ(1 − ρ2 ) + O(Z 2 ), ρ̇G = C

e ρ(1 − χ2 ) + O(Z 2 ). χ̇G = C

(164)

It remains to transfer the Gaussian drift estimates to the true population loss. By (24) and (25) from Lemma C.2, |ρ̇ − ρ̇G | ≤ Cη, |χ̇ − χ̇G | ≤ Cη. Combining this with (164), we get   e χ(1 − ρ2 ) − ρ(1 − χ2 ) + O(Z 2 + η). ρ̇ − χ̇ = C Finally,

χ(1 − ρ2 ) − ρ(1 − χ2 ) = (χ − ρ)(1 + ρχ) = −D(1 + ρχ).

Therefore

(165)

|E| ≤ CE (Z 2 + η),

e + ρχ)D + E, ρ̇ − χ̇ = −C(1

which proves (157) and (158). Since Z ≤ 2cbd , (159) follows immediately. D.2.2

Phase II Dynamics, Weak Recovery

Lemma D.16 (Phase II: signal-dominated amplification and balance). Assume Lemmas D.13, D.14, and D.15. Work on GTweak . There exist constants cZ ∈ (0, cbd ∧ 1/8),

C2 , Cη,II , CII > 0,

depending only on the model constants, such that the following holds for all sufficiently large d. If CII log d ≤ T δ ≤ CT log d, and then

t∗ := inf{t ≥ τA : Zt ≥ cZ },   t∗ ≤ τA + C2 δ −1 log d ,

Zt∗ ∈ [cZ , 2cZ ].

∗

Moreover, for every τA ≤ t ≤ t , log d ηt ≤ Cη,II √ , d

|Dt | ≤

1 Zt . 2

Consequently, ρt , χt ≥

1 Zt 4

(τA ≤ t ≤ t∗ ),

(166)

cZ . 4

(167)

ZτA ≥ cA d−1/4 ,

(168)

1 Zτ . 8 A

(169)

with t∗ < τB ,

min{ρt∗ , χt∗ } ≥

Proof. Work on GTweak . By Lemma D.13 and Lemma D.14, τA δ ≤ C log d,

τA < τB , and

Also, on GTweak , Choose

log d η τA ≤ C √ , d

ZτA = o(1),

|DτA | ≤

λZ , λD = O(d−1/2 ) = o(ZτA ). log d H := Cη,II √ , d

  T2 := τA + C2 δ −1 log d . 68

The constants Cη,II , cZ , and C2 will be fixed below, and CII is chosen large enough so that T2 ≤ T , using τA δ ≤ C log d. Define the bootstrap stopping time   1 σ := inf t ≥ τA : ηt > H or Zt > 2cZ or |Dt | > Zt or t = T2 . 2 At t = τA , none of the first three clauses holds for large d. For t < σ, 0 < Zt ≤ 2cZ ,

|Dt | ≤

1 Zt . 2

Hence

Zt − D t 1 3 Zt + D t , χt = , Zt ≤ ρt , χt ≤ Zt . (170) 2 2 4 4 Taking cZ < 1/3, this implies ρt , χt < 1/2, so all one-step expansions below are valid before σ. ρt =

Step 1: Z-growth and no-drop.

Let (Z)

εt

(Z)

(Z)

:= Nt+1 − Nt

.

For t < σ, the one-step overlap expansion and the aggregate drift bound give (Z)

Zt+1 ≥ Zt + δ(C9 Zt − C11 ηt ) − Cδ 2 d Zt + εt . Since δd = o(1), for large d, C9 − Cδd ≥ Set

3 C9 . 4



1 τ− := inf t ≥ τA : Zt < ZτA 2 For t < σ ∧ τ− , we have Zt ≥ ZτA /2 and ηt ≤ H. Therefore

 .

2H log d ηt ≤ ≤ C 1/4 = o(1). Zt Z τA d Thus, for large d, C11 ηt ≤

1 C9 Zt , 4

and hence

(Z)

Zt+1 ≥ (1 + βδ)Zt + εt , where β := C9 /2. Let γ := 1 + βδ.

t < σ ∧ τ− ,

(171)

Unrolling (171), for every s ≥ 0 with τA + s < σ ∧ τ− , ZτA +s ≥ γ

s

Z τA +

s−1 X

! (Z) γ −(k+1) ετA +k

.

k=0 (Z)

On ET , (102) gives sup |Nu(Z) | ≤ 2λZ .

u≤T

By Corollary E.9 applied on the shifted interval starting at τA , s−1 X

(Z)

γ −(k+1) ετA +k ≤ 4λZ .

k=0

Therefore

ZτA +s ≥ ZτA − 4λZ γ s .

Since λZ = o(ZτA ), the right side is at least (3/4)ZτA γ s . This rules out τ− ≤ σ. Hence Zt ≥

1 Zτ 2 A

(τA ≤ t < σ), 69

(172)

and (172) holds for every τA + s < σ. We shall use the following summability estimates. Summing (171) from τA to t − 1, for t < σ, gives t−1 X

βδ

(Z)

Zk ≤ Zt − ZτA − Nt

− Nτ(Z) . A

k=τA (Z)

Since Zt ≤ 2cZ and |Nt

(Z)

− NτA | ≤ 4λZ , t−1 X

δ

Zk ≤ CcZ .

(173)

Zk2 ≤ Cc2Z .

(174)

k=τA

Also, since Zk ≤ 2cZ , t−1 X

δ

k=τA

Finally, the backward form of (171) gives, for τA ≤ k < t < σ, Zk ≤ γ −(t−k) Zt + CλZ . Indeed, unroll (171) from k to t, divide by γ noise sum.

t−k

(175)

, and apply Corollary E.9 to the shifted weighted

(η)

Step 2: Closing the η-bootstrap. Since ET controls the cumulative coordinate martingales from time 0, for every t ≥ τA and every i ≤ d, δ

t−1 X

(θ)

ξk,i ≤ δ

k=τA

X (θ) X (θ) ξk,i + δ ξk,i ≤ 2Λη . k<t

k<τA

Given that the quadratic and cubic remainder sums over [τA , t] are bounded by their global good-event bounds, therefore, applying Lemma D.9 on the interval [τA , t], for every τA ≤ t ≤ σ, ! t−1 X (θ) ηt ≤ C ητA + δ ∥ḡk ∥∞ + 2Λη + λcub . k=τA

After increasing C, we write this as ηt ≤ C

η τA + δ

t−1 X

! (θ) ∥ḡk ∥∞ + Λη + λcub

.

k=τA

On the Phase-II band, the mean-field ℓ∞ -gradient bound gives (θ)

∥ḡk ∥∞ ≤ CZk ∥θ∗ ∥∞ + CZk2 ηk + Cηk2 . Using ∥θ∗ ∥∞ ≤ C

p log d/d and (173), δ

t−1 X

r ∗

Zk ∥θ ∥∞ ≤ CcZ

k=τA

log d = o(H). d

By (174) and the bootstrap bound ηk ≤ H, δ

t−1 X k=τA

Zk2 ηk ≤ Hδ

t−1 X

Zk2 ≤ Cc2Z H.

k=τA

Also, δ

t−1 X

ηk2 ≤ H 2 (t − τA )δ ≤ C

k=τA

70

(log d)3 = o(H). d

Together with log d ητA + λη (t) + λcub ≤ C √ , d we obtain

log d ηt ≤ C √ + Cc2Z H + o(H). d

√ Choose cZ small enough that Cc2Z ≤ 1/8, then choose Cη,II large enough that C(log d)/ d ≤ H/4. For large d, 1 ηt ≤ H (τA ≤ t ≤ σ). 2 Thus the stopping clause ηt > H cannot be the first active clause before t∗ . Step 3: Closing the D/Z-bootstrap. For t < σ, choose cZ ≤ cbd . Then (170) allows us to apply Lemma D.15: e + ρt χt )Dt + Et , Ḋt = −C(1 |Et | ≤ C(Zt2 + ηt ). The discrete one-step expansion for Dt has the form (D)

Dt+1 = (1 − αt δ)Dt + δet + εt where

(D)

(D)

,

(176)

(D)

:= Nt+1 − Nt ,   e + ρt χt ) + δ E[∥gt(β) ∥2 | Ft ] + E[∥gt(θ) ∥2 | Ft ] , αt = C(1 4 εt

and

|et | ≤ C(Zt2 + ηt ) + Cδd Zt . (177) Here the last term comes from the difference of the two normalization corrections. Since cZ is small, αt ≥ α0 > 0, and after decreasing the step-size constant, 0 ≤ 1 − αt δ ≤ 1 − α0 δ =: γD < 1. Unroll (176). For τA ≤ t ≤ σ, define Pk+1,t :=

t−1 Y

(1 − αℓ δ),

τA ≤ k ≤ t − 1,

ℓ=k+1

with Pt,t = 1. Then t−1 X

Dt = PτA ,t DτA +

Pk+1,t δek +

k=τA

t−1 X

(D)

Pk+1,t εk .

k=τA

The entry term is controlled by Lemma D.14 and the no-drop estimate: |PτA ,t DτA | ≤ |DτA | ≤

1 1 Z τA ≤ Z t . 8 4 (ρ,χ)

For the noise term, the weights Pk+1,t lie in [0, 1] and are nondecreasing in k. On ET sup |Nu(D) | ≤ 2λD .

u≤T

Therefore Lemma E.8 gives t−1 X

(D)

Pk+1,t εk

≤ 4λD = o(ZτA ) = o(Zt ).

k=τA

It remains to bound the forcing. Using (175), Zk ≤ γ −(t−k) Zt + CλZ . 71

,

t−1−k Since Pk+1,t ≤ γD , we have

δ

t−1 X

Pk+1,t Zk ≤ Cδ

X

j −j γD γ Zt + CλZ δ

j γD ≤ CZt + CλZ .

j≥0

j≥0

k=τA

X

Since λZ = o(Zt ), this is CZt + o(Zt ). Similarly, δ

t−1 X

Pk+1,t Zk2 ≤ CZt2 + Cλ2Z = CZt2 + o(Zt ).

k=τA

Also, by the η-bootstrap, δ

t−1 X

Pk+1,t ηk ≤ CH = o(ZτA ) = o(Zt ).

k=τA

Using (177), we conclude t−1 X

Pk+1,t δ|ek | ≤ CcZ Zt + o(Zt ) + Cδd Zt .

k=τA

Since δd = o(1), t−1 X

Pk+1,t δ|ek | ≤ CcZ Zt + o(Zt ).

k=τA

Thus 1 Zt + CcZ Zt + o(Zt ). 4 Choose cZ smaller, if necessary, so that CcZ ≤ 1/16, and then take d large enough so that the o(Zt )-term is at most Zt /16. Hence |Dt | ≤

|Dt | ≤

3 1 Zt < Zt 8 2

(τA ≤ t ≤ σ).

Therefore the stopping clause |Dt | > Zt /2 cannot be the first active clause before t∗ . Step 4: Hitting cZ before T2 .

From (172),
ZτA +s ≥ ZτA − 4λZ γ s

as long as τA + s < σ. Define  s∗ :=

 log(cZ /(ZτA − 4λZ )) . log γ

Since ZτA − 4λZ ≥ cd−1/4 and log γ ≥ cδ, we have s∗ ≤ Cδ −1 log d. Choose C2 larger than this constant. Suppose, for contradiction, that t∗ > τA + s∗ . Then Zt < cZ up to τA + s∗ , so the Zt > 2cZ clause cannot trigger before that time. By Steps 2 and 3, the η- and D/Z-clauses cannot trigger before t∗ . By the choice of C2 , the time clause t = T2 also cannot trigger before τA + s∗ . Hence τA + s∗ < σ, and
ZτA +s∗ ≥ ZτA − 4λZ γ s∗ ≥ cZ , contradicting the definition of t∗ . Therefore   t∗ ≤ τA + s∗ ≤ τA + C2 δ −1 log d . 72

Step 5: Overshoot and conclusion. Since ZτA = o(1), for large d, ZτA < cZ , so t∗ > τA and Zt∗ −1 < cZ . The one-step expansion gives (Z)

(Z)

Zt∗ = Zt∗ −1 + δFt∗ −1 + εt∗ −1 . By Lemma E.7, (Z)

|Ft Hence

(Z)

δ|Ft

√ | ≤ C d + Cδd.

√ | ≤ Cδ d + Cδ 2 d = o(1).

(Z)

Also, on ET , (Z)

(Z)

(Z)

|εt | ≤ |Nt+1 | + |Nt

| ≤ 4λZ = o(1).

Therefore Zt∗ ≤ cZ + o(1) ≤ 2cZ for all large d. By definition, Zt∗ ≥ cZ , so Zt∗ ∈ [cZ , 2cZ ]. The bootstrap conclusions give, for τA ≤ t ≤ t∗ , log d ηt ≤ H = Cη,II √ , d Thus ρt =

Zt + Dt 1 ≥ Zt , 2 4

|Dt | ≤

χt =

At t = t∗ , min{ρt∗ , χt∗ } ≥

1 Zt . 2

Zt − D t 1 ≥ Zt . 2 4

1 cZ Z t∗ ≥ . 4 4

Moreover, 3 3 1 Zt∗ + |Dt∗ | ≤ Zt∗ ≤ cZ < , 2 4 2 2 after choosing cZ < 1/3. Hence t∗ < τB . This completes the proof. max{ρt∗ , χt∗ } ≤

D.2.3

Weak Recovery Theorem

Now, we’re prepared to derive the final theorem for weak recovery, which states: Theorem D.17 (Escape from mediocrity under discrete-time SGD). Assume Assumption D.1 and the auxiliary estimates used in Lemmas D.13, D.14, and D.16. Let E0 be the favorable initialization event. Then there exist constants cesc , Ctime , K > 0 depending only on the model constants such that, for all sufficiently large d, the following holds. Let   T := Ctime δ −1 log d . Conditional on E0 , with probability at least 1 − Kε, there exists a time tesc ≤ Ctime δ −1 log d

(178)

min{ρtesc , χtesc } ≥ cesc .

(179)

such that

Moreover, throughout the escape window, log d ηt ≤ C √ , d

0 ≤ t ≤ tesc .

(180)

Equivalently, if P(E0 ) ≥ p0 > 0, then
P ∃ tesc ≤ Ctime δ −1 log d : min{ρtesc , χtesc } ≥ cesc ≥ p0 (1 − Kε). 73

(181)

Proof. Choose Ctime large enough so that Ctime > CI + CII + 2, where CI is the Phase-I time constant from Lemma D.13, and CII is the Phase-II time constant from Lemma D.16. Increasing Ctime if necessary, the horizon T satisfies all horizon assumptions required in the Phase-I, entry-balance, and Phase-II lemmas. Define the weak-recovery good event, excluding initialization, by (T )

HTweak := ET

(Z)

∩ ET

(η)

(ρ,χ)

∩ ET ∩ ET

(Ω)

∩ ET .

By bounds of the Taylor (89), accumulative (103) and difference (105) martingale and discretization control, η-control (135), (141) and third-moment estimates,
P (HTweak )c | E0 ≤ Kε. (182) Set

GTweak := E0 ∩ HTweak .

We now work on GTweak . By Lemma D.13, τA < τB , τA δ ≤ CI log d, ZτA ≥ cA d−1/4 . The same lemma also gives, for 0 ≤ t ≤ τA , √ log d + 1 + tδ √ . ηt ≤ Cη d

(183)

(184)

Next, Lemma D.14 gives |DτA | = o(ZτA ). Hence, for all sufficiently large d, |DτA | ≤

1 Zτ . 8 A

(185)

Now apply Lemma D.16. Define t∗ := inf{t ≥ τA : Zt ≥ cZ }, where cZ > 0 is the constant from Lemma D.16. The Phase-II lemma gives   t∗ ≤ τA + CII δ −1 log d , Zt∗ ∈ [cZ , 2cZ ], ∗

and, for every τA ≤ t ≤ t , log d ηt ≤ Cη,II √ , d Therefore,

|Dt | ≤

1 Zt . 2

t∗ ≤ Ctime δ −1 log d

for all sufficiently large d. Set tesc := t∗ ,

cesc :=

Zt + D t , 2

χt =

Since ρt =

cZ . 4

Zt − D t , 2

the balance estimate at t = t∗ gives min{ρtesc , χtesc } =

Ztesc − |Dtesc | 1 cZ ≥ Ztesc ≥ = cesc . 2 4 4

This proves weak recovery on GTweak . 74

(186)

It remains to verify the claimed ηt -bound. For 0 ≤ t ≤ τA , (184) and τA δ ≤ CI log d imply log d ηt ≤ C √ . d For τA ≤ t ≤ tesc , the same bound follows from (186). Hence log d ηt ≤ C √ , d

0 ≤ t ≤ tesc .

Finally, the conditional probability statement follows from (182). If P(E0 ) ≥ p0 , then P(E0 ∩ HTweak ) = P(HTweak | E0 )P(E0 ) ≥ p0 (1 − Kε), which gives the unconditional statement. D.3

Strong Recovery Setups

Remark D.18 (Local strong-recovery step-size regime). Entering strong recovery phase, we changed the step size to cloc δloc := , d where cloc > 0 is chosen sufficiently small so that Cret Csq δloc d ≤ µloc for the optimum convergence rate: our one-step recursion is Qt+1 ≤ Qt + δ⟨ḡt , Logt ⟩ + δΞt + Cret δ 2 Gt With ⟨ḡt , Logt ⟩ ≤ −2µloc Qt and E[Gt | Ft ] ≤ Csq dQt , we have in expectation E[Qt+1 | Ft ] ≤ (1 − 2µloc Qt + Cdδ 2 )Qt , which means by taking δ ≍ d1 , we will have the maximized contraction rate of O(d log 1ε ) samples. Population inputs reused from the continuous-time analysis. Throughout this subsection, let rloc > 0 and µloc > 0 be the constants from Lemma C.14. In particular, for every y ∈ Urloc ,

Urloc := {y ∈ M : distM (y, y ∗ ) ≤ rloc },

we have HessM L(y) ⪰ µloc I. Discrete-time basin-entry. Mirroring Lemma C.9, the following lemmas establish that after an additional O(d) iterations beyond weak recovery, the trajectory enters the local strongly-convex cap in discrete-time: Lemma D.19 (Entry-window delocalization and overlap-noise control). Work on the weak-recovery event at tesc . Assume the incoherence bound (A4). After tesc , run SGD with the local step size δloc = cloc constants, Cin , and the d , where cloc > 0 is chosen sufficiently small depending only onthe model  −1 desired entry-window failure level. Fix Cin < ∞, and set Nin := Cin δloc . Then, conditionally on Ftesc , there exists an event Ein , measurable with respect to the samples in [tesc , tesc + Nin ], such that P(Ein | Ftesc ) ≥ 1 − νin .

(187)

ηt ≤ η̄d ,

(188)

On Ein , sup tesc ≤t≤tesc +Nin

where, with Lν,d := 1 + log(d/νin ), η̄d ≤ C

log d √ + d

r

log d p 2 + δloc Lν,d + Lν,d δloc d3/2 d

In particular, for fixed or polynomially small νin ,   log d η̄d = O √ = od (1). d 75

! .

(189)

(190)

Moreover, for j ∈ {ρ, χ}, u−1 X

sup tesc ≤s≤u≤tesc +Nin

(j)

∆Nk+1 ≤ Λd ,

(191)

k=s

where Λd ≤ C

p

 2 δloc log(1/νin ) + Lν,d δloc d3/2 = od (1)

(192)

for fixed or polynomially small νin . Proof. Write t0 := tesc ,

  N := Nin = Cin δ −1 .

δ := δloc ,

Throughout the proof we condition on Ft0 and on the weak-recovery event stated in the lemma. Since the samples after t0 are fresh and independent of Ft0 , the martingale and moment estimates from the weak-recovery section apply conditionally on Ft0 , with the same constants. We shall use the normalized spherical one-step expansion (ρ)

+ ∆Nt+1 ,

(j)

(j)

ρt+1 = ρt + δFt where

(ρ)

∆Nt+1 = δMt

(χ)

χt+1 = χt + δFt (j)

(χ)

+ ∆Nt+1 ,

j ∈ {ρ, χ}.

+ Rt ,

(j) (j) Here Mt is a scalar martingale increment and Rt is the normalization/retraction remainder.

Step 1: overlap-noise control. By the scalar martingale moment bounds used in the weak-recovery proof, there is KM < ∞, independent of d, such that h i (j) E (Mt )2 | Ft ≤ KM (1 + δ 2 d2 ), j ∈ {ρ, χ}. ′ Since δd = cloc , this is bounded by another constant KM .

For u ∈ [t0 , t0 + N ], define Su(j) := δ

u−1 X

(j)

Mk .

k=t0

Then ′ ⟨S (j) ⟩t0 +N ≤ KM δ 2 N ≤ Cδ,

because N ≤ Cδ −1 . By the same conditional Freedman/truncation argument used for the weakrecovery overlap martingales, with conditional probability at least 1 − νin /8, sup

δ

t0 ≤u≤t0 +N

u−1 X

(j)

≤C

(j)

≤C

Mk

p

δ log(1/νin ).

k=t0

Consequently, sup t0 ≤s≤u≤t0 +N

δ

u−1 X

Mk

p

δ log(1/νin ),

k=s

after increasing C. Next consider the remainders. On the local Taylor event for the normalized spherical step, (ρ)

(β)

(χ)

|Rt | ≤ Cδ 3 ∥gt ∥3 ,

(θ)

|Rt | ≤ Cδ 3 ∥gt ∥3 .

Moreover, the gradient moment bounds give (β)

(θ)

E[∥gt ∥3 | Ft ] + E[∥gt ∥3 | Ft ] ≤ Cd3/2 . 76

Hence, for j ∈ {ρ, χ}, # "t +N −1 0X (j) |Rk | Ft0 ≤ CN δ 3 d3/2 ≤ Cδ 2 d3/2 . E k=t0

Using the same high-moment/truncation estimate as in the weak-recovery discretization control, there is an event of conditional probability at least 1 − νin /8 on which u−1 X

sup t0 ≤s≤u≤t0 +N

(j)

≤ CLν,d δ 2 d3/2 ,

Rk

Lν,d := 1 + log(d/νin ).

k=s

Therefore, on an event Eov of conditional probability at least 1 − νin /2, for both j ∈ {ρ, χ}, u−1 X

sup t0 ≤s≤u≤t0 +N

where Λd := C

(j)

∆Nk+1 ≤ Λd ,

k=s

p  δ log(1/νin ) + Lν,d δ 2 d3/2 .

Since δ = cloc /d, r Λd ≤ C

log(1/νin ) 1 + CLν,d √ = od (1) d d

for fixed or polynomially small νin . Step 2: ηt -control on the entry window. To control ηt = ∥θt ∥∞ + ∥θ∗ ∥∞ , we use the same ℓ∞ inequality as in the weak-recovery analysis: ∥θt ∥∞ ≤ ∥θt0 ∥∞ + δ

t−1 X k=t0

(θ) ∥ḡk ∥∞ + max i≤d

where

(θ)

δ

t−1 X

(θ) ξk,i

k=t0

(θ)

t−1

t−1

k=t0

k=t0

δ 3 X (θ) 3 δ 2 X (θ) 2 ∥gk ∥ ∥θk ∥∞ + ∥gk ∥ , + 2 2 (θ)

ḡk := E[gk | Fk ],

(θ)

(θ)

ξk := gk − ḡk .

Set t0 ≤ u ≤ t0 + N.

Θu := sup ηt , t0 ≤t≤u

By the conditional coordinate Freedman bound, with conditional probability at least 1 − νin /8, max

sup

i≤d t0 ≤u≤t0 +N

δ

u−1 X

(θ)

ξk,i ≤ C

p

δLν,d .

k=t0

The gradient moment estimates imply " t +N −1 # 0X (θ) 2 2 E δ ∥gk ∥ Ft0 ≤ Cδ 2 N d ≤ Cδd = Ccloc . k=t0

Choosing cloc > 0 sufficiently small, the same concentration/Markov argument used in the weak phase gives, with conditional probability at least 1 − νin /8, t +N −1

δ2 0 X 1 (θ) ∥gk ∥2 ≤ . 2 4 k=t0

Similarly, " E δ

3

t0 +N X−1

# (θ) ∥gk ∥3

Ft0 ≤ Cδ 3 N d3/2 ≤ Cδ 2 d3/2 ,

k=t0

77

and hence, on an event of conditional probability at least 1 − νin /8, δ3

t0 +N X−1

(θ)

∥gk ∥3 ≤ CLν,d δ 2 d3/2 .

k=t0

Let Eη,noise be the intersection of these three η-control events. On Eη,noise , for every u ∈ [t0 , t0 +N ], Θu ≤ Cηt0 + Cδ

u−1 X

(θ)

∥ḡk ∥∞ + C

p

k=t0

1 δLν,d + CLν,d δ 2 d3/2 + Θu . 4

Absorbing the last term, Θu ≤ Cηt0 + Cδ

u−1 X

(θ)

∥ḡk ∥∞ + C

p δLν,d + CLν,d δ 2 d3/2 .

(193)

k=t0

It remains to bound the mean-field term. The population ℓ∞ gradient estimate gives, uniformly along the entry window, (θ) ∥ḡk ∥∞ ≤ C∥θ∗ ∥∞ + Cηk + Cηk2 . Therefore, since N δ ≤ Cin + 1, (193) implies Θu ≤ Ad + Cδ

u−1 X

Θk + Cδ

k=t0

where

u−1 X

Θ2k ,

k=t0

  p Ad := C ηt0 + ∥θ∗ ∥∞ + δLν,d + Lν,d δ 2 d3/2 .

Now insert the scales. On the weak-recovery event, log d ηt0 ≤ Cη,esc √ , d and by the teacher incoherence assumption, r ∗

∥θ ∥∞ ≤ C∗

log d . d

Also δ = cloc /d. Hence, for fixed or polynomially small νin , ! r r   log d log d Lν,d Lν,d log d Ad ≤ C √ + + + √ =O √ = od (1). d d d d d We close the bootstrap. Define τη := inf{u ∈ [t0 , t0 + N ] : Θu ≥ 1} ∧ (t0 + N ). For u ≤ τη , one has Θ2k ≤ Θk , and therefore Θu ≤ Ad + Cδ

u−1 X

Θk .

k=t0

The discrete Gronwall inequality gives ′ Θu ≤ Ad exp(C(u − t0 )δ) ≤ Ad exp(C(Cin + 1)) ≤ Cin Ad ,

u ≤ τη .

Since Ad = od (1), for all sufficiently large d the right-hand side is strictly smaller than 1. Hence τη = t0 + N , and sup ηt = Θt0 +N ≤ η̄d , t0 ≤t≤t0 +N

78

where log d √ + d

′ η̄d := Cin Ad ≤ C

r

log d p + δLν,d + Lν,d δ 2 d3/2 d

! .

This proves the claimed delocalization bound. conclusion. Define Ein := Eov ∩ Eη,noise . By the conditional probability bounds above and a union bound, after adjusting constants, P(Ein | Ft0 ) ≥ 1 − νin . On Ein , both conclusions hold: ηt ≤ η̄d

sup t0 ≤t≤t0 +N

and, for j ∈ {ρ, χ}, u−1 X

sup t0 ≤s≤u≤t0 +N

(j)

∆Nk+1 ≤ Λd .

k=s

Substituting back t0 = tesc , N = Nin , and δ = δloc completes the proof. Lemma D.20 (One-step overlap recursion on the entry corridor). Assume Lemma C.3. After tesc , run SGD with step size δ = δloc . On any time t in the entry window for which the one-step Taylor expansion is valid and ρt ≥ 0, χt ≥ 0, the overlap recursions can be written as (ρ)

ρt+1 = ρt + δFt where

(ρ)

+ ∆Nt+1 ,

(ρ)

χt+1 = χt + δFt

(ρ)

(ρ)

∆Nt+1 = δMt

∆Nt+1 = δMt (ρ)

(χ)

Here Mt , Mt

(χ)

(χ)

+ Rt ,

(χ)

(χ)

+ ∆Nt+1 , (χ)

+ Rt .

are martingale increments and the remainders satisfy (ρ)

(β)

(χ)

|Rt | ≤ Cδ 3 ∥gt ∥3 ,

(θ)

|Rt | ≤ Cδ 3 ∥gt ∥3 .

Moreover, the predictable drifts satisfy (ρ)

Ft and

≥ aρ χt (1 − ρ2t ) − bρ ηt − Cδd,

(χ)

Ft ≥ aχ ρt (1 − χ2t ) − bχ ηt − Cδd. Consequently, on the event Ein from Lemma D.19, (ρ)

Ft

(χ)

≥ aρ χt (1 − ρ2t ) − Bd ,

Ft

≥ aχ ρt (1 − χ2t ) − Bd ,

where Bd := C(η̄d + δd). In particular, for δ = cloc /d,  Bd ≤ C

log d √ + cloc d



for fixed or polynomially small νin . Proof. We prove the ρ-recursion; the proof for χ is identical. By the normalized spherical update, (β)

βt+1 =

βt − δgt

(β)

∥βt − δgt ∥ 79

.

(β)

Since gt

(β)

⊥ βt , on the one-step Taylor event δ∥gt ∥ ≤ 1/2, we have (β)

βt+1 = βt − δgt

−

δ 2 (β) 2 (β) ∥g ∥ βt + O(δ 3 ∥gt ∥3 ). 2 t

Taking inner product with β ∗ , and using ρt = ⟨βt , β ∗ ⟩, gives (β)

ρt+1 = ρt − δ⟨gt , β ∗ ⟩ − where

δ2 (β) (ρ) ρt ∥gt ∥2 + Rt , 2

(ρ)

(β)

|Rt | ≤ Cδ 3 ∥gt ∥3 . Equivalently,   δ (β) (ρ) (β) ρt+1 = ρt − δ ⟨gt , β ∗ ⟩ + ρt ∥gt ∥2 + Rt . 2 Define

  δ (β) (β) 2 ∗ ⟨gt , β ⟩ + ρt ∥gt ∥ Ft , 2

(ρ) Ft := −E

and (ρ)

Mt

   δ δ (β) (β) (β) (β) := − ⟨gt , β ∗ ⟩ + ρt ∥gt ∥2 − E ⟨gt , β ∗ ⟩ + ρt ∥gt ∥2 Ft . 2 2

Then

(ρ)

E[Mt and

| Ft ] = 0,

(ρ)

ρt+1 = ρt + δFt Thus

(ρ)

(ρ)

+ δMt

+ Rt .

(ρ)

(ρ)

(ρ)

∆Nt+1 := δMt

+ Rt .

(ρ)

It remains to lower-bound Ft . By definition, (ρ)

Ft

δ (β) (β) = −E[⟨gt , β ∗ ⟩ | Ft ] − ρt E[∥gt ∥2 | Ft ]. 2

The first term is the population spherical-gradient drift: (β)

−E[⟨gt , β ∗ ⟩ | Ft ] = − ⟨∇S,β L(θt , βt ), β ∗ ⟩ . Since ρt , χt ≥ 0, Lemma C.3 implies − ⟨∇S,β L(θt , βt ), β ∗ ⟩ ≥ aρ χt (1 − ρ2t ) − bρ ηt . For the normalization correction, use |ρt | ≤ 1 and the conditional second-moment gradient bound: δ δ (β) (β) − ρt E[∥gt ∥2 | Ft ] ≥ − E[∥gt ∥2 | Ft ] ≥ −Cδd. 2 2 Combining the two estimates gives (ρ)

Ft

≥ aρ χt (1 − ρ2t ) − bρ ηt − Cδd.

The χ-recursion follows from the same argument applied to (θ)

θt+1 =

θt − δgt

(θ)

∥θt − δgt ∥

Thus

(χ)

χt+1 = χt + δFt with

(χ)

Ft

.

(χ)

+ ∆Nt+1 ,

≥ aχ ρt (1 − χ2t ) − bχ ηt − Cδd. 80

Finally, on Ein , ηt ≤ η̄d throughout the entry window. Increasing C to absorb bρ , bχ , the two drift bounds become (ρ)

Ft and

≥ aρ χt (1 − ρ2t ) − C(η̄d + δd),

(χ)

Ft ≥ aχ ρt (1 − χ2t ) − C(η̄d + δd). This proves the stated bounds with Bd := C(η̄d + δd). √ Since Lemma D.19 gives η̄d = O(log d/ d) for fixed or polynomially small νin , and since δ = cloc /d, we obtain   log d Bd ≤ C √ + cloc . d The proof is complete. Proposition D.21 (Discrete entry into the local product basin). Let rloc be the radius from Lemma C.14. After possibly shrinking rloc , assume 0 < rloc < π. Define 1 + sb . 2 Apply Lemma D.19 with a constant Cin at least as large as the fixed constant Cent constructed in the proof below. sb := cos(rloc /8),

s+ :=

Then there exist constants cloc,0 > 0, d0 < ∞, and Cent < ∞, depending only on rloc , cesc , and the model constants, such that the following holds. If cloc δloc = , 0 < cloc ≤ cloc,0 , d ≥ d0 , d −1 then on the event Ein from Lemma D.19, there exists a time t1 ≤ tesc + Cent δloc such that

ρt1 ≥ sb ,

χt1 ≥ sb .

Consequently, Qt1 =

1 1 r2 r2 distS (θt1 , θ∗ )2 + distS (βt1 , β ∗ )2 ≤ loc ≤ loc . 2 2 64 32

Proof. Write

cesc . 2 We work deterministically on the event Ein , together with the weak recovery event at tesc . Thus t0 := tesc ,

δ := δloc ,

ρt0 ≥ cesc ,

c− :=

χt0 ≥ cesc ,

and, throughout the entry window, ηt ≤ η̄d ,

u−1 X

sup t0 ≤s≤u≤t0 +Cin

δ −1

(j)

∆Nk+1 ≤ Λd ,

j ∈ {ρ, χ}.

k=s

By Lemma D.20, on Ein , (ρ)

Ft

(χ)

≥ aρ χt (1 − ρ2t ) − Bd ,

Ft

≥ aχ ρt (1 − χ2t ) − Bd ,

where Bd := C(η̄d + δd). Since

 log d √ , δd = cloc , η̄d = O d we may choose cloc > 0 sufficiently small and then d sufficiently large so that all smallness conditions below hold. 

81

Define m+ := 1 − s2+ > 0, and set

1 1 aρ c− m+ , γχ := aχ c− m+ . 2 2 These are fixed positive constants independent of d. Also define γρ :=

α∗ := Let Tρ :=

1 min {cesc , s+ − sb , γρ , γχ } > 0. 4

(s+ − cesc )+ + α∗ , γρ

Tχ :=

(s+ − cesc )+ + α∗ , γχ

and choose Cent := 2 + max{Tρ , Tχ }. We choose cloc small enough and then d large enough so that Bd ≤ min{γρ , γχ }, Λd ≤ α∗ , Cent Bd + Λd ≤ α∗ . √ These conditions are consistent because Bd ≤ C(log d/ d + cloc ) and Λd = od (1). Let

(194) (195) (196)

  Nent := max{Tρ , Tχ }δ −1 .

Since δ ≤ 1, we have Nent δ ≤ Cent . By the assumption Cin ≥ Cent , all estimates from Ein hold on [t0 , t0 + Nent ]. Step 1: the overlaps stay above the lower barrier. We claim that ρt ≥ c− ,

χt ≥ c− ,

t0 ≤ t ≤ t0 + Nent .

Suppose not, and define the first lower-barrier exit time τ− := inf {t ∈ [t0 , t0 + Nent ] : min{ρt , χt } < c− } . Since ρt0 , χt0 ≥ cesc > c− , one has τ− > t0 if τ− < ∞. For every k < τ− , both overlaps are at least c− , and hence nonnegative. Therefore the crude drift bounds (ρ) (χ) Fk ≥ −Bd , Fk ≥ −Bd hold. Summing the ρ-recursion from t0 to τ− , and using the cumulative noise bound from Ein , gives ρτ− ≥ ρt0 − (τ− − t0 )δBd − Λd . Since (τ− − t0 )δ ≤ Cent , (196) implies 3 cesc > c− . 4 The same argument gives χτ− > c− . This contradicts the definition of τ− . Hence no such exit occurs. ρτ− ≥ cesc − Cent Bd − Λd ≥ cesc − α∗ ≥

Step 2: each overlap reaches the buffer level s+ . First consider ρt . If ρt0 ≥ s+ , set τρ := t0 . Otherwise, suppose for contradiction that   ρt < s+ for all t0 ≤ t ≤ t0 + Tρ δ −1 . By Step 1, χt ≥ c− on this interval. Hence, for every such k, (ρ)

Fk

≥ aρ c− (1 − ρ2k ) − Bd . 82

Since ρk < s+ , we have 1 − ρ2k ≥ 1 − s2+ = m+ , and therefore (ρ)

Fk

≥ aρ c− m+ − Bd = 2γρ − Bd ≥ γρ

by (194). Summing the recursion gives ρt0 +⌈Tρ δ−1 ⌉ ≥ ρt0 + γρ Tρ − Λd . Using ρt0 ≥ cesc , the definition of Tρ , and (195), ρt0 +⌈Tρ δ−1 ⌉ ≥ cesc + (s+ − cesc )+ + α∗ − Λd ≥ s+ . This contradicts the assumption that ρt < s+ throughout the interval. Thus there exists   τρ ≤ t0 + Tρ δ −1 such that ρτρ ≥ s+ . The proof for χt is identical. If χt0 ≥ s+ , set τχ := t0 . Otherwise, using Step 1 and the lower bound (χ) for Ft , one obtains a time   τχ ≤ t0 + Tχ δ −1 such that χτχ ≥ s+ . Step 3: retain both overlaps above sb . Define t1 := max{τρ , τχ }. Then

t1 ≤ t0 + Nent ≤ tesc + Cent δ −1 .

We show that both overlaps are at least sb at time t1 . Since ρτρ ≥ s+ , and since Step 1 gives the (ρ) lower barrier on the whole window, the crude drift bound Fk ≥ −Bd holds on [τρ , t1 ]. Therefore ρt1 ≥ ρτρ − (t1 − τρ )δBd − Λd . As (t1 − τρ )δ ≤ Cent , (196) gives ρt1 ≥ s+ − Cent Bd − Λd ≥ s+ − α∗ . By the definition of α∗ , α∗ ≤

1 (s+ − sb ), 4

so ρt1 ≥ sb . The same argument gives χt1 ≥ sb . Step 4: convert overlap entry into product-distance entry. Since

ρt1 = ⟨βt1 , β ∗ ⟩ ≥ sb = cos(rloc /8),

and 0 < rloc < π, distS (βt1 , β ∗ ) = arccos(ρt1 ) ≤ Similarly, distS (θt1 , θ∗ ) ≤

rloc . 8

rloc . 8

Therefore 1 1 1  rloc 2 1  rloc 2 r2 distS (θt1 , θ∗ )2 + distS (βt1 , β ∗ )2 ≤ + = loc . 2 2 2 8 2 8 64 This proves the proposition. Qt1 =

83

Deriving Qt -iteration and its controlling bounds. Lemma D.22 (One-step local distance recursion). Let y ∗ = (θ∗ , β ∗ ),

yt = (θt , βt ),

(θ)

(β)

gt = (gt , gt ),

and define 1 1 1 distM (yt , y ∗ )2 = distS (θt , θ∗ )2 + distS (βt , β ∗ )2 . 2 2 2 Assume rloc < π/4. There exists a constant Cret < ∞, depending only on rloc , such that whenever yt ∈ Urloc , the normalized SGD step satisfies Qt :=

Qt+1 ≤ Qt + δ⟨gt , Lt ⟩ + Cret δ 2 ∥gt ∥2 , where

Lt := Logθt (θ∗ ), Logβt (β ∗ ) .

Consequently, if ḡt := E[gt | Ft ],

Ξt := ⟨gt − ḡt , Lt ⟩,

Gt := ∥gt ∥2 ,

then Qt+1 ≤ Qt + δ⟨ḡt , Lt ⟩ + δΞt + Cret δ 2 Gt .

(197)

Proof. We use the product metric on M = Sd−1 × Sd−1 , namely ⟨(aθ , aβ ), (bθ , bβ )⟩ := ⟨aθ , bθ ⟩ + ⟨aβ , bβ ⟩, and ∥(aθ , aβ )∥2 := ∥aθ ∥2 + ∥aβ ∥2 . Fix t such that yt ∈ Urloc . We prove the one-step estimate on one sphere and then add the two block estimates. Let u ∈ Sd−1 , let v ∈ Sd−1 , and let h ∈ Tu Sd−1 . Consider the normalized spherical step u b :=

u − δh . ∥u − δh∥

Since h ⊥ u, ∥u − δh∥2 = 1 + δ 2 ∥h∥2 . Define

  arctan(δ∥h∥) , τ (h) := ∥h∥  0,

h ̸= 0, h = 0.

Then u b = Expu (−τ (h)h). Indeed, when h ̸= 0, the normalized point lies on the great circle through u in direction −h/∥h∥, with angular displacement arctan(δ∥h∥). The case h = 0 is immediate. Now set

1 distS (u, v)2 . 2 Since u lies in the local cap around v, and since the normalized step has spherical length at most π/2, after shrinking rloc if necessary the geodesic segment used above remains in a compact subset away from the cut locus of v. Therefore qv has uniformly bounded Hessian on this region: qv (u) :=

∥ HessS qv ∥ ≤ Csph . Taylor expansion along s 7→ Expu (−sτ (h)h) 84

gives qv (b u) ≤ qv (u) + ⟨∇S qv (u), −τ (h)h⟩ +

Csph τ (h)2 ∥h∥2 . 2

Since ∇S qv (u) = − Logu (v), we get qv (b u) ≤ qv (u) + τ (h)⟨h, Logu (v)⟩ +

Csph τ (h)2 ∥h∥2 . 2

We compare τ (h) with δ. Since 0 ≤ τ (h) ≤ δ, and, writing x = δ∥h∥, 0 ≤ δ∥h∥ − arctan(δ∥h∥) = x − arctan x ≤ Cx2 , we have

|δ − τ (h)| ∥h∥ ≤ Cδ 2 ∥h∥2 . Also u is within distance rloc of v, so ∥ Logu (v)∥ = distS (u, v) ≤ rloc . Therefore

τ (h)⟨h, Logu (v)⟩ ≤ δ⟨h, Logu (v)⟩ + Crloc δ 2 ∥h∥2 .

Moreover,

τ (h)2 ∥h∥2 ≤ δ 2 ∥h∥2 .

Thus

(1)

qv (b u) ≤ qv (u) + δ⟨h, Logu (v)⟩ + Cret δ 2 ∥h∥2 , (1)

where Cret depends only on rloc and Csph . Apply this estimate first with (θ)

u = θt ,

v = θ∗ ,

h = gt ,

u = βt , The normalized SGD updates are

v = β∗,

h = gt .

and then with

(θ)

θt+1 =

(β)

(β)

θt − δgt

, (θ) ∥θt − δgt ∥

βt+1 =

βt − δgt

(β)

∥βt − δgt ∥

.

Hence E D 1 1 (θ) (1) (θ) distS (θt+1 , θ∗ )2 ≤ distS (θt , θ∗ )2 + δ gt , Logθt (θ∗ ) + Cret δ 2 ∥gt ∥2 , 2 2 and

D E 1 1 (β) (1) (β) distS (βt+1 , β ∗ )2 ≤ distS (βt , β ∗ )2 + δ gt , Logβt (β ∗ ) + Cret δ 2 ∥gt ∥2 . 2 2 Adding the two estimates gives Qt+1 ≤ Qt + δ⟨gt , Lt ⟩ + Cret δ 2 ∥gt ∥2 , after increasing the constant Cret . Finally, gt = ḡt + (gt − ḡt ), so ⟨gt , Lt ⟩ = ⟨ḡt , Lt ⟩ + ⟨gt − ḡt , Lt ⟩ = ⟨ḡt , Lt ⟩ + Ξt . With Gt = ∥gt ∥ , this yields 2

Qt+1 ≤ Qt + δ⟨ḡt , Lt ⟩ + δΞt + Cret δ 2 Gt , which proves (197). 85

Lemma D.23 (Sharp local self-quenching). Inside Urloc , there exists Csq > 0, independent of d, such that E[Gt | Ft ] ≤ Csq d Qt , E[Ξ2t | Ft ] ≤ Csq Q2t , and

h i 2 E (Gt − E[Gt | Ft ]) | Ft ≤ Csq d2 Q2t .

Proof. Condition on Ft , so that (θt , βt ) are deterministic. To lighten notation, write θ := θt , and set

β := βt ,

dθ := distS (θ, θ∗ ),

Since Qt = we have

dβ := distS (β, β ∗ ).
1 2 dθ + d2β , 2

d2θ + d2β = 2Qt .

Let

a := Logθ (θ∗ ) ∈ Tθ Sd−1 ,

b := Logβ (β ∗ ) ∈ Tβ Sd−1 ,

so that ∥a∥ = dθ ,

∥b∥ = dβ .

Write the one-sample prediction as s := zβ (x) · θ =

fβ,θ (x) = ψ(s),

d X

θi ϕ(ui ),

ui := x(i) · β,

i=1

and the teacher prediction as fβ ∗ ,θ∗ (x) = ψ(s∗ ),

s∗ := zβ ∗ (x) · θ∗ =

d X

θi∗ ϕ(u∗i ),

u∗i := x(i) · β ∗ .

i=1

Define the residual

R := fβ,θ (x) − fβ ∗ ,θ∗ (x) = ψ(s) − ψ(s∗ ).

The spherical one-sample gradients are (θ)

gt

= 2R ψ ′ (s) Pθ zβ (x),

and (β) gt = 2R ψ ′ (s) Pβ

d X

 θi ϕ′ (ui )x(i) .

i=1

Set

(θ)

g (θ) := gt ,

and

(β)

g (β) := gt ,

gt = (g (θ) , g (β) ),

Gt := ∥gt ∥2 = ∥g (θ) ∥2 + ∥g (β) ∥2 .

Also recall Ξt = ⟨gt − ḡt , (a, b)⟩,

ḡt := E[gt | Ft ].

All moment estimates below follow from Assumption (A3) and Remark B.1. residual self-quenching. We claim that for m = 2, 4, 8, m/2

E[|R|m | Ft ] ≤ Cm Qt 86

.

Since ψ ′ is bounded,

|R| ≤ ∥ψ ′ ∥∞ |s − s∗ |.

Decompose

s − s∗ = A + B,

where

d X A := (θi − θi∗ )ϕ(ui ), i=1

and B :=

d X

θi∗ ϕ(ui ) − ϕ(u∗i ) .

i=1

For A: the variables ϕ(ui ) are iid, centered, and have finite moments up to order 8. Hence by Rosenthal’s inequality, for m ∈ {2, 4, 8}, ∥A∥Lm (P(·|Ft )) ≤ Cm ∥θ − θ∗ ∥2 . Since Euclidean distance is bounded by geodesic distance on the sphere, ∥θ − θ∗ ∥2 ≤ dθ , we get ∥A∥Lm (P(·|Ft )) ≤ Cm dθ . For B: the blocks x(i) are independent, so the summands
Yi := θi∗ ϕ(ui ) − ϕ(u∗i ) are independent across i. They are centered because ui ∼ N (0, 1) and u∗i ∼ N (0, 1), hence E[ϕ(ui ) − ϕ(u∗i )] = 0. Using the Lipschitzness of ϕ, |Yi | ≤ |θi∗ | Lip(ϕ) |ui − u∗i | = |θi∗ | Lip(ϕ) |x(i) · (β − β ∗ )|. Thus for m ∈ {2, 4, 8}, d X

E[|Yi |m | Ft ] ≤ Cm ∥β − β ∗ ∥m 2

i=1

since

P

d X

|θi∗ |m ≤ Cm dm β ,

i=1

∗ m ∗ m i |θi | ≤ ∥θ ∥2 = 1 for m ≥ 2. Applying Rosenthal again,

∥B∥Lm (P(·|Ft )) ≤ Cm dβ ,

m = 2, 4, 8.

Combining the bounds for A and B, ∥s − s∗ ∥Lm (P(·|Ft )) ≤ Cm (dθ + dβ ) ≤ Cm Hence ∥R∥Lm (P(·|Ft )) ≤ Cm

p

Qt ,

q p d2θ + d2β = Cm 2Qt . m = 2, 4, 8,

which is equivalent to m/2

E[|R|m | Ft ] ≤ Cm Qt

,

m = 2, 4, 8.

feature-side moment bounds. Define z := zβ (x) = (ϕ(u1 ), . . . , ϕ(ud )),

h := Pβ

d X

θi ϕ′ (ui )x(i) ,

i=1

so that

g (θ) = A · Pθ z,

g (β) = A · h, 87

A := 2ψ ′ (s)R,

and Gt ≤ CR2 W where W := ∥z∥2 + ∥h∥2 . The same Gaussian moment computations used in the proof of Lemma E.7 give, for k = 1, 2, 4, E[∥z∥2k | Ft ] ≤ Cdk ,

E[∥h∥2k | Ft ] ≤ Cdk .

Hence E[W k | Ft ] ≤ Cdk for k = 1, 2, 4. bound on E[Gt | Ft ]. Since Gt ≤ CR2 W , Cauchy–Schwarz gives  1/2  1/2 E[Gt | Ft ] ≤ C E[R4 | Ft ] E[W 2 | Ft ] . Using Step 1 with m = 4 and Step 2, E[Gt | Ft ] ≤ C (Q2t )1/2 (d2 )1/2 = Cd Qt . bound on E[Ξ2t | Ft ]. Since conditional variance is bounded by the second moment, h
2 i E[Ξ2t | Ft ] ≤ E ⟨g (θ) , a⟩ + ⟨g (β) , b⟩ Ft . Hence

E[Ξ2t | Ft ] ≤ 2 E[⟨g (θ) , a⟩2 | Ft ] + 2 E[⟨g (β) , b⟩2 | Ft ].

For the θ-block, a ∈ Tθ Sd−1 , so Thus

⟨Pθ z, a⟩ = ⟨z, a⟩. ⟨g (θ) , a⟩ = 2R ψ ′ (s) ⟨z, a⟩.

Therefore

 1/2  1/2 E[⟨g (θ) , a⟩2 | Ft ] ≤ C E[R4 | Ft ] E[⟨z, a⟩4 | Ft ] .

P Now ⟨z, a⟩ = i ai ϕ(ui ) is a linear statistic of iid centered variables with finite fourth moment, so Rosenthal yields E[⟨z, a⟩4 | Ft ] ≤ C∥a∥4 . Hence E[⟨g (θ) , a⟩2 | Ft ] ≤ C Qt ∥a∥2 ≤ C Q2t . For the β-block, ⟨g (β) , b⟩ = 2R ψ ′ (s) ⟨h, b⟩. Conditionally on (u1 , . . . , ud ), the scalar ⟨h, b⟩ is Gaussian with variance ∥b∥2 σ 2 ,

σ2 =

d X

θi2 ϕ′ (ui )2 .

i=1

Therefore

E[⟨h, b⟩4 | u, Ft ] ≤ C∥b∥4 σ 4 . Averaging and using E[σ 4 | Ft ] ≤ C, E[⟨h, b⟩4 | Ft ] ≤ C∥b∥4 . Thus  1/2  1/2 E[⟨g (β) , b⟩2 | Ft ] ≤ C E[R4 | Ft ] E[⟨h, b⟩4 | Ft ] ≤ C Qt ∥b∥2 ≤ C Q2t . Combining the two block estimates, E[Ξ2t | Ft ] ≤ CQ2t . 88

centered normalization fluctuation. Since

i h 2 E (Gt − E[Gt | Ft ]) Ft ≤ E[G2t | Ft ],

it suffices to bound E[G2t | Ft ]. Using Gt ≤ CR2 W , G2t ≤ CR4 W 2 . Hence, by Cauchy–Schwarz,  1/2  1/2 E[G2t | Ft ] ≤ C E[R8 | Ft ] E[W 4 | Ft ] . Applying Step 1 with m = 8 and Step 2, E[G2t | Ft ] ≤ C (Q4t )1/2 (d4 )1/2 = Cd2 Q2t . Therefore

i h 2 E (Gt − E[Gt | Ft ]) Ft ≤ Cd2 Q2t .

This proves the lemma. D.4

Strong Recovery Dynamics

With all the controlling lemmas needed, we now move to prove the strong recovery dynamics. Proposition D.24 (One local SGD epoch contracts). Assume Lemma C.15, Lemma D.22, and Lemma D.23. Let the local step size satisfy cloc δ = δloc = , Cret Csq δd ≤ µloc . d Let s be a stopping time, and suppose that, conditionally on Fs , Qs ≤ B,

0≤B≤

2 rloc . 32

Define τB := inf{t ≥ s : Qt > 2B}. Then there exist constants Cep < ∞, C < ∞, and d0 < ∞, depending only on the model constants, such that for all d ≥ d0 , with   Cep Lep := , δ one has

  B P Qs+Lep ≤ and τB > s + Lep Fs ≥ 1 − Cδ. 2 Equivalently, with probability at least 1 − Cδ, Qs+Lep ≤

B 2

and

Qt ≤ 2B

for every s ≤ t ≤ s + Lep .

Proof. If B = 0, then Qs = 0, so ys = y ∗ . By Lemma D.23, E[Gs | Fs ] ≤ Csq dQs = 0, hence Gs = 0 almost surely conditionally on Fs . Thus the SGD step is trivial at y ∗ , and the same argument iterates. The conclusion is therefore immediate. Hence assume B > 0. Throughout the proof we condition on Fs . Since the samples after time s are fresh, all martingale and conditional moment estimates hold with the same constants. Step 1: stopped multiplicative recursion. For every t < τB , one has Qt ≤ 2B. Hence distM (yt , y ∗ )2 = 2Qt ≤ 4B ≤ 89

2 rloc 2 < rloc . 8

Thus yt ∈ Urloc , and the local one-step recursion is valid. By Lemma D.22, Qt+1 ≤ Qt + δ⟨ḡt , Lt ⟩ + δΞt + Cret δ 2 Gt . Since ḡt = ∇M L(yt ), Lemma C.15 gives ⟨ḡt , Lt ⟩ = ⟨∇M L(yt ), Lt ⟩ ≤ −2µloc Qt . Define ζt := Gt − E[Gt | Ft ]. Then, for t < τB , Qt+1 ≤ Qt − 2µloc δQt + δΞt + Cret δ 2 E[Gt | Ft ] + Cret δ 2 ζt . By Lemma D.23, E[Gt | Ft ] ≤ Csq dQt . Therefore

Qt+1 ≤ 1 − 2µloc − Cret Csq δd δ Qt + δΞt + Cret δ 2 ζt . The local step-size condition implies 2µloc − Cret Csq δd ≥ µloc . Set λ := µloc , α := 1 − λδ. For d large enough, α ∈ (0, 1). Thus, for every t < τB , Qt+1 ≤ αQt + ∆Nt+1 , where

∆Nt+1 := δΞt + Cret δ 2 ζt .

Moreover, E[∆Nt+1 | Ft ] = 0. Step 2: martingale variance on the stopped epoch. By Lemma D.23, and

E[Ξ2t | Ft ] ≤ Csq Q2t , h i 2 E[ζt2 | Ft ] = E (Gt − E[Gt | Ft ]) Ft ≤ Csq d2 Q2t .

Hence, using (a + b)2 ≤ 2a2 + 2b2 , E[(∆Nt+1 )2 | Ft ] ≤ Cδ 2 Q2t + Cδ 4 d2 Q2t . Since δd is bounded under the local step-size regime, E[(∆Nt+1 )2 | Ft ] ≤ Cδ 2 Q2t . On {t < τB }, Qt ≤ 2B, and therefore E[(∆Nt+1 )2 | Ft ] ≤ Cδ 2 B 2 . Define the stopped martingale increments et+1 := 1{t<τ } ∆Nt+1 , ∆N B and for u = s, . . . , s + Lep , eu := N

u−1 X

ek+1 , ∆N

es := 0. N

k=s ep eu )s+L Then (N is a martingale, and u=s

  s+Lep −1 h i X es+L |2 Fs = E ek+1 )2 | Fk ] Fs  ≤ CLep δ 2 B 2 . E |N E[(∆N ep k=s

90

(198)

Since

 Lep =

Cep δ



C δ

≤

for d large enough, i h es+L |2 Fs ≤ CδB 2 . E |N ep By Doob’s L2 maximal inequality, "

# eu |2 Fs ≤ CδB 2 . |N

sup

E

s≤u≤s+Lep

Therefore, by Markov’s inequality, !

eu | > B P sup |N 8 s≤u≤s+Lep

Fs

≤ Cδ.

Step 3: weighted martingale control. For 1 ≤ ℓ ≤ Lep , define the stopped geometric convolution Seℓ :=

ℓ−1 X

es+j+1 . αℓ−1−j ∆N

j=0

We claim that sup

|Seℓ | ≤ 2

sup

eu |. |N

s≤u≤s+Lep

1≤ℓ≤Lep

Indeed, for fixed ℓ, write es+j , Aj := N

j = 0, . . . , ℓ.

Then A0 = 0 and es+j+1 = Aj+1 − Aj . ∆N By summation by parts, Seℓ = Aℓ − (1 − α)

ℓ−1 X

αℓ−1−j Aj .

j=1

Thus |Seℓ | ≤ sup |Aj | + (1 − α) 0≤j≤ℓ

ℓ−1 X

αℓ−1−j sup |Aj | ≤ 2 sup |Aj |. 0≤j≤ℓ

j=1

0≤j≤ℓ

Taking the supremum over ℓ proves the claim. Define the epoch good event ( Eep :=

B sup |Seℓ | ≤ 4 1≤ℓ≤Lep

) .

By the preceding maximal bound, P(Eep | Fs ) ≥ 1 − Cδ. Step 4: contraction on Eep . We prove by induction that, on Eep , Qs+ℓ ≤ αℓ B +

B 4

for every 0 ≤ ℓ ≤ Lep .

For ℓ = 0, this follows from Qs ≤ B. 91

Assume the claim has been proved up to time s + ℓ − 1, where 1 ≤ ℓ ≤ Lep . Then, for every 0 ≤ j ≤ ℓ − 1, B B Qs+j ≤ αj B + ≤B+ < 2B. 4 4 Hence τB > s + ℓ − 1, so the original increments and the stopped increments agree on the interval s, . . . , s + ℓ − 1. Iterating (198) gives Qs+ℓ ≤ αℓ Qs +

ℓ−1 X

αℓ−1−j ∆Ns+j+1 = αℓ Qs + Seℓ .

j=0

On Eep , this yields Qs+ℓ ≤ αℓ B +

B . 4

This closes the induction. In particular, Qs+ℓ < 2B

for every 0 ≤ ℓ ≤ Lep ,

and therefore τB > s + Lep . It remains to choose Cep so that the endpoint is at most B/2. Choose Cep large enough that e−λCep ≤ Since

1 . 4



 Cep Lep = , δ

α = 1 − λδ,

we have, for d large enough, αLep ≤ e−λδLep ≤ e−λCep ≤

1 . 4

Thus, on Eep , Qs+Lep ≤ αLep B +

B B B B ≤ + = . 4 4 4 2

Together with τB > s + Lep , this gives   B Eep ⊆ Qs+Lep ≤ and τB > s + Lep . 2 Since P(Eep | Fs ) ≥ 1 − Cδ, the proposition follows. Proposition D.25 (Iterated local contraction). Assume Proposition D.24. Let cloc δ = δloc = d satisfy the local step-size condition, and let   Cep Lep := δ be the epoch length from Proposition D.24. Let s be a stopping time and suppose that, conditionally on Fs , r2 Qs ≤ B0 , 0 < B0 ≤ loc . 32 For ε ∈ (0, B0 ], define   B0 Jε := log2 , Nε := Jε Lep . ε 92

Then P (Qs+Nε ≤ ε and Qt ≤ 2B0 for every s ≤ t ≤ s + Nε | Fs ) ≥ 1 − CδJε .

(199)

In particular, since δ = cloc /d, 

B0 Nε ≤ Cd 1 + log ε

 ,

(200)

and P (Qs+Nε ≤ ε | Fs ) ≥ 1 − C

log(2B0 /ε) . d

(201)

Proof. If ε = B0 , then Jε = 0, Nε = 0, and the claim follows from Qs ≤ B0 = ε. Hence assume ε < B0 , so Jε ≥ 1. Set

sj := s + jLep , Bj := 2−j B0 , j = 0, 1, . . . , Jε . Since s is a stopping time, each sj is also a stopping time. Moreover, 2 rloc , 32 so Proposition D.24 is applicable at every epoch as soon as Qsj ≤ Bj .

Bj ≤ B0 ≤

For j = 0, . . . , Jε − 1, define the epoch-success event  Aj := Qsj+1 ≤ Bj+1 and Qt ≤ 2Bj for every sj ≤ t ≤ sj+1 . Also define Gm :=

m−1 \

Aj ,

m = 1, . . . , Jε ,

j=0

with G0 equal to the whole probability space. We claim that, for every m ≤ Jε , on Gm , Qsm ≤ Bm and Qt ≤ 2B0 for every s ≤ t ≤ sm . Indeed, the first statement follows by induction from the endpoint conclusion in each Aj . The second follows because, on the j-th epoch, Qt ≤ 2Bj ≤ 2B0 . Now fix j ∈ {0, . . . , Jε − 1}. On Gj , we have Qsj ≤ Bj . Applying Proposition D.24 conditionally on Fsj gives
P Acj | Fsj ≤ Cδ on Gj . Therefore,

 P Gj ∩ Acj | Fs = E 1Gj P Acj | Fsj Fs ≤ Cδ. Taking a union bound over the Jε epochs yields
P GJcε | Fs ≤ CδJε . Hence P (GJε | Fs ) ≥ 1 − CδJε . On GJε , we have

Qs+Nε = QsJε ≤ BJε = 2−Jε B0 ≤ ε, by the definition of Jε . Also, as shown above, Qt ≤ 2B0

for every s ≤ t ≤ s + Nε .

This proves the probability estimate. 93

It remains only to record the sample complexity. Since   Cep C Lep = ≤ δ δ for d sufficiently large, and     B0 B0 Jε = log2 ≤ C 1 + log , ε ε we get Nε = Jε Lep ≤ Cδ

−1



B0 1 + log ε

 .

With δ = cloc /d, this becomes 

B0 Nε ≤ Cd 1 + log ε

 .

Finally, δJε ≤

2B0 C log , d ε

after increasing C. This gives P (Qs+Nε ≤ ε | Fs ) ≥ 1 − C

log(2B0 /ε) . d

The proof is complete. Theorem D.26 (Discrete-time weak-to-strong recovery). Assume Theorem D.17 holds with event Gweak , stopping time tesc , and constants Tesc , cesc , Cη,esc , such that on Gweak , T T tesc ≤ Tesc ,

min{ρtesc , χtesc } ≥ cesc ,

log d ηtesc ≤ Cη,esc √ . d

Assume (A4). After tesc , switch to the local step size δloc = cloc d , where cloc > 0 is chosen sufficiently small so that the assumptions of Lemma D.19, Proposition D.21, and Proposition D.24 hold. Let 2 rloc . 64 Then, for every fixed or polynomially small νin , there exists an event Estr satisfying

Qin :=

c

P(Estr ) ≥ 1 − P(Gweak ) − νin , T

(202)

on which there is a stopping time t1 such that −1 −1 t1 ≤ tesc + Cent δloc ≤ Tesc + Cent δloc ,

(203)

and Qt1 ≤ Qin ≤ Moreover, for every ε ∈ (0, Qin ], define   Qin , Jε := log2 ε

 Lep :=

2 rloc . 32

 Cep , δloc

(204)

Nε := Jε Lep .

Then c

P (Qt1 +Nε ≤ ε) ≥ 1 − P(Gweak ) − νin − Cδloc Jε . T Equivalently, since δloc = cloc /d, c

) − νin − C P (Qt1 +Nε ≤ ε) ≥ 1 − P(Gweak T

log(2Qin /ε) . d

On the same event, 1 − ρt1 +Nε ≤ ε,

1 − χt1 +Nε ≤ ε, 94

(205)

and hence min{ρt1 +Nε , χt1 +Nε } ≥ 1 − ε.

(206)

In particular,   Qin Nε ≤ Cd 1 + log . ε With the weak phase satisfies Tesc = O(d log2 d), the total number of SGD samples needed to reach Q ≤ ε is   1 O(d log2 d) + O d log . (207) ε Proof. Let t0 := tesc , δ := δloc . Work first on the weak-recovery event Gweak . On this event, T log d ηt0 ≤ Cη,esc √ . d

min{ρt0 , χt0 } ≥ cesc , Together with the incoherence bound

r ∗

∥θ ∥∞ ≤ C∗

log d , d

these are exactly the assumptions of Lemma D.19. Apply Lemma D.19 with Cin ≥ Cent . Conditionally on Ft0 , there is an event Ein such that P(Ein | Ft0 ) ≥ 1 − νin . Hence, by the tower property, h i P(Gweak ∩ E ) = E 1 P(E | F ) ≥ (1 − νin )P(Gweak ). weak in in t GT T T 0 Therefore

c

P(Gweak ∩ Ein ) ≥ 1 − P(Gweak ) − νin . T T

On Gweak ∩ Ein , Proposition D.21 gives a time T t1 ≤ tesc + Cent δ −1 such that ρt1 ≥ cos(rloc /8),

χt1 ≥ cos(rloc /8),

and consequently 2 rloc = Qin . 64 Choosing t1 as the stopping time constructed in Proposition D.21, this proves the first part. Since tesc ≤ Tesc on Gweak , we also have T

Qt1 ≤

t1 ≤ Tesc + Cent δ −1 . Set

(0)

Estr := Gweak ∩ Ein . T

Then

c

(0)

P(Estr ) ≥ 1 − P(Gweak ) − νin . T (0)

Now fix ε ∈ (0, Qin ]. On Estr , we may apply Proposition D.25 conditionally on Ft1 , with s = t1 ,

B0 = Qin .

Indeed, Qt1 ≤ Qin ≤ 95

2 rloc . 32

Thus, conditionally on Ft1 , P (Qt1 +Nε ≤ ε | Ft1 ) ≥ 1 − CδJε . Using the tower property again,   (0) P Estr ∩ {Qt1 +Nε > ε} ≤ CδJε . Therefore

(0)

P (Qt1 +Nε ≤ ε) ≥ P(Estr ) − CδJε . (0)

Combining this with the lower bound on P(Estr ) gives c

P (Qt1 +Nε ≤ ε) ≥ 1 − P(Gweak ) − νin − CδJε . T Since

cloc d   Qin 2Qin Jε = log2 ≤ C log , ε ε δ=

and

we obtain CδJε ≤ C

log(2Qin /ε) . d

It remains to convert the Q-bound into the overlap bound. If Qt1 +Nε ≤ ε, then

distS (θt1 +Nε , θ∗ )2 + distS (βt1 +Nε , β ∗ )2 = 2Qt1 +Nε ≤ 2ε.

In particular,

distS (θt1 +Nε , θ∗ )2 ≤ 2ε, Using 1 − cos x ≤ x /2, we get

distS (βt1 +Nε , β ∗ )2 ≤ 2ε.

2

1 − χt1 +Nε = 1 − cos(distS (θt1 +Nε , θ∗ )) ≤ ε, and similarly 1 − ρt1 +Nε ≤ ε. Thus min{ρt1 +Nε , χt1 +Nε } ≥ 1 − ε. Finally, Proposition D.25 gives   Qin . Nε = Jε Lep ≤ Cd 1 + log ε Since

t1 ≤ tesc + Cent δ −1 = tesc + O(d), the total number of samples after initialization is bounded by   Qin tesc + O(d) + O d log . ε 2 Because Qin = rloc /64 is a fixed constant independent of d, this is   1 tesc + O d log . ε

If Tesc = O(d log2 d), then on Gweak , T 

1 t1 + Nε ≤ O(d log d) + O d log ε 2

The proof is complete. 96

 .

E

Auxiliary Lemmas

E.1

Stein Bounds

Here, we present the comparison bounds between the Riemannian Hessian of the true loss and that of the Gaussian model via a Stein-type approximation, for θ-Hessian and β-Hessian respectively. They are used in the Lemma C.12. We prove our results based on this theorem from [8]: Theorem E.1 (Stein coupling for product measures [8, Thm. D.1]). Let X = (X1 , . . . , Xd ) be a random vector with independent coordinates such that E[X] = 0, E[XX ⊤ ] = Id , and E|Xi |3 = τi3 < ∞. Let G ∼ N (0, Id ) be standard Gaussian. If h : Rd → R is three-times differentiable, then d

5X 3 3 τ ∂xi h ∞ . 6 i=1 i

E[h(X)] − E[h(G)] ≤

(208)

Stein Bound for θ-Hessian Corollary E.2 (Block Stein coupling with moment remainder). Let (Xi , Xi∗ )di=1 be independent R2 -valued random vectors with E[Xi ] = E[Xi∗ ] = 0,

E[Xi2 ] = E[(Xi∗ )2 ] = 1,

E[Xi Xi∗ ] = m,

1 m . If and τ := maxi (E|Xi |5 + E|Xi∗ |5 ) < ∞. Let (Gi , G∗i )di=1 be iid N (0, Σm ) with Σm = m 1 h : Rd × Rd → R is three-times differentiable in each block (xi , x∗i ), then   d X X ∂ a+b h (j) ∗(j) (W , W ) , (209) E[h(X, X ∗ )] − E[h(G, G∗ )] ≤ C(τ ) sup E ∂xai ∂(x∗i )b j≤d i=1 a,b≥0 a+b=3

where (W (j) , W ∗(j) ) is the hybrid that agrees with (G, G∗ ) on blocks 1, . . . , j and with (X, X ∗ ) on blocks j+1, . . . , d. Proof. We have (W (0) , W ∗(0) ) = (X, X ∗ ) and (W (d) , W ∗(d) ) = (G, G∗ ), so d  X     E[h(X, X ∗ )] − E[h(G, G∗ )] = E h(W (j−1) , W ∗(j−1) ) − E h(W (j) , W ∗(j) ) .

(210)

j=1

Each summand differs only in block j. With all other blocks frozen, define φj (s, t) :=
h W (j) (s,t) in block j , so the j-th summand equals E[φj (Xj , Xj∗ ) − φj (Gj , G∗j )]. Taylor-expand φj around (0, 0) to second order with integral remainder: X 1 ∂ a+b φj φj (s, t) = sa tb + Rj (s, t), a! b! ∂sa ∂tb (0,0) a,b≥0 a+b≤2

where the integral form of the remainder gives |Rj (s, t)| ≤ C0 (|s| + |t|)3

X a,b≥0 a+b=3

∂ a+b φj (ξ, ζ) . a b |ξ|≤|s|, |ζ|≤|t| ∂s ∂t sup

Taking E[φj (Xj , Xj∗ )] − E[φj (Gj , G∗j )]: the zeroth-order terms cancel trivially; the first-order terms cancel since E[Xj ] = E[Gj ] = 0 and E[Xj∗ ] = E[G∗j ] = 0; the second-order terms cancel since E[Xj2 ] = E[G2j ] = 1,

E[(Xj∗ )2 ] = E[(G∗j )2 ] = 1,

E[Xj Xj∗ ] = E[Gj G∗j ] = m.

Hence only the remainder contributes:  
X h ∂ a+b φj i E φj (Xj , Xj∗ ) − φj (Gj , G∗j ) ≤ C0 E(|Xj |+|Xj∗ |)3 + E(|Gj |+|G∗j |)3 E , ∂sa ∂tb a,b≥0 a+b=3

where the moment prefactors are bounded by C ′ (τ ), and the expectations of the third-order partials use Hölder to separate the (s, t)-dependence from the frozen blocks. Summing over j and noting ∂ a+b φj /∂sa ∂tb = ∂ a+b h/∂xaj ∂(x∗j )b evaluated at the hybrid yields (209). 97

Corollary E.3 (Stein estimate for a weighted quadratic statistic). Let X1 , . . . , Xd be independent, centered, unit-variance random variables with supi E|Xi |8 < ∞, and let Γ ∼ N (0, Id ). Let a, u ∈ Rd satisfy ∥a∥2 = ∥u∥2 = 1. If F ∈ C 3 (R) and ∥F (k) ∥∞ < ∞ for 1 ≤ k ≤ 3, then ! d d d X X X     2 2 2 2 3 E F (a · X)(u · X) − E F (a · Γ)(u · Γ) ≤ CF |ai |ui + ai |ui | + |ai | , i=1

i=1

i=1

where CF depends only on F and the law of X1 , but not on d, a, u. In particular,     E F (a · X)(u · X)2 − E F (a · Γ)(u · Γ)2 ≤ CF ∥a∥∞ . Proof. Apply the scalar-block moment version of the Lindeberg–Stein comparison, as in Corollary E.2, to h(x) := F (a · x)(u · x)2 . For each coordinate i, ∂i3 h(x) = a3i F (3) (a · x)(u · x)2 + 6a2i ui F ′′ (a · x)(u · x) + 6ai u2i F ′ (a · x). At every hybrid vector appearing in the Lindeberg interpolation, E|u · W |2 ≤ C∥u∥22 = C,

E|u · W | ≤ C,

because the coordinates have uniformly bounded second moments. Hence
E|∂i3 h(W )| ≤ CF |ai |3 + a2i |ui | + |ai |u2i . Summing over i gives the first estimate. For the final bound, use ∥a∥2 = ∥u∥2 = 1: !1/2 X

X

|ai |u2i ≤ ∥a∥∞ ,

i

X

a2i |ui | ≤

i

and

X

a4i

∥u∥2 ≤ ∥a∥∞ ,

i

|ai |3 ≤ ∥a∥∞

X

i

a2i = ∥a∥∞ .

i

This proves the claim. Lemma E.4 (Teacher-level θ-Hessian comparison). Let Aψ := E[ψ ′ (G)2 ], 2

G ∼ N (0, 1).

8

Assume E[ϕ(G)] = 0, E[ϕ(G) ] = 1, E|ϕ(G)| < ∞, and ψ ∈ C 4 (R) with bounded derivatives up to order 4. Then there exists Cθ > 0, independent of d, such that for every u ∈ Tθ∗ Sd−1 , HessS,θ L(y ∗ )[u, u] − 2Aψ ∥u∥2 ≤ Cθ ∥θ∗ ∥∞ ∥u∥2 . Equivalently, |HessS,θ L(y ∗ )[u, u] − HessS,θ LG (y ∗ )[u, u]| ≤ Cθ ∥θ∗ ∥∞ ∥u∥2 . Proof. By quadratic scaling, it is enough to prove the result for ∥u∥ = 1. Set a := θ∗ ,

Xi := ϕ(Gi ),

S := a · X =

d X i=1

ai Xi ,

Au := u · X =

d X

ui Xi ,

i=1

iid

where Gi ∼ N (0, 1). Also define

F := ψ ′2 .

Since u ∈ Tθ∗ Sd−1 , we have a · u = 0. At the teacher, the residual vanishes pointwise. Therefore, along a spherical geodesic in the θ-variable with initial velocity u,   HessS,θ L(y ∗ )[u, u] = 2E F (S)A2u . 98

Let Γ ∼ N (0, Id ). By Corollary E.3,     E F (a · X)(u · X)2 − E F (a · Γ)(u · Γ)2 ≤ CF ∥a∥∞ = CF ∥θ∗ ∥∞ . Since a · u = 0 and ∥a∥ = ∥u∥ = 1, the Gaussian random variables a · Γ and u · Γ are independent standard normals. Hence   E F (a · Γ)(u · Γ)2 = E[F (G)]E[G2 ] = Aψ . Thus

|HessS,θ L(y ∗ )[u, u] − 2Aψ | ≤ Cθ ∥θ∗ ∥∞

for every unit u ∈ Tθ∗ Sd−1 . Homogeneity gives the claim for arbitrary u. Finally, the Gaussian teacher Hessian satisfies HessS,θ LG (y ∗ )[u, u] = 2Aψ ∥u∥2 , because in the Gaussian model the same computation gives 2E[F (a · Γ)(u · Γ)2 ], and a · Γ is independent of u · Γ. This proves the comparison with LG . Stein Bound for β-Hessian To give a Stein bound for β-Hessian, we start by deriving this corollary from Theorem E.1. Corollary E.5 (One-dimensional linear statistic). Assume the setting of Theorem E.1. Let a ∈ Rd and f ∈ C 3 (R) with ∥f (3) ∥∞ < ∞. Define h(x) := f (⟨a, x⟩). Then E[f (⟨a, X⟩)] − E[f (⟨a, G⟩)] ≤

d X 5 (3) ∥f ∥∞ τi3 |ai |3 . 6 i=1

In particular, if supi τi ≤ τ and ∥a∥2 = 1, then

since

E[f (⟨a, X⟩)] − E[f (⟨a, G⟩)] ≤ C(τ, f ) ∥a∥∞ , P 3 2 i |ai | ≤ ∥a∥∞ i ai = ∥a∥∞ .

(211)

P

Proof. Define h(x) := f (⟨a, x⟩). By the chain rule, all partial derivatives act only through ⟨a, x⟩: ∂3h (x) = a3i f ′′′ (⟨a, x⟩). ∂x3i Therefore

∂3h = |ai |3 ∥f (3) ∥∞ . ∂x3i ∞

Applying Theorem E.1 directly gives E[f (⟨a, X⟩)] − E[f (⟨a, G⟩)] ≤

d d X 5 5X 3 3 τi ∂xi h ∞ = ∥f (3) ∥∞ τi3 |ai |3 . 6 i=1 6 i=1

If supi τi ≤ τ and ∥a∥2 = 1, then d X

τi3 |ai |3 ≤ τ 3

i=1

d X

3

3

|ai | ≤ τ ∥a∥∞

i=1

d X

a2i = τ 3 ∥a∥∞ ,

i=1

giving C(τ, f ) = 65 τ 3 |f (3) |∞ , so we obtained (211). Lemma E.6 (Teacher-level β-Hessian comparison). Let Aψ := E[ψ ′ (G)2 ],

Bϕ := E[ϕ′ (G)2 ] = κ′ (1),

G ∼ N (0, 1).

Assume (A3). Then there exists a constant Cβ > 0, independent of d, such that for every v ∈ Tβ ∗ Sd−1 , HL,ββ (y ∗ )[v, v] − 2Aψ Bϕ ∥v∥2 ≤ Cβ ∥θ∗ ∥∞ ∥v∥2 . Equivalently,

|HL,ββ (y ∗ )[v, v] − HLG ,ββ (y ∗ )[v, v]| ≤ Cβ ∥θ∗ ∥∞ ∥v∥2 . 99

Proof. By quadratic scaling, it is enough to consider ∥v∥ = 1. Set ai := θi∗ ,

qi := ϕ′ (Gi )2 ,

Xi := ϕ(Gi ),

S :=

d X

S−i := S − ai Xi ,

ai Xi ,

i=1

and define

F (s) := ψ ′ (s)2 .

By Lemma C.11, d X 1 HL,ββ (y ∗ )[v, v] = a2i E[F (S)qi ]. 2 i=1

We first decouple the i-th coordinate. Since S−i is independent of (Xi , qi ), E[F (S−i )qi ] = E[F (S−i )]E[qi ] = Bϕ E[F (S−i )]. Moreover, by the boundedness of F ′ and the moment assumption, |E[F (S)qi ] − E[F (S−i )qi ]| ≤ ∥F ′ ∥∞ |ai |E[|Xi |qi ] ≤ C|ai |. Multiplying by a2i and summing gives d X

a2i E[F (S)qi ] = Bϕ

i=1

Since

P

d X

a2i E[F (S−i )] + O

i=1

2 i ai = 1 and

we also have

d X

! |ai |

3

.

i=1

|E[F (S)] − E[F (S−i )]| ≤ ∥F ′ ∥∞ |ai |E|Xi |, d X

d X

a2i E[F (S−i )] = E[F (S)] + O

i=1

! |ai |

3

.

i=1

Therefore, d X

a2i E[F (S)qi ] = Bϕ E[F (S)] + O

i=1

d X

! |ai |3

.

i=1

Now apply the one-dimensional Stein comparison in Corollary E.5 with ai := θi∗ ,

Xi := ϕ(Gi ),

f := F = ψ ′2 .

By the assumptions on ϕ, E[Xi2 ] = 1,

E[Xi ] = 0,

τϕ3 := E|Xi |3 < ∞,

and by the regularity assumptions on ψ, F ∈ C 3 (R) with ∥F (3) ∥∞ < ∞. Moreover ∥a∥2 = ∥θ∗ ∥2 = 1. Let Γ = (Γ1 , . . . , Γd ) ∼ N (0, Id ). Then ⟨a, Γ⟩ ∼ N (0, 1). Hence Corollary E.5 gives ! d d X X 5 EF ai Xi − EF (⟨a, Γ⟩) ≤ ∥F (3) ∥∞ τϕ3 |ai |3 . 6 i=1 i=1 Since ⟨a, Γ⟩ ∼ G ∼ N (0, 1), this becomes |E[F (S)] − E[F (G)]| ≤ CF

d X

|ai |3 ,

i=1

where CF := 65 τϕ3 ∥F (3) ∥∞ . Finally, d X i=1

Therefore

|ai |3 ≤ ∥a∥∞

d X

a2i = ∥θ∗ ∥∞ .

i=1

|E[F (S)] − E[F (G)]| ≤ CF ∥θ∗ ∥∞ . 100

Combining this with the preceding leave-one-out estimate yields d X

a2i E[F (S)qi ] = Bϕ E[F (G)] + O(∥θ∗ ∥∞ ) = Bϕ Aψ + O(∥θ∗ ∥∞ ).

i=1

Multiplying by 2∥v∥2 proves the true ββ-block estimate. For the Gaussian loss, since we have (61) in Lemma C.10, X
LG (θ, β) = 2 αk2 1 − (κ(ρ)χ)k k≥1 ∗

, gives, at y ,  HLG ,ββ (y ∗ )[v, v] = 2κ′ (1) 

 X

kαk2  ∥v∥2 .

k≥1

But

X

kαk2 = E[ψ ′ (G)2 ] = Aψ ,

κ′ (1) = Bϕ .

k≥1

Thus

HLG ,ββ (y ∗ )[v, v] = 2Aψ Bϕ ∥v∥2 ,

and the comparison follows. E.2

Gradient Bounds

In the lemma below, we provide bounds for gradients g (θ) , g (β) to help us to control the evolution of Zt and ηt in weak recovery, and that of Qt in strong recovery. Lemma E.7 (Technical gradient bounds for blockwise spherical SGD). Assume (A3). Then there exists a constant K < ∞, depending only on Lϕ , Lψ , ϕ(0), such that if we draw a fresh sample xt then with g (β) (β, θ; xt ) := ∇Sβ ℓ(β, θ; xt ), g (θ) (β, θ; xt ) := ∇Sθ ℓ(β, θ; xt ), we have the moment bounds   sup E ∥g (β) (β, θ; xt )∥2 ≤ Kd, β,θ

  sup E ∥g (θ) (β, θ; xt )∥2 ≤ Kd, β,θ

  sup E ∥g (β) (β, θ; xt )∥4 ≤ Kd2 , β,θ

  sup E ∥g (θ) (β, θ; xt )∥4 ≤ Kd2 . β,θ

Moreover, for the centered gradient-noise fields defined as ∆g (β) (β, θ; xt ) := g (β) (β, θ; xt ) − E[g (β) (β, θ; xt )], ∆g (θ) (β, θ; xt ) := g (θ) (β, θ; xt ) − E[g (θ) (β, θ; xt )], we have the directional variance bounds   sup sup E (∆g (β) (β, θ; xt ) · u)2 ≤ K, β,θ ∥u∥=1

  sup sup E (∆g (θ) (β, θ; xt ) · u)2 ≤ K. β,θ ∥u∥=1

Consequently, along the SGD trajectory (βt , θt ), since xt is independent of Ft , the same bounds hold conditionally on Ft :  (β)   (β)    (β) E ∥gt ∥2 | Ft ≤ Kd, E ∥gt ∥4 | Ft ≤ Kd2 , sup E ((∆gt ) · u)2 | Ft ≤ K, ∥u∥=1

and similarly for the θ-block. 101

Proof. Since spherical gradient is the orthogonal projection of the Euclidean gradient onto the tangent space via ∇S F (x) = ∇F (x) − (∇F (x) · x)x, we have ∥∇S F (x)∥ ≤ ∥∇F (x)∥ and |u · ∇S F (x)| ≤ |u · ∇F (x)| ≤ ∥∇F (x)∥ for |u| = 1, it suffices to bound the Euclidean gradients. Based on the model definition in Section 3.1 and define the per-sample loss as l = (f − f ∗ )2 ,

(212)

we have the per-sample Euclidean gradients w.r.t. each layer as: ∇θ l = 2(f − f ∗ ) · ψ ′ (s) · z

(213)

d X

∇β l = 2(f − f ∗ ) · ψ ′ (s) ·

θi ϕ′ (ui )x(i) .

(214)

i=1

We simplify the expression by defining A(x; β, θ) := 2(f − f ∗ ) · ψ ′ (s) ∈ R.

(215)

θ bounds E[∥g (θ) ∥2 ] ≤ E[∥∇θ l∥2 ] 2

(216)

2

= E[|A| ∥z∥ ] 4 1/2

≤ (E|A| )

(217) 4 1/2

· (E∥z∥ )

by Cauchy-Schwarz

(218)

′

For the first term of (218), with the bounded ψ, ψ assumption, we have |A| = 2|ψ(s) − ψ(s∗ )||ψ ′ (s)| ≤ 2(|ψ(s)| + |ψ(s∗ )|)|ψ ′ (s)| ≤ 4Cψ Cψ′ .

(219)

For the second term, given z has i.i.d. entries, we have E[∥z∥4 ] = E

d 2 X zi2

(220)

i=1 d d  X
X
=E zi2 zi2 ] i=1

(221)

i=1

d X  X zi4 + zi2 zj2 =E i=1

=

d X

(222)

i̸=j

E[zi4 ] +

i=1

X

E[zi2 ]E[zj2 ]

(223)

i̸=j

= dµ4 + d(d − 1)µ22

(224)

2

= Θ(d ).

(225)

Therefore, there exists a unifying constant K such that  2  1/2 E[∥g (θ) ∥2 ]| ≤ 4Cψ Cψ′ Θ(d2 ) = Kd.

(226) 4

Similarly, to bound the forth moment, using the Minkowski inequality in L and the i.i.d. of {zi }, we have d d d X X 4
1/4 X
1/4 ∥ zi2 ∥L4 = E zi2 ≤ E[zi8 ] = d(E[zi8 ])1/4 i=1

i=1

i=1

Raise both sides to the 4th power, we have 8

E[∥z∥ ] = E

d X

zi2

4

≤ d4 E[z18 ] = O(d4 ).

i=1

Thus, with the same, finite A, there exists a unifying constant K such that E[∥g (θ) ∥4 ] ≤ E[|A|4 ∥z∥4 ] ≤ E[|A|8 ]1/2 E[∥z∥8 ]1/2 ≤ Kd2 102

(227)

θ-block directional variance. Fix (β, θ) ∈ Sd−1 × Sd−1 , and condition on Ft , so that (β, θ) are deterministic. For a fresh sample xt , write (i)

ui := xt · β,

zi := ϕ(ui ),

Also write

z := (z1 , . . . , zd ),

s := z · θ.

At := 2(ft − ft∗ )ψ ′ (s).

By the boundedness of ψ and ψ ′ , |At | ≤ CA . The Euclidean θ-gradient is At z, hence the spherical θ-gradient is (θ)

gt

= Pθ⊥ (At z) = At Pθ⊥ z,

Pθ⊥ := I − θθ⊤ .

Let v ∈ Rd be any unit vector. Since Pθ⊥ is self-adjoint, (θ)

⟨gt , v⟩ = At ⟨z, Pθ⊥ v⟩. Set ve := Pθ⊥ v. Then ∥e v ∥2 ≤ 1. Since E[ϕ(G)] = 0 and E[ϕ(G)2 ] < ∞, the coordinates zi = ϕ(ui ) are iid centered variables with bounded second moment. Therefore   E ⟨z, ve⟩2 Ft = E[ϕ(G)2 ]∥e v ∥22 ≤ C. Using |At | ≤ CA , h i (θ) E ⟨gt , v⟩2 Ft ≤ C. Hence

h i h i
2 (θ) (θ) (θ) E ⟨gt , v⟩ − E[⟨gt , v⟩ | Ft ] Ft ≤ E ⟨gt , v⟩2 Ft ≤ C.

Equivalently, h i
2 (θ) sup E (∆gt ) · v Ft ≤ C.

∥v∥2 =1

β bounds

Recall

∇β l = 2(f − f ∗ ) ψ ′ (s)

d X

θi ϕ′ (ui ) x(i) ,

ui := x(i) · β,

s := z · θ, zi = ϕ(ui ),

i=1

and define the scalar

A(x; β, θ) := 2(f − f ∗ )ψ ′ (s).

By boundedness of ψ, ψ ′ , we have the uniform bound
|A| ≤ 2 |ψ(s)| + |ψ(s∗ )| |ψ ′ (s)| ≤ 4Cψ Cψ′ =: CA . Let Pβ := I − ββ ⊤ be the orthogonal projector onto β ⊥ . The β-spherical gradient is g (β) (β, θ; x) = ∇Sβ l(β, θ; x) = Pβ ∇β l = A

d X

θi ϕ′ (ui ) Pβ x(i) .

i=1

Now exploit the Gaussian block structure: each block x(i) ∼ N (0, Id ) independently. Decompose x(i) = ui β + vi ,

ui := x(i) · β ∼ N (0, 1),

vi := Pβ x(i) ∼ N (0, Pβ ),

and for each i, ui is independent of vi (and the pairs (ui , vi ) are i.i.d. over i). Hence g

(β)

=A

d X

θi ϕ′ (ui ) vi .

i=1

103

Conditioned on (u1 , . . . , ud ), the vector S :=

d X

θi ϕ′ (ui ) vi

i=1

is Gaussian with mean 0 and covariance Cov(S | u) =

d X

 θi2 ϕ′ (ui )2 Pβ .

i=1

Write σ 2 (u) :=

d X

θi2 ϕ′ (ui )2 .

i=1

Since rank(Pβ ) = d − 1, standard Gaussian moment formulas give     E ∥S∥2 | u = (d − 1)σ 2 (u), E ∥S∥4 | u = (d − 1)(d + 1)σ 4 (u). At this point we use the regularity of ϕ. For instance, if ϕ is Lϕ -Lipschitz (as in our regularity assumptions), then |ϕ′ | ≤ Lϕ a.e., so σ 2 (u) ≤ L2ϕ

d X

θi2 = L2ϕ ,

σ 4 (u) ≤ L4ϕ .

i=1

Therefore,

  2 2 E∥g (β) ∥2 = E |A|2 ∥S∥2 ≤ CA E∥S∥2 ≤ CA (d − 1)L2ϕ ≤ Kd,

and similarly   4 4 E∥g (β) ∥4 = E |A|4 ∥S∥4 ≤ CA E∥S∥4 ≤ CA (d − 1)(d + 1)L4ϕ ≤ Kd2 . Finally, we bound the directional variance of the centered noise field. Fix any unit vector q ∈ Rd . Conditioned on u, (q · S) is Gaussian with variance Var(q · S | u) = q ⊤ Cov(S | u) q = σ 2 (u) q ⊤ Pβ q ≤ σ 2 (u). Thus,

  2 2 2 E[(q · g (β) )2 ] = E |A|2 (q · S)2 ≤ CA E[σ 2 (u)] ≤ CA Lϕ .

Since Var(X) ≤ E[X 2 ], we get     2 2 sup E ((∆g (β) ) · q)2 ≤ sup E (g (β) · q)2 ≤ CA Lϕ ≤ K, ∥q∥=1

∥q∥=1

which proves the claimed directional variance bound (absorbing constants into K). Along the SGD trajectory (βt , θt ) is Ft -measurable and xt is independent of Ft , hence the above bounds apply conditionally on Ft by evaluation at (β, θ) = (βt , θt ). E.3

Summation by Parts

The following estimates are deterministic and pathwise; no martingale property is used. They are stated on subintervals because the SGD proof applies them after stopping times. Lemma E.8 (Summation by parts for monotone weights). Let (Nt )0≤t≤T be a real-valued sequence and set εt := Nt+1 − Nt . Fix integers 0 ≤ r < t ≤ T , and let (bk )t−1 k=r be monotone with 0 ≤ bk ≤ 1. If sup |Nu | ≤ M, 0≤u≤T

then

t−1 X

bk εk ≤ 2M.

k=r

104

Proof. Abel summation gives t−1 X

bk (Nk+1 − Nk ) = bt−1 Nt − br Nr −

k=r

t−2 X

Nk+1 (bk+1 − bk ).

k=r

Hence

t−1 X

bk ε k ≤ M

bt−1 + br +

k=r

t−2 X

! |bk+1 − bk | .

k=r

If (bk ) is increasing, then t−2 X

|bk+1 − bk | = bt−1 − br ,

k=r

so the coefficient is 2bt−1 ≤ 2. If (bk ) is decreasing, then t−2 X

|bk+1 − bk | = br − bt−1 ,

k=r

so the coefficient is 2br ≤ 2. This proves the claim. Corollary E.9 (Summation by parts for decreasing geometric weights). Let (Nt )0≤t≤T be a realvalued sequence and set εt := Nt+1 − Nt . Let γ > 1. If sup |Nu | ≤ M, 0≤u≤T

then for every 0 ≤ r < T and every s ∈ {1, . . . , T − r}, s−1 X

γ −(k+1) εr+k ≤ 2M.

k=0

In particular, if M = 2λ, the bound is 4λ. If moreover r = 0 and N0 = 0, then the bound improves to M . Proof. Apply Lemma E.8 with br+k := γ −(k+1) ,

0 ≤ k ≤ s − 1.

These weights are decreasing and lie in [0, 1], so the first bound follows. If r = 0 and N0 = 0, then the initial boundary term in Abel summation vanishes: s−1 s−2   X X γ −(k+1) (Nk+1 − Nk ) = γ −s Ns − Nk+1 γ −(k+2) − γ −(k+1) . k=0

k=0

Therefore s−1 X

γ

−(k+1)

εk ≤ M

γ

−s

! s−2   X −(k+1) −(k+2) + γ −γ = M γ −1 ≤ M.

k=0

k=0

Corollary E.10 (Summation by parts for increasing geometric weights). Let (Nt )0≤t≤T be a realvalued sequence and set εt := Nt+1 − Nt . Let α ∈ (0, 1). If sup |Nu | ≤ M, 0≤u≤T

then for every 0 ≤ r < t ≤ T , t−1 X

αt−1−k εk ≤ 2M.

k=r

In particular, if M = λ, the bound is 2λ, and if M = 2λ, the bound is 4λ. 105

Proof. Apply Lemma E.8 with bk := αt−1−k ,

r ≤ k ≤ t − 1.

Since α ∈ (0, 1), these weights are increasing in k and lie in [0, 1]. The claim follows.

106

NeurIPS Paper Checklist 1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? Answer: [Yes] Justification: We provide proofs of the main claims and clearly distinguish between claims and open questions. Guidelines: • The answer [N/A] means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A [No] or [N/A] answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 2. Limitations Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: In the final conclusion section Guidelines: • The answer [N/A] means that the paper has no limitation while the answer [No] means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate “Limitations” section in their paper. • The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren’t acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 3. Theory assumptions and proofs Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] 107

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

Answer: [Yes] Justification: We provide the code used to run the experiments. Guidelines: • The answer [N/A] means that paper does not include experiments requiring code. • Please see the NeurIPS code and data submission guidelines (https://neurips.cc/ public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so [No] is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark). • The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //neurips.cc/public/guides/CodeSubmissionPolicy) for more details. • The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc. • The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable). • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 6. Experimental setting/details Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer) necessary to understand the results? Answer: [Yes] Justification: Guidelines: • The answer [N/A] means that the paper does not include experiments. • The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 7. Experiment statistical significance Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: Guidelines: • The answer [N/A] means that the paper does not include experiments. • The authors should answer [Yes] if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. 109

• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g., negative error rates). • If error bars are reported in tables or plots, the authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 8. Experiments compute resources Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] Justification: Guidelines: • The answer [N/A] means that the paper does not include experiments. • The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn’t make it into the paper). 9. Code of ethics Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] Justification: Guidelines: • The answer [N/A] means that the authors have not reviewed the NeurIPS Code of Ethics. • If the authors answer [No], they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [N/A] Justification: Guidelines: • The answer [N/A] means that there is no societal impact of the work performed. • If the authors answer [N/A] or [No], they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. 110

• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate Deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 11. Safeguards Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pre-trained language models, image generators, or scraped datasets)? Answer: [N/A] Justification: Guidelines: • The answer [N/A] means that the paper poses no such risks. • Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [N/A] Justification: Guidelines: • The answer [N/A] means that the paper does not use existing assets. • The authors should cite the original paper that produced the code package or dataset. • The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset. • For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 111

• If this information is not available online, the authors are encouraged to reach out to the asset’s creators. 13. New assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [N/A] Justification: Guidelines: • The answer [N/A] means that the paper does not release new assets. • Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 14. Crowdsourcing and research with human subjects Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)? Answer: [N/A] Justification: Guidelines: • The answer [N/A] means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 15. Institutional review board (IRB) approvals or equivalent for research with human subjects Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [N/A] Justification: Guidelines: • The answer [N/A] means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 16. Declaration of LLM usage 112

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigor, or originality of the research, declaration is not required. Answer: [N/A] Justification: Guidelines: • The answer [N/A] means that the core method development in this research does not involve LLMs as any important, original, or non-standard components. • Please refer to our LLM policy in the NeurIPS handbook for what should or should not be described.

113

Related documents

• Localized Diffusion Models#1063752
• On the independence of the slow and fast scales in multiple-scale expansions, with application to Van der Pol's equation#1076914
• Do Flatter Minima Drive Better Generalization? An Algorithmic Separation in Grokking#1248133
• Broken scale symmetries in undercomplete linear autoencoders#1188351
• Broken scale symmetries in undercomplete linear autoencoders#1181633