On-Average Stability of Multipass Preconditioned SGD and Effective Dimension

Simon Vary; Tyler Farghly; Ilja Kuzborskij; Patrick Rebeschini

cs.LGMar 12, 2026

On-Average Stability of Multipass Preconditioned SGD and Effective Dimension

Authors: Simon Vary, Tyler Farghly, Ilja Kuzborskij, Patrick Rebeschini

Organizations: Department of Statistics, University of Oxford · Google DeepMind

Abstract

We study trade-offs between the population risk curvature, geometry of the noise, and preconditioning on the generalisation ability of the multipass Preconditioned Stochastic Gradient Descent (PSGD). Many practical optimisation heuristics implicitly navigate this trade-off in different ways -- for instance, some aim to whiten gradient noise, while others aim to align updates with expected loss curvature. When the geometry of the population risk curvature and the geometry of the gradient noise do not match, an aggressive choice that improves one aspect can amplify instability along the other, leading to suboptimal statistical behavior. In this paper we employ on-average algorithmic stability to connect generalisation of PSGD to the effective dimension that depends on these sources of curvature. While existing techniques for on-average stability of SGD are limited to a single pass, as first contribution we develop a new on-average stability analysis for multipass SGD that handles the correlations induced by data reuse. This allows us to derive excess risk bounds that depend on the effective dimension. In particular, we show that an improperly chosen preconditioner can yield suboptimal effective dimension dependence in both optimisation and generalisation. Finally, we complement our upper bounds with matching, instance-dependent lower bounds.

Figures & tables

Panel of Figure 1 — Figure 1: Illustration of model misspecification. The geometry of the expected loss curvature $\nabla^{2}f$ differs from the geometry of the gradient noise ( $\Sigma$ ). While setting $P\approx\Sigma^{-1}$ whitens the noise, it may result in unstable updates along high-curvature directions.

Stochastic Gradient Descent

Stochastic Gradient Descent (SGD) uses noisy, data-sampled gradients to train models and solve large-scale optimization problems.

See all 221 papers

Uniform Stability and Generalization Error of GD and SGD on Fixed-Point Parameters
Jun 5, 2026Jonghyun Shin, Sejun ParkAlgorithmic Stability Gradient Descent
Department of Artificial Intelligence Korea University Seoul, Korea
Exact information accounting for SGD methods
Sep 30, 2026Akshay BalsubramaniNeural Network Generalization Stochastic Optimization Convergence
Stochastic Gradient Descent with Momentum is Algorithmically Stable
May 27, 2026Yunwen Lei, Zimeng Wang, Xiaoming YuanGeneralization Bounds Algorithmic Stability
Department of Mathematics, The University of Hong Kong, Hong Kong, China · Department of Mathematics and Mathematical Statistics, Umeå University, Sweden

Jun 5, 2026J/K move · Enter open · S save

Uniform Stability and Generalization Error of GD and SGD on Fixed-Point Parameters

Jonghyun Shin, Sejun Park

Department of Artificial Intelligence Korea University Seoul, Korea

Algorithmic Stability Gradient Descent Stochastic Gradient Descent Generalization Bounds

Open paper AskPDF arXiv

We analyze generalization error, uniform stability, and uniform argument stability of gradient descent (GD) and stochastic gradient descent (SGD) over discrete parameter spaces, where each update involves deterministic or stochastic rounding. We show that deterministic rounding degrades the generalization error of GD on convex, Lipschitz, and smooth loss functions, increasing the rate from

O(T/n)

O(T/\sqrt{n})

, and establish matching lower bounds. We further prove that uniform stability of GD becomes

Ω(T)

, showing that stability-based generalization bounds are vacuous in this setting. In contrast, for the same losses, stochastic gradient descent with deterministic rounding admits nontrivial uniform stability guarantees, which differ qualitatively from the real-valued case and exhibit distinct dependencies on the number of iterations and the dimension: we prove tight bounds

O(T/n)

for one dimension and

O(T^2/n)

for higher dimensions. We also show that stochastic rounding can introduce generalization error that increases with the dimension; such a phenomenon is absent in standard real-valued optimization and in the deterministic rounding case. Finally, we provide upper bounds on uniform argument stability for stochastic rounding schemes and show that these bounds are tight when the loss can be represented as a sum of coordinate-wise functions.

Abstract

Stochastic Gradient Descent

Explore similar work

Uniform Stability and Generalization Error of GD and SGD on Fixed-Point Parameters

Exact information accounting for SGD methods

Stochastic Gradient Descent with Momentum is Algorithmically Stable