🔢 mathematics

Design Criteria for SGD Preconditioners: Local Conditioning, Noise Floors, and Basin Stability

This paper establishes a theoretical framework for designing SGD preconditioners that simultaneously optimize convergence rates and minimize stochastic noise floors by improving local conditioning in the preconditioner-induced metric, a principle validated through experiments on scientific machine learning benchmarks.

Original authors: Mitchell Scott, Tianshi Xu, Ziyuan Tang, Alexandra Pichette-Emmons, Qiang Ye, Yousef Saad, Yuanzhe Xi

Published 2026-06-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mitchell Scott, Tianshi Xu, Ziyuan Tang, Alexandra Pichette-Emmons, Qiang Ye, Yousef Saad, Yuanzhe Xi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to roll a heavy boulder down a mountain to reach the very bottom of a valley (the "perfect solution"). This is what machine learning models do when they train: they try to minimize errors to find the best possible answer.

The paper you provided is about a specific tool called Stochastic Gradient Descent (SGD), which is like a hiker taking steps down the mountain. Usually, this hiker moves fast at first, but as they get closer to the bottom, they start to stumble, slow down, and jitter around the true bottom without ever quite settling there. This happens because the mountain has weird, bumpy shapes (curvature) and the ground is slippery and noisy (random data errors).

The authors of this paper are asking: "How can we give the hiker better shoes or a better map so they can reach the bottom faster and stop jittering?"

Here is the breakdown of their findings in simple terms:

1. The Problem: The "Bumpy Valley" and the "Jitter"

In the late stages of training, the hiker (the algorithm) faces two main problems:

Anisotropic Curvature: The valley isn't a smooth bowl. It's shaped like a long, narrow canyon. If you try to walk straight down, you might hit the walls. You need to zig-zag, which is slow.
Gradient Noise: The hiker is wearing foggy glasses. They can't see the exact slope; they only get a blurry, noisy guess of which way is down. This causes them to jitter around the bottom instead of stopping exactly at the lowest point.

2. The Solution: "Preconditioning" (The Magic Map)

The paper studies a technique called Preconditioning. Think of this as giving the hiker a special, stretchy map (a matrix called M) that reshapes the world in their mind.

On this new map, the long, narrow canyon looks like a perfect, round circle.
The hiker can now walk straight down without zig-zagging.
Crucially, this map also helps filter out the "fog," making the noisy steps more stable.

3. The Two Golden Rules for the Map

The authors discovered that for this "Magic Map" to work well, it needs to do two specific things at the same time:

Rule A: Smooth out the bumps (Improve Conditioning). The map should stretch the narrow parts of the valley so the hiker doesn't have to take tiny, inefficient steps. This makes the path to the bottom straighter.
Rule B: Dampen the noise (Attenuate Noise). The map should also act like noise-canceling headphones. It needs to reduce the impact of the random "jitter" caused by the foggy glasses.

The Catch: You can't just focus on one. If you make the valley perfectly round but leave the foggy glasses, you still jitter. If you clear the fog but the valley is still a narrow canyon, you still move slowly. You need a map that does both.

4. The "Basin Stability" (Staying in the Neighborhood)

The paper also looks at a safety guarantee. Imagine the bottom of the valley is a small, safe room. If the hiker takes a step that is too big or too shaky, they might accidentally kick the door open and fall out of the room (diverge).

The authors proved that if you choose the right map, you can calculate the probability that the hiker will stay inside this safe room for a long time. A good map doesn't just help you move fast; it keeps you from wandering off the cliff.

5. Why This Matters for Science (SciML)

The authors tested this on "Scientific Machine Learning" problems (like predicting weather patterns or how fluids move).

In normal video games or cat-photo apps, a little bit of error at the end doesn't matter much.
But in science, if your math is slightly off, your prediction might break the laws of physics (e.g., creating energy out of nothing).
The paper shows that using the right "Magic Map" allows scientists to get the error down to a tiny, precise level where the laws of physics are actually respected.

6. The Experiments

They tested their theory on:

