Exact information accounting for SGD methods

Akshay Balsubramani

cs.LGSep 30, 2026

Exact information accounting for SGD methods

Authors: Akshay Balsubramani

Abstract

As an alternative to the standard geometric analyses, we give an exact, information-theoretic analysis of stochastic gradient descent (SGD) and its variants. We show that a preconditioned SGD step is the posterior-mean update of a Gaussian Bayes model, and that its one-step regret splits into an intrinsic-time cost and a change in comparator information. The split extends to an identity for the objective itself. Convex convergence, strict-saddle-point escape, the link between flatness and generalization, the standard learning-rate schedules, adaptive optimizers, and the noisy, momentum, heavy-tailed, and gradient-free variants of SGD each correspond to a term or a special case of this identity. We measure its terms on synthetic and real training runs. On real networks it attributes the slack of classical convergence bounds to the terms their derivations drop and separates optimizers that reach the same training loss. That separation follows the number and consistency of their steps. Its relation to which of them generalizes better differs between networks. For gradient-free SGD the identity determines how a curvature preconditioner should enter the update. The sharpness-based generalization certificate it yields, with a data-independent isotropic prior, is vacuous at network scale unless the curvature spectrum is nearly flat across all parameters.

Figures & tables

Figure 1: At a matched training loss the information ledger separates five optimizers by

31\times

where their losses differ by less than

1.6\times

. In the learning-rate and width sweeps the intrinsic-time fraction follows the shape of the path. (a) Five optimizers stopped at a common training loss, with a

1.04\times

spread in training loss and a

1.53\times

spread in test error. Each marker is one seed. (b) A learning-rate sweep at a fixed budget: the test error (blue) is smallest at

\eta=0.1

, while the fraction (vermillion) rises with

\eta

as larger steps make the path less consistent. (c) A width sweep under

15\%

label noise: the test error’s double-descent peak sits at the interpolation threshold (dotted,

1257

training points), where the fraction has no feature.

Optimizer	Choice of $(\Sigma_{t},m_{t},\eta_{t})$	Information-ledger term it incurs
Plain SGD	fixed $\Sigma$ , no forecast	$Q_{t}=\frac{1}{2}g_{t}^{\top}\Sigma g_{t}$ in fixed $\Sigma$ -geometry
Natural gradient † [ 2 ]	$\Sigma_{t}=\mathcal{F}(\theta_{t})^{-1}$ , $\mathcal{F}$ the Fisher matrix	Fisher-metric intrinsic time $\frac{1}{2}\\|g_{t}\\|_{\mathcal{F}^{-1}}^{2}$
K-FAC † [ 68 ] , Shampoo † [ 33 ]	block-structured $\Sigma_{t}=P_{t}$	$\frac{1}{2}\\|g_{t}\\|_{P_{t}}^{2}$ per round
Muon † [ 46 ]	momentum buffer $B_{t}$ , orthogonalized (Newton–Schulz)	$\frac{1}{2}\langle B_{t},(B_{t}B_{t}^{\top})^{-1/2}B_{t}\rangle=\frac{1}{2}\\|B_{t}\\|_{*}$ , the nuclear norm; the Newton–Schulz iterate is a deliberate approximation to this
AdaGrad-Norm [ 26 ]	$\eta_{t}\propto V_{t}^{-1/2}$ (the scalar variant; per-coordinate AdaGrad instead sets $\Sigma_{t}$ )	terminal-remainder-saturating square-root clock
RMSProp † [ 98 ] , Adam † [ 51 ]	$\Sigma_{t}=\mathrm{diag}(\sqrt{\hat{v}_{t}}+\varepsilon)^{-1}$ , score $\hat{m}_{t}$ , fixed $\eta$	$\frac{1}{2}\\|\hat{m}_{t}\\|^{2}_{\Sigma_{t}}$

Table 1: Every SGD-class optimizer is one reading of the same information ledger. Each row gives an optimizer’s choice of preconditioner

\Sigma_{t}

, forecast

m_{t}

, and schedule

\eta_{t}

, and the term incurred by that choice. Rows marked

\dagger

have a time-varying

\Sigma_{t}

, for which the pathwise account is exact only with the metric-drift term

\Delta_{t}^{\mathrm{metric}}(u)

of ( 87 ) included alongside

Q_{t}

. Derivations are in § 5.2 , § 5.3 , § 5.4 , and § E.3 .

Figure 2: On real networks the information ledger attributes each classical envelope’s slack to the terms dropped by its derivation. It also separates optimizers that reach the same training loss: their intrinsic-time fractions differ by more than an order of magnitude. (a) Comparator information

I_{n}(u)

still to be transported at prefix

n

, as a share of the linearized envelope

I_{0}(u)+\sum_{t}\eta_{t}^{2}Q_{t}

