Natural gradient: a covariant derivative for distributions

In 1998 Amari posed a strange question: what does it mean to step by $\epsilon$ in distribution space? If parameterizations $\sigma$, $\log\sigma$, and $\sigma^2$ are all valid, three different SGDs trace three different trajectories. Natural gradient kills that arbitrariness: one geometry, one direction, invariant under reparameterization. The same gap separates Euclidean and Riemannian geometry - and for neural networks it is real, not academic.

• **TRPO and PPO in RL.** OpenAI Five (Dota 2) and DeepMind AlphaStar trained on natural policy gradient inside a trust region. Without a KL constraint, policy collapse happens within tens of iterations.
• **K-FAC at DeepMind.** AlphaGo Zero's policy network was trained with K-FAC. Sophia (Liu 2023) sped up GPT-2-scale pretraining 2x relative to Adam at the same FLOPs.
• **RLHF in Claude/ChatGPT.** PPO clipping is a surrogate for the natural-gradient KL constraint: clipping the ratio $\frac{\pi_\theta}{\pi_{\text{old}}}$ is a coarse approximation of TRPO's trust region.

Предварительные знания

• Fisher information metric
• KL and Bregman divergences

Definition: a covariant derivative for distributions

Vanilla gradient depends on parameterization

The ordinary gradient $\nabla L(\theta)$ is just the vector of partial derivatives in some coordinate system $\theta$. The problem: a change of parameterization changes the gradient. Pick $\theta=(\mu, \sigma)$ for a Gaussian or $\theta=(\mu, \log\sigma)$ - SGD takes different steps, follows different trajectories, converges at different speeds. Geometrically meaningless: the distribution is the same.

The natural gradient $\tilde{\nabla} L$ is the vanilla gradient pre-multiplied by the inverse of the Fisher information matrix $F(\theta)$. The key property: under reparameterization $F$ transforms as a metric tensor and $\nabla L$ as a covector, so the product is invariant. The step direction in distribution space is the same regardless of which coordinates do the bookkeeping.

**Fisher information matrix:** $F_{ij}(\theta) = \mathbb{E}_{x\sim p_\theta}\big[\partial_i \log p_\theta(x)\, \partial_j \log p_\theta(x)\big]$. It is the metric on the manifold of distributions, also called the Fisher-Rao metric.

Geometrically: natural gradient is steepest descent in KL divergence, not in Euclidean parameter norm. The step $\theta \leftarrow \theta - \eta \tilde{\nabla} L$ minimises loss among all steps with a fixed KL radius. Vanilla SGD minimises among steps with a fixed Euclidean radius - which is metrically arbitrary.

Riemannian geometry and the Fisher inverse

The manifold of distributions and the Fisher metric

A parametric family $\{p(\cdot|\theta) : \theta \in \Theta\}$ is a multidimensional manifold: each point $\theta$ corresponds to a distribution. The Fisher metric $F(\theta)$ defines an inner product on the tangent space: the distance between nearby distributions is measured by KL, and to leading order KL is quadratic in $F$.

This identity is the heart of natural gradient. Minimising $L(\theta + \epsilon) - L(\theta) \approx \nabla L^T \epsilon$ subject to $D_{KL} \le \delta$ gives the direction $\epsilon^* \propto -F^{-1} \nabla L$. So natural gradient solves a trust-region problem in KL geometry.

**Geodesics and parallel transport.** SGD in the Euclidean metric follows straight lines in $\theta$-space - but a straight line in $\theta$ is not a straight line on the manifold of distributions. Natural gradient follows Fisher-Rao geodesics: shortest paths in KL distance. Christoffel symbols describe how tangent vectors are transported along curves; for an exponential family they have a closed form.

**Connection to TRPO in RL.** Trust Region Policy Optimization (Schulman 2015) solves $\max_\theta \mathbb{E}[A(s,a)]$ subject to $\mathbb{E}[D_{KL}(\pi_{\theta_{\text{old}}} \| \pi_\theta)] \le \delta$. Linearising the objective and quadratising the constraint produces the natural gradient. PPO is a cheap heuristic replacement of that constraint via clipping.

K-FAC and Shampoo: practical approximations

Why direct $F^{-1}$ is impossible

For a model with $d$ parameters the Fisher matrix has size $d \times d$. Storage costs $O(d^2)$, inversion costs $O(d^3)$. For GPT-2 small ($d \approx 10^8$) that is $10^{16}$ FLOPs per inversion every step - unfeasible. Every production-grade natural-gradient method approximates $F^{-1}$, trading accuracy against compute.

**Answer to 'which approximation does Adam use?'.** Adam keeps a diagonal empirical Fisher (through $v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2$). It is natural gradient at its crudest. That is why Adam is often called a 'cheap natural gradient' - empirically close, but ignoring off-diagonal Fisher.

**Where K-FAC has actually shipped.** DeepMind used K-FAC for policy training in AlphaGo and AlphaZero. Sophia (Liu et al. 2023) sped up GPT-2-scale pretraining 2x compared to Adam at the same FLOPs. Shampoo was adopted by Google for PaLM. The trade-off: compiler and memory complexity grow, but wall-clock convergence improves at large batch sizes.

Key ideas

• Natural gradient $\tilde{\nabla}L = F^{-1}\nabla L$ is invariant under reparameterization - the step direction is identical in any coordinates
• Geometrically it is steepest descent in KL divergence: $D_{KL}(p_\theta \| p_{\theta+\epsilon}) \approx \tfrac12 \epsilon^T F \epsilon$, and natural gradient solves the trust-region problem
• Direct $F^{-1}$ costs $O(d^3)$ - neural nets need approximations: diagonal Fisher (Adam), K-FAC, Shampoo, Sophia, conjugate gradient
• Natural gradient is the principled foundation of TRPO, PPO, natural policy gradient in RL and K-FAC/Shampoo in deep learning

Connections to other topics

Natural gradient is the central hub of information geometry: it links the Fisher metric to practical optimisation.

• Fisher information metric — The metric in which natural gradient is steepest descent. Without Fisher, the formula $F^{-1}\nabla L$ has no meaning.
• RL and bandits — TRPO, PPO, natural Q-learning - all are built on natural gradient. The trust-region KL constraint is exactly the Fisher metric.
• Neural network training — K-FAC, Shampoo, Sophia are production-grade approximations of natural gradient for deep nets. AlphaGo and LLM pretraining use them at scale.

Вопросы для размышления

• If parameterization does not affect natural gradient, why is vanilla SGD still the default - and in which scenarios is the parameterization arbitrariness harmless in practice?
• Adam = diagonal empirical Fisher: is this a principled approximation or a coincidence? What is preserved and what is lost relative to full Fisher?
• TRPO solves a trust-region problem explicitly; PPO replaces it with clipping. What geometric intuition underlies the clip ratio - and where does the approximation break?

Связанные уроки

• ig-02-fisher-metric — Natural gradient is the gradient measured in the Fisher metric; without it the formula $F^{-1}\nabla L$ is meaningless
• ig-04-kl-bregman — Local KL is quadratic in Fisher: $D_{KL}(p_\theta \| p_{\theta+\epsilon}) \approx \tfrac12 \epsilon^T F \epsilon$ - this is the justification for natural gradient
• ig-05-dual-flat — Dually flat structure explains why natural gradient is invariant under reparameterization
• ig-12-rl-bandits — TRPO, PPO, and natural policy gradient are direct applications of natural gradient in RL
• ig-14-neural — K-FAC, Shampoo, and Sophia are practical approximations of natural gradient for deep nets
• stat-01-sampling