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

Thomas Dybdahl Ahle; Yaroslav Bulatov; Christopher De Sa; Christopher Ré

cs.LGSep 28, 2026

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

Authors: Thomas Dybdahl Ahle, Yaroslav Bulatov, Christopher De Sa, Christopher Ré

Organizations: Normal Computing · Cornell University · Stanford University

Abstract

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.

Figures & tables

Figure 1: Validation of IID SGD theory across data sources. Rows: Gaussian Data, Resnet CIFAR-10 embeddings, Transformer Wikipedia embeddings. Columns: (1) empirical eigenvalue CDF of

H=X^{\top}X/m

with a left-edge power-law fit

F(\lambda)\approx Q(\lambda-\lambda_{0})^{\beta}

; (2) test loss,

\|w-w^{*}\|^{2}

; (3) training loss,

\|Xw\|^{2}/m

. Blue curves show empirical results (5 runs); orange dashed curves show theoretical predictions using the fitted spectral constants. Training loss curves are visibly noisier than test loss because they weight eigendirections by

\lambda_{j}

, amplifying fluctuations in high-eigenvalue directions. See Section 6 for experimental details and reproducibility information.

Panel of Figure 2 — Figure 2: Test loss $c_{n}(\eta)$ vs learning rate $\eta\in(0,2)$ for $d=512$ , $\rho=1$ . Solid curves show theory; hollow circles show Monte Carlo validation ( $40$ probes). (a) IID uses ODE-based theory from the $d\to\infty$ limit. (b) Shuffle-once and (c) flip-flop have exact finite- $n$ formulas; for $n{=}20$ we use the asymptotic formula. At $\eta=1$ , shuffle-once achieves a $\sqrt{2}$ factor improvement over IID. See Appendix F for derivations.

Figure 3: Effect of shuffling method on convergence. Test loss (normalized by $\sqrt{n}$ ) for four shuffling methods on isotropic Gaussian data with $d=m=512$ . Single Shuffle (green) consistently outperforms other methods. The theoretical predictions (dashed lines) from Section 4 match the empirical curves precisely.

Figure 4: Validation of Theorem 5 . Expected trace

\frac{1}{d}\mathbb{E}[\Tr(S_{n})]

for IID SGD on Gaussian data with various step sizes

\eta

. All curves exhibit

n^{-1/2}

decay (the predicted rate for

\beta=1/2

), with step size affecting only the constant.

Method	Test Constant	Training Constant
IID Sampling	$\displaystyle\frac{1}{\sqrt{2\pi}}$	$\displaystyle\frac{1}{2\sqrt{2\pi}}$
Shuffle-Once	$\displaystyle\frac{1}{2\sqrt{\pi}}$	$\displaystyle\frac{1}{8\sqrt{\pi}}$
FlipFlop	$\displaystyle\frac{1}{\sqrt{2\pi}}$	$\displaystyle\frac{\sqrt{2}}{3\sqrt{\pi}}$

Table 1: Exact leading constants for test and training loss under different shuffling methods on isotropic Gaussian data with

m=d

Figure 5: Long-horizon run showing the third (final) exponential regime. Population loss for randomized Kaczmarz on isotropic Gaussian data with

d=m=256

, plotted against the number of row updates

k

(log-log). The curve exhibits three regimes: an initial exponential transient (up to

k\asymp d

), an extended power-law regime with slope

k^{-1/2}

(for

d\ll k\ll d^{3}

, equivalently

1\ll n\ll d^{2}

epochs), and a final exponential regime when spectral extremes dominate.

Dataset (preprocessed, $m=d=512$ )	$\hat{\beta}$	effective rank	$q_{0.95}/q_{0.05}$
Gaussian ( $X_{ij}\sim\mathcal{N}(0,1)/\sqrt{d}$ )	0.50	257	$4.5\times 10^{2}$
CIFAR-10 ResNet-50 embeddings	0.46	40	$3.7\times 10^{3}$
Wikipedia OpenAI embeddings	0.47	98	$3.7\times 10^{3}$

Table 2: Slope fits and anisotropy diagnostics.

\hat{\beta}

is the fitted left-edge spectral exponent from the eigenvalue CDF of

H=X^{\top}X/m

. The effective rank is

\mathrm{erank}(H):=\Tr(H)^{2}/\Tr(H^{2})

(

R_{k}

in Definition 3 of Bartlett et al. (2020) ), and

q_{p}

denotes the

p

-quantile of eigenvalues. CIFAR embeddings are highly anisotropic in the bulk (small effective rank), yet their left edge still yields a

\hat{\beta}

close to

1/2

, consistent with the power-law slopes observed in Figure 1 .

Figure 6: Transformer head-only training exhibits shifted-edge power-law behavior. A tiny Transformer LM is frozen and only the final linear head is trained (MSE), reducing the dynamics to a quadratic problem with curvature $\Cov_{\mathrm{train}}$ . Left: shifted-edge fit to the spectrum of $\Cov_{\mathrm{train}}$ , used to estimate $(\lambda_{0},\beta,Q)$ . Right: theory-aligned parameter error and residual compared with the finite- $d$ shifted-edge prediction (formula shown in the plot; dashed curves shown only once the Laplace window enters the fitted edge regime).

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Training loss for shuffle-once at $\rho=1$ , showing excellent agreement between empirical results and the theoretical prediction $\frac{1}{8\sqrt{\pi}}n^{-3/2}$ .