. (b) Cumulative Bregman remainder

\sum_{t}\eta_{t}B_{F}(u,\theta_{t})

over the same envelope; below zero the objective-level envelope does not apply. AdamW stays within

10^{-3}

of zero on both networks. (c) Intrinsic-time fraction at the deepest training loss reached by every optimizer. Solid lines ResNet-20, dashed lines the character-level language model (char LM); each line is a median over seeds.

Figure 3: The intrinsic-time inflation and the strict-saddle-point escape times track their closed forms. (a) Realized inflation

\hat{V}_{T}-V_{T}^{\mathrm{det}}

against inverse batch size on linear regression with mini-batches drawn without replacement from

N=5000

examples; the fit proportional to

(N-b)/(b(N-1))

solid and the plain

1/b

law dashed. (b) Mean escape iteration against the leading-order time ( 37 ) on the strict saddle point of

F=\frac{1}{2}(x^{2}+y^{2})-\frac{1}{2}z^{2}

(

\lambda_{\min}=-1

), started from

\theta_{0}=(1,1,\varepsilon_{\mathrm{init}})

; deterministic runs filled, their constant ratio solid, and equality dotted. Here

\varepsilon_{\mathrm{eff}}

in ( 37 ) is

\max(\varepsilon_{\mathrm{init}},\sigma_{\xi}\sqrt{\eta/2})

, so runs with the same

\varepsilon_{\mathrm{eff}}

share a leading-order time. Error bars are standard errors over seeds.

Figure 4: On a real objective with a wide curvature spread, a curvature probe drawn into the sampling covariance speeds zeroth-order search more than 50-fold, while applied to the step it makes no progress. (a) Function queries to close

90\%

of the initial suboptimality of seven-class logistic regression on

49

principal components of Covertype, at three spreads of the Hessian diagonal; small markers are seeds, bars are medians, and open markers on the dashed line did not reach the target within the budget. (b) Top Hessian eigenvalue against the train–test gap at the end of training, Wide-ResNet-28-4 on CIFAR-10, one marker per run; circles batch

128

, squares batch

8192

. Sharpness tracks the generalization gap only along the batch axis. (c) Final test accuracy of ResNet-20 on CIFAR-10 against the share of logged steps on which the pressure-cycled controller’s stall signal is on, one marker per seed. The stall signal responds to mini-batch noise.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Queries to the target grow far more slowly with the condition number when the probe captures the curvature. (a) Probes on the axis-aligned quadratic drawn from the isotropic covariance, from an AdaGrad accumulator and from the oracle covariance

\mathrm{diag}(H)^{-1}

. The accumulator does not identify the curvature. (b) Probes that apply the preconditioner to the step. The oracle is

\mathrm{diag}(H)^{-1}

on the axis-aligned quadratics and

(|H|+\lambda I)^{-1}

on the rotated ones. The probes’ curvature sweeps count toward their queries. On rotated quadratics only the full-matrix probe captures the curvature. Bars are medians over seeds of the function-evaluation pairs in (a) and the function queries in (b) to

f<10^{-3}

, or to

f<1

on the Rosenbrock valley. A pale dashed bar marks a sampler that reached the target on no seed and ends at the budget. A fraction at the foot of a bar is the share of seeds that reached the target, where that share is below one.

Figure 6: At fine-tuning scale the probe allocation that follows the curvature ends at the lowest validation loss. Validation cross-entropy of the fine-tuned encoder after

700

queries for four per-group probe scales that share one learning rate and one total probe energy. Open circles are seeds, filled circles their mean, and bars the

95\%

percentile-bootstrap interval of the mean. The dashed line is the cross-entropy before fine-tuning.

Noise or momentum class	How it enters the information ledger
Mini-batch score noise, post-step Gaussian, Langevin thermal	one additive quadratic entry each. The expected ledger depends on the perturbation only through its covariance, but the three attach at different places and scale differently in $\eta_{t}$ : score noise inflates $Q_{t}$ , post-step injection adds a channel term, and the Langevin thermal entry is $\beta^{-1}\eta_{t}d$ (§ F.3 )
Finite-cumulant non-Gaussian (Rademacher, Laplace, mixtures)	the Gaussian quadratic is replaced by the cumulant-generating function $\Lambda_{t}(v)$ (Proposition F.8 ), of which it is the second-order term
Heavy-tailed and $\alpha$ -stable	the pathwise and channel relative-entropy identities still hold. The conditional-expectation entry is finite only after truncation, robustification, or a fractional or jump regime
State-dependent (Riemannian) preconditioning	a metric-drift term with an Itô correction
Phase-space momentum: stochastic-gradient Hamiltonian Monte Carlo (SGHMC), underdamped Langevin, thermostats	conservative Hamiltonian transport plus an Ornstein–Uhlenbeck thermostat channel; SGHMC’s friction is the covariance-balance condition. Heavy-ball is instead a linear phase-space recursion (Proposition F.11 ) and Nesterov a lookahead-transport correction

