Information Geometry & Natural Gradients

The coordinate dependence of gradient descent

Let $(\Omega, \mathcal{F}, \mu)$ be a measure space. We have a parametric family $\mathcal{S} = \{ P_\xi \mid \xi \in \Xi \}$ dominated by $\mu$, $\Xi \subseteq \mathbb{R}^d$ open, with densities:

We want a Riemannian metric $g$ and a connection $\nabla$ on $\mathcal{S}$ that sit inside the family itself and do not care about how we happen to write the parameters down.

The trouble is that gradient descent takes the distance between $P_\xi$ and $P_{\xi + \delta \xi}$ to be $\|\delta \xi\|_2$, and under a reparameterization $\xi \to \phi(\xi)$ this Euclidean distance shifts around even though the distributions have not moved at all. So we need a divergence $D[P : Q]$ whose induced metric is

1. Reparameterization-invariant, so that $ds^2(\xi, \xi + d\xi)$ is a scalar invariant.
2. Sufficient-statistic-invariant, so that if $T$ from $\Omega$ to $\Omega'$ is sufficient for $\xi$, the geometry on $\mathcal{S}$ matches the geometry on the induced family over $\Omega'$.

Regularity conditions and their failures

The standard theory sits on top of the following assumptions and each one quietly kicks out a bunch of cases we actually care about.

A1 (Identifiability). $\xi \mapsto P_\xi$ is injective. Neural networks break this all over the place through permutation symmetry and overparameterization and entire manifolds of solutions that all give the same distribution.

A2 (Common support). $\text{supp}(P_\xi)$ does not depend on $\xi$. The uniform distribution $U[0, \xi]$ blows this up right away and the likelihood stops being differentiable at the boundary.

A3 (Smoothness). $\ell(\xi; x) = \log p(x; \xi)$ is at least $C^3$ in $\xi$.

A4 (Interchange of derivative and integral).

This needs the score to be uniformly integrable and it gets assumed way more often than anyone bothers to check it.

Deriving the metric from KL divergence

KL divergence is asymmetric and it breaks the triangle inequality so it is not a metric, but if you zoom in and do a second-order expansion you get back a real Riemannian metric.

Set $\xi' = \xi + \delta \xi$ and Taylor-expand $\ell(x; \xi')$ around $\xi$ to see what comes out.

Plugging in gives

The linear term drops out and the calculation below shows why.

By A4,

So $D_{KL}$ sits at a local minimum at $\delta \xi = 0$ like you would expect and what is left over is the quadratic form

Now for the Hessian-to-outer-product identity. Differentiating $\int (\nabla \ell) p dx = 0$ one more time gives

and so

The local distance $ds^2 = \delta \xi^T G(\xi) \delta \xi$ gives a Riemannian metric on $\mathcal{S}$ and the matrix $G(\xi)$ is the Fisher Information Matrix.

Chentsov’s uniqueness theorem

Chentsov (1972) showed that the Fisher metric is the only Riemannian metric (up to scale) that stays invariant under congruent embeddings by Markov morphisms.

The idea is that a measurable map $F$ from $\Omega$ to $\Omega'$ pushes measures forward and if $F$ is a sufficient statistic then no information is lost, so distances have to be preserved. The Fisher metric does this on its own and no other Riemannian metric does.

Dual connections

A metric gives lengths and angles but if you want geodesics and flatness you also need a connection.

In ordinary Riemannian geometry the obvious pick is the Levi-Civita connection which is metric-compatible and torsion-free and unique. Statistical manifolds carry a whole one-parameter family of connections $\nabla^{(\alpha)}$ with $\alpha \in \mathbb{R}$ and this extra structure is where a lot of the action is.

The $\alpha$-connection has Christoffel symbols

Using the skewness tensor $T_{ijk} = \mathbb{E}[(\partial_i \ell)(\partial_j \ell)(\partial_k \ell)]$ this cleans up to

$\nabla^{(\alpha)}$ and $\nabla^{(-\alpha)}$ are dual with respect to $g$, and so for any vector fields $X$ and $Y$ and $Z$ we have

