Optimizing Neural Networks with Kronecker-factored Approximate Curvature

By James Martens, Roger Baker Grosse. 2015

1,395 citations recorded by the source corpus.

We propose an efficient method for approximating natural gradient descent in neural networks which we call Kronecker-Factored Approximate Curvature (K-FAC). K-FAC is based on an efficiently invertible approximation of a neural network's Fisher information matrix which is neither diagonal nor low-rank, and in some cases is completely non-sparse. It is derived by approximating various large blocks of the Fisher (corresponding to entire layers) as being the Kronecker product of two much smaller matrices. While only several times more expensive to compute than the plain stochastic gradient, the updates produced by K-FAC make much more progress optimizing the objective, which results in an algorithm that can be much faster than stochastic gradient descent with momentum in practice. And unlike some previously proposed approximate natural-gradient/Newton methods which use high-quality non-diagonal curvature matrices (such as Hessian-free optimization), K-FAC works very well in highly stochastic optimization regimes. This is because the cost of storing and inverting K-FAC's approximation to the curvature matrix does not depend on the amount of data used to estimate it, which is a feature typically associated only with diagonal or low-rank approximations to the curvature matrix.

Subfields: Artificial Intelligence, Statistical and Nonlinear Physics, 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 Statistical and Nonlinear Physics

In Optimization & Training

Citations in the corpus

Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model (2019, methodology): “Through large scale experiments with Adam [Kingma and Ba, 2014] and K-FAC [Martens and Grosse, 2015], we confirm that, as predicted by the NQM, preconditioning extends perfect batch size scaling to larger batch sizes than are possible with momentum SGD alone.”
Don't Decay the Learning Rate, Increase the Batch Size (2017, background): “K-FAC (Martens & Grosse, 2015) is also gaining popularity as an efﬁcient alternative to SGD. Wilson et al. (2017) argued that adaptive optimization methods tend to generalize less well than SGD and SGD with momentum (although they did not include K-FAC in their study), while our work reduces the…”
Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks (2016, methodology): “Such an approximate inverse can for example be obtained by using a Kronecker factored approximation to the Fisher matrix and inverting it (KFAC, [19]), by using an approximate Cholesky factorization of the inverse Fisher matrix (FANG, [8]), or by whitening the input of each layer in the neural…”
On Empirical Comparisons of Optimizers for Deep Learning (2019, background): “We can also view (cid:15) as setting a trust region radius (Martens and Grosse, 2015; Adolphs et al., 2019) and controlling an interpolation between momentum and diagonal natural gradient descent, by either diminishing or increasing the effect of v t on the update direction.”
Sharpness-Aware Minimization for Efficiently Improving Generalization (2020, background): “, and others (Duchi et al., 2011; Dozat, 2016; Martens & Grosse, 2015)) has become an important design choice, though understanding of its relationship to model generalization remains nascent (Shirish Keskar et al., 2016; Wilson et al., 2017; Shirish Keskar & Socher, 2017; Agarwal et al., 2020;…”
Fast and Furious Convergence: Stochastic Second Order Methods under Interpolation (2019, methodology): “In order to further scale these methods to high-dimensional datasets or models with a large number of parameters, there has been a rich body of research in developing computationally efficient approximations to the Hessian matrix (Becker et al., 1988; Liu and Nocedal, 1989; Martens and Grosse, 2015;”
Trust-PCL: An Off-Policy Trust Region Method for Continuous Control (2017, background): “It is often more eﬀective to deﬁne an appropriate Riemannian metric that respects the loss surface (Amari, 2012), which allows much steeper descent directions to be identiﬁed within a local neighborhood ( e.g., Amari (1998); Martens & Grosse (2015)).”
Large Batch Optimization for Deep Learning: Training BERT in 76 minutes (2019, background): “…- 50 training on ImageNet in a few minutes (You et al., 2018; Iandola et al., 2016; Codreanu et al., 2017; Akiba et al., 2017; Jia et al., 2018; Smith et al., 2017; Martens & Grosse, 2015; Devarakonda et al., 2017; Mikami et al., 2018; Osawa et al., 2018; You et al., 2019; Yamazaki et al., 2019).”
A Comprehensive Survey of Continual Learning: Theory, Method and Application (2023, methodology): “To overcome these issues, the Hessian can be approximated by the Fisher information matrix (FIM) ) For ease of computation, the FIM can be further simpliﬁed with a diagonal approximation [49], [53] or a Kronecker-factored approximation [50], [54].”
Composing graphical models with neural networks for structured representations and fast inference (2016, background): “Natural gradients are invariant to smooth invertible reparameterizations of the variational family [16, 17] and provide effective second-order optimization updates [18, 11].”
BayesNAS: A Bayesian Approach for Neural Architecture Search (2019, methodology): “…of weight matrix has been widely used in second-order optimization techniques to speed up the training process (LeCun et al., 1990; Amari, 1998), it still remains infeasible to calculate explicit Hessian directly due to the intensive computation burden (Martens & Grosse, 2015; Botev et al., 2017).”
Recasting Gradient-Based Meta-Learning as Hierarchical Bayes (2018, methodology): “Martens & Grosse (2015) developed Kronecker-factored approximate curvature (K-FAC), a scheme for approximating the curvature of the objective function of a neural network with a block-diagonal approximation to the Fisher information matrix.”
Bayesian Model-Agnostic Meta-Learning (2018, methodology): “However, unlike our proposed model, the posterior on validation set is approximated by local Laplace approximation and used a relatively complex 2nd-order optimization using K-FAC [Martens and Grosse, 2015].”
Incorporating neuro-inspired adaptability for continual learning in artificial intelligence (2023, background): “Although p(θ|DA) is generally intractable, we can locally approximate it with a second-order Taylor expansion around θ∗ A = argmaxθ log p(θ|DA), resulting in a Gaussian distribution whose mean is θ∗ A and precision matrix is the Hessian of the negative log posterior [7, 11, 69].”
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.”
Recent Advances of Continual Learning in Computer Vision: An Overview (2021, neutral): “[16] also proposed an efficient alternative to EWC called EWC++ which maintains a single Fisher information matrix for all the previously learned tasks and updates the matrix using moving average [23].”
Natural continual learning: success is a journey, not (just) a destination (2021, methodology): “In this work, we use a Kronecker-factored approximation (Martens and Grosse, 2015; Ritter et al., 2018).”