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

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 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.

It is increasingly clear that models known as “artificial neural networks” are effective at learning patterns arising in words, images, and many other forms of information. A natural question is whether this technology can also be used to learn solutions to the mathematical equations that govern molecules, materials, fluids, and other physical systems. A major challenge is that scientific applications require precise quantitative predictions, whereas mainstream machine learning methods are designed to produce only qualitatively correct predictions. Recently, a class of learning methods called “subsampled natural gradient algorithms” has emerged as a promising way to improve the quantitative precision of artificial neural networks for scientific applications. However, it has not been explained why these methods work so well. Our work gives a new explanation for this phenomenon. Our main idea is to connect subsampled natural gradient algorithms to a better-understood class of methods called "sketch-and-project algorithms," which tackle large problems by repeatedly solving smaller, randomly sampled subproblems. This connection explains a number of recent observations about subsampled natural gradient algorithms, and it suggests new directions for future work aimed at developing even better algorithms.