Natural gradient descent

Climbing slower on the tricky bits

July 18, 2019 — May 26, 2020

functional analysis
geometry
optimization
statistics

A placeholder.

Figure 1: The misalignment of the magnetic and geographic poles is not an ideal metaphor for the misalignment of natural and Euclidean gradients (because magnetic field descent leads to the wrong optimum, not the wrong rate) but I can’t be searching online for illustrations all day there is work to do.

Gradient descent with the natural gradient, or a close approximation thereto. Related: Information geometry, which formalises and generalises this. What is that natural gradient then?

Salimbeni, Eleftheriadis, and Hensman (2018):

The ordinary gradient turns out to be an unnatural direction to follow for variational inference since we are optimising a distribution, rather than a set of parameters directly. One way to define the gradient is the direction that achieves maximum change subject to a perturbation within a small Euclidean ball. To see why the Euclidean distance is an unnatural metric for probability distributions, consider the two Gaussians \(\mathcal{N}(0, 0.1)\) and \(\mathcal{N}(0, 0.2)\), compared to \(\mathcal{N} (0, 1000.1)\) and \(\mathcal{N}(0, 1000.2)\). The former pair are different and the latter similar, yet in Euclidean distance they are equally far apart in the mean and variance. Using the precision in place of the variance gives the opposite result, yet the distributions are unchanged. There is a fundamental mismatch between the ordinary gradient and the objective function: the gradient is dependent on parameterisation whereas the objective function is not.

Fortunately there is a way to solve the disparity: the natural gradient. The natural gradient can be defined as the direction that achieves maximum change in KL divergence. It is well known that paths following the natural gradient are invariant to reparameterisation (see e.g., Martens (2020)), and that the natural gradient direction is the ordinary gradient rescaled by the inverse Fisher information matrix (Amari 1998).

Martens (2020) says

Natural gradient descent is an optimisation method traditionally motivated from the perspective of information geometry, and works well for many applications as an alternative to stochastic gradient descent. In this paper we critically analyse this method and its properties, and show how it can be viewed as a type of approximate 2nd-order optimisation method, where the Fisher information matrix can be viewed as an approximation of the Hessian. This perspective turns out to have significant implications for how to design a practical and robust version of the method. Additionally, we make the following contributions to the understanding of natural gradient and 2nd-order methods: a thorough analysis of the convergence speed of stochastic natural gradient descent (and more general stochastic 2nd-order methods) as applied to convex quadratics, a critical examination of the oft-used “empirical” approximation of the Fisher matrix, and an analysis of the (approximate) parameterisation invariance property possessed by natural gradient methods, which we show still holds for certain choices of the curvature matrix other than the Fisher, but notably not the Hessian.

Returning to Salimbeni, Eleftheriadis, and Hensman (2018):

Our fundamental problem is to minimise \(-\mathcal{L}(\xi)\). All the approaches we consider find a sequence of parameters \(\left\{\boldsymbol{\xi}_{t}\right\}_{t=0}^{T}\) using the iterative update

\[\boldsymbol{\xi}_{t+1}=\boldsymbol{\xi}_{t}-\gamma_{t} \mathrm{P}_{t}^{-1} \mathbf{g}_{t}, \quad \mathbf{g}_{t}=\left.\nabla_{\xi}^{\top} \mathcal{L}\right|_{\xi=\xi_{t}}\] where \(\gamma_{t}\) denotes the step size and \(\mathrm{P}_{t}^{-1} \mathbf{g}_{t}\) the direction

