Natural Gradient & Riemannian Optimization

Why It Matters for Modern Models

The ordinary gradient depends on how the model is parameterized; the natural gradient follows the geometry of the distributions instead, so in the limit of small steps it traces the same path under any smooth reparameterization, and approximately so for finite steps (Amari, 1998; Martens, 2020).
Trust-region policy optimization (TRPO) applies this idea to reinforcement learning: it computes a natural-gradient direction with the conjugate-gradient method and then searches along it for a step that satisfies a KL constraint (Schulman et al., 2015). Proximal policy optimization (PPO) keeps the aim of limiting each policy change but replaces the constraint with a clipped first-order objective (Schulman et al., 2017).
In many important cases, including the usual losses for regression and classification, the Fisher matrix equals the generalized Gauss–Newton matrix, so natural gradient descent can be read as a second-order method with the Fisher standing in for the Hessian (Martens, 2020).
Adam divides by the square root of a running average of squared gradients, which resembles a diagonal of the "empirical Fisher"; Kunstner et al. (2019) argue that the empirical Fisher does not generally capture curvature, so the link between Adam and natural gradient is loose.

Questions to Think About

If one parameter is rescaled, say θ₁ → 10θ₁, how does an ordinary gradient step change, and why does the natural gradient step describe the same change in the distribution?
Why might natural gradient avoid or shorten the plateaus that slow backpropagation in multilayer perceptrons, as Amari (1998) suggested?

Core Math

Key Equation

\tilde{\nabla} L = F(\theta)^{-1} \nabla_\theta L

The natural gradient preconditions the ordinary gradient with the inverse of the Fisher information matrix $F(\theta)$ :

\tilde{\nabla} L = F(\theta)^{-1} \nabla_\theta L

and natural gradient descent takes the step

\theta_{t+1} = \theta_t - \eta\, F(\theta_t)^{-1} \nabla_\theta L(\theta_t)

Why it is called natural. Among small steps $\delta$ that change the model's distribution by a fixed amount, measured by the KL divergence, the natural gradient step decreases the loss fastest:

\delta^* = \arg\min_\delta \, \langle \nabla L, \delta \rangle \quad \text{s.t.} \quad \tfrac{1}{2}\, \delta^\top F \delta \le \epsilon

where $\tfrac{1}{2}\delta^\top F \delta$ is the second-order approximation of $\text{KL}(p_\theta \,\|\, p_{\theta+\delta})$ . Setting the gradient of the Lagrangian to zero gives

\delta^* \propto -F^{-1} \nabla L

and the minus sign makes it a descent direction. For a network with millions of parameters $F$ cannot be stored or inverted, so practical methods approximate it, for example with Kronecker-factored blocks (K-FAC; Martens & Grosse, 2015).