Improved high-dimensional estimation with Langevin dynamics and stochastic weight averaging

Stanley Wei; Alex Damian; Jason D. Lee

cs.LGMar 6, 2026

Improved high-dimensional estimation with Langevin dynamics and stochastic weight averaging

Authors: Stanley Wei, Alex Damian, Jason D. Lee

Organizations: Princeton University · Harvard University · University of California, Berkeley

Abstract

Significant recent work has studied the ability of gradient descent to recover a hidden planted direction $θ^\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 $θ^\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.

Figures & tables

\displaystyle d\theta=\quantity(-\frac{d-1}{2}\theta+ \epsilon b(\theta) )dt+P_{\theta}^{\perp}dW_{t},\quad b(\theta):=-\nabla_{\theta}L_{n}(\theta)

(1)

Algorithm 1 Learning algorithm

Panel of Figure 1 — Figure 1: We run with $d=100$ with $n=10d^{\lceil k^{\star}/2\rceil}$ samples, using various learning rates. Here, the dark curves correspond to the correlation of the time average as a function of iteration, in which it indeed converges to the direction of $\theta^{\star}$ . The light curves correspond to the actual iterate as a function of time, which can be seen to stay near the equator over the entire training process.

Panel of Figure 2 — Figure 2: Simulations for $k^{\star}=4$ , run with $d=100$ with $n=10d^{2}$ samples.

Explore similar work

Jun 24, 2026

High-Probability PL-SGD with Markovian Noise: Optimal Mixing and Tail Dependence

We study first-order methods for smooth objectives satisfying the Polyak-Łojasiewicz (PL) condition when gradient samples are generated by an exogenous Markov chain. In the light-tailed setting, prior uniform-in-time high-probability bounds for ordinary Stochastic Gradient Descent (SGD) under a standard growth envelope scale as

\widetilde{O}(t_{mix}^2/k)

, leaving a gap with the

\widetilde{O}(t_{mix}/k)

expectation bounds. We close this gap using a lag-blocking argument to establish a uniform high-probability guarantee with a leading stochastic term of

\widetilde{O}(t_{mix}/(k+K_0))

under geometric mixing. We prove this linear dependence on the mixing time is optimal via a matching

Ω(σ^2 t_{mix}/k)

lower bound on a quadratic objective driven by a persistent two-state chain. We then extend this framework to heavy-tailed Markovian gradients satisfying a stationary finite-

p

-moment condition,

p \in (1,2]

. We design an all-samples clipped block method that uses every Markov transition while mitigating Markovian bias. Under a transition budget

T

, this algorithm achieves a high-probability stochastic error of

\widetilde{O}(σ_p^2(t_{mix}/T)^{2(p-1)/p})

. We establish a matching lower bound by reducing PL optimization to heavy-tailed mean estimation for a sticky Markov chain. Ultimately, this work tightly characterizes the optimal polynomial dependence on mixing time for light-tailed PL-SGD, and the optimal heavy-tail exponent and effective-sample-size dependence in the robust regime.

Dhruv Sarkar, Aprameyo Chakrabartty, Vaneet Aggarwal

Indian Institute of Technology Kharagpur · Mohamed bin Zayed University of Artificial Intelligence · Purdue University

Stochastic Optimization Optimization Convergence Analysis Heavy-Tailed Noise Stochastic Optimization Convergence Heavy-Tailed Distributions Heavy-Tailed Stochastic Optimization+1

Ask PDF

Mar 15, 2026

High-Probability Bounds for SGD under the Polyak-Lojasiewicz Condition with Markovian Noise

We present the first uniform-in-time high-probability bound for SGD under the PL condition, where the gradient noise contains both Markovian and martingale difference components. This significantly broadens the scope of finite-time guarantees, as the PL condition arises in many machine learning and deep learning models while Markovian noise naturally arises in decentralized optimization and online system identification problems. We further allow the magnitude of noise to grow with the function value, enabling the analysis of many practical sampling strategies. In addition to the high-probability guarantee, we establish a matching

1/k

decay rate for the expected suboptimality. Our proof technique relies on the Poisson equation to handle the Markovian noise and a probabilistic induction argument to address the lack of almost-sure bounds on the objective. Finally, we demonstrate the applicability of our framework by analyzing three practical optimization problems: token-based decentralized linear regression, supervised learning with subsampling for privacy amplification, and online system identification.

Avik Kar, Siddharth Chandak, Rahul Singh +3

Department of Computer Science and Automation, Indian Institute of Science, Bengaluru, India. · Department of Electrical Engineering, Stanford University, Stanford, CA, USA. · CMAP, CNRS, École polytechnique, Institut Polytechnique ´de Paris, Palaiseau, France.

Optimization Convergence Analysis Stochastic Optimization Convergence Stochastic Gradient Descent

Ask PDF

Sep 28, 2026

Stochastic Optimization Under Power-Law Spectra: Tight Bounds and Shuffling Analysis

Recent work has established that power-law spectral conditions on data enable tight convergence bounds for deterministic gradient descent, resolving the conflict between classical exponential bounds and observed power-law learning curves. In this work, we extend this result to the stochastic regime of high-dimensional machine learning. We provide two main contributions: (1) We generalize the power-law spectral theory to Stochastic Gradient Descent (SGD), showing that the same spectral exponents govern stochastic dynamics; (2) For the fundamental case of isotropic Gaussian data, we provide a precise analysis of data shuffling, deriving exact constants that prove Single Shuffle is strictly superior to Flip-Flop and IID sampling. Our results bridge the gap between abstract spectral theory and practical stochastic training choices, offering a unified picture of how data geometry drives optimization speed.

Thomas Dybdahl Ahle, Yaroslav Bulatov, Christopher De Sa +1

Normal Computing · Cornell University · Stanford University

Stochastic Optimization Convergence Stochastic Gradient Descent

Ask PDF