A simple math puzzle: Where they could prove the map worked exactly as predicted.
Three real-world science problems:
1. Fitting a noisy curve (Franke surface).
2. Solving a physics equation with a neural network (PINN).
3. Learning how a fluid spreads (Green's function).

The Result: In all cases, the methods that used "curvature-aware" maps (maps that understood the shape of the valley) reached the bottom faster and stopped with much less jitter than the standard methods. Specifically, a method called CG-GGN (which uses a specific type of map based on how the data changes) performed the best.

Summary

The paper says: To get the best results when training AI models, especially for science, don't just pick a random step size. You need a preconditioner (a smart way to reshape the problem) that flattens the difficult curves of the problem AND quiets the random noise. If you do both, you get a faster, more stable, and more accurate result.

Technical Summary: Design Criteria for SGD Preconditioners

Problem Statement

Stochastic Gradient Descent (SGD) frequently exhibits slowed convergence in the late stages of training, particularly in Scientific Machine Learning (SciML) contexts where reaching small training losses is critical for physical fidelity, numerical stability, and constraint satisfaction. This slowdown is driven by two primary factors: anisotropic curvature (ill-conditioning) and persistent gradient noise. While various preconditioned optimizers (e.g., Adam, K-FAC, L-BFGS) are empirically successful, there lacks a unified theoretical framework that explicitly identifies which properties of a preconditioner determine the late-stage convergence rate and the attainable noise floor.

Methodology

The paper analyzes preconditioned SGD updates of the form $w_{k+1} = w_k - \alpha_k M^{-1} g(w_k, \xi_k)$ , where $M \succ 0$ is a symmetric positive definite (SPD) matrix defining the geometry in which curvature and noise are measured. The analysis proceeds in two regimes:

Globally Strongly Convex Baseline: The authors extend classical convergence theory to the $M$ -induced geometry. They define effective smoothness ( $\hat{L}$ ) and strong convexity ( $\hat{c}$ ) constants relative to the $M$ -norm and analyze the first and second moments of the stochastic gradient in the $M^{-1}$ -norm.
Local Nonconvex Regime: Recognizing that deep learning objectives are typically nonconvex, the authors establish convergence guarantees within a local basin around a minimizer set. This analysis relies on:
- A local $M$ -Polyak–Łojasiewicz (PL) condition to handle flat or weakly curved directions.
- Local $M$ -Lipschitz gradient assumptions.
- Local stochastic gradient moment bounds in the $M^{-1}$ -norm.
- A local quadratic growth condition and a "controlled one-step overshoot" assumption to derive basin-stability guarantees.

The theoretical framework derives explicit bounds for both fixed and diminishing learning rates, characterizing the convergence behavior in terms of the preconditioned condition number and the preconditioned noise level.

Key Contributions

1. Theoretical Bounds for Preconditioned SGD

The paper derives explicit convergence bounds that factorize the late-stage behavior into two components:

Effective Conditioning: The convergence rate is governed by the condition number in the $M$ -geometry ( $\hat{L}/\hat{c}$ for strongly convex, $\hat{L}/\hat{\mu}_{PL}$ for local nonconvex). Improved conditioning allows for larger admissible step sizes and faster contraction.
Noise Floor: The attainable error floor (for fixed step sizes) or the leading constant (for diminishing step sizes) scales with the product of the effective condition number and the preconditioned noise level $K$ . Here, $K$ is defined as an upper bound on the trace of the preconditioned noise covariance, specifically $K \approx \text{tr}(M^{-1}\Sigma(w))$ .

2. Basin-Stability Guarantee

For nonconvex objectives, the authors provide a probabilistic lower bound on the likelihood that iterates remain within a well-behaved local basin up to a finite time horizon. This bound explicitly accounts for:

The objective barrier height (determined by local quadratic growth).
The probability of rare one-step overshoots out of the basin, which depends on the conditional second moment of the gradient in the $M^{-1}$ -norm.

3. Design Criteria

Synthesizing the theory, the paper proposes a practical design principle for preconditioners: Choose $M$ to improve local conditioning while attenuating noise in the $M^{-1}$ -norm.

Curvature-aware methods (e.g., Fisher, Gauss-Newton, Hessian) primarily improve conditioning ( $\hat{L}/\hat{c}$ ).
Methods aligned with gradient noise structure (e.g., Fisher-based) effectively reduce the preconditioned noise level $K$ .
The optimal preconditioner balances these two effects to minimize the product governing the noise floor.

Experimental Results

The authors validate their theory through two types of experiments:

Diagnostic Quadratic Model: Using a synthetic quadratic objective with controlled eigenvalues and noise, the authors demonstrate that:
- Deflating large eigenvalues reduces the smoothness constant $\hat{L}$ and the noise level $K$ , lowering the floor.
- Deflating small eigenvalues increases the PL constant but simultaneously increases $K$ , resulting in only modest gains due to the counteracting noise effect.
- The steady-state loss strictly tracks the theoretical noise-floor scaling.
SciML Benchmarks: Experiments on three tasks (Noisy Franke surface regression, Physics-Informed Neural Networks for Poisson equations, and Green's function learning for convection-diffusion) reveal:
- Curvature-aware preconditioners (specifically Conjugate Gradient with Gauss-Newton/Fisher approximations, CG-GGN) consistently outperform vanilla SGD, Momentum, Adam, and L-BFGS in the late stage.
- CG-GGN achieves the lowest training losses and fastest wall-clock convergence. The authors attribute this to the Gauss-Newton matrix's alignment with the gradient covariance structure (providing effective noise whitening) and its positive semi-definiteness (avoiding negative curvature issues present in full Hessians).
- Quantitative analysis on the PINN and Green's function tasks confirms that CG-GGN reduces the effective smoothness constant $\hat{L}$ by orders of magnitude (e.g., 3710x) and significantly reduces the estimated trace of the preconditioned noise covariance (e.g., 1505x).

Significance and Claims

The paper claims to provide a principled, geometry-aware framework for understanding late-stage SGD optimization. Its significance lies in:

Unifying Theory: It connects the empirical success of diverse preconditioners (adaptive, second-order, quasi-Newton) to a single theoretical mechanism: the trade-off between local conditioning and preconditioned noise attenuation.
SciML Relevance: It highlights that in Scientific Machine Learning, where small reductions in loss are tied to physical constraints and stability, the asymptotic noise floor is a critical metric often overlooked by standard convergence rates.
Design Guidance: It offers a concrete criterion for selecting preconditioners: one must not only align with curvature but also explicitly dampen gradient noise in the metric defined by the preconditioner.

The authors maintain modesty by noting that their theory assumes a fixed SPD metric $M$ and does not fully analyze fully time-varying natural gradient methods, though it provides a local lens for them. They identify covariance-aware preconditioning and online diagnostics for local constants as important future directions.