aiqbench analysis
This research reframes subsampled natural gradient descent as a sketch-and-project method, resolving theoretical gaps in small-sample regimes by replacing standard stochastic proxies with squared volume sampling. The analysis provides global convergence guarantees for single mini-batch settings and clarifies how SNG leverages spectral decay in the model Jacobian more effectively than SGD. Additionally, the work offers a rigorous derivation of the SPRING momentum scheme as a natural outcome of accelerated sketch-and-project techniques.
A Sketch-and-Project Analysis of Subsampled Natural Gradient Algorithms
arXiv:2508.21022v3 Announce Type: replace 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.
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.
Machine Learning (cs.LG); Optimization and Control (math.OC); Machine Learning (stat.ML)
arXiv:2508.21022 [cs.LG]
(or arXiv:2508.21022v3 [cs.LG] for this version)
Full text under the publisher's open licence · 211 words