📊 statistics

A Sketch-and-Project Analysis of Subsampled Natural Gradient Algorithms

This paper advances the analysis of subsampled natural gradient descent by reframing it as a sketch-and-project method with a squared volume sampling proxy, thereby establishing global convergence guarantees for single mini-batches and providing new insights into its spectral advantages over SGD and its connection to the SPRING momentum scheme.

Original authors: Gil Goldshlager, Jiang Hu, Lin Lin

Published 2026-06-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Gil Goldshlager, Jiang Hu, Lin Lin

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 find the perfect recipe for a cake (the "optimal solution") by tasting small samples of batter. You have a massive kitchen with millions of ingredients (parameters), but you can only taste a tiny spoonful at a time (a "mini-batch").

This paper is about a specific cooking technique called Subsampled Natural Gradient Descent (SNG). In the world of scientific machine learning (like simulating quantum physics or solving complex equations), this technique is a superstar because it finds the perfect recipe much faster and more accurately than standard methods. However, until now, scientists didn't fully understand why it worked so well when they only had a tiny spoonful of batter to taste.

Here is the breakdown of the paper's findings using simple analogies:

1. The Old Way: The "Two-Spoon" Mistake

For a long time, mathematicians tried to explain why SNG works by pretending you use two different spoons for every step:

Spoon A tastes the batter to figure out which direction to move.
Spoon B tastes the batter to figure out how to adjust your sensitivity (the "preconditioner").

They assumed these two spoons were independent. The problem? In real life, you only have one spoon. You use the same small sample to do both jobs. When the sample is tiny (which is common in scientific problems), the "two-spoon" theory breaks down. It's like trying to predict the weather by looking at two different, unrelated clouds; it doesn't tell you what's actually happening in the sky.

2. The New Lens: "Sketch-and-Project"

The authors propose a new way to look at the problem. Instead of thinking of SNG as a fancy gradient descent, they view it as a Sketch-and-Project method.

The Sketch: Imagine you have a giant, complex 3D sculpture (the problem). You take a quick, blurry photo of it from one angle (the "sketch"). This photo is incomplete, but it captures the most important features.
The Project: You then try to adjust your current model to fit that blurry photo.

The paper argues that SNG is essentially taking a blurry photo of the problem and projecting your current guess onto the solution space defined by that photo. This perspective is powerful because it works perfectly even when your photo is tiny and blurry.

3. The Secret Sauce: "Squared Volume Sampling"

To prove their theory, the authors used a mathematical trick called Squared Volume Sampling (SVS).

The Analogy: Imagine you are trying to pick a few photos to represent a whole album. A normal method picks photos randomly. SVS is smarter: it picks photos that are as different from each other as possible. If you pick two photos that look exactly the same, you learn nothing new. SVS ensures you pick a diverse set of "angles" that cover the most ground.

The authors showed that if you analyze SNG as if it were using this "smart diversity" sampling (even if real-world SNG uses simple random sampling), the math suddenly makes sense. It reveals that the expected direction of the algorithm is actually a perfect, pre-conditioned step toward the solution.

4. What This Tells Us About Speed

The paper reveals two big insights about why SNG is so fast in small-sample settings:

The "Spectral Decay" Advantage: In many scientific problems, the "information" in the data is concentrated in a few key directions (like a few dominant flavors in a soup), while the rest is just noise. Standard methods (like SGD) treat all directions equally. SNG, however, is like a chef who instantly recognizes the dominant flavors and focuses only on them. The paper proves mathematically that SNG can exploit this "spectral decay" to converge super-fast, even with tiny samples.
The "SPRING" Algorithm: There is a popular, fast version of SNG called SPRING (which adds momentum, like a rolling ball). For years, people used it because it worked, but didn't know why. The paper shows that SPRING is actually just a natural result of applying "acceleration" techniques to the Sketch-and-Project method. It's not magic; it's just the math of projecting and accelerating working together.

Summary

The paper says: "Stop trying to analyze this complex algorithm by pretending we have two separate data sets. Instead, look at it as a method that takes a quick sketch of the problem and projects the solution onto it. When you do this, you realize that SNG is incredibly efficient at finding solutions in small-sample settings because it naturally focuses on the most important parts of the data, and popular acceleration tricks like SPRING are just the logical next step in this process."

This helps scientists trust these algorithms more when they are solving high-precision problems in physics and chemistry, where they often can't afford to use massive amounts of data.

Technical Summary: A Sketch-and-Project Analysis of Subsampled Natural Gradient Algorithms

1. Problem Statement

Subsampled Natural Gradient (SNG) algorithms have emerged as a critical tool for high-precision scientific machine learning (SciML), particularly in applications like Neural Network Wavefunctions (NNWs) and Physics-Informed Neural Networks (PINNs). In these domains, precision is prioritized over generalization, and SNG has demonstrated superior performance compared to standard optimizers like Adam or SGD.

However, the theoretical understanding of SNG remains limited, especially in small-sample settings where the number of samples ( $k$ ) is orders of magnitude smaller than the number of parameters ( $n$ ). Standard analyses rely on stochastic preconditioning, which typically requires a theoretical proxy using two independent mini-batches to decouple the stochastic gradient from the stochastic preconditioner. The authors argue that this decoupling tactic fails to provide meaningful insights into realistic small-sample regimes. In such settings, the stochastic preconditioner captures only a tiny subset of the parameter space, making it unlikely to be informative for an independent stochastic gradient. Consequently, existing analyses often yield only generic global convergence guarantees rather than detailed convergence rates that explain the empirical success of SNG.

2. Methodology