Appendix

Table 2: Every noise class and every momentum scheme resolves to one explicit entry of the information ledger.

Figure 7: The information ledger separates the optimizers that keep a momentum buffer from those that do not. (a) Intrinsic-time fraction ( 13 ) at the end of the run, in the Euclidean metric against the final iterate. (b) Share of the logged grid points at which the cumulative Bregman remainder is negative. There the objective-level envelope does not apply. (c) Max/min ratio across optimizers at the deepest training loss reached by every optimizer of a network. The fraction spreads by an order of magnitude on ResNet-20 and the char LM, where the training and test losses spread by less than

1.5\times

. Markers are medians over seeds. The vision transformer runs only AdamW, Muon, SOAP and Sophia-G.

Figure 8: Further terms of the information ledger behave as its identities imply. (a) Realized regret over

I_{1}+\sum_{t}\eta_{t}^{2}Q_{t}

over the late half of each run, for noisy SGD on quadratics with

d=100

and

\sigma_{g}=0.5

, by condition number and step size. The dashed line is one half. Error bars are standard errors over seeds. (b) The unstable clock

V_{T}^{-}

and the total clock

V_{T}

at escape from strict saddle points in

d=50

with

\lambda_{\min}\in\{-0.1,-0.3,-1\}

, three noise scales and two initializations, each divided by its mean over the runs. CV is the coefficient of variation. (c) Realized over deterministic intrinsic time for SVRG on logistic regression (UCI Adult and Covertype) against its inner-loop length, with SGD dashed. Error bars are

95\%

bootstrap intervals. (d) Per-step intrinsic time for logistic regression under Student-

t_{3}

gradient noise over

T=20{,}000

steps, for plain SGD, clipping and Catoni’s score. Circles are seeds, bars their mean, and each percentage is the change of the mean against SGD. (e) The realized noise-energy entry on a

d=10

quadratic against the stability index of

\alpha

-stable noise: its 99th percentile, and its mean without robustification, under Catoni’s score and truncated at

R=2

Figure 9: Scaled coordinate by coordinate, the gradient noise of a character-level transformer at batch 32 is not heavy-tailed. (a) Hill tail index at a

5\%

tail fraction for three architectures, one marker per seed. Open circles pool the raw noise over coordinates, and crosses pool Gaussian noise with the measured per-coordinate scales. Filled circles pool the noise after each coordinate is divided by its own robust scale, with bars over tail fractions from

1\%

10\%

. Lines mark the standardized Gaussian and Student-

t

(1.5) controls and index

2

. (b) Ratio of the noise variance along the mean-gradient direction to its diagonal approximation, for the real noise and for Gaussian noise with the same covariance. The dashed line is equality.

Figure 10: The intrinsic clock and its step-weighted form seldom fit or forecast the loss better than the step counter. Each forecast

\hat{N}

estimates the step

N^{\star}

at which the run reaches a target loss. In (b) and (c) the clock is the step-weighted sum

\sum_{s\leq t}\eta_{s}^{2}Q_{s}

. (a) Completed runs of the curriculum sweep:

\Delta R^{2}=R^{2}_{V}-R^{2}_{T}

of the two log-loss fits versus the defect

1-r^{2}

\log V_{T}

regressed on

\log T

. Open markers are runs that neither fit describes with

R^{2}\geq 0.9

. (b) Held-out prefix instances: the disagreement

|\log(\hat{N}_{V}/\hat{N}_{t})|

of the two forecasts at a common window versus the defect on the prefix. (c) The difference in median forecast error

|\log(\hat{N}/N^{\star})|

, clock minus step counter, with each forecast at its own fit fraction; below zero the clock is closer to the truth. Bars are

95\%

paired intervals, hidden where narrower than the marker.

Figure 11: The pressure-cycled controller finds the deep basin of a synthetic potential on every seed but slows or degrades training on real networks. (a) Best-iterate

f-f^{\star}

on the two-basin potential for fixed-step and constant-Langevin SGD and for the controller. The dashed line is the target. (b) Steps to grok for AdamW with and without the escape, and for SGD at a fixed and at a pressure-target step. (c) Best validation accuracy of the vision transformer on CIFAR-10 for the same four. (d) Steps to grok for AdamW without kicks and with isotropic or curvature-targeted kicks. Markers are seeds and bars their means. Color gives the base optimizer, and an open triangle marks a run with the escape phase.