The exponential connection at $\alpha=1$ and the mixture connection at $\alpha=-1$ are the pair that does most of the work.

The hyperbolic geometry of Gaussians

The univariate Gaussian makes the theory concrete. Take $\mathcal{S} = \{ N(\mu, \sigma^2) \mid \mu \in \mathbb{R}, \sigma > 0 \}$.

The scores are

We work out the Fisher matrix one entry at a time.

$g_{\mu\mu} = \frac{1}{\sigma^4} \mathbb{E}[(x-\mu)^2] = \frac{1}{\sigma^2}$.

$g_{\mu\sigma} = 0$ because the odd central moments of a Gaussian wipe out, so $\mu$ and $\sigma$ sit orthogonal to each other in the Riemannian sense.

For $g_{\sigma\sigma}$ we expand and use $\mathbb{E}[(x-\mu)^4] = 3\sigma^4$ which gives

The Fisher matrix is

and the line element is

Put this next to the Poincare upper half-plane where $ds^2 = \frac{dx^2 + dy^2}{y^2}$. The factor of 2 just rescales the curvature but the shape is the same and the manifold of univariate Gaussians has constant negative curvature.

Geodesics between $N(\mu_1, \sigma_1)$ and $N(\mu_2, \sigma_2)$ come from solving the Euler-Lagrange equations for $L = \int \frac{\sqrt{\dot{\mu}^2 + 2\dot{\sigma}^2}}{\sigma} dt$. When $\mu_1 = \mu_2$ the geodesic is a vertical line that just rescales the variance and otherwise the geodesics are semi-ellipses sitting on the $\mu$-axis.

Because of this negative curvature, averaging parameters coordinate by coordinate does not land you on the geometric center and the Frechet mean on this manifold can sit pretty far from the naive average.

Exponential families and e-flatness

Take the exponential family in natural parameters $\theta$ and write it out as

The second derivative of the log-likelihood is $\partial_i \partial_j \ell = -\partial_i \partial_j \psi(\theta)$ and this has no $x$ in it at all, so the Hessian is just a deterministic function of $\theta$.

So the e-connection Christoffel symbols collapse to

because the score has zero mean. So the manifold is flat under the e-connection and the natural parameters are affine coordinates and geodesics are straight lines given by $\theta(t) = (1-t)\theta_1 + t\theta_2$.

Mixture families and the Pythagorean theorem

And on the other side, the mixture family

is flat under the m-connection at $\alpha=-1$ and the expectation parameters $\eta$ are affine and geodesics are linear mixtures given by $P(t) = (1-t)P_1 + tP_2$.

Now that we have dual flat structures sitting on the same space, a Pythagorean theorem falls out. When the e-geodesic $PQ$ hits the m-geodesic $QR$ at right angles at $Q$ we get

The KL divergence between members of an exponential family is a Bregman divergence on the cumulant potential $\psi$, and it reads

where $\eta_P = \nabla \psi(\theta_P)$.

If you write out $D_{KL}(P\|Q) + D_{KL}(Q\|R)$ and subtract $D_{KL}(P\|R)$, a bunch of things cancel and what is left is

The vector $\eta_P - \eta_Q$ is the tangent in dual coordinates along $P \to Q$ and $\theta_R - \theta_Q$ is the tangent in primal coordinates along $Q \to R$. Orthogonality says their inner product wipes out and so $\Delta = 0$.

This gives you the projection theorem, which says that the m-projection of $P$ onto an e-flat submanifold is unique and satisfies the Pythagorean relation, and the MLE is just the m-projection of the empirical distribution onto the model manifold.

The uniform distribution and how regularity breaks

Breaking A2 changes the theory in a qualitative way. Take $p(x; \theta) = U[0, \theta] = \frac{1}{\theta} \mathbb{I}(0 \le x \le \theta)$.

The score does not have zero mean anymore and the whole derivation from the KL section falls apart, because $\frac{d}{d\theta} \int_0^\theta p\, dx \neq \int_0^\theta \nabla p\, dx$ and Leibniz’s rule picks up a boundary term $p(\theta; \theta) = 1/\theta$.

