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

Mar 31, 2026

Learning to Shuffle: Block Reshuffling and Reversal Schemes for Stochastic Optimization

Shuffling strategies for stochastic gradient descent (SGD), including incremental gradient, shuffle-once, and random reshuffling, are supported by rigorous convergence analyses for arbitrary within-epoch permutations. In particular, random reshuffling is known to improve optimization constants relative to cyclic and shuffle-once schemes. However, existing theory offers limited guidance on how to design new data-ordering schemes that further improve optimization constants or stability beyond random reshuffling. In this paper, we design a pipeline using a large language model (LLM)-guided program evolution framework to discover an effective shuffling rule for without-replacement SGD. Abstracting from this instance, we identify two fundamental structural components: block reshuffling and paired reversal. We analyze these components separately and show that block reshuffling strictly reduces prefix-gradient variance constants within the unified shuffling framework, yielding provable improvements over random reshuffling under mild conditions. Separately, we show that paired reversal symmetrizes the epoch map and cancels the leading order-dependent second-order term, reducing order sensitivity from quadratic to cubic in the step size. Numerical experiments with the discovered algorithm validate the theory and demonstrate consistent gains over standard shuffling schemes across convex and nonconvex benchmarks.

Lam M. Nguyen, Dzung T. Phan, Jayant Kalagnanam

IBM Research, Thomas J. Watson Research Center, Yorktown Heights, NY, USA

Stochastic Optimization Convergence Stochastic Gradient Descent

Ask PDF