Natural gradient descent (NGD). Another way of interpreting the update \((\sqrt{5})\) is to use the fact that the direction of steepest descent with respect to a norm \(\|\boldsymbol{\delta}\|_{\mathrm{A}}=\boldsymbol{\delta}^{T} \mathrm{A} \boldsymbol{\delta}\) is given by \(\mathrm{A}^{-1} \nabla_{\boldsymbol{\xi}}+\mathcal{L}\). (This can be seen by minimising \(\frac{1}{\epsilon} \mathcal{L}(\xi+\delta)\) subject to the constraint that \(\|\boldsymbol{\delta}\|_{\mathrm{A}}=\epsilon\) and letting \(\epsilon \rightarrow 0\).) Identifying \(\mathrm{P}\) with \(\mathrm{A},\) the update corresponds to the steepest descent with respect to the norm induced by the matrix \(\mathrm{P}\). Gradient descent (where \(\mathrm{P}\) is the identity and the induced metric is Euclidean) can therefore be seen as moving in the direction that maximises the change in objective with respect to the Euclidean norm of the parameters. The Euclidean norm is an unnatural way to compare two parameter vectors if the parameters correspond to distributions, however. If instead we consider the KL divergence between two distributions and take the small perturbation limit, we obtain \(\mathrm{KL}[q(\mathbf{u} ; \boldsymbol{\xi}), q(\mathbf{u} ; \boldsymbol{\xi}+\boldsymbol{\delta})]=\frac{1}{2} \delta^{\top}\left[\mathbb{E}_{q(\mathbf{u} ; \xi)} \nabla_{\xi}^{2} \log q(\mathbf{u} ; \boldsymbol{\xi})\right] \boldsymbol{\delta}+\mathcal{O}\left(\|\boldsymbol{\delta}\|^{3}\right) .\) Therefore, in a sufficiently small neighbourhood the KL divergence induces a quadratic norm with curvature given by the expected Hessian of the log density. This matrix is known as the Fisher information \(\mathbf{F}_{\xi}\).

\[ \mathbf{F}_{\xi}=-\mathbb{E}_{q(\mathbf{u} ; \xi)} \nabla_{\xi}^{2} \log q(\mathbf{u} ; \boldsymbol{\xi})\] The direction of steepest descent with respect to this norm is called the natural gradient \(\tilde{\nabla}_{\xi} \mathcal{L},\) given by the gradient scaled by the inverse Fisher information: \(\tilde{\nabla}_{\boldsymbol{\xi}} \mathcal{L}=\left(\nabla_{\boldsymbol{\xi}} \mathcal{L}\right) \mathbf{F}_{\boldsymbol{\xi}}^{-1}\) (Amari 1998)

Ahhhh.

Figure 2

1 Connection to Bayesian inference

One interesting family of methods the “Bayesian Learning Rule” tweaks Adam to more directly approximate Bayesian inference (M. E. Khan and Rue 2024) via natural gradients. See Bayes by Backprop.

2 Natural Policy Gradient

What the reinforcement learning people do? A brutally short explanation here, and a longer informal one here, and a downright lengthy one here.

3 Incoming

4 References