And the fallout is real. The Fisher information $G = 1/\theta^2$ looks finite on paper but the Cramer-Rao bound does not apply because A2 has failed. The MLE $\hat{\theta} = \max(X_i)$ has variance $O(n^{-2})$ and this beats the $O(n^{-1})$ Fisher rate, because the support boundary holds information that a score-based story just cannot see.

Natural gradient descent

Say we want to minimize a loss $\mathcal{L}(\theta)$ on $\mathcal{S}$. The update $\theta_{new} = \theta - \eta \nabla \mathcal{L}$ mashes covectors and vectors together because $\nabla \mathcal{L}$ is a covector and $\Delta \theta$ is a vector, and they do not live in the same space.

So instead we solve

and plugging in the quadratic approximation $D_{KL} \approx \frac{1}{2} \delta \theta^T G \delta \theta$ gives

Cramer-Rao as Cauchy-Schwarz on the tangent space

The Cramer-Rao bound has a clean geometric proof that is really just Cauchy-Schwarz on the tangent space.

The score $S_\theta(x) = \nabla_\theta \ell(x; \theta)$ sits in $L^2(P_\theta)$ and the Fisher information is its squared norm $\|S_\theta\|^2 = G(\theta)$.

For an unbiased estimator $\hat{\theta}$ we work out $\text{Cov}(\hat{\theta}, S_\theta)$. Using $\nabla p = p S_\theta$ and A4 we get

And matrix Cauchy-Schwarz on this gives

High curvature means large $G$ and that says the distributions are easy to tell apart so estimators can be precise, and low curvature means small $G$ and the distributions look locally the same so no estimator can do well.

Singular models and Watanabe’s RLCT

When identifiability fails and A1 goes out the window, the Fisher matrix goes singular and we get

The manifold dimension collapses at these points and the tangent space drops to a tangent cone, and in deep learning singular Fisher matrices are just the default rather than the exception.

Watanabe (2009) showed that near singularities the Bayesian posterior does not settle down to $N(0, G^{-1})$ and the standard BIC complexity term $\frac{d}{2} \log n$ gets replaced by $\lambda \log n$ where $\lambda$ is the Real Log Canonical Threshold (RLCT) and satisfies

Singular models are simpler than their parameter count makes them look. The posterior volume is set by resolution of singularities and that is a question for algebraic geometry and not Riemannian geometry. Standard information geometry stops working here and you have to pick up Watanabe’s singular learning theory instead.

Implementation (JAX)

We check e- and m-geodesic orthogonality and put Natural Gradient up against SGD on a warped Gaussian to see how they converge.

import jax
import jax.numpy as jnp
from jax import random, grad, jit, vmap, lax
from jax.scipy.stats import multivariate_normal
from typing import NamedTuple, Tuple
 
# ------------------------------------------------------------------
# SYSTEM CONFIGURATION
# ------------------------------------------------------------------
SEED = 42
LEARNING_RATE = 0.1
NUM_STEPS = 100
DAMPING = 1e-4
 
# ------------------------------------------------------------------
# 1. Manifold Definition: Warped Gaussian
# ------------------------------------------------------------------
def get_sigma(theta: jax.Array) -> jax.Array:
    """
    Constructs the covariance matrix Sigma(theta) = diag(theta^2 + 1).
    Ensures positive definiteness everywhere.
    """
    return jnp.diag(theta**2 + 1.0)
 
def log_likelihood(theta: jax.Array, x: jax.Array) -> jax.Array:
    """ Computes sum of log-likelihoods for data x given theta. """
    mu = theta
    cov = get_sigma(theta)
    return jnp.sum(multivariate_normal.logpdf(x, mu, cov))
 
# ------------------------------------------------------------------
# 2. Fisher Information Computation
# ------------------------------------------------------------------
@jit
def compute_fisher_mc(theta: jax.Array, key: jax.Array, num_samples: int = 1000) -> jax.Array:
    """
    Approximates the Fisher Information Matrix using Monte Carlo integration.
    G(theta) = E[score * score^T]
    """
    cov = get_sigma(theta)
    # Sampling from the model distribution at theta
    samples = random.multivariate_normal(key, theta, cov, shape=(num_samples,))
    
    def score_fn(t, x_single):
        return grad(lambda p: multivariate_normal.logpdf(x_single, p, get_sigma(p)))(t)
    
    # Vectorized score computation
    scores = vmap(lambda x: score_fn(theta, x))(samples)
    
    # Outer product expectation
    outer_products = vmap(lambda s: jnp.outer(s, s))(scores)
    return jnp.mean(outer_products, axis=0)
 
