Improved high-dimensional estimation with Langevin dynamics and stochastic weight averaging
Authors: Stanley Wei
Alex Damian
Jason D. Lee
Venue: ICLR 2026 Poster (2026)
Keywords: Learning theory
high-dimensional statistics
non-convex optimization
Primary area: learning theory

Abstract: Significant recent work has studied the ability of gradient descent to recover a hidden planted direction $\theta^\star \in S^{d-1}$ in different high-dimensional settings, including tensor PCA and single-index models. The key quantity that governs the ability of gradient descent to traverse these landscapes is the information exponent $k^\star$ (Ben Arous et al., (2021)), which corresponds to the order of the saddle at initialization in the population landscape. Ben Arous et al., (2021) showed that $n \gtrsim d^{\max(1, k^\star-1)}$ samples were necessary and sufficient for online SGD to recover $\theta^\star$, and Ben Arous et al., (2020) proved a similar lower bound for Langevin dynamics. More recently, Damian et al., (2023) showed it was possible to circumvent these lower bounds by running gradient descent on a smoothed landscape, and that this algorithm succeeds with $n \gtrsim d^{\max(1, k^\star/2)}$ samples, which is optimal in the worst case. This raises the question of whether it is possible to achieve the same rate without explicit smoothing. In this paper, we show that Langevin dynamics can succeed with $n \gtrsim d^{ k^\star/2 }$ samples if one considers the average iterate, rather than the last iterate. The key idea is that the combination of noise-injection and iterate averaging is able to emulate the effect of landscape smoothing. We apply this result to both the tensor PCA and single-index model settings. Finally, we conjecture that minibatch SGD can also achieve the same rate without adding any additional noise.