Abt, and Welch. 1998. Fisher Information and Maximum-Likelihood Estimation of Covariance Parameters in Gaussian Stochastic Processes.” Canadian Journal of Statistics.
Amari. 1998. Natural Gradient Works Efficiently in Learning.” Neural Computation.
Amari, Karakida, and Oizumi. 2018. Fisher Information and Natural Gradient Learning of Random Deep Networks.” arXiv:1808.07172 [Cond-Mat, Stat].
Amari, Park, and Fukumizu. 2000. Adaptive Method of Realizing Natural Gradient Learning for Multilayer Perceptrons.” Neural Computation.
Arbel, Gretton, Li, et al. 2020. Kernelized Wasserstein Natural Gradient.”
Botev, Ritter, and Barber. 2017. Practical Gauss-Newton Optimisation for Deep Learning.” In Proceedings of the 34th International Conference on Machine Learning.
Chen, Fox, and Guestrin. 2014. Stochastic Gradient Hamiltonian Monte Carlo.” In Proceedings of the 31st International Conference on Machine Learning.
Csiszár. 1975. I-Divergence Geometry of Probability Distributions and Minimization Problems.” The Annals of Probability.
Durmus, and Moulines. 2016. High-Dimensional Bayesian Inference via the Unadjusted Langevin Algorithm.” arXiv:1605.01559 [Math, Stat].
Efron, and Hinkley. 1978. Assessing the Accuracy of the Maximum Likelihood Estimator: Observed Versus Expected Fisher Information.” Biometrika.
Ge, Lee, and Risteski. 2020. Simulated Tempering Langevin Monte Carlo II: An Improved Proof Using Soft Markov Chain Decomposition.” arXiv:1812.00793 [Cs, Math, Stat].
Grosse. 2021. Metrics.” In CSC2541 Winter 2021.
Grosse, and Martens. 2016. A Kronecker-Factored Approximate Fisher Matrix for Convolution Layers.” In Proceedings of The 33rd International Conference on Machine Learning.
Hensman, Rattray, and Lawrence. 2012. Fast Variational Inference in the Conjugate Exponential Family.” In Advances in Neural Information Processing Systems.
Honkela, Tornio, Raiko, et al. 2008. Natural Conjugate Gradient in Variational Inference.” In Neural Information Processing. Lecture Notes in Computer Science.
Kakade. 2002. “A Natural Policy Gradient.” In Advances In Neural Information Processing Systems.
Karakida, and Osawa. 2020. Understanding Approximate Fisher Information for Fast Convergence of Natural Gradient Descent in Wide Neural Networks.” Advances in Neural Information Processing Systems.
Khan, Mohammad Emtiyaz, and Lin. 2017. Conjugate-Computation Variational Inference : Converting Variational Inference in Non-Conjugate Models to Inferences in Conjugate Models.” In Artificial Intelligence and Statistics.
Khan, Mohammad, Nielsen, Tangkaratt, et al. 2018. Fast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adam.” In Proceedings of the 35th International Conference on Machine Learning.
Khan, Mohammad Emtiyaz, and Rue. 2024. The Bayesian Learning Rule.”
Ly, Marsman, Verhagen, et al. 2017. A Tutorial on Fisher Information.” Journal of Mathematical Psychology.
Martens. 2016. Second-Order Optimization for Neural Networks.”
———. 2020. New Insights and Perspectives on the Natural Gradient Method.” Journal of Machine Learning Research.
Martens, and Grosse. 2015. Optimizing Neural Networks with Kronecker-Factored Approximate Curvature.” In Proceedings of the 32nd International Conference on Machine Learning.
Mosegaard, and Tarantola. 2002. Probabilistic Approach to Inverse Problems.” In International Geophysics. International Handbook of Earthquake and Engineering Seismology, Part A.
Nielsen. 2018. An Elementary Introduction to Information Geometry.” arXiv:1808.08271 [Cs, Math, Stat].
Norton, and Fox. 2016. Tuning of MCMC with Langevin, Hamiltonian, and Other Stochastic Autoregressive Proposals.” arXiv:1610.00781 [Math, Stat].
Nurbekyan, Lei, and Yang. 2022. Efficient Natural Gradient Descent Methods for Large-Scale Optimization Problems.”
Ollivier. 2017. Online Natural Gradient as a Kalman Filter.” arXiv:1703.00209 [Math, Stat].
Osawa, Ishikawa, Yokota, et al. 2023. ASDL: A Unified Interface for Gradient Preconditioning in PyTorch.”
Osawa, Swaroop, Khan, et al. 2019. Practical Deep Learning with Bayesian Principles.” In Advances in Neural Information Processing Systems.
Salimbeni, Eleftheriadis, and Hensman. 2018. Natural Gradients in Practice: Non-Conjugate Variational Inference in Gaussian Process Models.” In International Conference on Artificial Intelligence and Statistics.
Sato. 2001. Online Model Selection Based on the Variational Bayes.” Neural Computation.
Schraudolph. 2002. Fast Curvature Matrix-Vector Products for Second-Order Gradient Descent.” Neural Computation.
Welling, and Teh. 2011. Bayesian Learning via Stochastic Gradient Langevin Dynamics.” In Proceedings of the 28th International Conference on International Conference on Machine Learning. ICML’11.
Wilkinson, Särkkä, and Solin. 2021. Bayes-Newton Methods for Approximate Bayesian Inference with PSD Guarantees.”
Zellner. 1988. Optimal Information Processing and Bayes’s Theorem.” The American Statistician.
Zhang, Sun, Duvenaud, et al. 2018. Noisy Natural Gradient as Variational Inference.” In Proceedings of the 35th International Conference on Machine Learning.