# ------------------------------------------------------------------
# 3. Optimization Loop (JIT-Compiled Scan)
# ------------------------------------------------------------------
def loss_fn(theta: jax.Array, batch: jax.Array) -> jax.Array:
    """ Negative Log Likelihood Loss. """
    return -log_likelihood(theta, batch) / batch.shape[0]
 
class OptState(NamedTuple):
    theta_sgd: jax.Array
    theta_ngd: jax.Array
    key: jax.Array
 
@jit
def update_step_sgd(theta: jax.Array, batch: jax.Array) -> jax.Array:
    grads = grad(loss_fn)(theta, batch)
    return theta - LEARNING_RATE * grads
 
@jit
def update_step_ngd(theta: jax.Array, batch: jax.Array, key: jax.Array) -> jax.Array:
    grads = grad(loss_fn)(theta, batch)
    fisher = compute_fisher_mc(theta, key)
    
    # Natural Gradient: G^-1 * grad
    # Numerically stable solve: (G + damping * I) * update = grad
    regularized_fisher = fisher + DAMPING * jnp.eye(fisher.shape[0])
    nat_grad = jnp.linalg.solve(regularized_fisher, grads)
    
    return theta - LEARNING_RATE * nat_grad
 
@jit
def run_experiment() -> Tuple[jax.Array, jax.Array]:
    """ Fully compiled training loop using lax.scan. """
    key = random.PRNGKey(SEED)
    key, subkey_data, subkey_init = random.split(key, 3)
    
    # Ground Truth
    true_theta = jnp.array([2.0, 3.0])
    data = random.multivariate_normal(
        subkey_data, true_theta, get_sigma(true_theta), shape=(500,)
    )
    
    # Initialization
    theta_0 = jnp.array([0.5, 0.5])
    init_state = OptState(theta_sgd=theta_0, theta_ngd=theta_0, key=subkey_init)
    
    def step_fn(state: OptState, _):
        key, subkey_ngd = random.split(state.key)
        
        # Parallel updates
        next_sgd = update_step_sgd(state.theta_sgd, data)
        next_ngd = update_step_ngd(state.theta_ngd, data, subkey_ngd)
        
        new_state = OptState(theta_sgd=next_sgd, theta_ngd=next_ngd, key=key)
        # Record trajectories
        return new_state, (next_sgd, next_ngd)
 
    # Execute simulation
    final_state, (path_sgd, path_ngd) = lax.scan(step_fn, init_state, None, length=NUM_STEPS)
    
    return path_sgd, path_ngd

Implications for optimization

Pulling all the theory together we land on a handful of concrete takeaways.

1. The Fisher metric is the unique valid choice (Chentsov).
2. Exponential families are both e-flat and m-flat (Amari).
3. The Cramer-Rao bound is Cauchy-Schwarz on $T_\theta \mathcal{S}$.
4. Deep learning lives in the singular regime where $G$ is rank-deficient (Watanabe).
5. The natural gradient $G^{-1} \nabla \mathcal{L}$ is the unique first-order update that respects the geometry.

A quick word on Adam. It approximates $G$ with a diagonal $\text{diag}(\sqrt{\mathbb{E}[g^2]})$ and that has no clean geometric story behind it. Natural gradient scales by $G^{-1}$ with units $1/g^2$ and Adam scales by $G^{-1/2}$ with units $1/g$. In practice Adam is much closer to sign descent where you just normalize gradient magnitudes per coordinate, and not really Riemannian steepest descent at all. Adam works, but for reasons that are mostly beside the curvature correction story.

Without information geometry, optimization is running around in a coordinate system that has no real meaning attached to it.

Timeline

Legendre Duality

The exponential and mixture duality is really just a Legendre transform.