The paper proposes a paradigm shift in analyzing SNG: instead of viewing it as a stochastic preconditioning scheme, the authors analyze it as a sketch-and-project method. This perspective treats the SNG update as the minimal-norm solution to a subsampled linear system, drawing connections to randomized block Kaczmarz methods.

Key Technical Shifts:

Rejection of Two-Mini-Batch Proxies: The authors discard the standard theoretical proxy that uses two independent mini-batches to decouple gradients and preconditioners.
Adoption of Squared Volume Sampling (SVS): To rigorously analyze the single mini-batch algorithm without decoupling, the authors introduce a new theoretical proxy based on Squared Volume Sampling (SVS). SVS is a sampling strategy where rows are selected with probability proportional to the squared volume of the parallelepiped formed by the sampled rows (regularized by $\lambda$ ).
Theoretical Insight via SVS: While SVS is not proposed as a practical implementation strategy (due to computational costs), it serves as a powerful analytical tool. It enables exact computation of expectation values involving the inverse of subsampled matrices, a task that is otherwise intractable under general sampling distributions.

3. Key Contributions

A. Analysis Without Decoupling

The authors prove that under SVS, the expectation of the SNG direction becomes equal to a preconditioned gradient descent step, even without decoupling the gradient and preconditioner.

Lemma 4.1: Shows that $\mathbb{E}_S [J_S^{+(\lambda)} r_S] = \tilde{W} J^\top r$ , where $\tilde{W}$ is a specific preconditioner derived from the expected projector $P$ and the natural gradient preconditioner $J^\top J$ .
Theorem 4.2: Establishes global convergence guarantees for SNG using a single mini-batch of any size, provided the sampling follows SVS. This removes the need for the unrealistic two-mini-batch assumption.

B. Characterization of Convergence Rates in Small-Sample Settings

Focusing on the Linear Least-Quadratics (LLQ) problem (a second-order approximation of consistent parametric optimization), the authors derive an explicit convergence rate.

Theorem 5.1: The convergence rate is governed by the quantity $\alpha/\gamma$ $α / γ$ .
- $\alpha$ : The standard sketch-and-project convergence rate (minimal eigenvalue of the expected projector).
- $\gamma$ : A new quantity related to the second moment of the sketch-and-project step.
Insight on Spectral Decay: The appearance of $\alpha$ suggests that SNG's advantage over SGD lies in its ability to exploit spectral decay in the model Jacobian. Unlike SGD, where the convergence rate scales at most linearly with sample size, $\alpha$ can scale superlinearly (or even exponentially) with sample size when the Jacobian exhibits fast spectral decay. This allows SNG to compensate for its higher iteration cost ( $O(k^2)$ vs $O(k)$ ).
Role of $\gamma$ : The quantity $\gamma$ modulates the benefits of the sketch-and-project structure. Preliminary analysis and numerical evidence suggest $\gamma$ behaves benignly (often $O(1)$ ), meaning it does not significantly suppress the superlinear scaling of $\alpha$ .

C. Explanation of the SPRING Algorithm

The paper provides a theoretical foundation for SPRING (Subsampled Projected-Increment Natural Gradient), a popular structured momentum scheme for SNG.

Theorem 6.1: Demonstrates that SPRING arises naturally as an application of accelerated sketch-and-project methods (specifically Nesterov-accelerated regularized sketch-and-project) to the natural gradient subproblem.
Implication: This explains why SPRING is most effective when the underlying sketch-and-project steps converge slowly (i.e., in small-sample settings where $\alpha$ is small), a phenomenon confirmed empirically.

4. Results and Empirical Validation

The authors validate their theoretical framework through numerical experiments:

Proxy Validity (Figure 2): Experiments on a discrete Poisson problem show that the SVS-based proxy matches the behavior of the realistic single-mini-batch algorithm closely, whereas the traditional two-mini-batch proxy exhibits drastically different behavior.
Spectral Decay (Figure 4): Empirical results confirm that SNG exhibits superlinear convergence rates exactly when the Jacobian has fast spectral decay, validating the theoretical prediction regarding $\alpha$ .
SPRING Acceleration (Figure 5): Experiments show that SPRING provides substantial acceleration for small sample sizes, but the gap narrows as the sample size increases (and $\alpha$ grows), consistent with the conjecture that acceleration benefits are greatest when base convergence is slow.
Scaling of $\alpha$ and $\gamma$ (Figure 3): Numerical estimation shows that $\alpha$ grows superlinearly with sample size, while $\gamma$ exerts only a mild damping effect.

5. Significance and Claims

The paper claims that the sketch-and-project perspective offers a more fruitful understanding of SNG than stochastic preconditioning.

Shift in Focus: The authors argue that when using SNG for SciML, attention should shift from statistical estimation properties (e.g., variance of gradient/preconditioner estimates) to sketch-and-project properties (specifically $\alpha$ and $\gamma$ ).
Theoretical Proxy: SVS-SNG is established as a new, rigorous theoretical proxy for SNG that is faithful to realistic single-mini-batch algorithms in small-sample regimes.
Mechanism of Success: The work provides the first detailed characterization of why SNG succeeds in small-sample settings: its ability to leverage the spectral decay of the Jacobian through the sketch-and-project structure, a mechanism inaccessible to standard SGD analysis.
Foundation for Future Work: The paper sets the stage for leveraging randomized linear algebra to develop and analyze subsampled optimizers, suggesting future directions such as tighter bounds on $\gamma$ , extensions to other sampling distributions, and practical adaptations of SVS-inspired sampling (e.g., introducing negative correlations in mini-batches).

The authors remain modest regarding the practical implementation of SVS, noting that direct implementation is computationally expensive. Instead, they position SVS as a theoretical tool to derive insights and motivate practical heuristics (like structured momentum or specific sampling correlations) that approximate the favorable properties of SVS.