Published as a conference paper at ICLR 2026
IMPROVED HIGH-DIMENSIONAL ESTIMATION WITH LANGEVIN DYNAMICS AND STOCHASTIC WEIGHT AV- ERAGING
Stanley Wei Princeton University
Alex Damian
Jason D. Lee UC Berkeley
Harvard University
ABSTRACT
Significant recent work has studied the ability of gradient descent recover a hidden planted direction 0* E 5d-1 in different high-dimensional settings, including tensor PCA and single-index models. The key quantity that governs the ability of gradient descent to traverse these landscapes is the information exponent k* (Ben Arous et al. (202I). which corresponds to the order of the saddle at initialization in the population landscape. [Ben Arous et al. (2021) showed that n 2 d'max(1,k" -1) samples were necessary and sufficient for online SGD recover 8", and Ben Arous et al. (2020) proved similar lower bound for Langevin dynamics. More recently, Damian et al. (203)) showed it was possible eircumvent these lower bounds by running gradient descent on a smoothed landscape, and that this algorithm succeeds with 2 d'maa(1,k" /2) samples, which is optimal in the worst case. This raises the question of whether it is possible show that
achieve the same rate without explicit smoothing. In this paper, succeed with 2 d 12 samples if one considers the
Langevin dynamics average iterate, rather than the last iterate. The key idea is that the combination of noise-injection and iterate averaging is able emulate the effect of landscape smoothing. We apply this result to both the tensor PCA and single-index model settings. Finally, we conjecture that minibatch SGD can also achieve the same rate without adding any additional noise.
INTRODUCTION
In many learning settings, gradient descent is the default algorithm, and recent years have significant progress in understanding its theoretical properties and learnability guarantees different feature learning settings (Damian et al.) 2022 Mei et al., 2022). While the optimization process is non-convex in general, there are many settings in which we can nonetheless tractably give learning guarantees. Single index models, or functions of the form (0*-2)) provide one such sandbox; here, the goal is to recover this planted direction 8" E gd-1 through which the target depends on the imput. In the statisties literature, single index models have been studied for decades (Hristache et al.) [000T] Hirdle et al.] 2004). and are also known as generalized linear models. In the special case where monotonic, the information- thhooretic sample complexity of n > learn
the link function 0* is achieved via perceptron-like algorithms (Kalai and Sastry 2009. Kakade et al., 2011) For non non-monotonic link functions, one classic example is the phase-retricval problem where (t) = |t]. which has been well-studied (Chen al. 2019 Maillard et al. (020).
For the case of Gaussian input data, the information exponent A of the link function tells us the sample complexity needed to learn 0* with "correlational learners" (Ben Arous et al. 2021). This
be extended allow for label preprocessing (Mondelli and Montanari 2018
Maillard et al. 2020; Chen et al., 2025 Dandi et al. 2024 Troiani et al.] 2024 Lee et al. 2024
Arnaboldi et 2024) and the resulting exponent becomes the "generative exponent" (Damian et al.
2024). Ben Arous al. (2021) shows that using dk -1 samples is necessary and sufficient for certain class of online stochastic gradient descent (SGD) algorithms. Damian et al. (2023) improves this to n 2 (max(11," /2) samples by running online SGD on a smoothed loss, and they provide a
matching correlational statistical query (CSQ) lower bound. Key their analysis is the fact that the

---

Published as a conference paper at ICLR 2026
smoothed loss boosts the signal-to-noke ratio in the region near initialization (i.e. when the current iterate lies in the equatorial region with respect to 8*).
Overall, the information exponent has been shown to determine the sample complexity in many settings (Ben Arous et al. 202T, Damian et al. 2023 Bietti al. [022] Abbe et al. 2023 Dandi et al. 2023). A recent work Joshi et al. (2025) analyzes the spherical symmetric distribution case, which slightly relaxes the Gaussian data assumption. In particular, the work by Abbe al. (2023) provides a generalization of the information exponent to the multi index setting, in which the target depends on low dimensional subspace of the input instead of just a single direction (Ren and Lee 2024 Damian et al. 2025). We would also like to note the connection of learning information exponent & single index models the order k tensor PCA problem (Montanari and Richard 2014). In both problems, it turns out that the partial trace estimator returns the direction of the planted spike with optimal sample complexity of dd//2 in the CSQ framework, and similar smoothing- bassdd approaches there (Anandkumar et 2017 Biroli et al. 2020) have been proposed to return this estimator.
Notably, along this line of work, Ben Arous et al. (020) conjectures that Langevin dynamics in the tensor PCA setting does not work due to the divergence of the computational-statitti setting. In our work, we surprisingly show that Langevin dynamics can still be used to recover the planted direction of the single index model. To achieve this, run Langevin dynamics, but take the time average of all the iterates. Our analysis reveals that with Z dlk /21 samples, we are able to recover the direction of the partial trace estimator and hence 0* The key insight is that this Langevin dynamics process closely tracks the Brownian motion on the sphere, and averaging out the iterates roughly corresponds to an ergodicity concentration argument on the sphere. Our main theorem the following.
gap in this
Theorem 1 (Main theorem (informal)) Consider link function with information exponent k*. Then, with n 2 d samples drawn i.i.d. from the standard d-dimensional Gaussian, rumning Algorithm 7 recovers the ground truth direction H*
We can also shave off a factor vd to improve the sample complexity to n Z dk /2 by running Algorithm| I and running online SGD on the returned time averaged estimator. This corresponds to the warm start in Damian et al. 2023) for the odd case.
2 SETUP AND MAIN CONTRIBUTIONS
2.1 NOTATION
We use ||- IIp to denote the vector fn-norm; furthermore, when p === 2, often drop the subscript and write -1. Given a probability measure 7 over IRd, we denote (IR, 7) the space of meeasurabll and square-integral functions; we shorthand this to L²())) when the domain is clear. For f E (7),,, we denote | (=)}). We also denote μ to be the uniform measure on 5d-1
2.2 SETTING
We consider in this paper tensor PCA (Montanari and Richard 2014 and single-index models.
2.2.1 TENSOR PCA
For tensor PCA, we will assume there is planted direction 0' E 5d-1 and we observe the k-tensor T defined by:
T = 0* n---- -1/2 z where
Zi.... ltd. N(0,
We consider optimizing the negative log-likelihood:
L(0) -
Information theoretically, 8* is possible to recover whenever n Z d. However, common techniques like approximate message passing (AMP), tensor power method, and online SGD require 2 k-1
2

---

Published as a conference paper at ICLR 2026
to recover 0* [Montanari and Richard 2014 Ben Arous al. 2021). Nevertheless, it possible
to recover /* with 2 dh/²2 samples using tensor unfolding (Montanari and Richard, 2014), the partial-trace estimator (Hopkins et al. 2016), and landscape smoothing (Anandkumar a 2017
Biroli et al., 2020) Damian et al., 2023 J. In our paper, show Langevin dynamics combined with iterate averaging can recover 0* with dl without explicit unfolding smoothing.
2.2.2 SINGLE-INDEX MODELS
We mostly follow the setting of Damian et al. (2023). Let (x44 R Re be the set of training data. The input data are drawn i.i.d. from standard d-dimensional Gaussian N0, Ia), and the labels yi are generated through target or teacher function f*. In particular, we consider the setting where f is a single index model, in which the label only depends on the input through planted direction E Sd-1, Formally, we have for each
V1 (x,) + o00 xx4) Ess
24
Lid.
((0,1a),C iid.
(00,1)
where is a known link function. We will consider the setting where our learner is f(0, a:) == 00.x.)
where / 5d-1 the learnable parameter.
Assumption 1. We will assume the following regarding the link function a.
E (==) (Normalization)
[x)) < C for k = 0, 1, and for all z. (Lipschitzness)
We note the assumption the boundedness of (*) can be relaxed to it having polynomial tails Damian al. 2023), but at the cost of increasing the complexity of the proof.
We consider training via the correlation loss; the loss on specific sample (ar, y) is:
L(0; x. y) = 1 f(0, x)y
The empirical loss our training set is therefore:
L, (0) = Σ L(0:
iE(n)
We also denote the population loss over (x, y) from the data distribution be L(0) (L) y)]-
In this setting, Ben Arous al. (20TT) showed that the sample complexity for learning depends on quantity called the information exponent ke of the link function a. To motivate this definition, consider first the probabilist's s Hermite polynomials.
Definition (Probabilist's Hermite polynomials). For k 2 0, the kth normalized probabilist Hermite polynomial h4 R + R is:
(-1)4 (x)
(x) = Vk! drl (x)
(x))
where (xx) = =
is the probability density function of a standard univariate Gaussian.
Of importance is that the Hermite polynomials form an orthogonal basis in I,²(y) (i.e. the space of gquare-integrable functions with respect to the standard Gaussian measure). Henceforth, for link function E (7))), let {cx J22o denote the Hermite coefficients of a:
Definition (Hermite coefficients). Let the Hermite coefficients of a E L²(>>) be (cu In other words.
2C
(x) = Σ (xh)(x), Ca = (==x
k=
This leads to the key quantity, the information exponent.

---

Published as a conference paper at ICLR 2026
Definition 3 (Information exponent). We define the information exponent be:
ky = = 2 1 C 0)
In other words, this the first Hermite coefficient with positive index that is nonzero. Some examples of information exponents below:
Example 1. (Link functions and their information exponents)
o (t) = and (t) (t)U((t) have information exponent 1.
σ (t) ||] and (t)) (2 have information exponent 2.
a (t) = t t² has information exponent 4.
o (t) (t) has information exponent k.
Ben Arous al. 2021) showed that d_aax(1, -1) samples were necessary and sufficient for
online SGD to recover H*, mirroring the tensor PCA setting. Damian et al. (202)) showed that
this rate could be improved to n Z d_aax(1,k' ka /2) by running online SGD on a smoothed landscape. number of papers have managed circumvent the information exponent by applying a label
transformation before running SGD Mondelli and Montanari (2018) 8); Maillard et al. (2020); Chen
al. (2025); Dandi et al. (2024); Troiani
Damian et al. (2024); Lee et al. (2024). These
results apply a transformation T to the labels (2/ 1.- to derive samples from the single index model defined by T This link function can have smaller information exponent than Ø, and the smallest
exponent such a transformation can achieve is called the "generative exponent" Damian et al. (2024).
For the purposes of this paper, we can assume that such label transformation has already been applied so that the information exponent and the generative exponent coincide.
2.3 THE LEARNING ALGORITHM
Definition 4 (Spherical gradient operator). For 0 E Sch and function R + R, define the
spherical gradient operator to be Vog(f) P. Vg(=)lz=o- where Pd := I- 0g projection operator with respect to 0 and the standard Fuclidean gradient.
is the orthogonal
T
We now formally define our learning algorithm; here, (W.)220 is the standard Wiener process in TRd.
Algorithm Learning algorithm
Input: Inverse temperature parameter e, number of time steps T, data points HL) Initialize Oo ~ y (e.g. uniform over 5d-1)
Run the following SDE up to time T:
dl (- 1--2
1 0 + eb(() )at Pe dW,, b(0) = - (0)
(I)
2
0 ;= 1 Jo
T 0,tt
M 1. J 0,0 dt
If 1:- is odd, return 0/|@@]
Otherwise if 1et is even, return the top eigenvector of MM
can be shown that when 0g follows the SDE in Equation [ remains on the sphere for all time t. Thus, this SDE is the natural analogue of the standard Langevin dynamics on the sphere. A discussion regarding this is deferred the appendix.
2.4 MAIN CONTRIBUTIONS
We now highlight our main contributions in this work.
4

---

Published as a conference paper at ICLR 2026
We show that by combining Langevin dynamics with weight averaging, we recover #* in both the tensor PCA and single-index model settings with 2 d /2] samples, which nearly matches the optimal
peapuutatiooalstttttiiiin tradeoff for these problems (Damian 2024 Hopkins et al. 2015).
ct al.
In contrast with previous work (Damian et al. 2023 Biroli et al., 2020 Anandkumar et al. 2017). which attain the sample complexity guarantee via smoothing the existing loss landscape create high signal-to-noise ratio regime, we utilize the other end of the spectrum a low signal-to-nok toisee ratio setting. Our method of uniform averaging takes advantage of the noise, and allows us learn the estimator that one would obtain by running landscape smoothing.
One other feature of our algorithm is that it does not see the data in an online manner, unlike previous works (Damian al., 2023 Ben Arous et al. 202T We use the empirical risk minimization (ERM) loss to obtain our results.
2020) shows that Langevin dynamics struggles to escape the "equator" {0 10 0*1 S d 172} without n Z d samples. Surprisingly, we show that it is not necessary to escape the equator to get good estimate of #* our process #(t) indeed lies on the equator throughout the training process so that its correlation with 0 remains small, but the time-averann averggde iterate still converge 8*
(Ben Arous et al.
MAIN RESULTS
Our high level framework to show ergodic concentration to an estimator that recovers the planted direction with enough samples. We will state our results for both the odd and even algorithm.
Theorem 2 (Odd k²). Let o(d--k"-3)/?) and T 2 d /e2. Then, Algorithm succeeds in estimating 2-1 Bz-n [b=)] up to O(e) relative error, Moreover, for A > if 2 /21 AA we recover the ground truth 0* up to error A with probability least 1 - edd
Consider first the setting where c + 0; this corresponds to a convergence to the pure Brownian motion on 5d-1, which has Itô SDE
d 3 dt + P dW.
dB
2
In the regime of in Theorem 2. it turns out that at time we can write 0, 3, Fe where E, is an error term of order c, and we couple the processes 0 and / with the same noise process W, We 0n = fo, and Es O, with the former being drawn from the uniform distribution on the sphere. Then, time averaging allows us to obtain:
1 + 1/, Ordl - - [ Brdt 1
0rde=
Eddt
By ergodicity of Brownian motion, we can prove that the first term concentrates to zero. For the second term E,, we show that the time average of it converges to the direction of Eg-p L, (2)]. In both the tensor PCA and single-index model settings, this estimator
be shown to recover the
planted direction 6* with 2 d'kk/2 samples. Moreover, it is possible to this estimator warm start before rumning online SGD. This idea was also used by Hopkins et al. (2066): Anandkumar ct al. 22017; : Damian et al. (2023) to boost this estimator, and allow it to recover 0* with n 2 d= d='/3
samples:
Corollary Using the same and T in the setting of Theorem 2 and Oidk (((kk"2),, (2), we can run Algorithm 1 followed by online SGD with (2dh'11) samples to recover the ground truth 0* to arbitrary accuracy.
The idea here is with 22(ak"/2 /2) samples (which is multiple of /d less than in Theorem 2). the averaging estimator gives warm start that obtains correlation ((d 1/4) with 0* From here, can run online SGD using the result from Ben Arous et al. (202T) to recover the ground truth. We now proceed to state our result for the even case.
5

---

Published as a conference paper at ICLR 2026
Theorem (Even k*). Let (k--(k"2))/) and let T 2 d 12-+122 Then, Algorithm succeeds {Bmnp(2((2) in estimating Ez~ulzz T E (2b(2)] + b(z)z'] up to O(c) relative ervor in operator norm. Moreover, for A > 0, if 2 d 12 /2/ then the top eigenvector of our estimator recovers the ground truth 8r up to enror A with probability at least 1 e
Intuitively, the algorithm for the odd case does not work here because of the first order terms vanish upon taking time average, due to the symmetry of the uniform distribution/Browninn motion. More specifically, L.. Rs 0 and does not have any meaningful correlation with 0*. On the other hand, when we consider the time average of the second order information given by 00T, we can precisely recover the planted direction 0" by taking the top eigendirection of our estimator. More formally, time averaging gives US:
1 T J. 0,0, dt
1 J T L'
I T 1." (B.ET + E, jtt + T / E,E,
B,
We prove concentration of each of these terms to the stationary average via the ergodicity of the spheri- cal Brownian motion, which leads to a final quantity of approximately Bz~uiza I _
b( (==)=]] The first term converges to //d, and the final term is a negligible error term. When Z d /2, the middle term converges to matrix with a rank-one spike 0~0*T
OVERVIEW OF PROOF IDEAS
4.1 ERGODIC CONCENTRATION
In showing general ergodic concentration result, we first give some preliminaries Markov processes on compact Riemannian manifolds.
Definition 5 (Markov semigroup). Let (Xt)zo be time time-homogeneous Markov process. Then, its associated Markov semigroup (Pi)czo the family of operators acting on bounded measurable functions f through:
P.f(z) = E()())) x]
At this point, it is useful define the infinitesimal generator of a Markov process. Definition 6 (Infinitesimal generator). Let (Pt) be the associated Markov semigroup for a Markov process. Then, the infinitesimal generator L associated with this semigroup defined as:
P
Cf lim c-0
for all functions f for which this limit exists.
Having these definitions introduced, consider the Brownian motion on Sd that we defined earlier:
d (1 ) dt
d3 = (0---c0
Note that by rotational invariance, the stationary distribution is μ.. Moreover, by classic results (Saloff-Costt 1994), we know that the infinitesimal generator of this process is L = A4-11, where Agu is the Laplace-Bei Laplace-Beltrami operator on 5dl-1 We now give general lemma for ergodic averages of functions of a Brownian motion over the sphere.
Lemma 1. Let f ddd + R such that L²))), where μ is the stationary uniform measure the sphere for the Brownian motion, and Sga-s fdj = 0. Then, we have:
1 1/' [(3.)dt= o(Bb) o(Brr
Mr
T
T
T
where
6(B) 1 P.f(
and My = fo VØ(3)) (o dW, is martingale.

---

Published as a conference paper at ICLR 2026
The proof is deferred to the appendix, and it now remains to bound these terms, which depends our choice of f. Recall that we need to make this ergodicity argument for 3n and ((3t) (defined in Section 4.2). well as (,b(3,) for the even case. For the sake of exposition. we look at this function coordinate-wist wise the main text; our full proofs in the appendix directly handle the tensorized version.
The bounds these quantities are given by the following lemma, with full proof in the appendix. Lemma 2. In the setting of Lemmal 7 the following holds:
(
((3r)
2 2sup /vl
T
(d 2)T
Mr T
27
sup flll
(d 2)
The d - term comes from the Ricci curvature of 5d being = d - 2, which leads to a bound on the gradient decay in the sense that P: e e- V fll (Bakry et al. 2016). A detailed discussion of this is included in the appendix. We now sketch the remainder of the ergodicity arguments in the main result. The previous lemmas tell us that the concentration happens at time T that depends on the function f.
4.2 ANALYZING THE ERROR COMPONENT E
Recall in the previous section that the time average consists of a Brownian component that is averaged out to zero, and an eπror component # Jod E,dt. First, let us recall our definition b(0) := LL (0) _ P Evenn 10 x, vr)).; By decomposing the time average of E, even further, it turns out write the above roughly:
e T % Fadl dT
1
( b(,)) [MM0)dd
From here, we derive the following:
T %
1
T
T T / b(3.)dt +
1 T T //" (b(0,) b(B)))tt
b(6, (t.)dd=
The first term concentrates to b = R.-p [b=)] using the ergodicity arguments from the previous section, and the second term can be controlled via upper bound on || ||0 - Be due to Lipschitzness. Indeed, in the regime of that we work in, can further argue that with high probability, ||0 - 31 remains order (e)) over all time, which we outline below. Recall the SDE's for the coupled processes 0, 8:
d 0 + eb(0) de + p
d0 =
2
d3 d
1
3dt + P3 dW,
This tells that:
dE (0----111 E+bb()) dt + (Pf p
(---1
The key observation here that the noise matrix //2 1= Py. - Pj satisfies the property that
tr Σ < 2| Intuitively, this means that the size of the noise scales with the norm of E, and this allows us to get high probability uniform bound on || over all time. The following lemma makes this rigorous.
Lemma (High probability uniform bound of sup E|). With probability at least 1 - dl' -d there exists an absolute constant C' such that:
csup |b||
sup T
[((t]
S C'
7

---

Published as a conference paper at ICLR 2026
Hes
less
Cocvelation of SWW estiintto with
1.0
Cocrdaaionn SWA extimatoo with
1.0
6.5-
6.6-
6.4-
04-
6.2-
02-
6.0-
0.
0.0001
0.01
0.
-0.2
-0.4
10
-0.-
0.25
1.00
1.25
3.00 -00P
0.5
1.0
300
Eteracion
Ineratiion
4007
Figure 1: We with d 100 with 10d /21 samples, using various learning rates. Here, the dark curves correspond to the correlation of the time average as a function of iteration, which indeed converges to the direction of 0* The light curves correspond the actual iterate function of time, which can be seen to stay near the equator over the entire training process.
The key idea of this uniform bound lies in bijection between this Oristein-Uhlebeec--ii process and suitable subgaussian process. From there, we can apply the chaining method to obtain uniform bound of sup || over time. Indeed, the fact that || = O(c) throughout training is key to both the proof's of odd and even k*, since it heuristically reduces our process to a Brownian component plus an signal component that can leverage the randomness in the Brownian component
4.3 RECOVERY OF 0*
Let Ó()) hide non-c terms. In the odd case, our estimator converges the direction of Ez~,, [b(2)) with a magnitude of Ō(c). We prove in Section G that for the tensor PCA setting, this recovers 0* with n 2 d[k*/2 /2) and we prove in Section that for the single-index model setting, it recovers #* with n 2 df**/2] as well. Moreover, we prove that when n 2 d''/2, /2, we obtain nontrivial correlation with 0*, from which we can then run online SGD to get total sample complexity of d' /2, For the case, full proof's are included in Section Gand and Section G as well; also leverage the uniform bound on sup |E] prove convergence of our estimator M approximately & plus Ö(e) spike in 0-0-T From here, we can perform PCA similar algorithm recover 0*
5 DISCUSSION
5.1 EXPERIMENTS
We sanity check our findings experimentally via different choices of link functions which correspond to different k** For k* 3, 4, 5, we let ott) ha- (t), as defined in Definition] [ Specifically, we run the minibatch update defined in Section 5.2 with batch size Our findings are included in Figure T and Figure for the odd and even cases, respectively. For k = 3, 5, our first-order estimator indeed recovers 0* even though the iterates stay near the equator throughout training. For k 4, this same estimator does not recover 0* but the second-order estimator's top eigendirection does, with the iterates once again staying near the equator. Our experiments are run with different learning rates, and observe that smaller learning rates behave more and more like gradient flow, whereas larger ones behave more like Brownian motion and stay near the equator, as we would predict with Langevin dynamics. However, there are some nuances to this, describe in the next section.
As an aside, our technique is one way to prove convergence to the stationary Gibbs distribution exp( (ap(-2eT...) and helieve this useful way to approach minibatch conjecture in Scction 5.2
8

---

Published as a conference paper at ICLR 2026
lle
lle
1.0
Corvection of SWA estinator with
Correlation of SWA estimater with
0,8
0.2
0.1
0.0
0.0
0.01
-0.2
0,0001
-0.4
0.001
-0.1
10
0.01
1.0
Iteration
Ieraaion
(a) For various learning rate choices, track the time average (e.g. the first order estimator) as a function of iteration, which can seen not have any meaningful correlation with 0" This is due to the o' being an odd function, causing the first order estimator to vanish.
(b) The solid curves correspond to the correlation of 6" with the top cigenvector of the time average 00 and dotted lines are for the correlation hetween actual iterate 0 and 6*. Indeed, the actual iterate itself remains near the equator over all time.
Figure 2: Simulations for ks = 4, run with d = 100 with n = 10d2 samples.
5.2 EXTENSION TO MINI-BATCH SGD
Our experimental results suggest that pure mini-batch SGD should have theoretically guarantees too. Consider mini-batch SGD with learning rate 77 and batch size 1:
02 799t
0
VoL(6;; xi, Me.), i, (((m))) St
10:
ng.
Di is approximately standard Gaussian, since VL x, v) = yyo x) and 0 x is O(1) for the most part, and hence |e xu O(v)). For < 1/2, we have the following approximation:
0: 19t
02 79 ~ (0, -1
0.41
2" (d
10,
19:
7 g
Let 2, : Sr b(0)) be the mini-batch noise Because we are in noise-dominated regime, z; is approximately isotropic so we approximate this process by an SDE, we would heuristically get:
0.+1 = 0 -9t
2" nd 1)0,
0
v.. TT. 17 2 nd(d d 1)0 + r)b(O.)
d - 1
d0
110 b(0) dt + vnp dW,
2
d -
1 b(0) dt P dW.
d#
2
" We therefore conjecture that there exists learning rate regime for which this SGD argument holds even without the noise boosting that is present in Langevin dynamics. The main technical challenge in extending our results in this direction is not just controlling the discretization error, but also the dependencies that arise between the noise covariance and the smoothing estimator. In particular, the stationary distribution for the pure-noise process will longer be isotropic over the sphere and will have data-dependent stationary distribution, which introduces additional complications. However, extending our results and techniques the minibatch SGD setting is promising direction for future work.
which roughly recovers Langevin setting with =
By choosing batch size maximize the scale of the noise without explicit noise hoosting.

