Information Geometry & Natural Gradient
Information geometry applies differential geometry to probability theory, revealing the geometric structure of statistical models. The natural gradient respects this geometry, leading to more efficient optimization in parameter spaces.
1. Statistical Manifolds
1.1 Definition
A statistical manifold is a Riemannian manifold where:
- is the space of probability distributions
- is a Riemannian metric tensor (the Fisher information metric)
1.2 Exponential Families
An exponential family has the form:
where:
- : natural parameters
- : transformed parameters
- : sufficient statistics
- : log-partition function
Example: Gaussian with :
1.3 Dual Coordinates
Exponential families have two natural coordinate systems:
Expectation parameters:
Natural parameters:
where is the convex conjugate: .
Duality: The two coordinate systems are related by:
2. Fisher Information Metric
2.1 Definition
The Fisher information matrix defines the Riemannian metric:
2.2 Properties
Positive definiteness: for regular statistical models.
Invariance: Under reparameterization :
where is the Jacobian.
Chain rule: For independent observations:
2.3 Fisher Information for Common Families
Gaussian :
Bernoulli :
Poisson :
2.4 Connection to Cramér-Rao Bound
The Fisher information matrix determines the minimum variance of unbiased estimators:
This is the Cramér-Rao bound, with equality for efficient estimators.
3. Geodesics and Divergences
3.1 Riemannian Distance
The geodesic distance between two distributions:
3.2 KL Divergence as Divergence
The KL divergence is not a metric (asymmetric), but defines a divergence:
Taylor expansion:
The Fisher information is the Hessian of the KL divergence at zero.
3.3 α-Divergences
A family of divergences parameterized by :
- : KL divergence
- : reverse KL
- : Hellinger distance
3.4 f-Divergences
The f-divergence is defined by a convex function :
- : KL divergence
- : Pearson divergence
- : reverse KL
4. Natural Gradient Descent
4.1 Motivation
Standard gradient descent updates parameters in the Euclidean direction:
This is not invariant to reparameterization. The natural gradient respects the geometry:
4.2 Information-Geometric Interpretation
The natural gradient is the steepest descent direction in the statistical manifold:
This is the direction that maximizes the decrease in per unit of KL divergence (not Euclidean distance).
4.3 Properties
Invariance: Under reparameterization :
where is the Jacobian. The natural gradient transforms correctly.
Fisher information as preconditioner: The natural gradient preconditions the gradient by the inverse Fisher information, adapting to the local geometry.
4.4 Amari's Theorem
Theorem (Amari, 1998): The natural gradient is the unique gradient that is:
- Invariant under reparameterization
- Equal to the standard gradient in canonical coordinates
- Multiplicative invariant for scale parameters
5. Practical Natural Gradient Methods
5.1 Empirical Fisher
Instead of the true Fisher, use the empirical Fisher:
Limitation: Not positive definite in general; can be rank-deficient.
5.2 KFAC (Kronecker-Factored Approximate Curvature)
For neural networks, approximate the Fisher as a Kronecker product:
where is the input covariance and is the gradient covariance for layer .
Inversion:
5.3 Natural Gradient Descent for Neural Networks
# Pseudocode for KFAC natural gradient
for each batch:
# Forward pass
activations = forward(batch)
# Backward pass
gradients = backward(batch)
# Compute Fisher factors
for each layer l:
A_l = cov(activations[l]) # Input covariance
G_l = cov(gradients[l]) # Gradient covariance
# Kronecker approximation
F_l = kron(A_l, G_l)
# Natural gradient
natural_grad_l = solve(F_l, gradients[l])
# Update parameters
params -= lr * natural_grad
5.4 SFO (Sum-of-Functions Optimizer)
Decompose the objective into sum of per-sample functions:
Estimate the Fisher using only a subset of samples.
6. Mirror Descent
6.1 Definition
Mirror descent generalizes natural gradient by using a Bregman divergence instead of the Fisher information:
where is the Bregman divergence generated by the convex function .
6.2 Connection to Natural Gradient
When (local approximation), mirror descent reduces to natural gradient.
6.3 Examples
Entropy mirror: (for probability simplices)
Quadratic mirror:
(standard gradient descent)
6.4 Regret Bounds
Mirror descent achieves regret:
The constant depends on the diameter of the constraint set under the Bregman divergence.
7. Connections to Deep Learning
7.1 Loss Surface Geometry
The Fisher information defines the local geometry of the loss surface:
The second-order term is the Fisher information, not the Hessian. This is the expected curvature, not the observed curvature.
7.2 Flat Minima and Generalization
Hypothesis: Flat minima (small eigenvalues of ) generalize better.
Natural gradient tends to find flatter minima because it penalizes directions with high curvature.
7.3 Information Geometry of Dropout
Dropout can be interpreted as sampling from a distribution over networks. The variational dropout objective:
The KL term is an information-geometric regularization.
8. Code: Information Geometry
import numpy as np
from scipy.optimize import minimize
from scipy.linalg import inv, det
class StatisticalManifold:
"""Statistical manifold for exponential families."""
def __init__(self, dim):
self.dim = dim
def log_partition(self, theta):
"""Log-partition function (to be implemented by subclasses)."""
raise NotImplementedError
def natural_to_expectation(self, theta):
"""Convert natural to expectation parameters."""
from scipy.optimize import fsolve
def equation(eta):
return self.grad_log_partition(eta) - theta
eta0 = np.zeros(self.dim)
return fsolve(equation, eta0)
def expectation_to_natural(self, eta):
"""Convert expectation to natural parameters."""
from scipy.optimize import fsolve
def equation(theta):
return self.grad_log_partition(theta) - eta
theta0 = np.zeros(self.dim)
return fsolve(equation, theta0)
def fisher_information(self, theta):
"""Compute Fisher information matrix."""
return self.hessian_log_partition(theta)
def kl_divergence(self, theta1, theta2):
"""Compute KL divergence D(p_theta1 || p_theta2)."""
A1 = self.log_partition(theta1)
A2 = self.log_partition(theta2)
eta1 = self.grad_log_partition(theta1)
return A2 - A1 - eta1 @ (theta2 - theta1)
def geodesic(self, theta0, theta1, t):
"""Approximate geodesic (for visualization)."""
# Simple linear interpolation (exact for flat manifolds)
return (1 - t) * theta0 + t * theta1
class GaussianManifold(StatisticalManifold):
"""Statistical manifold for Gaussian distributions."""
def __init__(self):
super().__init__(dim=2) # (mu, sigma^2)
def log_partition(self, theta):
mu, sigma2 = theta
return mu**2 / (2 * sigma2) + 0.5 * np.log(2 * np.pi * sigma2)
def grad_log_partition(self, theta):
mu, sigma2 = theta
return np.array([mu / sigma2, -mu**2 / (2 * sigma2**2) + 1 / (2 * sigma2)])
def hessian_log_partition(self, theta):
mu, sigma2 = theta
return np.array([
[1 / sigma2, -mu / sigma2**2],
[-mu / sigma2**2, mu**2 / sigma2**3 - 1 / (2 * sigma2**2)]
])
def fisher_information(self, theta):
mu, sigma2 = theta
return np.array([
[1 / sigma2, 0],
[0, 1 / (2 * sigma2**2)]
])
def sample(self, theta, n=1):
mu, sigma2 = theta
return np.random.normal(mu, np.sqrt(sigma2), n)
class BernoulliManifold(StatisticalManifold):
"""Statistical manifold for Bernoulli distributions."""
def __init__(self):
super().__init__(dim=1)
def log_partition(self, theta):
return np.log(1 + np.exp(theta))
def grad_log_partition(self, theta):
return np.array([1 / (1 + np.exp(-theta))])
def hessian_log_partition(self, theta):
p = 1 / (1 + np.exp(-theta))
return np.array([[p * (1 - p)]])
def fisher_information(self, theta):
p = 1 / (1 + np.exp(-theta))
return np.array([[1 / (p * (1 - p))]])
def sample(self, theta, n=1):
p = 1 / (1 + np.exp(-theta))
return np.random.binomial(1, p, n)
class NaturalGradientDescent:
"""Natural gradient descent optimizer."""
def __init__(self, manifold, lr=0.01):
self.manifold = manifold
self.lr = lr
def step(self, theta, grad):
"""Perform one natural gradient step."""
# Compute Fisher information
F = self.manifold.fisher_information(theta)
# Natural gradient: F^{-1} * grad
natural_grad = np.linalg.solve(F, grad)
# Update
return theta - self.lr * natural_grad
class KFACOptimizer:
"""Kronecker-Factored Approximate Curvature optimizer."""
def __init__(self, layers, lr=0.01, damping=0.01):
self.layers = layers
self.lr = lr
self.damping = damping
self.A = {} # Input covariances
self.G = {} # Gradient covariances
self.step_count = 0
def _compute_factors(self, layer_idx, activations, gradients):
"""Compute Kronecker factors for a layer."""
# Input activation covariance
A = activations.T @ activations / len(activations)
# Gradient covariance
G = gradients.T @ gradients / len(gradients)
return A, G
def _natural_gradient(self, layer_idx, gradient, A, G):
"""Compute natural gradient using Kronecker factorization."""
# Add damping for numerical stability
A_damped = A + self.damping * np.eye(A.shape[0])
G_damped = G + self.damping * np.eye(G.shape[0])
# Invert factors
A_inv = np.linalg.inv(A_damped)
G_inv = np.linalg.inv(G_damped)
# Natural gradient: (A^{-1} ⊗ G^{-1}) * vec(grad)
# For matrix gradient: G^{-1} * grad * A^{-1}
if gradient.ndim == 2:
natural_grad = G_inv @ gradient @ A_inv
else:
# For vector gradient: (A^{-1} ⊗ G^{-1}) * grad
natural_grad = np.kron(A_inv, G_inv) @ gradient
return natural_grad
def step(self, activations_dict, gradients_dict):
"""Perform one KFAC step."""
updates = {}
for layer_idx in self.layers:
if layer_idx in activations_dict and layer_idx in gradients_dict:
# Compute Kronecker factors
A, G = self._compute_factors(
layer_idx,
activations_dict[layer_idx],
gradients_dict[layer_idx]
)
# Update running estimates
self.A[layer_idx] = 0.9 * self.A.get(layer_idx, A) + 0.1 * A
self.G[layer_idx] = 0.9 * self.G.get(layer_idx, G) + 0.1 * G
# Compute natural gradient
updates[layer_idx] = self._natural_gradient(
layer_idx,
gradients_dict[layer_idx],
self.A[layer_idx],
self.G[layer_idx]
)
self.step_count += 1
return updates
class MirrorDescent:
"""Mirror descent optimizer with various mirror maps."""
def __init__(self, mirror='quadratic', lr=0.01):
self.mirror = mirror
self.lr = lr
def _mirror_map(self, theta):
"""Compute mirror map h(theta)."""
if self.mirror == 'quadratic':
return 0.5 * np.sum(theta**2)
elif self.mirror == 'entropy':
# Negative entropy for probability simplex
return np.sum(theta * np.log(theta + 1e-10))
elif self.mirror == 'log':
return np.sum(np.log(theta + 1e-10))
else:
raise ValueError(f"Unknown mirror: {self.mirror}")
def _bregman_divergence(self, theta, theta0):
"""Compute Bregman divergence D_h(theta, theta0)."""
h_theta = self._mirror_map(theta)
h_theta0 = self._mirror_map(theta0)
# Gradient of h at theta0
grad = self._grad_mirror(theta0)
return h_theta - h_theta0 - grad @ (theta - theta0)
def _grad_mirror(self, theta):
"""Gradient of mirror map."""
if self.mirror == 'quadratic':
return theta
elif self.mirror == 'entropy':
return np.log(theta + 1e-10) + 1
elif self.mirror == 'log':
return 1 / (theta + 1e-10)
def _proximal_mirror(self, grad, theta0):
"""Solve proximal step for mirror descent."""
def objective(theta):
return grad @ theta + self._bregman_divergence(theta, theta0) / self.lr
# For quadratic mirror, closed-form solution
if self.mirror == 'quadratic':
return theta0 - self.lr * grad
# For other mirrors, use optimization
result = minimize(objective, theta0, method='L-BFGS-B')
return result.x
def step(self, theta, grad):
"""Perform one mirror descent step."""
return self._proximal_mirror(grad, theta)
class InformationGeometryNeuralNet:
"""Neural network with natural gradient updates."""
def __init__(self, input_dim, hidden_dim, output_dim, lr=0.01):
self.lr = lr
# Initialize weights
self.W1 = np.random.randn(input_dim, hidden_dim) * 0.01
self.b1 = np.zeros(hidden_dim)
self.W2 = np.random.randn(hidden_dim, output_dim) * 0.01
self.b2 = np.zeros(output_dim)
# KFAC factors
self.A1 = None
self.G1 = None
self.A2 = None
self.G2 = None
# Damping
self.damping = 0.01
def forward(self, X):
"""Forward pass."""
self.z1 = X @ self.W1 + self.b1
self.a1 = np.maximum(0, self.z1) # ReLU
self.z2 = self.a1 @ self.W2 + self.b2
# Softmax for classification
exp_z = np.exp(self.z2 - np.max(self.z2, axis=1, keepdims=True))
self.probs = exp_z / np.sum(exp_z, axis=1, keepdims=True)
return self.probs
def backward(self, X, y):
"""Backward pass with KFAC factor computation."""
n = len(y)
# One-hot encoding
y_onehot = np.zeros_like(self.probs)
y_onehot[np.arange(n), y] = 1
# Output gradient
dz2 = self.probs - y_onehot
dW2 = self.a1.T @ dz2 / n
db2 = np.mean(dz2, axis=0)
# Hidden gradient
da1 = dz2 @ self.W2.T
dz1 = da1 * (self.z1 > 0) # ReLU derivative
dW1 = X.T @ dz1 / n
db1 = np.mean(dz1, axis=0)
# Compute KFAC factors
self.A1 = X.T @ X / n
self.G1 = dz1.T @ dz1 / n
self.A2 = self.a1.T @ self.a1 / n
self.G2 = dz2.T @ dz2 / n
return dW1, db1, dW2, db2
def natural_gradient_step(self, X, y):
"""Update using KFAC natural gradient."""
# Forward and backward
self.forward(X)
dW1, db1, dW2, db2 = self.backward(X, y)
# Natural gradient for layer 2
A2_damped = self.A2 + self.damping * np.eye(self.A2.shape[0])
G2_damped = self.G2 + self.damping * np.eye(self.G2.shape[0])
A2_inv = np.linalg.inv(A2_damped)
G2_inv = np.linalg.inv(G2_damped)
nat_grad_W2 = G2_inv @ dW2 @ A2_inv
nat_grad_b2 = G2_inv @ db2
# Natural gradient for layer 1
A1_damped = self.A1 + self.damping * np.eye(self.A1.shape[0])
G1_damped = self.G1 + self.damping * np.eye(self.G1.shape[0])
A1_inv = np.linalg.inv(A1_damped)
G1_inv = np.linalg.inv(G1_damped)
nat_grad_W1 = G1_inv @ dW1 @ A1_inv
nat_grad_b1 = G1_inv @ db1
# Update parameters
self.W2 -= self.lr * nat_grad_W2
self.b2 -= self.lr * nat_grad_b2
self.W1 -= self.lr * nat_grad_W1
self.b1 -= self.lr * nat_grad_b1
# Compute loss
loss = -np.mean(np.log(self.probs[np.arange(len(y)), y] + 1e-10))
return loss
def train(self, X, y, epochs=100, batch_size=32):
"""Train with natural gradient."""
losses = []
for epoch in range(epochs):
# Shuffle data
idx = np.random.permutation(len(y))
X_shuffled = X[idx]
y_shuffled = y[idx]
epoch_loss = 0
n_batches = 0
for i in range(0, len(y), batch_size):
X_batch = X_shuffled[i:i+batch_size]
y_batch = y_shuffled[i:i+batch_size]
loss = self.natural_gradient_step(X_batch, y_batch)
epoch_loss += loss
n_batches += 1
losses.append(epoch_loss / n_batches)
if (epoch + 1) % 10 == 0:
print(f"Epoch {epoch+1}: loss = {losses[-1]:.4f}")
return losses
# Example: Visualize Gaussian manifold
def visualize_gaussian_manifold():
"""Visualize the statistical manifold of Gaussians."""
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
manifold = GaussianManifold()
# Create grid of parameters
mu_range = np.linspace(-2, 2, 20)
sigma2_range = np.linspace(0.1, 2, 20)
MU, SIGMA2 = np.meshgrid(mu_range, sigma2_range)
# Compute Fisher information determinant
DET = np.zeros_like(MU)
for i in range(len(mu_range)):
for j in range(len(sigma2_range)):
theta = np.array([MU[i, j], SIGMA2[i, j]])
F = manifold.fisher_information(theta)
DET[i, j] = np.sqrt(np.linalg.det(F))
# Plot
fig = plt.figure(figsize=(12, 4))
# 3D surface
ax1 = fig.add_subplot(131, projection='3d')
ax1.plot_surface(MU, SIGMA2, DET, cmap='viridis', alpha=0.8)
ax1.set_xlabel('μ')
ax1.set_ylabel('σ²')
ax1.set_zlabel('√det(F)')
ax1.set_title('Fisher Information Determinant')
# Contour plot
ax2 = fig.add_subplot(132)
contour = ax2.contourf(MU, SIGMA2, DET, levels=20, cmap='viridis')
plt.colorbar(contour, ax=ax2)
ax2.set_xlabel('μ')
ax2.set_ylabel('σ²')
ax2.set_title('Fisher Information Contours')
# KL divergence contours
ax3 = fig.add_subplot(133)
theta0 = np.array([0, 1]) # Reference distribution
KL = np.zeros_like(MU)
for i in range(len(mu_range)):
for j in range(len(sigma2_range)):
theta = np.array([MU[i, j], SIGMA2[i, j]])
KL[i, j] = manifold.kl_divergence(theta0, theta)
contour = ax3.contourf(MU, SIGMA2, KL, levels=20, cmap='hot')
plt.colorbar(contour, ax=ax3)
ax3.set_xlabel('μ')
ax3.set_ylabel('σ²')
ax3.set_title('KL Divergence from N(0,1)')
plt.tight_layout()
plt.show()
# Compare natural gradient vs standard gradient
def compare_optimizers():
"""Compare natural gradient vs standard gradient on Gaussian estimation."""
manifold = GaussianManifold()
# True parameters
theta_true = np.array([1.0, 0.5])
# Generate data
np.random.seed(42)
data = manifold.sample(theta_true, n=100)
# Loss function: negative log-likelihood
def loss(theta):
mu, sigma2 = theta
return 0.5 * np.log(2 * np.pi * sigma2) + np.mean((data - mu)**2) / (2 * sigma2)
def grad_loss(theta):
mu, sigma2 = theta
dmu = -np.mean(data - mu) / sigma2
dsigma2 = -0.5 / sigma2 + np.mean((data - mu)**2) / (2 * sigma2**2)
return np.array([dmu, dsigma2])
# Standard gradient descent
theta_gd = np.array([0.0, 1.0])
lr_gd = 0.1
history_gd = [loss(theta_gd)]
for _ in range(50):
grad = grad_loss(theta_gd)
theta_gd = theta_gd - lr_gd * grad
history_gd.append(loss(theta_gd))
# Natural gradient descent
ngd = NaturalGradientDescent(manifold, lr=0.1)
theta_ngd = np.array([0.0, 1.0])
history_ngd = [loss(theta_ngd)]
for _ in range(50):
grad = grad_loss(theta_ngd)
theta_ngd = ngd.step(theta_ngd, grad)
history_ngd.append(loss(theta_ngd))
# Plot
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 4))
plt.subplot(121)
plt.plot(history_gd, label='Gradient Descent')
plt.plot(history_ngd, label='Natural Gradient')
plt.xlabel('Iteration')
plt.ylabel('Negative Log-Likelihood')
plt.legend()
plt.title('Convergence Comparison')
plt.subplot(122)
plt.plot([h[0] for h in [theta_gd]], [h[1] for h in [theta_gd]], 'ro',
label='GD final')
plt.plot([h[0] for h in [theta_ngd]], [h[1] for h in [theta_ngd]], 'bo',
label='NGD final')
plt.plot(theta_true[0], theta_true[1], 'k*', markersize=15, label='True')
plt.xlabel('μ')
plt.ylabel('σ²')
plt.legend()
plt.title('Parameter Space')
plt.tight_layout()
plt.show()
if __name__ == "__main__":
# Visualize Gaussian manifold
visualize_gaussian_manifold()
# Compare optimizers
compare_optimizers()
9. Summary
Information geometry reveals the rich mathematical structure underlying probability distributions:
- Statistical manifolds treat distributions as points in a geometric space
- Fisher information defines the natural metric on this space
- Natural gradient respects the geometry, leading to invariant and efficient optimization
- Mirror descent generalizes natural gradient using Bregman divergences
- Connections to Cramér-Rao bound link estimation theory to geometry
The key insight is that the parameter space of statistical models has a natural geometry defined by the Fisher information, and algorithms that respect this geometry (natural gradient, mirror descent) are more efficient and invariant to reparameterization.