$\psi(\theta)$ is the cumulant generating function and it reads $\psi(\theta) = \log \int \exp(\theta \cdot F(x)) d\mu(x)$

and its Legendre conjugate is $\phi(\eta) = \sup_{\theta} \{ \theta \cdot \eta - \psi(\theta) \}$

The sup gets hit where $\eta = \nabla \psi(\theta)$, which says $\eta_i = \mathbb{E}[F_i(x)]$, and this is the map that takes you from natural parameters to expectation parameters.

$\phi(\eta)$ is negative entropy up to constants and the inverse map goes the other way as $\theta = \nabla \phi(\eta)$.

The Hessians give you the metric in each coordinate system, $G_{ij}(\theta) = \frac{\partial^2 \psi}{\partial \theta_i \partial \theta_j}, \qquad G^{ij}(\eta) = \frac{\partial^2 \phi}{\partial \eta_i \partial \eta_j}$

and these two are inverses of each other up to coordinate-change Jacobians, and that confirms the Riemannian structure is consistent across the dual representations.

Fisher vs. Wasserstein

People mix these two up all the time but they actually live on different spaces.

Fisher geometry sits on parameter space $\mathcal{S}$ and is about densities, and Wasserstein geometry sits on the sample space $\Omega$ lifted up to measures and it needs a ground metric $d_\Omega(x, y)$ that Fisher geometry does not need at all.

Their geodesics are fundamentally different. The Fisher geodesic from the e-connection interpolates multiplicatively as $\log p_t(x) = (1-t) \log p_0(x) + t \log p_1(x) - \psi(t)$

and the Wasserstein geodesic just slides mass around horizontally as $T(x) = x + t \nabla \phi(x)$

Try interpolating $N(0, 1)$ and $N(10, 1)$ and see what happens. The Fisher m-geodesic runs through a bimodal mixture in the middle and the Wasserstein geodesic just slides the bump over smoothly and hits $N(5, 1)$ at the halfway point.

Which one you reach for depends on the question you are asking. Fisher geometry fits inference and asks how much a sample tells you about $\theta$, and Wasserstein fits transport and asks what it costs to push one distribution into another. Fisher comes out of the entropy Hessian and Wasserstein comes out of kinetic energy minimization from Benamou-Brenier.

Vocabulary

An affine connection sets up parallel transport and covariant derivatives on a manifold and tells you how to compare tangent vectors at different points. The Fisher Information Metric is the only Riemannian metric on families of probability distributions (up to scale) that stays invariant under sufficient statistics and it falls out of the Hessian of KL divergence. The natural gradient $G^{-1} \nabla \mathcal{L}$ is the steepest descent direction that actually accounts for the curvature of the statistical manifold instead of pretending parameter space is flat. A statistical manifold is a family of probability distributions equipped with the Fisher metric and $\alpha$-connections. Dual connections are a pair $(\nabla^{(\alpha)}, \nabla^{(-\alpha)})$ that satisfy the compatibility condition $X \langle Y, Z \rangle = \langle \nabla_X Y, Z \rangle + \langle Y, \nabla_X^* Z \rangle$ with respect to the metric. The Kullback-Leibler divergence is the asymmetric divergence whose second-order Taylor expansion spits out the Fisher metric, and even though it is not a distance in the topological sense its local geometry is Riemannian.

References

1. Amari, S., & Nagaoka, H. (2000). “Methods of Information Geometry”. The standard reference. It sets up $\alpha$-connections and dually flat spaces and applications to estimation.

2. Chentsov, N. N. (1972). “Statistical Decision Rules and Optimal Inference”. Proves the uniqueness of the Fisher metric from Markov invariance.

3. Rao, C. R. (1945). “Information and the accuracy attainable in the estimation of statistical parameters”. The original paper that proposed the Riemannian metric.

4. Watanabe, S. (2009). “Algebraic Geometry and Statistical Learning Theory”. The foundation of Singular Learning Theory and it handles the place where regular information geometry falls over in neural networks.

5. Martens, J. (2014). “New insights and perspectives on the natural gradient method”. A modern take on why Natural Gradient works for deep learning (K-FAC).