Figure 8: Single-shuffle test loss with learning rate $\eta\in[0,2]$ at $\rho=1$ . Interestingly, the optimal $\eta$ is not at 1 when $n\geq 2$ . For $n=20$ we use the asymptotic formula derived above. Note: Legend label $\alpha$ corresponds to step size $\eta$ used in the text.

Figure 9: Flip-flop test loss with learning rate $\eta\in[0,2]$ at $\rho=1$ .

Explore similar work

Apr 11, 2026

Shuffling the Data, Stretching the Step-size: Sharper Bias in constant step-size SGD

From adversarial robustness to multi-agent learning, many machine learning tasks can be cast as finite-sum min-max optimization or, more generally, as variational inequality problems (VIPs). Owing to their simplicity and scalability, stochastic gradient methods with constant step size are widely used, despite the fact that they converge only up to a constant term. Among the many heuristics adopted in practice, two classical techniques have recently attracted attention to mitigate this issue: Random Reshuffling of data and Richardson--Romberg extrapolation across iterates. Random Reshuffling sharpens the mean-squared error (MSE) of the estimated solution, while Richardson-Romberg extrapolation acts orthogonally, providing a second-order reduction in its bias. In this work, we show that their composition is strictly better than both, not only maintaining the enhanced MSE guarantees but also yielding an even greater cubic refinement in the bias. To the best of our knowledge, our work provides the first theoretical guarantees for such a synergy in structured non-monotone VIPs. Our analysis proceeds in two steps: (i) we smooth the discrete noise induced by reshuffling and leverage tools from continuous-state Markov chain theory to establish a novel law of large numbers and a central limit theorem for its iterates; and (ii) we employ spectral tensor techniques to prove that extrapolation debiases and sharpens the asymptotic behavior even under the biased gradient oracle induced by reshuffling. Finally, extensive experiments validate our theory, consistently demonstrating substantial speedups in practice.

Konstantinos Emmanouilidis, Emmanouil-Vasileios Vlatakis-Gkaragkounis, Rene Vidal

Department of Computer Science, University of Pennsylvania · University of Wisconsin-Madison, Archimedes Research Unit

Stochastic Optimization Convergence Stochastic Gradient Descent

Ask PDF

Jun 30, 2026

Random Reshuffling Dominates Stochastic Gradient Descent

Stochastic Gradient Descent (

\textsf{SGD}

) is one of the most classical optimization algorithms with favorable theoretical guarantees, yet the practical implementation of

\textsf{SGD}

differs subtly from its well-known form and is often referred to as Shuffling Stochastic Gradient Descent (

\textsf{Shuffling SGD}

). A particularly popular strategy in

\textsf{Shuffling SGD}

is Random Reshuffling (

\textsf{RR}

), which has achieved great empirical success across numerous experiments. Despite its strong performance,

\textsf{RR}

has long been considered a heuristic due to a lack of theoretical support. Over the last decade, people have finally established provable convergence rates for

\textsf{RR}

, thus justifying its observed superiority. However, for smooth convex optimization, two clouds over the convergence theory of

\textsf{RR}

remain to this day. More precisely, according to the current theory,

\textsf{Shuffling SGD}

under

\textsf{RR}

converges only when the stepsize is smaller than a threshold proportional to

1/n

, where

n

is the number of summands in the objective (or the number of data points). Consequently, the optimally tuned theoretical rate of

\textsf{Shuffling SGD}

under

\textsf{RR}

is strictly worse than that of

\textsf{SGD}

when the number of epochs is smaller than another threshold proportional to

n

. These two restrictions heavily limit the applicability of existing theories and leave a critical mismatch with practice. In this work, for the first time, we prove that

\textsf{RR}

dominates

\textsf{SGD}

in smooth convex optimization under any reasonable stepsize after any finite number of epochs, thereby addressing a longstanding open question.

Zijian Liu

Stern School of Business, New York University

Stochastic Optimization Convergence Stochastic Gradient Descent

Ask PDF

Oct 14, 2026

Large-Step Last-Iterate Convergence of Random-Reshuffling Stochastic Gradient Descent

Random reshuffling stochastic gradient descent (RR-SGD) is widely used for finite-sum optimization in machine learning. For convex

L

-smooth components, recent work establishes convergence guarantees for an averaged iterate with component stepsizes as large as order

1/L

. However, the final iterate is a natural output in practice, and existing last-iterate guarantees in this setting require stepsizes that scale inversely with the number of components

n

. We establish an expected objective-gap bound for the last iterate of standard RR-SGD on unconstrained problems with a constant component stepsize satisfying

0<η\lesssim 1/(L(1 + \log(nK)))

, where

K

is the number of epochs. By retaining the sampling bias at random comparison points, we remove the explicit

1/n

factor from the stepsize restriction. Combined with existing small-step bounds, our result matches the corresponding averaged-iterate convergence rate up to logarithmic factors.

Zijun Chen, Zihan Zhang

Department of Computer Science and Engineering, Hong Kong University of Science and Technology

Stochastic Optimization Last-Iterate Convergence Stochastic Optimization Convergence Stochastic Gradient Descent

Ask PDF