Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis

By Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, Pascal Vincent. 2018

210 citations recorded by the source corpus.

Optimization algorithms that leverage gradient covariance information, such as variants of natural gradient descent (Amari, 1998), offer the prospect of yielding more effective descent directions. For models with many parameters, the covariance matrix they are based on becomes gigantic, making them inapplicable in their original form. This has motivated research into both simple diagonal approximations and more sophisticated factored approximations such as KFAC (Heskes, 2000; Martens & Grosse, 2015; Grosse & Martens, 2016). In the present work we draw inspiration from both to propose a novel approximation that is provably better than KFAC and amendable to cheap partial updates. It consists in tracking a diagonal variance, not in parameter coordinates, but in a Kronecker-factored eigenbasis, in which the diagonal approximation is likely to be more effective. Experiments show improvements over KFAC in optimization speed for several deep network architectures.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Optimization & Training

Citations in the corpus

DLCFT: Deep Linear Continual Fine-Tuning for General Incremental Learning (2022, methodology): “With neural networks having intra-batch dependency due to batch normalization, Extended K-FAC (XK-FAC) [22] generalized the method to take intra-batch dependency into account.”
Continual Learning With Extended Kronecker-Factored Approximate Curvature (2020, methodology): “[8] introduces Eigenvaluecorrected Kronecker factorization, an approximate factorization of the Fisher that is computationally manageable, accurate, and amendable to cheap partial updates.”
EigenDamage: Structured Pruning in the Kronecker-Factored Eigenbasis (2019, methodology): “It shows that even C-OBD and C-OBS can almost match NN slimming on CIFAR10 and CIFAR100 with VGGNet, if trained with L 1 sparsity on BatchNorm, and outperform when trained without it. signiﬁcant when the pruning ratio is high, e.g. 90% , especially on more complicated networks, e.g. (Pre)ResNet, whi”
LQF: Linear Quadratic Fine-Tuning (2020, methodology): “We use Kronecker-factorized (KFAC) approximation of the Fisher [35, 8, 15] to efficiently approximate the curvature of the loss function.”
Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (2023, unknown): “Martens & Grosse (2015) and follow-up works (Ba et al., 2017; George et al., 2018; Martens et al., 2018) proposed to approximate the Hessian based on the structure of neural networks.”
Structured Pruning for Deep Convolutional Neural Networks: A Survey (2023, methodology): “This second algorithm adopts K-FAC [242] and projects weight space to a Kronecker-Factored eigenspace (KFE) [243] where there is little correlation.”
BackPACK: Packing more into backprop (2019, background): “This work sparked a new line of research to improve the second-order approximation (Grosse & Martens, 2016; Botev et al., 2017; Martens et al., 2018; George et al., 2018).”
Limitations of the Empirical Fisher Approximation (2019, background): “There are also several works building upon KFAC which substitute the EF for the Fisher [George et al., 2018, Osawa et al., 2018].”
A Trace-restricted Kronecker-Factored Approximation to Natural Gradient (2020, unknown): “…considered the eﬃcient Kronecker-factored approximation to the FIM, such as KFAC (Martens and Grosse, 2015; Grosse and Martens, 2016; Martens et al., 2018,Martens and Grosse, 2015; Zhang et al., 2018; Bae et al., 2018) and Eigenvalue-corrected Kronecker Factorization (EKFAC) (George et al., 2018).”
A survey of uncertainty in deep neural networks (2021, background): “A recent extension can be found in [201] where the authors propose to re-scale the eigen-values of the Kronecker factored matrices so that the diagonal variance in its eigenbasis is accurate.”
Eigenvalue Corrected Noisy Natural Gradient (2018, unknown): “Eigenvalue corrected K-FAC (EK-FAC) [George et al., 2018] extends on K-FAC to compute a more accurate diagonal re-scaling factor in K-FAC eigenbasis.”
Laplace Redux - Effortless Bayesian Deep Learning (2021, methodology): “Diagonal approximations perform significantly worse than KFAC and are therefore not suggested.”