---

Published as a conference paper at ICLR 2026
REFERENCES
Emmanuel Abbe, Enric Boix-Adserà, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics, 2023.
Radoslaw Adamezak. tail inequality for suprema of empii processes with applica- tions to markov chains, 2008. http:pss/aaaa...ooooo/oa////////
Anima Anandkumar, Yuan Deng, Rong Ge, and Hossein Mobahi. Homotopy analysis for tensor pea, 2017.
Luca Arnaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iu- vant: Data repetition allows sgd to learn high-dimensional multi-index functions. arXiv preprint arXiv:2405. 15459, 2024.
D. Bakry, I. Gentil, and M. Ledoux. Analysis and Geometry of Markov Diffusson Operators. Grundlehren der mathematischen Wissenschaften. Springer International Publishing, 2016. ISBN 9783319343235.
gowgle http:://ooooooooooooooooooo//eee
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent non-convex losses from high-dimensional inference, 2021.
Gérard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Algorithmic thresholds for tensor pca.
The Annals of Probability, 48(4), July 2020. ISSN 0091-1798. doi: 10.1214/19-aop1455. URL http: dx. doi www.oogg/o.o..o..oo http:///...111111111 12111A1A-11111
Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks, 2022.
Giulio Biroli, Chiara Cammarota, and Federico Ricci- Tersengiis How to iron out rough landscapes and get optimal performances: averaged gradient descent and its application tensor pca. Journal of Physics A: Mathematical and Theoretical,, 53(17) April 2020. ISSN 1751-8121. 11188811751888 10881111--81111111333
doi: 10.
1088/1751-8121/ab7b6 URL htRttttp::// htt.tttpp//.xxdeee doi wrg.
Siyu Chen, Beining Wu, Miao Lu, Zhuoran Yang, and Tianhao Wang. Can neural networks achieve optimal computational staittati tradeoff? an analysis on single-index model. In The Thirteenth International Conference on Learning Representations, 2025. URL
net./ met/fooumm: forum? ? i 4CCCKSFA
Yuxin Chen, Yuejie Chi, Jianqing Fan, and Cong Ma. Gradient descent with random initialization: fast global convergence for nonconvex phase retrieval. Mathematical Programming. 176(1-2): 5-37, February 2019. ISSN 1436-4646. doi: 10. 10.100771100007009--11336 URL http:
//dx.doi.orgg 10. 1007 00t.1100011110001111111 00 01363-6
Alex Damian, Jason D. Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent, 2022.
Alex Damian, Eshaan Nichani, Rong Ge, and Jason D. Lee. Smoothing the landscape boosts the signal for sgd: Optimal sample complexity for learning single index models, 2023.
Alex Damian, Loucas Fillaud-Vivien, Jason D. Lee, and Joan Bruna. caapuutational-stttiiiie gaps in gaussian single-index models, 2024.
Alex Damian, Jason D. Lee, and Joan Bruna. The generative leap: Sharp sample complexity for efficiently learning gaussian multi-index models, 2025. URL https::/aazziv.orr//aa/ 2256..08 05500
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time, 2023. URL. httpp:://axxiv.org//
2305. 82700
Yatin Dandi, Emanuele Troiani, Luca Amaaolli, Luca Pesce, Lenka Zdeborová. and Florent Krzakala. The benefits reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents, 2024. URL httpps//aaaiiv.o/////a//////////
10

---

