A Sketch-and-Project Analysis of Subsampled Natural Gradient Algorithms
SVS-SNG analyzes single-mini-batch natural gradient through sketch-and-project geometry, yielding convergence rates governed by α/γ.
Key Findings
Methodology
The paper reinterprets subsampled natural gradient (SNG) as a regularized randomized block Kaczmarz, or sketch-and-project, method rather than as an estimator of a deterministic preconditioner. For linear least-quadratics, it uses squared volume sampling (SVS), with p(S) proportional to det(J_SJ_S^⊤+λI). This preserves the coupling between the stochastic gradient and stochastic preconditioner and makes the coupled expectation analytically tractable.
Key Results
- Lemma 4.1 proves that for S∼SVS(J,k,λ), E[J_S^{+(λ)}r_S]=f_WJ^⊤r, where f_W=(J^⊤J)^{-1/2}P(J^⊤J)^{-1/2} and P=E[P(S)]. Thus a single mini-batch has an expected preconditioned-gradient form without independent gradient and preconditioner batches.
- Theorem 4.2 establishes global convergence for any batch size k under λ>0 and diminishing step sizes. Theorem 5.1 gives an LLQ rate governed by α/γ, where α is the smallest positive eigenvalue of the expected projector and γ captures a second moment of the sketch-and-project step.
- Figure 2 uses a discrete Poisson problem with m=100, n=7801, k=10 and 10^3 iterations. SVS-SNG closely tracks realistic uniform-sampling SNG, whereas the two-batch proxy behaves very differently. Figure 4 supports the spectral-decay explanation; Figure 5 shows stronger SPRING benefits at small batch sizes.
Significance
The work addresses a central mismatch in SNG theory: practical algorithms use one small batch for both the gradient and Jacobian preconditioner, while conventional analyses decouple them with two independent batches. When k is far smaller than n, that proxy can be qualitatively misleading. By shifting attention from estimator variance to Jacobian row spaces, spectral decay, and projection geometry, the paper offers a more realistic explanation for high-precision scientific machine learning in neural-network wavefunctions and PINNs.
Technical Contribution
The paper contributes an exactly analyzable coupled expectation under SVS, a single-mini-batch global convergence theorem, an α/γ rate characterization for consistent LLQ, and a derivation of SPRING from accelerated sketch-and-project methods. The practical SNG step is θ_{t+1}=θ_t−ηJ_S^{+(λ)}r_S, requiring inversion of a k×k kernel rather than an n×n Fisher matrix.
Novelty
The novelty is not a claim that full SVS should replace uniform sampling in production. Instead, SVS is introduced as a principled theoretical proxy that connects SNG, randomized block Kaczmarz, and accelerated sketch-and-project analysis. Unlike independent-mini-batch theory, it retains the realistic coupling and exposes why very small batches can still benefit from Jacobian spectral structure.
Limitations
- SVS is mainly an analytical device. Sampling from it generally requires access to all of J or many of its rows, with costs described as O(m·poly(k)) in general, so the theory does not directly provide a cheap production sampler.
- The sharp rate analysis focuses on consistent linear least-quadratics. Nonlinear, inconsistent, and stochastically normalized settings receive extensions or assumptions, while the benign behavior of γ remains supported mainly by preliminary analysis and numerics.
Future Work
The authors point toward approximate SVS, negatively correlated mini-batches, MCMC, and oversample-then-subselect schemes. Future work should characterize γ under realistic Jacobian spectra, test the α/γ prediction on large neural-network wavefunction and PINN systems, build samplers that access only a few Jacobian rows, and optimize SPRING’s regularization and momentum parameters.
AI Executive Summary
Scientific machine learning often values precision over ordinary predictive performance. Neural-network wavefunctions and physics-informed neural networks therefore use natural-gradient methods, which account for how parameter changes affect the represented function. Subsampled natural gradient (SNG) makes these methods affordable, but creates a theoretical problem: the same small mini-batch determines both the stochastic gradient and the stochastic preconditioner. Standard analyses avoid the coupling by using two independent batches, an approximation that can become misleading when k is much smaller than the parameter dimension n.
Goldshlager, Hu, and Lin propose a sketch-and-project analysis centered on squared volume sampling (SVS), with probability proportional to det(J_SJ_S^⊤+λI). They reinterpret SNG as a regularized randomized block Kaczmarz step. Under SVS, Lemma 4.1 shows that the expected coupled direction equals f_WJ^⊤r, a preconditioned gradient step. This yields single-mini-batch global convergence for arbitrary batch size and, in the linear least-quadratics model, a rate controlled by α/γ: α measures expected projection quality, while γ captures second-moment effects.
The evidence is deliberately diagnostic. In Figure 2, a discrete Poisson problem uses m=100, n=7801, k=10 and 10^3 iterations. SVS-SNG closely matches realistic uniform-sampling SNG as step size and regularization vary, whereas the independent two-batch proxy differs dramatically. Figure 4 supports the claim that SNG exploits Jacobian spectral decay more effectively than SGD; Figure 5 indicates that SPRING, derived from accelerated sketch-and-project methods, helps most when batches are small. SVS remains expensive, and γ lacks a complete general theory, but the framework redirects attention from generic variance estimates to the geometry of the model Jacobian.
Deep Analysis
Background
Natural gradient descent seeks to imitate gradient descent in function space and is used in neural-network wavefunctions, variational Monte Carlo, and PINNs. Its full update is θ_{t+1}=θ_t−η(J^⊤J+λI)^+J^⊤r. SNG replaces J and r by sampled J_S and r_S, reducing the linear algebra to a k×k kernel system and enabling regimes such as k=10^3 with n=10^6.
Core Problem
The update J_S^{+(λ)}r_S contains an inverse random matrix and a gradient generated by the same samples. Independent two-batch analyses remove this difficulty but fail to describe realistic k≪n behavior. They also obscure whether SNG gains arise from variance reduction or from the geometry and spectrum of the Jacobian.
Innovation
First, SNG is identified with regularized randomized block Kaczmarz and sketch-and-project iteration. Second, SVS provides an analyzable proxy that preserves gradient–preconditioner coupling. Third, LLQ convergence is expressed through α/γ, combining expected projection quality with step second moments. Fourth, SPRING is explained as an accelerated sketch-and-project method rather than an ad hoc momentum modification.
Methodology
- �� Model the local problem as v_θ=Jθ with L(v)=1/2v^⊤Hv−v^⊤b.
- �� Sample rows to form J_S and r_S, then apply θ_{t+1}=θ_t−ηJ_S^{+(λ)}r_S.
- �� Define P(S)=J_S^{+(λ)}J_S, P=E[P(S)], and α=λ_min^+(P).
- �� Draw S with p(S)=det(J_SJ_S^⊤+λI)/Σ_{|S'|=k}det(J_{S'}J_{S'}^⊤+λI).
- �� Use SVS expectation identities to obtain f_WJ^⊤r, then combine stochastic-optimization and sketch-and-project arguments for convergence and rate bounds.
Experiments
Figure 2 compares realistic uniform-sampling SNG, SVS-SNG, and the independent two-batch proxy on a discrete Poisson problem with m=100, n=7801, and k=10. It runs 10^3 iterations and five repetitions, varying λ at fixed η and η at fixed λ. Figure 3 studies γ, Figure 4 examines SNG versus SGD under spectral decay, and Figure 5 evaluates SPRING; appendices add LLQ experiments and implementation details.
Results
The SVS proxy closely reproduces realistic SNG in Figure 2, while the two-batch proxy shows markedly different error behavior. The theory allows global convergence with any batch size under stated smoothness, boundedness, regularization, and step-size conditions. The α/γ rate explains why projection geometry matters, and the experiments support both spectral-decay benefits for SNG and stronger SPRING gains when the batch is small.
Applications
The framework applies directly to stochastic reconfiguration for neural-network wavefunctions, variational quantum ground-state search, PINN optimization, and subsampled Gauss–Newton for nonlinear least squares. Users need efficient Jacobian-row and gradient evaluation and stable regularization. In practice, the paper suggests designing diverse or negatively correlated batches rather than implementing full SVS.
Limitations & Outlook
SVS may require global Jacobian access or many row queries, making it expensive. The strongest analysis is local and based on consistent LLQ; nonlinear global behavior, stochastic normalization, and approximate-SVS bias remain open. The quantity γ has only preliminary characterization, while SPRING parameter selection and large-scale scientific benchmarks need systematic study.
Plain Language Accessible to non-experts
Imagine a factory with thousands of control knobs and a product that must be made extremely precise. Ordinary adjustment looks at the product’s errors and changes knobs in the indicated direction. Natural adjustment is smarter: it also notices that several knobs may produce almost the same effect, so it avoids wasting effort by turning all of them equally.
To save time, the factory inspects only a few product features at each round. But those same features are used both to identify the defect and to decide which knobs matter. Older mathematical analyses pretend that two unrelated inspection teams perform these jobs. With very few inspected features, that is like using one customer order to diagnose a problem and a completely different order to choose the repair—it may describe the wrong factory.
This paper studies a selection rule called squared volume sampling, which favors inspection sets containing complementary information. It proves that, on average, the resulting adjustment still behaves like a well-calibrated correction. In the Poisson experiment, there are 7801 parameters but only 10 samples per step; the SVS proxy follows the real algorithm much more closely than the two-team proxy. The broader lesson is that SNG may win not merely because it is less noisy, but because it recognizes that only a few important patterns dominate the system.
ELI14 Explained like you're 14
Imagine a video game with 7801 control buttons, but each turn you can inspect only 10 spots on the track. A basic strategy sees a mistake and changes nearby buttons. Natural gradient is more thoughtful: it checks whether several buttons actually do the same job, then chooses a cleaner combination of changes.
Here is the tricky part: the same 10 spots tell you both what went wrong and which buttons seem useful. Older analyses pretend those jobs use two different groups of spots. That is like using one exam to find your mistakes and a completely different exam to decide what to study. With tiny samples, the prediction can be wildly off!
The paper uses squared volume sampling, or SVS, to choose spots that contain different, complementary information. It also views each update as moving the current answer back toward a set of rules it should satisfy. The math then shows that, on average, the one-batch method still moves in a reliable direction.
In a discrete Poisson test, the system has 7801 parameters, uses only 10 samples at a time, and runs for 1000 rounds. The SVS version behaves much more like the real algorithm than the old two-batch version. SPRING adds acceleration, especially when batches are tiny. The surprising idea is that the advantage is not just “less random noise”; SNG can also exploit the fact that a few important patterns dominate the model!
Glossary
Subsampled Natural Gradient
A natural-gradient method that estimates the function gradient and Jacobian preconditioner from a small sample. It replaces an n×n solve with a k×k kernel solve.
The paper’s main algorithm, θ_{t+1}=θ_t−ηJ_S^{+(λ)}r_S.
Sketch-and-Project
A method that compresses a linear system with a low-dimensional sketch and projects the current iterate onto the sketched solution space. Its theory supports arbitrary sketch sizes.
The central lens used to analyze SNG and SPRING.
Squared Volume Sampling
A sampling rule selecting row subsets with probability proportional to det(J_SJ_S^⊤+λI), favoring complementary information. It is related to determinantal point processes.
A theoretical proxy that preserves gradient–preconditioner coupling.
Randomized Block Kaczmarz
An iterative solver that selects a random block of equations and projects the current point onto the block’s solution set. Regularization improves stability.
SNG is obtained by applying its regularized form to the natural-gradient subproblem.
SPRING
Subsampled projected-increment natural gradient, a structured momentum scheme for SNG. The paper interprets it as accelerated sketch-and-project iteration.
Used to explain why acceleration is most helpful for slowly converging small-batch projections.
Open Questions Unanswered questions from this research
- 1 γ is linked to the second moment of the sketch-and-project step, but the paper provides only preliminary bounds and numerical evidence. A general characterization under realistic Jacobian spectra remains open.
- 2 SVS is expensive because it may require broad Jacobian access. An approximate sampler that uses only a few rows while preserving negative dependence is still needed.
- 3 The predictive value of LLQ theory for nonlinear wavefunction and PINN training, including stochastic normalization errors, requires larger and more diverse benchmarks.
Applications
Immediate Applications
High-precision PINN training
PINN researchers can use SNG with residual Jacobian rows and a k×k kernel solve to reduce iteration cost. Regularization and sample coverage must be monitored; complementary, negatively correlated batches are a practical direction suggested by the theory.
Neural-network wavefunction optimization
Neural-network wavefunction and variational Monte Carlo workflows can use SNG or SPRING when precision is more important than generic prediction. The paper recommends examining Jacobian spectral decay, not only gradient variance, when selecting batch sizes.
Long-term Vision
Low-cost approximate SVS optimizers
MCMC, oversample-then-subselect, or other diversity-aware procedures could approximate SVS without constructing the full Jacobian. Such methods may retain its geometric benefits while approaching the cost of ordinary mini-batching.
Abstract
Subsampled natural gradient descent (SNG) has been used to enable high-precision scientific machine learning, but standard analyses based on stochastic preconditioning fail to provide insight into realistic small-sample settings. We overcome this limitation by instead analyzing SNG as a sketch-and-project method. Motivated by this lens, we discard the usual theoretical proxy which decouples gradients and preconditioners using two independent mini-batches, and we replace it with a new proxy based on squared volume sampling. Under this new proxy we show that the expectation of the SNG direction becomes equal to a preconditioned gradient descent step even in the presence of coupling, leading to (i) global convergence guarantees when using a single mini-batch of any size, and (ii) an explicit characterization of the convergence rate in terms of quantities related to the sketch-and-project structure. These findings in turn yield new insights into small-sample settings, for example by suggesting that the advantage of SNG over SGD is that it can more effectively exploit spectral decay in the model Jacobian. We also extend these ideas to explain a popular structured momentum scheme for SNG, known as SPRING, by showing that it arises naturally from accelerated sketch-and-project methods.