Figure 12: Sharpness tracks the generalization gap with the sign the flatness bound implies along the batch axis, but not across optimizers. Each entry is the Spearman correlation of a sharpness statistic with the train–test gap or with the test error over the $n$ endpoints of one sweep. Its color is that correlation times the sign the bound implies ( $+$ for $\lambda_{\max}$ and $\mathrm{tr}H$ , $-$ for $-\log\det\widetilde{H}$ ), so blue agrees with the bound and red opposes it. Gray entries were not measured. On ResNet-20 at $50$ epochs the gap is the final test loss minus the last logged training loss. $K$ is the Lanczos rank.

Explore similar work

Apr 30, 2026cs.LG

Information-Theoretic Generalization Bounds for Stochastic Gradient Descent with Predictable Virtual Noise

Information-theoretic generalization bounds analyze stochastic optimization by relating expected generalization error to the mutual information between learned parameters and training data. Virtual perturbation analyses of SGD add auxiliary Gaussian noise only in the proof, making mutual information tractable while leaving the actual SGD trajectory unchanged. Existing bounds, however, typically require perturbation covariances to be fixed independently of the optimization history, limiting their ability to represent geometries induced by moving gradient statistics, preconditioners, curvature proxies, and other pathwise information. We introduce predictable history-adaptive virtual perturbations, where the perturbation covariance at each iteration may depend on the past real SGD history but not on current or future randomness. This predictability enables a conditional Gaussian relative-entropy argument and yields generalization bounds for SGD with adaptive virtual-noise geometry. The bounds replace fixed sensitivity and gradient-deviation terms with conditional adaptive counterparts, include an output-sensitivity penalty from accumulated perturbation covariance, and reduce the deviation term to a conditional variance only under conditional unbiasedness. Since adaptive covariances may be data-dependent, we separate local Gaussian smoothing from global reference-kernel comparison. The resulting bound includes a covariance-comparison cost measuring the KL price of using an admissible reference geometry different from the actual adaptive covariance. Fixed-noise-style bounds are recovered under admissible synchronization, such as deterministic, public, or prefix-observable covariance rules. The framework recovers fixed isotropic and geometry-aware bounds as special cases while extending virtual perturbation analysis to history-dependent SGD without modifying the algorithm.

Mohammad Partohaghighi

Electrical Engineering and Computer Science University of California, Merced Merced, CA, USA

Stochastic Gradient Descent Generalization Bounds Stochastic Optimization Covariance

Ask PDF

Jul 6, 2026stat.ML

Non-asymptotic Convergence of Stochastic Gradient Descent in Score-based Generative Models

Score-based Generative Models (SGMs) have achieved impressive performance in data generation across a wide range of applications. While the statistical properties of their sampling procedures are increasingly well understood, the optimization dynamics underlying their training remain less explored. SGMs are typically trained by minimizing a weighted denoising scorematching objective, yet optimization guarantees with stochastic gradients remain limited. In this work, we study Stochastic Gradient Descent (SGD) for SGMs, contributing results in two complementary regimes. First, for general score parameterizations, we establish a non-convex convergence rate for SGD on the weighted denoising score-matching objective, with explicit dependence on the schedule-dependent weighting factors. Second, for overparameterized two-layer ReLU networks, we develop a Neural Tangent Kernel analysis tailored to diffusion training with stochastic gradients, yielding score-approximation error bounds along the SGD trajectory. Finally, our analysis quantifies the role of the reweighting factor in the score approximation error, providing theoretical guidance for weighting choices used in practice.

Stanislas Strasman, Sobihan Surendran, Sylvain Le Corff

SU, LPSM · Sorbonne Université and Université Paris Cité, CNRS, LPSM, F-75005 Paris, France · LOPF, Califrais’ Machine Learning Lab, Paris, France

Stochastic Gradient Descent Generative Models Finite Time Analysis Rectified Linear Unit Networks

Ask PDF

Nov 14, 2025math.OC

Non-Euclidean SGD for Structured Optimization: Unified Analysis and Improved Rates

Recently, several instances of non-Euclidean SGD, including SignSGD, Lion, and Muon, have attracted significant interest from the optimization community due to their practical success in training deep neural networks. Consequently, a number of works have attempted to explain this success by developing theoretical convergence analyses. Unfortunately, these results cannot properly justify the superior performance of these methods, as they could not beat the convergence rate of vanilla Euclidean SGD. We resolve this important open problem by developing a new unified convergence analysis under the structured smoothness and gradient noise assumption. In particular, our results indicate that non-Euclidean SGD (i) can exploit the sparsity or low-rank structure of the upper bounds on the Hessian and gradient noise, (ii) can provably benefit from popular algorithmic tools such as extrapolation or momentum variance reduction, and (iii) can match the state-of-the-art convergence rates of adaptive and more complex optimization algorithms such as AdaGrad and Shampoo.

Dmitry Kovalev, Ekaterina Borodich

Yandex Research

Stochastic Gradient Descent Convergence Improvement

Ask PDF