Published as a conference paper at ICLR 2026
Samuel B. Hopkins, Jonathan Shi, and David Steurer. Tensor principal component analysis via sum-of-squares proofs, 2015.
Samuel B. Hopkins, Tselil Schramm, Jonathan Shi, and David Steurer. Fast spectral algorithms from sum-of-squares proofs: tensor decomposition and planted sparse vectors, 2016.
Marian Hristache, Anatoli Juditsky, and Vladimir Spokoiny. Direct estimation of the index coefficient in a single-index model. The Anmals of Statistics, 29(3):593 623, 2001. doi: 10. 10.12///aos/ 1009210682.
10. 121 1ww11/211110000 /ttp://...crr/.0022222/////12222
Wolfgang Karl Hirdle, Marlene Miller, Stefan Sperlich, and Axel Werwatz. Nonparametric and Semiparametric Models / Edition 1. Springer Berlin Heidelberg. 2004.
Nirmit Joshi, Hugo Koubbi, Theodor Misiakiewicz, and Nathan Srebro. Learning single-index models via harmonic decomposition. arXiv preprint arXiv:2506. 09887, 2025.
Sham Kakade, Adam Tauman Kalai, Varun Kanade, and Ohad Shamir. Efficient learning of generalized linear and single index models with isotonic regression, 2011. URL https
/ a/arxa //rxiiv. www.//gg//b/bs/e
http/bbs/11000 /1104 2018
Adam Tauman Kalai and Ravi Sastry. The isotron algorithm: High-dimensional isotonic regres-
sion. In Annual Conference Computational Learning Theory, 2009. URL https://api
semanticscholes wwwpeaar.rr//oooooeeoeeee 1tt.a//......rooooor1.111..1.11.
orq
Jason D Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with sgd near the information-theoreiic limit. Advances in Neural Information Processing Systems, 11758711-555 16-5566,, 2024.
Antoine Maillard, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase retrieval in high dimensions: Statistical and computational phase transitions, 2020. httttss:://aaxi
org/abs/2006, 05228
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Generalization error of random feature and kernel methods: Hypercontractivity and kernel matrix concentration. Applied and Computational Harmonic Analysis, 59:3-84, 2022.
Marco Mondelli and Andrea Montanari, Fundamental limits of weak recovery with applications to phase retrieval, 2018. URL
http///xxiv.or//b////////s/ htarg//bb///i 05932
Andrea Montanari and Emile Richard. A statistical model for tensor pea, 2014.
Yunwei Ren and Jason D. Lee. Learning orthogonal multi-index models: A fine ame-grained information exponent analysis, 2024. URLhttps://aaxxv org/abs. ht/////s2/2211 09678
L. Saloff-Coste. Precise estimates on the rate at which certain diffusions tend equilibrium. Mathematische Zeitschrift, 94, 1994.
Emanuele Troiani, Yatin Dandi, Leonardo Defilippis, Lenka Zdeborová, Bruno Loureiro, and Florent Krzakala. Fundamental computational limits of weak learnability high-dimensional multi-index models, 2024. URL htt.s: http/arxivv.ooo abs/2400 15480
Ramon van Handel. Probability in high dimension, 2016.
https web. math. h. princetonn
edu/=rvan/APCCI pdf
11

---

