Information Geometry & Natural Gradient

Fisher information matrix, statistical manifold, natural gradient descent, K-FAC, connection to Adam

Standard gradient descent treats all parameter directions as equally meaningful. But two parameters that differ by the same Euclidean distance in parameter space may correspond to distributions that are almost identical or drastically different — the Euclidean metric ignores how sensitive the model's outputs are to perturbations in each direction. This produces the ravine problem: gradient descent zigzags across high-curvature directions while inching along low-curvature directions, wasting steps.

Information geometry gives the right metric: the space of probability distributions is a Riemannian manifold where distances are measured by the Fisher information matrix (FIM), which quantifies how much the output distribution changes per unit perturbation of parameters. Natural gradient descent premultiplies the gradient by the inverse FIM, taking steps that are equal-sized in distribution space.

The result is faster convergence per step. The catch: computing and inverting the full FIM is O(p²) in storage and O(p³) in compute — prohibitive at any useful scale. Adam's per-parameter learning rate scaling is a diagonal FIM approximation, which explains both why Adam works and why it fails when parameters are highly correlated.

Key points

• The Fisher information matrix: F = E_{x~p_θ}[∇_θ log p(x|θ) ∇_θ log p(x|θ)ᵀ]. The FIM measures how much the output distribution changes when θ is perturbed: KL[p_θ ‖ p_{θ+δ}] ≈ ½ δᵀFδ locally. So the FIM is the curvature of the KL divergence landscape — the metric tensor for distribution space. Flat directions in F (zero eigenvalues) are parameter directions that do not change the output distribution at all — overparameterisation, symmetries, dead neurons.
• Standard gradient descent is not covariant under reparameterisation: if you apply a bijective transformation φ = h(θ), the gradient direction in φ-space is different from the gradient direction in θ-space (after accounting for the Jacobian). This means the solution gradient descent finds depends on how you chose to parameterise the model — different implementations of the same model (batch norm vs weight normalisation) converge to different solutions. Natural gradient is covariant: the update corresponds to the same distribution-space step regardless of parameterisation.
• Natural gradient update: θ_{t+1} = θ_t - η F(θ_t)⁻¹ ∇_θ L. In a loss landscape with a ravine (high curvature in one direction, low in another), standard SGD zigzags — large steps in the steep direction oscillate, small steps in the shallow direction barely move. Natural gradient rescales: small steps in the steep direction, large steps in the shallow direction. Per distribution-space step, you get more loss reduction than any Euclidean-metric gradient step.
• FIM = second derivative of KL: KL[p_θ ‖ p_{θ+δ}] ≈ ½δᵀFδ. This is the operational definition that connects information geometry to practical optimisation. Directions with large Fisher eigenvalues change the model's output distribution a lot per unit parameter change — natural gradient takes small steps there. Directions with small eigenvalues change the distribution very little — natural gradient takes large steps there. Euclidean gradient ignores all of this.
• K-FAC (Kronecker-Factored Approximate Curvature) makes natural gradient tractable for neural networks. For a layer with weight matrix W, the FIM admits a Kronecker product approximation F ≈ A ⊗ G where A = E[a_ta_tᵀ] (input activation covariance) and G = E[g_tg_tᵀ] (output gradient covariance). Storage drops from O(p²) to O(d_in² + d_out²) per layer. Inversion is separable: (A⊗G)⁻¹ = A⁻¹⊗G⁻¹. K-FAC captures within-layer input-output correlations — the structure that Adam's diagonal approximation misses.
• Adam is a diagonal FIM approximation. The second moment v_t ≈ diag(F) estimates only the diagonal of the Fisher. Dividing the gradient by √v_t + ε approximates rescaling by the diagonal Fisher — each parameter gets an independent learning rate based on its gradient variance. This works well when parameters are approximately uncorrelated. When parameters are highly correlated (collinear features, attention across similar tokens), the off-diagonal terms of F matter and Adam's approximation fails.
• K-FAC vs Adam: Adam is the practical choice for most deep learning because K-FAC's per-step cost is much higher. K-FAC wins when data is small, per-step compute is affordable, and the correlations between parameters are strong — some supervised learning benchmarks and RL settings. The convergence advantage is real: K-FAC typically converges in fewer steps, but each step costs more than Adam.
• TRPO and PPO formalise the natural gradient idea in RL. TRPO explicitly constrains KL[π_old ‖ π_new] ≤ δ at each update — this is a trust region in distribution space, exactly what natural gradient descent respects. Euclidean constraints on weight updates do not prevent large changes in the policy distribution. KL constraints do — and large distribution changes destabilise RL training. PPO approximates the KL constraint with a clipped surrogate, trading theoretical precision for engineering simplicity.
• The FIM connects to confidence intervals in MLE. By the Cramér-Rao bound, Var(θ̂) ≥ F(θ)⁻¹ for any unbiased estimator. The MLE achieves this bound asymptotically. The inverse FIM is the asymptotic covariance of the MLE — this is where frequentist confidence intervals for point estimates come from. In continual learning, Elastic Weight Consolidation (EWC) penalises changes to parameters with high Fisher information, preserving knowledge from previous tasks by anchoring the highest-curvature directions.

Recap

• Parameter space is the wrong metric — what matters is how much the output distribution moves per step, measured by the FIM.
• FIM = curvature of KL: $KL[p_θ ‖ p_{θ+δ}] ≈ frac12 δ^T F δ$ — the metric tensor for distribution space.
• Natural gradient: $θ_{t+1} = θ_t - η F^{-1}∇L$ — small steps in steep (high-Fisher) directions, large in flat ones; fixes ravine zigzag.
• Natural gradient is covariant to reparameterisation; plain SGD finds different optima depending on parameterisation (batch vs weight norm).
• Full FIM is $O(p^2)$ store / $O(p^3)$ invert — infeasible; K-FAC uses Kronecker factoring $F ≈ A ⊗ G$ per layer.
• Adam is a diagonal FIM approximation ($v_t ≈ ext{diag}(F)$) — works when parameters are uncorrelated, fails when off-diagonals matter.
• TRPO/PPO = natural gradient in RL: constrain $KL[π_{old} ‖ π_{new}]$, not Euclidean weight change — Euclidean constraints don't bound policy shift.

Check your understanding