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

Shuffling the Data, Stretching the Step-size: Sharper Bias in constant step-size SGD
Apr 11, 2026Konstantinos Emmanouilidis, Emmanouil-Vasileios Vlatakis-Gkaragkounis, Rene VidalStochastic Optimization Convergence Stochastic Gradient Descent
Department of Computer Science, University of Pennsylvania · University of Wisconsin-Madison, Archimedes Research Unit
Random Reshuffling Dominates Stochastic Gradient Descent
Jun 30, 2026Zijian LiuStochastic Optimization Convergence Stochastic Gradient Descent
Stern School of Business, New York University
Large-Step Last-Iterate Convergence of Random-Reshuffling Stochastic Gradient Descent
Oct 14, 2026Zijun Chen, Zihan ZhangStochastic Optimization Last-Iterate Convergence
Department of Computer Science and Engineering, Hong Kong University of Science and Technology

Apr 11, 2026J/K move · Enter open · S save

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

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Random Reshuffling Dominates Stochastic Gradient Descent

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