In plain words: Instead of solving a costly matrix equation each step, this loop lets a physical system settle to equilibrium to do the math, making a second-order update cost like plain gradient descent. It beat top digital first- and second-order trainers on classification and language-model tuning.
Abstract
Second-order training methods have better convergence properties than gradient descent but are rarely used in practice for large-scale training due to their computational overhead. This can be viewed as a hardware limitation (imposed by digital computers). Here we show that natural gradient descent (NGD), a second-order method, can have a similar computational complexity per iteration to a first-order method, when employing appropriate hardware. We present a new hybrid digital-analog algorithm for training neural networks that is equivalent to NGD in a certain parameter regime but avoids prohibitively costly linear system solves. Our algorithm exploits the thermodynamic properties of an analog system at equilibrium, and hence requires an analog thermodynamic computer. The training occurs in a hybrid digital-analog loop, where the gradient and Fisher information matrix (or any other positive semi-definite curvature matrix) are calculated at given time intervals while the analog dynamics take place. We numerically demonstrate the superiority of this approach over state-of-the-art digital first- and second-order training methods on classification tasks and language model fine-tuning tasks.
Kaelan Donatella, Samuel Duffield, Maxwell Aifer, Denis Melanson, Gavin Crooks, Patrick J. Coles
arXiv:2405.13817 · cs.LG, cs.ET · submitted May 22, 2024
abstract · pdf · html · 17 pages, 7 figures
∇̃L(θ) = F⁻¹∇L(θ)
which requires solving a linear system. For this, you can use the methods from the author's previous paper [Thermodynamic Linear Algebra](https://arxiv.org/abs/2308.05660).
Since it's hard to implement a full neural network on a thermodynamic computer, the paper suggests running one in parallel to a normal GPU. The GPU computes F and ∇L(θ), but offloads the linear system to the thermo computer, which runs in parallel to the digital system (Figure 1).
It is important to note that the "Runtime vs Accuracy" plot in Figure 3 uses a "timing model" for the TNGD algorithm, since the computer necessary to run the algorithm still doesn't exist.