Published as a conference paper at ICLR 2026
A PRELIMINARIES
Definition 7. Let i. = C. (og(()))fF for a sufficiently large constant C.- We define high probability events be events that happen with probability at least 1 poly(d)e where poly(d) does not depend C.
Note that high probability events are closed under polynomial number of union bounds.
Lemma 4. The Itô stochastic differential equations for 3 and remain Sd-1 for all time.
Proof. This follows by Itô's lemma on f(X) = 11 More concretely,
d- d( (- 0) P eb(0)0 + tr P di + 0 PW =
The derivation for proceeds similarly.
We proceed by applying Lemma 24 that gives high probability control of E over all time.
Lemma (High probability uniform bound of sup E|], With probability at least dTe-d, there exists an absolute constant C" such that:
csup |b||
sup [(t) S C² :
d
Proof. Recall the SDE for E(I):
dE (- d d
E cb(0) dt (P/ Pt)) dWW
2
By Lemma 24 can apply the result with C' d-1 x d. G ssup ||b and B 2.
B ERGODIC CONCENTRATION
Lemma 6 (Lemmal) I restated). Let f: :R - Rm such that f E [2)),,, where is the stationary uniform measure over the sphere for the Brownian motion, and Jga-1 fdp = Then, we have:
1
l f(B,)dt x(B6) -(BT)
MT
T
T
T
where
か(B) b(B)= S.a Pif(3)dt
and My = fo Vo(3.) (6(3))PP).. dW, is martingale.
Proof.
To begin, observe that @ satisfies -Lo = To see why, note that:
CŌ(z)= poo [-
d( ((P.())) fxx
where in the second equality we used Kolmogorov's backward equation:
d - P,Cf CPff. Pof =
di
Applying Itô's to ((3)), we obtain:
dø(3) Vo(g) ( +3 Co(3)dt
= -Vo(8) dB + Lo(3)dt
-VO(B) P) dW, - L)
12

---

Published as a conference paper at ICLR 2026
where the second line follows from that fact that 3 (d3) 0 (i.e. Brownian motion stays the sphere). Therefore, it holds that by integrating from to T.
-13r)-((33)) = L' Vo(Be) P
P PW, + I'- T
Lo(3,)dt
MM [" f(Be)dt
Rearranging gives the desired result.
Lemma 7, In the setting of Lemma 6 with f Rd TR. the following holds:
((50) ((3r)
ssup V
T
2)T
Proof. We recall that V (3))ll2 can be interpreted as the Lipschitz constant of ʃ with respect to the Euclidean norm. First, note that two points on 5d can differ by at most in Euclidean norm. Therefore, have:
(1)()) ((3)) & suu Vollz
We then bound the supremum follows:
sup [ sup
(B) dt
sup
(d-2)) sup
V (f))
sup [
2-2
where the second last inequality follows from the Ricci curvature of S being - 2 and the gradient bound of Theorem 3.2.3 in Bakry et al. (2016), and the first result follows upon division by T.
Lemma 8. In the setting of Lemma with f Rd RT, the following holds with probability 1 c m.
MT T
suup [V
T(d 2)2 2)
Proof. Recall that Mr == 5od vo(B.) Pj. dW, We consider the predictable quadratic variation
matrix (Mi) 1oo0 Vob(B.) PH Vo(3.) p dt. Then, we have that:
[Vv((3) P/] 1 < [o(B)]22 sup |
d-2
Since have operator control here (rather than Frobenius), applying Lemma 20 yields that with probability
M S sup |[]l d
VT(m + log(1/8))
from which the desired result follows upon division by T.
Corollary 2. In the setting of Lemma b with Rd mm it holds with probability 1 cm that:
1 f(3. dt
sup V (B)1l2
m msup
T
Td
Td²
13

---

Published as a conference paper at ICLR 2026
C PROOF OF THE ODD k* CASE
We now show that after sufficiently long running time, the time average roughly approximates the time average of the Brownian motion, which in expectation over the stationary measure μ should converge to the partial trace estimator for k* odd (i.e. (==)(6(=)))
Proposition 1 (Decomposition E). At time 2 O, it holds that:
E(t)= = J/ e 2b( + [ ---00 - )dW
Proof. Recall the SDE's for the coupled processes 1 and B3
d -
d0
eb(00 di Po dW,
2
d d-1 1
d3
3dl + P3 dW.
This implies that:
dE = - )110))a) + (Po) -P3
PP)) )dW,
Integrating this gives the desired expression.
We now give the ergodic concentration results for the relevant functions. Lemma 9 (Ergodic concentration of b). Suppose T Z d=1 With probability at least 1 - e have:
I' T
1 T b(3)dss - b sup /bbl2
(2)
VTd
Td
Proof. This follows directly from Corollary 2 setting f(3) = b(B) - b, and using the fact that b is (ipssiips
Lemma 10 (Ergodic concentration of 3). Suppose T d-1 holds that:
With probability at least 1 e-d,
1
B,ds S 1
T
VTd
Proof. This follows directly from Corollary
2 setting f(ß) = B.
We now prove the main theorem.
Theorem 4 (Theorem restated), Let c= od- and Z d le. Then for 6, A >
if n 2 d /21/22. Algorithm succeeds in recovering the ground truth 6* up error A with probability at least ed
Proof. The time average of the E up to time the sum of the time averages of the two terms in For the second term, which the noise term, we have the following:
Proposition I
(' (1-80) (P - JdW,t
Mr = T 0
( (Pr P+)
T
T-
c - dtdW,
T /
[" (P)
2
1=1 (T-)) dW.
1
14

---

Published as a conference paper at ICLR 2026
Note that PP Pa || 2 sup sup < 2. Therefore, by Lemma 21 we have that with probability e_d
MT
T
For the first term in Proposition I we have
1 I
{' l0 e -
T-s
eb(0,,
T eb(O,)
c
- dtds
T
J
/ eb(O.))
T
2
(1-e - (1' ds
T
d
T /'" eb(0.)
d ds T {'" eb(Os) d-1
1 (T - ds
We analyze these two terms separately. For the second term, note that:
T cb(O.)
1 ds
x esup |[6(0)||
Esup [lb(0)
Td
Td²
For the first term, we decompose follows isolate the Brownian motion:
T L' cb(O.)
1
2 L'
2
/ c(b(0.) - b(3.))ds
d-1 ds T(d- 1) cb(3.)ds + T(d 1)
Once again, the second term can be bounded by the Lipschitz constant of b:
2
1' e(b(0.) d(
2e sesup Vbble
||0, Ballds
T(d 1)
T(d
2e sup /bble (d
Tesup ||6||
d
The remaining term is the main term stationary average for in Lemma 9 Therefore, the time average of E satisfies via triangle inequality:
2 T 0o' b(3,) which we proved concentration around the
2x
T
2c d-1
T E,ds
2E
esup ||6||
2ca
sup
V6ll2
sup |6|
(b(B) - b)ds
+
T J0
d
VTd
Tdl
d²
26 sup Vbl2
csup
||b||
2e sup
Vbll2
sup
|bb||
S
d 1-1 VTd
TTd
Td²
d²
/Td
de
Combining our results with Lemma 10
using triangle inequality, we obtain with probability at least
1 14
1 0,ds
2e
1
T
2
b
(B,
E,)ds
b
T J
d
T
d-1
1
1
26
8,ds
+
E.ds
T
T
d
1
2
VTd
VT
d2
15

---

With probability at Published as a conference paper at ICLR 2026
Let 2-1 and = + -o O,dl. Then, in our regime of ' and 6, the total error is bounded as:
è²
2x
||u u
+ <
d (k* 11/2
VTd
VTd
d
d-1
By our lemma, have that with probability
1 e=
b
E. Ey S Ad Ad 111-
We wish to analyze
w. Mo
which we calculate via triangle inequality as:
V. 0
2c 2-1
Ex
0*
d-1
2x b
2c 8-1
1
d-1
2x
|
E 6
+
b
8
d-1
d-1
(1-1-A
(1++)
2 21
as desired.
D PROOF OF THE EVEN As CASE
Lemma 11 (Ergodic concentration of (BT). Suppose T 2
d-2,
With probability least e-d
holds that:
1 T
B. ds
S
d
VT
Proof. This follows directly from Corollary into vector in Rdd
2 setting S(B) = =3BT 1. and flattening the matrix
Lemma 12 (Ergodic concentration of (b(B)T + b(B))T).
Suppose T Z d 2
least 1 e-a, we have that:
1 [' (8.D(3.) + b(B)8 ds Bz-n [=b(2) + b(2)=T] T
x sup (TVp|((B(((T] +1
VT
VT
Proof. This follows directly from Corollary setting f(8) Bb(B) b(8)8T Ezpp [2b(2) T + T]. and flattening the matrix into vector in I
Lemma 13. With probability 1 -d", it holds that:
2
T (E.b(O.) + b(0.)E
- a [2b(x)) + b())
S dvT
+ de
Proof. Recall the SDE's for E and B:
d3 d- 1
3dl + PJdW.
d
dE =
+ cb(0) dt + (P)) - -P) )dW)
2
By Itô's lemma, we calculate the SDE for E,3 as:
d(EBT) - 1-E3-1)E3 + Tb + (Pg) PT) (P)) JWW33 + +EdW, Pj
16

---

Published as a conference paper at ICLR 2026
The SDE of 3E is just the transpose of the above, so we have:
d(EB')= (-(d -
+ +Bb(0)) 1-(d-1)3E'
+ (Pa d 3dW, (P P+) + P3 dW
Let G = F.3 + + Then the SDE for G is:
d(G) = 1old-166-
(- c(b()) + (B(())))))))))))) + [Pp PP3Ppp + Ps - + (Pg - )aW,"" + Edw," P3 + 3dW, (Pg Pg) Pj dW,F"
where the first line is the drift term, and the second line is the noise term. Moreover, we can further simplify the final term in the drift:
(P& - P3) + Pp (Pg - Pt))
= --33E - EET + (ET ()83 + + (-E3" - EE + (ET 8)(B3 + +B)
-(BET E +
where is the remainder term satisfying | x ||E|| 2 2//6. The last line follows from Lemma [4] for simplification. Our SDE for G can therefore be rewritten as:
dG (-dG + =b(0)) + (b(0))))))))))))))))))()))))(( + =)
+ (P) PS + Edw, PS- 3dW, (Po - -P))+ P
This implies that:
G(I) [' d(t (e(b(0.)3 + (2.b(0.))++))
c-dle-)) P + Edw, P + 3dw, (Pz PP) + PH FF
We first analyze the time average of the second term, which is the noise term. Intuitively, the time average should concentrate around as time increases.
1
dil -dless) ((P P JaW . + EdW P + BdW. (Pø -P)) + P3 dW, J.W.ET E. Jdt
T 0
T S
1
TT)
(Pt
-Pf) / T-
e dtdW. ... / E
T-S
" Ps
1 P
T-
/-- ''
T-
+ T Jo
E dt -adddW. - P +
Ps
E d didW F.
T J
1 T
(Po
pb) pb) (1 -e ))) )dW
+ T J
T
E
7
(1 e - a(T -(r-ss) dd, Pg
+ T ["
1
1 8 (1 -e Ja, Pτ) + T Jo
/
Pj
1
d(T 5)) dW.F.T
It now suffices to bound the Frobenius norm of the time average of the top two terms of the last expression (since the latter two terms are just transposes). We again observe that P P l| S sup E S &. For the first term, we have that:
E
1
(P)
- pp) (
(1) ( (i--e-ddr-) -10 d(T-s)))
dW,3!
T
E (1 d d(T lidt -P 12 ds
<
T
sup tsT
Fr |2
d
d.T
17

---

Published as a conference paper at ICLR 2026
where the second to last inequality follows from Lemma 15 For the second term in the time average the noise component, we have:
E
2 j
e - dw, PJ
T
E
a(T (2
T
?
d
Combining all four noise terms together using Lemma 21 and triangle inequality, we have that with probability 1
1 / (dd(t - P9 33 + EdW. P3 + BdWT (Pa) - Pb) P3 dW E
c
di
S
VTd
F
We now analyze the drift term of G. First, to isolate the Brownian motion, we once again do another decomposition:
f e d(t d(t- (c(MO.)3 + +.b(0))) + E.
((((oo))c((((()))))))))))))))) + +J + +3.(6(33)) + +
+ E.)ds
c-d(t-o) E(b(8))B + +Bbb(B.) )ds f (-dl-s)(E(03 + +B. + =.) ds
where here we define v b(0)-b(3) which by Lipschitzness has norm bounded by 0(I x 5. Hence, for all < T', this second term satisfies:
% dit dit (E(DB] + +
+ =.))ds
1 S [e[B3 + B,u + E .5t
J
which means the time average over this component also has Frobenius norm O(c2/d3). For the time average the first term, we have the following:
1
(' e c(b33))) + (.b(33) )ds dss
T Jo
á - a(T a(T ((M(B.)8 + +BBb(3.) Jdds
T
{"
(b(333)) + +BMBB.
1
1 T -(1- -0-000)))((((3.0)) + )ds
L' -
T
T
For the second term, we can bound this in Frobenius norm by:
1
der (bbB.)) + BBb(3a) )ds
sup [b(33) + +666())T111 11 x e
T o J
< Td
Td
18

---

Published as a conference paper at ICLR 2026
Finally, for the first term, have shown concentration (Eg_.ge-i((())) + =b(2) in the previous lemma. Combining everything through triangle inequality, we have:
T I G(s)ds [=b2) + b(2)2
]
- d
B.b(3.) + ((.)33 ds -Ezn -Ezu [2b(2) + b(2)2
T
F
0
+ 2
+
Td
d2
VTd
?
+
dvT
Td
d
VTd
2²
+
dvT
d
and the result follows.
Theorem 5 (Theorem 3 restated). Let ((dd(kk-2)/))) and let T Z d= Then, for A >> 0.
if n d /2/ /2/2?, the algorithm succeeds in recovering f up error A with probability at least 1 1
Proof. Recall that 00T BB + EBT + TE + EET, In the previous lemmas, we have analyzed each of these terms separately, and our goal is to prove ergodic concentration to d + (zbb))) + =[2))
0,0 ds
و Eu 12bb2)) + b(())=
T JO
F
<
1
B.B
ds
+
1
(E3T + BB )s _ [Bess- + b()) TT)
+
T I
EE ds
T
T
F
F
2
22
VT
dvT
d
de
2
dvT
d
Consider the stationary average of Mn := 11 + $Ezng [zb(2)T + b(z)2 71. By Lemma 30 with probability - e ed, holds that:
ESWS [ab(z) b(2)2 Esngu 1 (==(=(()) b(2)2 11112 S dd
Therefore, we obtain via triangle inequality that:
1 T
T 0.0_ ds [M[]]]]
< T 0.0 ds -Ma
+ +IIM, - Ez [Mn]ll
è + d VVd-k12/1
+
d/T
de
k- 12/n
V
where the last inequality follows from our regime of and T. We now note that the eigengap for
E2 [MM]] is (O /22). Then, when n = 02dk'2//2) /2/ when applying Davis-Kahar we see that
the top eigenvector can be recovered up accuracy:
d /2 SA
Sin(w100) S
(0(d)k122
where иn denotes the top eigenvector of our time averaged matrix.
19

---

Published as a conference paper at ICLR 2026
E
USEFUL LEMMAS
Lemma 14. Let 8,8' E St-1,, and let E = -3--.. - Then, we have that
E 3' --2 E
Proof.
118' + E11 ||3|| - 2E 8' + +||| -0
since || ||8'1| 1. Rearranging gives the desired result.
Lemma 15. Let 8, E St-1 Then, we have that
((Pg - - E|I =11111*
where = = B
Proof.
ur ((P))) - (P) - = -r(p Psr
Note that
(B3 (B' + + EB + +EET)
= (E3 + )
=ETT + -.E - -' EB - BB F.F.T
and similarly
Pg ( -33 + EE + 33 EBT - EET Summing these, we get the trace to be
2 E - 1/2
Lemma 16. Let 2 ~ 8d-1, Then, for integers 2 0, it holds that:
(2k 1)!! O(d=k)
2k
11 (d 2j)
Lemma 17. Suppose (xx has information exponent k* 1. Then, the information exponent of g(z) == zz(z) has information exponent k* -.
Lemma 18. Let g = Ex Cyltg where by the k-th normalized Hermite polynomial and let be the index of the first nonzero even coefficient. Then,
E (EE.9( 2)) S (o==
Proof. Note that we can rearrange this as:
E a 1=(= 2)g(=' =)] Σ Be ((>-211 Σ c5 Be.r 22kk)
We can now upper bound this by:
2/) (2 z2)2
Eur.No..ylo(()
Σ (z. 2/)2
k28/2
The result now follows from Damian et al.
2023
Lemma 26).
20

---

Published as a conference paper at ICLR 2026
F MISCELLANEOUS CONCENTRATION INEQUALITIES
Lemma 19 (Concentration of norm), Let Z ~ (0, Ta). Then, it holds that:
Pr[1 ||| s] S exp(-s211 22)
Lemma 20, Suppose MT 10 A,dW, is a vector martingale in Rd, with ||Ar|12 S a for all t. Then, it holds that:
P || 2 aVT Vd y
2 <
Lemma 21. In the setting of Lemma 20
suppose we instead have Frobenius norm control (e.g.
| S a for all t). Then, holds that:
|| 2 T 1+ 2 log
26
Lemma 22. Let X : R - IR satisfy (0) 0 and
dX dx + dX).W,
(X) S for all X, then for all 0 < t, it holds that X(t) X(s) is = e-20((-s))-
subgaussian.
Proof. Let Y()) = e Then,
dY(() = o
Thus, (t) a martingale. Furthermore, the quadratic variation of Y satisfies
M- 2M ( so d² S' 2 = 0².
2At
< co
2A
Therefore, Novikov's condition tells us that
E(AY), == exp (A(()) - 2 (8))
is a martingale. Hence,
E(XY), = E( = exp XY(I) -// (Y), |Fa
12
Rearranging the above inequality gives
Elexp(AY
X202 2AL 22
S exp (Y(s) +
exp
FF
2
24
Now, converting back X and replacing +- de- At obtain
Elexp(A(X(() X())))|FF.
X202 -2((t=))
S S exp (x) A(C-8) -1) (e
FFF
2A
Applying this for (s, 0) instead (t, s) gives us
X202 1 e
2A6
1202
E(exp(AX(s))] exp
$ exp
2
2A
4A
21

---

Published as a conference paper at ICLR 2026
Plugging this in the previous equation upon taking expectation over F,, obtain
(e-2(e-Ait-- 1)2 + A2 e (A((-8)))
E(exp(A(X((t) X(s))) - exp
4A
4A
X²o02
S exp
(1
21((-)))
2A
where we substituted and used the fact that
AltA((->)___ - 1)2 < -2A(t-))
Lemma 23 (Chaining tail inequality (van Handel 2016)). Let {Xeher be separable subgaussian process the metric space (I', d). Then for all to E T and 0.
sup Xx+ x 2 C Jo Vlog (T, d, + x SCe Cidiman (T KET
Pr
where < 0o is a universal constant, and (T, d, E) denotes the covering number of an e-net for (T.d).
Corollary 3. In the setting of Lemma 22 there exists absolute constant C < oo such that for any 8 8>0
1 AT
Pr
sup 1<T
X2
C
log 5
VĀ
Proof. Define
02
d(s,t) = V
2-2((t-s)
A
Then, X, - -x, is t(s, subggausiia from the Lemma
22
When we invert this distance, we obtain
2AT
(),,)) -log(
Note that for < a/VA, this can be upper bounded by 1+ 2T'2 and the diameter upper bounded
by a/vA. Applying the chaining tail inequality Lemma 23 we have:
22
Pr
sup EST
Xell 2 Cx
log(1 + AT) + x Se
where we used the fact that:
R
log N( c)de S
/ /log(1 AT)
VA
Rearranging gives the desired result.
Lemma 24. Let (0) and suppose X satisfies the following SDE.
dxX AAAX + 2
and that uniformly for all X,
|((X) SG, tr (()) S BI|
Then, there exists an absolute constant C 0 such that for any 5, T 0, I= 1 log 1+4T
x- and
A 2 CBL, then with probability at least 8:
CG
sup t<
[ <
A
22

---

Published as a conference paper at ICLR 2026
Proof. We begin by decomposing X(t) X1(t) + (t) where X1 X2 follow:
dXi = |-AX, + b(X))dt, dX₂ - AX2dt + E (X)dW,
and X1(0) X2 = 0. Define R = a. Observe that for all
(1) = / - A(t b(X = (1x
S l E e-Alt- S A R.
For X2, note that:
x| = | +tr E(X) X2 /2/2(X)
We now decompose |2 = Yi + Y2 so that:
dY 2AY + tr E(X) dY₃2 2A + X2 (X)W() Define the stopping time : inf
[X2(t)] 2 R). Then
t=E(((tAT) - B ( T|] s 2B
+ R² 4BR2.
A
Therefore 2BR AA Next, the noise term in the SDE for Y2 can be bounded by:
xe(tAr)E((X(( A (tt S [2 tr E(X(tAT)) 4BR
Now, let C be sufficiently large constant. Substituting into Corollary 3 we have that with probability
at least 1 - 6,
BR A
2(1 + AT)
sup (2(LAT)] ST
SC
log
Under this event, we have that
B
B
22 AT)
sup
|2p(XX((AA)]I CR² ST
+
log
Now since A 2 C' V log(1 + AT)) where C' a sufficiently large constant then the right hand side is strictly less than R, which implies that with probability at least 8, T and supesT [t)l $ R.
We now give the following standard definition of the Orlicz norm, which will be used extensively for concentration results.
Definition 8. For > O, define the function (.((t) = exp((!) 1. Then, for random variable X, we define the a Orlicz norm of X be:
|| = inf(A 0 Eeja (I 1
In particular, zero random variable X is || -subexponential, and is || -ubgauusaaassssaa
We now give the following lemma, which is adapted from Theorem 4 in (Adamczak 2008 for our setting.
Lemma 25 (Adapted from Theorem 4, (Adamezak] 2008)). Suppose XI,..., X, are ii.d random variables in measureable space (S,B), and let F be a countable class of measureable function f: R. Assume for every f E F, holds that (X,) = O and that sup f(X))llly, co. JeF
Define Z = sup E f(Xi)) and o² = sup Σn 1 2x( Then, with probability at least
sup fEF
fef
1 - 6, it holds that:
1 max sup |X( log
S EZ o² log
fEF
23

---

Published as a conference paper at ICLR 2026
Lemma 26. Suppose X1.... Xn are i.i.d mean-zero random variables R such that for any
V E 5tt1,, holds that E (X; . v) S o² and X, is R-subgaussian. Then, it holds with probability 1 that
1 Σ X
(d + log(1/))) + Rlog(1/0)vl + log $
n
n
n
Proof. Consider a 1/4-net Nijs of Sa-1, Define Z supee /4 E, X, w/.
Then, by
Lemma 25 it holds that with probability least 1 d:
2 log
max; SupmEN X . w ll log 1
S EZ
71
n
By union bound over exp(d) terms with standard subgaussian tails, we have that:
max sup |X, ol llaa x R log
vEN1/4
Since the ei norm is upper bounded by the 122 norm, we have that the above is an upper bound of the 4h norm as well. For EZ, we have that:
o2d
EZ S E 22
("Cor
Cov
n Σ x,
T
n
where in the second incquality we used the fact that N1/4 C 5d-1 covering argument, we have that with probability at least 1 8,
Combining everything with the
2 (d + log )
R d + log n log 1
X
desired.
G TENSOR PCA
Let T= = (0*) + 1122 where every coordinate of drawn i.i.d. from (1, 1). We consider the negative log-likelihood:
L(0) - (0 T)
The spherical gradient is given by:
b(0) kPp T
G.1 ODD k
Lemma 27. Ez,zb(z) c0 where c (('-~"11)
Proof. A direct calculation shows:
Ez,zb(=) = E0 (0 2)1-11 (0 +1)
Note that 2 is equal in distribution to 21 so
kE= [0 2) -0+. 22
is of order O(d )
Next, will control concentrate the norm of the deviation from this population expectation.
24

---

Published as a conference paper at ICLR 2026
Lemma 28. With probability least 1 8, we have the following:
d 1-(2-1)/2(6 + log())))
[E,(2)) Ex,2b(2]l X
n
and in the 6* direction,
d- 11/2 log log (1/8)
0 (E,b(2) - -b2,2 2
Proof. We first note that:
- Ez,x)(), kn 1n/211111p 1/
P 2 2| 0k-1]
which can be seen to be a linear functional of the Gaussian tensor Z, as well as rotationally invariant by symmetry. Hence, we obtain that [E])(2) 2x,) _ x2 where 7² Ez,zb(z) 12 This can be calculated as:
Ez [[E,(2)) Ez, zb()l = Ex
km 1/ [
Be,s' 2 (P Z[ 112 11 Ps 0k-1| (='))
x ni
"g,s' (z x2) - (P, P
B (3.22') - 2 + (2.22)²))
x n d (k 3)/2
Therefore, d- Finally, by xa concentration, we have that with probability at least
1 - 6, the magnitude of xa random variable is bounded by O(d + /dlog(1/8) + log(1/5)) Old + (og( by the AM-GM inequality, and the result follows.
For the second equation, we note that 0* (Ezb(z) (bz))(z)) is also a linear functional of the Gaussian tensor Z, so we simply have to consider its variance for concentration. Using the previous calculation, we obtain:
Varz (E_b(=) E. (bz))) Ez km R. Z(2 (² -d
where the d factor follows from the covariance being isotropic. This gives that with probability 1 8,
- (E.((z) - Ee,zb(=))] d (k- 1)/2 (og(1//)
Proposition When n 2 d(b+1)///A2 / for A
(0, it holds with probability edd that:
E_b(z)
0* > A
[=)]
Moreover, when 2 d22,, it holds with probability 1 ed for e < 1/2 that:
E.b(2)
8* 8*2 1-14
[E,b(2)|
Proof. When d(k+1)/2/A3 we have with probability -d that:
E_b(2) 8*
b2) 0 (E.b(z) bz,)b(z)) l*1
[[,b())]|
|E, bb)) + [x)(2) Ez, Zb,)1(2)]]
d (k- ()/2 2/2)
d d (1 2/2)
> 1 A
25

---

Published as a conference paper at ICLR 2026
as desired. When d we have with probability 1 - e- df that:
E. (() 0
(z)z 0 (Ezb(2) - Ez,zb(2))- 1-1
[
|B,,b(22]| + [[zb(2)) B2,2b(z)1
d-
d d + dd//4)
2 d/1 1/4
where in the second inequality we use the fact that (E,)(z)) Ezzzb(z)) 0-1 < Ez,b((2) or due
to 1/2.
G.2 EVEN
In this section, our goal is to concentrate the self-adjoint random matrix G := E,(zb(2) + b(2)2 Lemma 29. Es O(d k/2)0-0-T - ( (k+2)/2)Pp
Proof. We have that:
( (0" 22)
By symmetry, it holds that:
E, [2[20T (0 2)) 2)h
(8* 2))222 0(d=k//20000 ((d Pa-
Similar calculations hold for E,, (bb2)2 it is just the transpose, and the result follows.
Lemma 30. With probubility at least 1 - 8, it holds that:
(k+2)/2 (d + log(110)) (1/8))
| - Ez
Proof. Note that G - EzG is self-adjoint. Therefore, it holds that:
||G EzG 2 sup
| 1 (G BzG)ul
vEN1/4
< 2 sup vEN1/4
|o (B=[=b()) [B,,,( Dol +
|u (B=(()))) - DDo
vEN1/4
where M1/4 denotes 1/ 1/4-net of 5 It now suffices to bound each of these two terms individually.
Let us start with the first term. Consider for fixed 5 the quantity
E. [) E. (=b((()(=)) []] kn 10--- 1/2 TT.. EEg [=(PP] Z[=0k-11)
Since this quantity a linear functional Gaussian tensor Z, suffices analyze just the variance to obtain a concentration.
Varz [u
E, [2((2)
- []
$ n Ez [
[1: 11-1) ]
n- EE [ (v z²) (PP, 7[2 0k- '1) (P T]]]
= n Ez,er (v.2)(u.=))B (P.- Z[2 00 11) Z12 11 To]]
= n E., (u.2)(v.2)) v P. Ez [2 0h (2
10 P
n E.. (v.2)(v p. (2 21)x k- 11
n ER (1)=)(
E (0- I11 -n "E,,'' [(v 2)(0 22)()22))(6))
22)3(2 2/)x-1]
n E (v. (2) 2//
26

---

Published as a conference paper at ICLR 2026
In the last expression, the first term is the main term, and the latter three due to the at least of the projection terms. For the main term, we can bound it by (dd-(k+2)/2 /n). For the latter three, we have that they are ((d-(k+4)/2 (0+44/2//)) /n). Hence, the entire variance expression is (dd /n) Therefore, we have that for an arbitrary S it holds with probabillity at least 8/94.
1 T [_,[=(=) (x < d 2k+2)/2 1og(99 /6)
d (242)/2 (d + log (1)))
We now consider 1/4-net N1/4 over Sd-1, which has size at most 9d. Union bounding over E N1/4, we have that with probability at least 1 - 8 that:
(k+22///2 log(1/6))
sup 10 T [=[=b(=) NEN1/4
Ex [
n1
By similar argument, we obtain that:
d d (lg((188)
[E -E,,,)(()))
sup vENiza
n
Adding these yields the desired result.
Proposition 3. When n 2 dd/2 A for A E (0, it holds with probability at least 1 e-d that the top eigenvector of Ex + b(z)= satisfies (v 0-)² 2 21
Proof. First, note that the eigengap of the expectation over is (dd From the previous lemma, we know that with probability - e e-d ||G - EzGl|2 I 4/2 for our choice of n Z //2 // (with an appropriately chosen constant). Hence, the eigengap of G is bounded below by k/2 A/2). From the Davis-Kahan theorem, we have that
/-//2
V - (u-0+)2- sin <(v,0)
S
(1-
Rearranging yields the corollary.
G.3 LIPSCHITZNESS OF
remains show that b is bounded and Lipschitz, which is formalized through the next two lemmas. Lemma 31. With probubility 1 e-d, it holds that:
sup [(b(0)] S1+ dd/n
Proof. Observe that with probability least ccd
sup [B(O)]| I 1 + 1/2 sup Z Sk S + n 111211/ S1 + vd/n
where the operator norm bound on Z follows from a standard covering argument.
Lemma 32. In the same setting as Lemma 37
[b6(0) - b(0')l s (1 /d/n)ll0 0'll
Proof.
[b(00) 6(0) k|| 1108k-11
$ (P - Pp 1 + Po (T (8) (0)00-111))
$ (1 /d/n)]l0 81
27

---

Published as a conference paper at ICLR 2026
where the inequality for the second term follows from the fact that if 0' = 0 E:
= Σ ( TI
k-1
(k- (k-1) & goh-1-oj - |T op
k
k-1
[ + E) @)) 10h-1111
Σ
||
|
B||
==1
==1
H SINGLE INDEX MODELS
Recall that by assumption, our activation satisfies sup, o (2) = (1) for k = 1, Define 6; (0) to be the negative spherical gradient on the ith datapoint:
(0) P. x;0' xi).
We will use E, to denote the expectation with respect to the data. We will also let ~ Uni((Sd-1).
H.1 ODD k*
Lemma 33. B,,bb(2)) = of where e = O )
Proof. First note that by Hermite expanding y and we have that:
E,yo(2 xz) [,(a(0" z)0(2 2) Σ (0)
Taking spherical gradient with respect to gives:
E2d,(2) 1
We can now average over the sphere. First by (Damian et al. 2023
Lemma 26),
E2 Σ
(0)=0-0)611 S K
In addition by isolating the ls 1e* term, it is at least order d
Next we handle the projection
term:
E2
Σ
kotz(z 0+)* O',,
Σ kcî(z 0+)h+1
k
k
and this upper bounded by θ #1
which completes the proof.
Lemma 34. With probability 1 8. we have that:
d d (b+ (11--1)66 1)/2(d + log 1)
vd + log log à
- - Ex,=b(2)11 $
(3)
and in the 6* direction,
(-(---11) 1)/2 122log?
Mog n log
10 - (E,((2) Be)) /
Proof. We can decompose:
E.d,(2) D,o, xx)) V.E2/2(2 22)00(2 (2. x.)].
28

---

Published as a conference paper at ICLR 2026
We first concentrate in the direction of 8*. We will analyze the main term and the projection term separately. For the main term, we have that:
Var Σ (is 0 (') (z. ms))
Var, farz((((: (000000)00 (2)
_ E. [v" (z. " (o' (2 (=' x)]]
Ex (ax Eg, (z z)0o (z' x)]
n
27
En
0-)
Fx || (o' (z x2' =)]]
||>|||
n1
1
d-(k--3//2
d (k* 11/2
2^
n
where in the third to last line we used the polar decomposition of æ, and in the second to last line we
used the fact that Exg
(1ना 0+)
) 21 = @
(1/d)) /d) and:
Ex [|= Ex (z x)) x]]]
= Ex scal==l Be. lo' (z- x)''('' x]]] + E. calll/ lo' x)o (=)+
2 dEE (o (z 2)d' x)]l + |[| 2 CdIE2 [1| (o' x)oo(s' =))}
S dx'
which follows from of having information exponent k- - 1, and P|[|0||3 2 Cal being exponentially small for C by standard x2 concentration.
We also note that (z, (2) is ((1)-subgaussian) Therefore, by Lemma 25 holds with probability that:
Σ (o; O') 2i) Ex
1 n1
Σ
u 0')
S
d (e+ 1)/2 log
Vlg n log(1/8)
For the projection term in the direction of 8", we have that:
Var Σ 1(= (=) xx)oo(= mz)
Varz [yE2 (2 0*)(z 2)) x)]]
x g
Var,
y/
||
E. (z - x) ) 20 mm]
E
y²
(xc
Eg [(
3) (z m
|| ||
E2 (2 x ) x;)]
E
I
0*
Ex
||=|||
| 1
d- (66' -3)/2
7
d (K* +1)/2
29

---

Published as a conference paper at ICLR 2026
where the second to last line follows from:
Eg (2
7 (2 2
x)]
E. (2-2) 20 xx]] x)|²
Ex (2 x)) (2. 22)
27
Ex
Ex 10
+ EE,
1 || ssca
|0
|0
|
27
1
Ezl(2 22)2o'( x²)²
Ex [ (z 2) x)]]]
P Eg
d
|
1 (k* 31/2 d
since (t) has information exponent k- by Lemma 17 and P[[N))PP is exponentially small for C' 1. Moreover, we can see that this is ((l/d)--ubgaussian
x 0"
ly (z) (2) x)]l
y/
||>||
Ez [= xx) 2))00(2))))
x
( 0 )E 2/
(z
|
(x. 8*),
Eg
E. lo' 20
x
(x. 0*)
Therefore by Lemma 25 with probability
holds that:
d +1/2 11/2+=///000////00
Vlgg n log(1/8)
Σ M. 8*) (z 2z)o (z- 24)1 Ez ((> 8*) x) z)]l 3 V
dn
Altogether combining the main and projection term in the 0* direction, we have that with probability 1 that:
10 - (E,b(()) E,, b()) d (6+ 1)/2 log (1/6) Vlog n l1g(1/6)
now concentrate the entire norm of of E2)()) - b(2)) and once again we will handle the main and projection term separately. By the same variance and subgaussian calculations before, we can apply Lemma 26 obtain that with probability 1-
d (k-k--1)/22 + log } + Vd + log log !
[B,b(=)- - b, S
The desired result follows.
Proposition 4. When 2 d d( +1)/2/22 for A (0, 1), it holds with probability 1 e that: E.b(2)
||,,((2]|| 2
Moreover, when n 2 d212, ,, it holds with probability - E dt" that:
E.b(z)
0* 1r200 1/14
[E_b(2)
Proof. When n 2 +11/22 /22 we have with probability 1 -d that:
E.b(2)
E,, b(2) 0 (E.((2) bx)) 0*
E. b[)](
E b())l + (() 14..)()))
(k* 11/ A/2)
d- 11) + A/2)
> 21
30

---

Published as a conference paper at ICLR 2026
as desired. When / have with probability 1 - c d² that:
E, b(2)
[(z) 0 (E.b(z) [z,) 8*1
E b())l
b())l + [|bzb(z) (0))==) bx,=b(2)] d -1)/2
(k- 11/221 (1 d1/4)
1/1//4
where the second inequality we use the fact that the first term of the numerator dominates for c < 1/4, and the second term in the denominator is of order d 11)22 d!/4 for this same choice of
H.2
EVEN A**
Lemma 35. [i,x<b((=)]] = c000+ - -- where c 12 and - Old ((d-(x++2)2)
Proof. We will fix z first and then take average over the sphere of z. First,
E, a0 2z)o'(2 (z P+1 2x, ( xz)o'(( æ;)] - 2E, [ar a 22)0' 22)
Let e be the Hermite coefficients for a. For the first term, we have by Stein's lemma that:
2 [a] (0 2x)oo(2 x.)) =o' =x))o
x)) tafoa 2x)0"(z 22))22
= Σ
( O + Σ
(k + 2)(k + 1)c2 2))²
42k-11
k
We now proceed to handle the projection term:
2IE [x] (0 2x)00(e (2 22)122 = Σ
2) 2) 22
Σ k2b
(k: 2) + +))xck22(8 2))²22 T2222 22 (k 2) 1)ck(4+2(00 22
12k--1
= Σ (0 2)
k TT
Σ
1
Therefore, after combining and before taking expectation over our expression is:
2 (0 =T
Σ
20 TT
We now take expectation of over the sphere. For the first term, we have that
E2
z Σ
(0
5 0- = -=>ee (d (k" +20/230-00 = O(d 02)
2)a
12k*
j20
For the second term, we have that
E,
c0 22 12
Σ
O(d -(2//0++2/))0
where the two è hide different absolute constants. Nonetheless, the main part of our desired expression is ((d-k' k² /2 02 and this gives the desired result.
Corollary 4. E,, (zb(z) + b(z)=T] = c + gPg. where c = 1(d-k1/2 and = ((d-()+2)/2)
Proof. This follows directly from the fact that we are just adding the transpose.
In the rest of the section, goal is to concentrate the self-adjoint random matrix G == [2b(2)) b(2)zT], T]: however, for the sake of exposition we will simply consider [=[z((>) TJ and we will note, when necessary, the properties that transition to the case of G.
31

---

Published as a conference paper at ICLR 2026
Lemma 36, With probability least 1 - 6, it holds that:
B2 [2y:0' (z x;)x, P. Ei, x2)o, P=l
(k 2) + log(1/6)) - log(1/6)
dn
Proof. We will now concentrate 1 E. Ez [>b(=) Σ E. P=-
in operator norm. By the epsilon-net bound on the operator norm, it suffices consider for arbitrary
v E Sd-1 the quantity v [ Σx= E, [z (2.2;)=; 2;)x; Ps E,, [zb(z) 'I First, note that the
projection term gives the following decomposition:
E2 x4) ] yz (= 2. x,)or.]]
ys (z( ax)(2 x,)=
We will handle the two terms separately. For the first term, the variance is bounded above by:
E2
E2[( z)yoo (z (x . v)] = E. [ (x v) BE 2)0'(= x)1²]
72
x
E, [(x v) 2E [V z)0 (z -)
(x v)
Ex[(x z)0 - x)}²²
E,
||=||
(x
(x-v)4
v)a
E,
EE, 2)0 (x - =]]²]
||=||
1
d
-2)/2
d²
+22/2
where in the third line we use the fact that by symmetry:
[ z)o' x)]
V.
Bel(z x)o' (z- x)]
|
In the fourth line we use the independence between the direction of a and the norm of a and in the second to last line the fact that the information exponent of t o'(t) is k 2.
In addition, we will also show the term itself is O expoonennll from which we will combine with the variance bound via Bernstein. Rewriting the term, we have:
(a U)E, [w 2)a'(= xx:)] = y/ (x; v)) v. 24 By((2 1x)oo(2 ax)]
|211|
Since |3:| = 01) and x4 v is O(1) subgaussian, it suffices to show that is ((1/d) subgaussian. This follows from:
| Ex(s x. (2 x.)]
|r || Ez[(2 2x)00(2 m3)l
v. Fi E.((2 (= (0)((((() + 2:) (2)=))))(((011
|
U E.((2 (o xx) o' (0))
||
[2 x2)
|
|
xzl d
Since this is upper bounded by O(1/d) times a half-Gaussian we have that this is O(1/d) subgaussian. Therefore, the entire term is O(1/d) subexponential. In particular, for a fixed v E -d Bernstein's 's inequality gives that with probability at least 8/9dd
d +2)/2 log (94 /8)
log(9d) (9/4/// /8)
E Ez (2 (x,,,,,.]] E ( (z x: 20)o:
S
1
n
dn
32

---

Published as a conference paper at ICLR 2026
We now handle the projection term. Here, the variance is bounded above by:
E, [E =)² yo' 22 x)]²
E,yy E, [(v: 2 2) - x)]²]
7
72
24
B, [E. 2' 2) 2)]
E, Br 1 2 a)(2 2eo x xx)]
E.. [wv 2) E lo (z z)(2 wjd ()(e x)]
n
1
E [[v 2 (v 2 z) 22
n d +2)/2
In addition, will also show that the projection term E. 2) z)(2 x)] is O(1/d) subexponential. This follows by:
Ez[(v- (e_((u=)²uo''' zo²yo'(z () x)] x B2 =)²''(z-))) x)]
x Bz((v- =)) | + 2))
d(d + 2)
By triangle incquality, this is just O(L/d) subexponential, since the chi-squared | is ((d)
subexponential. Therefore, Bernstein's inequality tells with probability least 1 8/94.
E [=yio''' (= (x) =)= ]
Er,, [2u:o (=: xxx 2)z
d (b+ +2)/2 log (9 /8) Log(9%/0)
dn
Combining the main term and the projection term, have that for arbitrary Sd-1, with probability at least 1 - 0/9d:
Eg 22yoo (z- xx)xx P - ley 2.)a, p-11 10
71
d (k* 22/2 19g(94/8)
log 1og(9%/5)
n
dn
d (k. +2)/2 log(1 + d log(1/5) dn
n
We now consider a 1/4-net N1/4 over 5d-1, which has size at most 9d. Union bounding over Ni we have that with probability at least 8 that:
sup vE/i/4
Eg (2yno' (z m.),
E=,==zyoo (2 mx)ar P.l
d +2) d + log(1/0)) d + log(1/0)
dn
Using the fact that the supremum over the 1 1/4-net upper bounds the operator up constant factors, we obtain:
B2 (zyid' (z _x)x]]] P: Ex,= (= xx)x P!
2
d +2)/2 (d log(1/0)) + log(1/0)
n +
dn
desired.
33

---

Published as a conference paper at ICLR 2026
Proposition 5. When 2 /2/12//// for A E (0.1). holds with probability at least 1 e that the top eigenvector of E2 [G] satisfies (v 8*)² 2 1 A.
Proof. This follows directly from the Davis-Kahan theorem since with probability 1- we have:
[R_ (2b(2) ] B., [=b((= 1112 X Ad-
The similar holds true for Eg [b(2)= ]. and hence it holds for the random matrix G well. Since the eigengap of E&G is is O(d /2), we have the desired result.
H.3 LIPSCHITZNESS OF
Lemma 37. With probability at least 1 - e -cod
sup [(((0)] S
Proof. Let X E IRaxd be the stacked matrix with all the data points. Then,
d
| Σ Po 2 x4)
$ |
Σ
vio'(0
x1)² 2 1+
i=1
Lemma 38. In the same setting as Lemma 37
sup (((() - b(0')) $ (1+ 0 0'
Proof. We have
(6(0)- < Σ yi [P 0'00000000000 xx)) Pooo'('' (0:))0=x:))
Now have that:
Ppl o' (0 xz) Pp. (8' xz))
Py Jo' xz) o' 00' xx)] +
(.) 1
1
For the first term, the same argument as above proves that the sum is bounded by:
X 10 01) $ (1 +1+ vd/n)]l0 0'I
For the second term, is bounded by:
||X 11
Ps Vn
Pot 112
S (1+ Vd/n)]10 8'l
which completes the proof.
34


Reviews:
[{"review_id": "xdLZZThnFn", "reviewer": "Reviewer_yNf2", "summary": "The paper proposes to solve the single index models by a Langevin algorithm and demonstrates that this achieves the state of the art theoretical scaling with respect to the information exponent of the link function, on par with smoothing the loss landscape.", "questions": "1. Compared to the landscape smoothing, what are the pros and cons of using the Langevin algorithm? Some clarification of how the landscape is smoothened would be useful. Can smoothing be applied in practice, if so, could it be shown in Figure 1?\n\n2. What is the practical relevance of the findings in the paper, which practical problems are modeled by the problems considered here (single index models with k > 2 and tensor PCA)? In other words, in which problems should the practitioner consider applying the Langevin algorithm?\n\n3. What is the intuition as to why the Hermite polynomials play an important role here and what is an intuitive meaning of the information exponent?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper addresses the problem of learning the implanted direction with a Langevin algorithm, even though this was previously considered inefficient and shows that it is not. The paper nicely explains that it does not solve the problem of escaping from the \"equator\" for individual samples, but only once the averaging over the samples is performed. It also demonstrates this with numerical simulations.", "weaknesses": "It is not clear what are the implications of this work and what should a practitioner gather from these findings.\nAs far as I understand, Langevin dynamics only boosts the performance when the information exponent k is larger than 2 (otherwise max(1, k-1) = 1 and so is max(1, k/2)). It is not clear to me how relevant is the k > 2 case, especially if the single index model is viewed as a toy model of the neural network. In Example 1, k = 1 and k = 2 cases are often encountered, while cases with k > 2 seem made up.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PN9v8XMJa6", "reviewer": "Reviewer_Vo4P", "summary": "The paper shows how combining Langevin dynamics with stochastic weight averaging can enhance signal recovery in high-dimensional estimation tasks by navigating complex loss landscapes more effectively.", "questions": "Concerning the general idea of dynamical smoothing can you comment on the relation between your proposal and https://www.pnas.org/doi/10.1073/pnas.1608103113 and other papers on the same lines such as https://arxiv.org/abs/1611.01838?\nThe finite learning rate could represent an effective correlation between the regions over which information is dynamically averaged.\nNote that an optimal protocol for the coupling is considered in the contexts of dynamical explorations inspired by robust ensemble.\n\nPossibly related to this discussion: in figure 1 and figure 2 is there a non monotonicity appearing in the behaviour when learning rate increases? Can you discuss why that is? Is there a best learning rate for the smoothing process? \nWhat are the drawbacks deteriorating the quality of the average information for large and small learning rates?\n\nAlso a significant chunk of literature about tensor PCA and its sample complexity for the success of Langevin dynamics was overlooked: https://arxiv.org/abs/1812.09066, https://arxiv.org/abs/1902.00139, https://arxiv.org/abs/1907.08226, please consider adding it.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "The paper conveys an interesting check about the possibility to dynamically smooth the landscape in a number of hard inference problems.\nIt is robust in terms of derivations and conclusions.", "weaknesses": "Not sure that it is relevant to Machine Learning applications where landscape is expected to be way smoother already thanks to overparametrization and large datasets. \nThe paper is not very innovative both concerning ideas and execution. I do not need to specify further the very rich literature on the same subject already cited in the paper and in other parts of this review. Also little parts of the proof are contained in other papers (as explicitly mentioned in the text). \nI suggest the authors to stress the added value of their contribution more explicitly, also with reference to possible ML applications especially as they intendo to discuss SGD implications in terms of stochastic averaging.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pxl0iJ8msT", "reviewer": "Reviewer_kQAb", "summary": "This paper investigates the ability of Langevin dynamics on the sphere to recover a latent signal $\\theta^\\star$ in the spiked PCA and single-index models. The key idea is to consider the time average of the iterates instead of the final one. The author shows that this averaging procedure allows recovery of $\\theta^\\star$ with a sample complexity matching that of smoothed online SGD, suggesting that iterate averaging has an implicit smoothing effect. The paper also provides a heuristic connection between mini-batch SGD and Langevin dynamics.", "questions": "(1) The role of the empirical risk minimization (ERM) formulation (as opposed to an online or streaming setting) is mentioned as a distinction from prior work. However, the algorithm used in experiments appears very close to the standard spherical SGD analyzed in previous works. Is this distinction mostly of conceptual interest? Furthermore, other works such as [1] already provided a similar ``ERM'' formulation, so the novelty of this contribution is not entirely clear to me.\n\n(2) The theoretical analysis is conducted in continuous time, but in practice, the Langevin dynamics must be discretized. The choice of step size can crucially affect both stability and sample complexity, as discussed for instance by Arous et al.(2021). Could the authors comment on how the results would translate to the discrete-time setting? In particular, what scaling of the discretization step would be required for the discrete algorithm to inherit the same statistical guarantees as the continuous process?\n\n\n[1] Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics, Mousavi-Hosseini et al. (2025)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The idea that averaging iterates can replace explicit loss smoothing to recover a planted signal with optimal sample complexity in the online regime is novel and conceptually appealing. The paper provides a clear theoretical treatment of this idea in the context of Langevin dynamics on the sphere, supported by sound derivations and proofs.", "weaknesses": "(1) The analysis is limited to the continuous-time setting, and it remains unclear how a discretized version of the proposed dynamics would perform in practice or whether it would preserve the same sample complexity guarantees.\n\n(2) While the idea of using iterate averaging could potentially lead to simplified algorithms, the specific procedure analyzed in this work is not practical. It requires oracle knowledge of the information exponent $k^\\star$, and depending on its parity, a different estimator must be used. \n\n(3) The use of iterate averaging is quite common in convex optimization, yet this connection to existing literature is not discussed.\n \n(4) The heuristic connection between mini-batch SGD and Langevin dynamics presented in Section~5.2 is somewhat confusing, as it requires $\\varepsilon = 1/\\eta$ while both quantities are assumed to be small. \n\n(5) The algorithm implemented in the experiments differs from the one analyzed theoretically, and no empirical comparison is provided with smoothing SGD or with the partial-trace method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)