Natural gradient descent

Climbing slower on the tricky bits

A placeholder.

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 optimizing 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 parameterization 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 reparameterization (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 optimization 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 analyze this method and its properties, and show how it can be viewed as a type of approximate 2nd-order optimization 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) parameterization 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 minimize \(-\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 minimizing \(\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 maximizes 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)


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.


Abt, Markus, and William J. Welch. 1998. β€œFisher Information and Maximum-Likelihood Estimation of Covariance Parameters in Gaussian Stochastic Processes.” Canadian Journal of Statistics 26 (1): 127–37.
Amari, Shun-ichi. 1998. β€œNatural Gradient Works Efficiently in Learning.” Neural Computation 10 (2): 251–76.
Amari, Shun-ichi, Ryo Karakida, and Masafumi Oizumi. 2018. β€œFisher Information and Natural Gradient Learning of Random Deep Networks.” arXiv:1808.07172 [Cond-Mat, Stat], August.
Amari, Shun-ichi, Hyeyoung Park, and Kenji Fukumizu. 2000. β€œAdaptive Method of Realizing Natural Gradient Learning for Multilayer Perceptrons.” Neural Computation 12 (6): 1399–409.
Arbel, Michael, Arthur Gretton, Wuchen Li, and Guido Montufar. 2020. β€œKernelized Wasserstein Natural Gradient.” arXiv.
Botev, Aleksandar, Hippolyt Ritter, and David Barber. 2017. β€œPractical Gauss-Newton Optimisation for Deep Learning.” In Proceedings of the 34th International Conference on Machine Learning, 557–65. PMLR.
Chen, Tianqi, Emily Fox, and Carlos Guestrin. 2014. β€œStochastic Gradient Hamiltonian Monte Carlo.” In Proceedings of the 31st International Conference on Machine Learning, 1683–91. Beijing, China: PMLR.
Csiszar, I. 1975. β€œI-Divergence Geometry of Probability Distributions and Minimization Problems.” Annals of Probability 3 (1): 146–58.
Durmus, Alain, and Eric Moulines. 2016. β€œHigh-Dimensional Bayesian Inference via the Unadjusted Langevin Algorithm.” arXiv:1605.01559 [Math, Stat], May.
Efron, Bradley, and David V. Hinkley. 1978. β€œAssessing the Accuracy of the Maximum Likelihood Estimator: Observed Versus Expected Fisher Information.” Biometrika 65 (3): 457–83.
Ge, Rong, Holden Lee, and Andrej Risteski. 2020. β€œSimulated Tempering Langevin Monte Carlo II: An Improved Proof Using Soft Markov Chain Decomposition.” arXiv:1812.00793 [Cs, Math, Stat], September.
Grosse, Roger. 2021. β€œMetrics.” In CSC2541 Winter 2021, Chapter 3.
Grosse, Roger, and James Martens. 2016. β€œA Kronecker-Factored Approximate Fisher Matrix for Convolution Layers.” In Proceedings of The 33rd International Conference on Machine Learning, 573–82. PMLR.
Hensman, James, Magnus Rattray, and Neil Lawrence. 2012. β€œFast Variational Inference in the Conjugate Exponential Family.” In Advances in Neural Information Processing Systems. Vol. 25. Curran Associates, Inc.
Kakade, Sham M. 2002. β€œA Natural Policy Gradient.” In Advances In Neural Information Processing Systems, 8.
Karakida, Ryo, and Kazuki Osawa. 2020. β€œUnderstanding Approximate Fisher Information for Fast Convergence of Natural Gradient Descent in Wide Neural Networks.” Advances in Neural Information Processing Systems 33.
Khan, Mohammad Emtiyaz, and HΓ₯vard Rue. 2022. β€œThe Bayesian Learning Rule.” arXiv.
Khan, Mohammad, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. 2018. β€œFast and Scalable Bayesian Deep Learning by Weight-Perturbation in Adam.” In Proceedings of the 35th International Conference on Machine Learning, 2611–20. PMLR.
Ly, Alexander, Maarten Marsman, Josine Verhagen, Raoul P. P. P. Grasman, and Eric-Jan Wagenmakers. 2017. β€œA Tutorial on Fisher Information.” Journal of Mathematical Psychology 80 (October): 40–55.
Martens, James. 2016. β€œSecond-Order Optimization for Neural Networks.” University of Toronto.
β€”β€”β€”. 2020. β€œNew Insights and Perspectives on the Natural Gradient Method.” Journal of Machine Learning Research 21 (146): 1–76.
Martens, James, and Roger Grosse. 2015. β€œOptimizing Neural Networks with Kronecker-Factored Approximate Curvature.” In Proceedings of the 32nd International Conference on Machine Learning, 2408–17. PMLR.
Mosegaard, Klaus, and Albert Tarantola. 2002. β€œProbabilistic Approach to Inverse Problems.” In International Geophysics, 81:237–65. Elsevier.
Nielsen, Frank. 2018. β€œAn Elementary Introduction to Information Geometry.” arXiv:1808.08271 [Cs, Math, Stat], August.
Norton, Richard A., and Colin Fox. 2016. β€œTuning of MCMC with Langevin, Hamiltonian, and Other Stochastic Autoregressive Proposals.” arXiv:1610.00781 [Math, Stat], October.
Nurbekyan, Levon, Wanzhou Lei, and Yunan Yang. 2022. β€œEfficient Natural Gradient Descent Methods for Large-Scale Optimization Problems.” arXiv.
Ollivier, Yann. 2017. β€œOnline Natural Gradient as a Kalman Filter.” arXiv:1703.00209 [Math, Stat], March.
Osawa, Kazuki, Satoki Ishikawa, Rio Yokota, Shigang Li, and Torsten Hoefler. 2023. β€œASDL: A Unified Interface for Gradient Preconditioning in PyTorch.” arXiv.
Osawa, Kazuki, Siddharth Swaroop, Mohammad Emtiyaz E Khan, Anirudh Jain, Runa Eschenhagen, Richard E Turner, and Rio Yokota. 2019. β€œPractical Deep Learning with Bayesian Principles.” In Advances in Neural Information Processing Systems. Vol. 32. Red Hook, NY, USA: Curran Associates, Inc.
Salimbeni, Hugh, Stefanos Eleftheriadis, and James Hensman. 2018. β€œNatural Gradients in Practice: Non-Conjugate Variational Inference in Gaussian Process Models.” In International Conference on Artificial Intelligence and Statistics, 689–97.
Schraudolph, Nicol N. 2002. β€œFast Curvature Matrix-Vector Products for Second-Order Gradient Descent.” Neural Computation 14 (7): 1723–38.
Welling, Max, and Yee Whye Teh. 2011. β€œBayesian Learning via Stochastic Gradient Langevin Dynamics.” In Proceedings of the 28th International Conference on International Conference on Machine Learning, 681–88. ICML’11. Madison, WI, USA: Omnipress.
Wilkinson, William J., Simo SΓ€rkkΓ€, and Arno Solin. 2021. β€œBayes-Newton Methods for Approximate Bayesian Inference with PSD Guarantees.” arXiv.
Zellner, Arnold. 1988. β€œOptimal Information Processing and Bayes’s Theorem.” The American Statistician 42 (4): 278–80.
Zhang, Guodong, Shengyang Sun, David Duvenaud, and Roger Grosse. 2018. β€œNoisy Natural Gradient as Variational Inference.” In Proceedings of the 35th International Conference on Machine Learning, 5852–61. PMLR.

No comments yet. Why not leave one?

GitHub-flavored Markdown & a sane subset of HTML is supported.