Module 15. Information Geometry: Fisher, Natural Gradient, and Distribution Flows — exercises and worked solutions · Information Theory for ML

Exercise 1. Bernoulli Fisher Geometry

For $X\sim\operatorname{Bernoulli}(p)$ :

derive $G_p(p)=1/[p(1-p)]$ ;
switch to the logit and obtain $G_z(z)=p(1-p)$ ;
show that $u(p)=2\arcsin\sqrt p$ transforms the line element into $ds^2=du^2$ ;
compute the Fisher–Rao distance between $p_1=0.1$ and $p_2=0.9$ .

Show answer and solution

For

\log p(X)=X\log p+(1-X)\log(1-p),

the score in coordinate $p$ is

\frac{\partial}{\partial p}\log p(X) = \frac{X-p}{p(1-p)}.

Since $\operatorname{Var}(X)=p(1-p)$ ,

G_p(p) = \mathbb E \left[ \frac{(X-p)^2}{p^2(1-p)^2} \right] = \boxed{\frac1{p(1-p)}}.

For logit $z$ , $dp/dz=p(1-p)$ , hence

G_z = \left(\frac{dp}{dz}\right)^2G_p = \boxed{p(1-p)}.

With

u(p)=2\arcsin\sqrt p,

we have

\frac{du}{dp}=\frac1{\sqrt{p(1-p)}}

and therefore

ds^2=\frac{dp^2}{p(1-p)}=du^2.

Thus

\boxed{ d_{\mathrm{FR}}(p_1,p_2) =2\left| \arcsin\sqrt{p_2}-\arcsin\sqrt{p_1} \right|. }

For $p_1=0.1$ , $p_2=0.9$ :

\boxed{d_{\mathrm{FR}}\approx1.854590436}.

Exercise 2. Natural Step from a Trust Region

Let $g=\nabla L(\theta)$ and $G\succ0$ . Solve

\min_\delta g^\top\delta \qquad \text{subject to} \qquad \frac12\delta^\top G\delta\le\varepsilon.

Find the direction and normalization factor. What changes when $G$ is singular?

Show answer and solution

The Lagrangian is

\mathcal L(\delta,\lambda) = g^\top\delta + \lambda \left( \frac12\delta^\top G\delta-\varepsilon \right).

Stationarity gives

g+\lambda G\delta=0,

so

\delta=-\lambda^{-1}G^{-1}g.

The constraint is active at the optimum:

\frac1{2\lambda^2}g^\top G^{-1}g=\varepsilon.

Therefore

\boxed{ \delta^* =- \sqrt{ \frac{2\varepsilon}{g^\top G^{-1}g} } G^{-1}g }.

For singular $G$ , use a pseudoinverse, restrict to the identifiable subspace, or add damping $G+\lambda I$ . If the objective has a component along the null space, the undamped local problem may be unbounded.

Exercise 3. Exponential Families and Bregman Divergence

For

p_\theta(x) =h(x)e^{\theta^\top T(x)-A(\theta)},

show that

\nabla A(\theta) =\mathbb E_\theta[T(X)], \qquad \nabla^2A(\theta) =G(\theta),

and

D_{\mathrm{KL}} (p_{\theta_1}\|p_{\theta_2}) =B_A(\theta_2,\theta_1).

Show answer and solution

Normalization gives

e^{A(\theta)} = \int h(x)e^{\theta^\top T(x)}dx.

Differentiating,

\boxed{ \nabla A(\theta)=\mathbb E_\theta[T(X)] }.

Differentiating once more,

\boxed{ \nabla^2A(\theta) = \operatorname{Cov}_\theta(T(X)) = G(\theta) }.

For KL,

\begin{aligned} D_{\mathrm{KL}}(p_{\theta_1}\|p_{\theta_2}) &= \mathbb E_{\theta_1} \left[ (\theta_1-\theta_2)^\top T(X) -A(\theta_1)+A(\theta_2) \right]\\ &= A(\theta_2)-A(\theta_1) -\nabla A(\theta_1)^\top(\theta_2-\theta_1)\\ &= \boxed{B_A(\theta_2,\theta_1)}. \end{aligned}

The Bregman arguments are reversed because the expectation and tangent point are at $\theta_1$ .

Exercise 4. Model and Empirical Fisher

For $Y\sim\operatorname{Bernoulli}(q)$ and a model $p=\sigma(z)$ :

derive the model Fisher $p(1-p)$ ;
compute $\mathbb E_{Y\sim q}[(Y-p)^2]$ ;
determine when the two quantities coincide;
explain why coincidence at an optimum does not make empirical Fisher a universal curvature matrix.

Show answer and solution

For Bernoulli log-likelihood in logit coordinate,

\frac{\partial}{\partial z}\log p(Y)=Y-p.

When $Y\sim\operatorname{Bernoulli}(p)$ ,

G_{\mathrm{model}} = \mathbb E_p[(Y-p)^2] = \boxed{p(1-p)}.

Under data $Y\sim\operatorname{Bernoulli}(q)$ ,

\boxed{ G_{\mathrm{data}} =q(1-p)^2+(1-q)p^2 }.

The quantities agree when $q=p$ , and also at some accidental points such as $p=1/2$ . Near a well-specified optimum $p\approx q$ , they may become close. This does not produce a universal equality throughout training.

Exercise 5. TRPO, PPO, and Reference KL

State the KL orientation used in practical TRPO.
Construct an example in which PPO ratios on the sample do not change while global KL becomes large.
Compare stepwise KL and fixed-reference KL in RLHF.
Explain why DPO is not a natural-gradient update.

Show answer and solution

TRPO

The standard practical KL direction is

\mathbb E_s D_{\mathrm{KL}} \left( \pi_{\mathrm{old}}(\cdot\mid s) \| \pi_{\mathrm{new}}(\cdot\mid s) \right) \le\delta.

A PPO counterexample

Let

\pi_{\mathrm{old}}=(0.50,0.49,0.01)

and suppose the batch contains only the first action, with positive advantage. Consider

\pi_q=(0.50,q,0.50-q), \qquad 0<q<0.50.

The probability ratio of the sampled action is always $1$ , so the sampled clipped surrogate is unchanged. The full KL, however, is

D_{\mathrm{KL}}(\pi_{\mathrm{old}}\|\pi_q) = 0.49\log\frac{0.49}{q} + 0.01\log\frac{0.01}{0.50-q}.

As $q\to0$ the first term diverges while the second tends to the finite value $0.01\log0.02$ , so

D_{\mathrm{KL}}(\pi_{\mathrm{old}}\|\pi_q)\to\infty.

Clipping is not a global constraint on KL.

Two different KL terms

stepwise KL to the old policy stabilizes one on-policy update;
KL to a fixed SFT or base policy limits accumulated drift.

DPO

DPO uses the idealized relation

\pi^*(y\mid x) \propto \pi_{\mathrm{ref}}(y\mid x)e^{r(x,y)/\beta}

together with a preference likelihood, but performs no explicit natural-gradient or TRPO update.

Exercise 6. Flow Matching and Couplings

Let

X_0\sim\mathcal N(\mu_0,\sigma_0^2), \qquad X_1\sim\mathcal N(\mu_1,\sigma_1^2).

Compare

shared noise: $X_0=\mu_0+\sigma_0Z$ , $X_1=\mu_1+\sigma_1Z$ ;
independent $X_0$ and $X_1$ .

For

X_t=(1-t)X_0+tX_1,

find the mean, variance, and kinetic action. Show that the shared-noise coupling gives the one-dimensional $W_2$ geodesic.

Show answer and solution

Shared-noise coupling

Let

X_0=\mu_0+\sigma_0Z, \qquad X_1=\mu_1+\sigma_1Z, \qquad Z\sim\mathcal N(0,1).

Then

X_t=(1-t)X_0+tX_1 = \mu_t+\sigma_tZ,

where

\mu_t=(1-t)\mu_0+t\mu_1, \qquad \sigma_t=(1-t)\sigma_0+t\sigma_1.

Therefore

X_t\sim\mathcal N(\mu_t,\sigma_t^2).

The velocity is constant:

\dot X_t =(\mu_1-\mu_0)+(\sigma_1-\sigma_0)Z.

Its kinetic action is

\mathbb E[\dot X_t^2] = (\mu_1-\mu_0)^2+(\sigma_1-\sigma_0)^2 = W_2^2(P_0,P_1).

This is the one-dimensional optimal coupling and the displacement geodesic.

Independent coupling

If $X_0$ and $X_1$ are independent,

\operatorname{Var}(X_t) =(1-t)^2\sigma_0^2+t^2\sigma_1^2.

The kinetic action is

\mathbb E[(X_1-X_0)^2] =(\mu_1-\mu_0)^2+\sigma_0^2+\sigma_1^2.

It exceeds the optimum by

2\sigma_0\sigma_1.

The endpoint marginals are the same, but the coupling and the whole intermediate path differ.

Exercise 7. Python

Run

text

M15_information_geometry_experiments.py

The script

compares exact Bernoulli KL with its quadratic Fisher approximation;
demonstrates coordinate dependence of ordinary gradient and consistency of the natural-gradient flow;
compares model Fisher with the data-gradient second moment;
constructs the PPO counterexample;
compares optimal and independent couplings for a Gaussian path.

Show answer and solution

For $p=0.2$ , the Fisher quadratic approximation approaches exact KL as $\delta$ decreases:

$\delta$	Exact KL	Quadratic	Ratio
0.1	0.02573209	0.03125000	0.8234
0.01	0.00030499	0.00031250	0.9760
0.001	0.0000031172	0.0000031250	0.9975

Ordinary-gradient trajectories in probability and logit coordinates differ substantially. Two Euler discretizations of the natural-gradient vector field differ only through finite-step numerical error; the discrepancy shrinks with the step size.

At $q=0.8$ , $p=0.2$ ,

G_{\mathrm{model}}=0.16, \qquad G_{\mathrm{data}}=0.52.

In the PPO example, the sampled ratio stays at 1 while

D_{\mathrm{KL}}(\pi_{\mathrm{old}}\|\pi_{\mathrm{new}})

reaches roughly $8.64$ nats. The Gaussian experiment verifies the action gap $2\sigma_0\sigma_1$ between independent and shared-noise couplings.

Code solution · the same source file for RU/EN

PythonM15_information_geometry_experiments.pyDownload .py

from __future__ import annotations

from pathlib import Path
import math

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

ROOT = Path(__file__).resolve().parent
OUT = ROOT / 'generated_experiment'
OUT.mkdir(exist_ok=True)


def bernoulli_kl(p: np.ndarray | float, q: np.ndarray | float) -> np.ndarray:
    p = np.asarray(p, dtype=float)
    q = np.asarray(q, dtype=float)
    eps = 1e-15
    p = np.clip(p, eps, 1 - eps)
    q = np.clip(q, eps, 1 - eps)
    return p*np.log(p/q) + (1-p)*np.log((1-p)/(1-q))


def local_kl_experiment() -> dict[str, float]:
    p0 = 0.2
    deltas = np.logspace(-5, -1, 160)
    exact = bernoulli_kl(p0, p0+deltas)
    fisher = 1/(p0*(1-p0))
    quadratic = 0.5*fisher*deltas**2
    fig, ax = plt.subplots(figsize=(8.8, 5.3))
    ax.loglog(deltas, exact, label='exact KL')
    ax.loglog(deltas, quadratic, '--', label=r'$\frac{1}{2}F(p)\delta^2$')
    ax.set_xlabel(r'$\delta$')
    ax.set_ylabel('nats')
    ax.set_title('Bernoulli: local KL and Fisher quadratic approximation')
    ax.legend()
    fig.tight_layout()
    fig.savefig(OUT/'M15_local_kl.png', dpi=190)
    plt.close(fig)
    print('LOCAL KL')
    for d in (0.1, 0.05, 0.01, 0.001):
        e=float(bernoulli_kl(p0,p0+d)); q=0.5*fisher*d*d
        print(f'delta={d:g} exact={e:.10f} quadratic={q:.10f} ratio={e/q:.6f}')
    return {'fisher': fisher, 'ratio_1e3': float(bernoulli_kl(p0,p0+1e-3)/(0.5*fisher*1e-6))}


def sigmoid(z: np.ndarray | float) -> np.ndarray:
    return 1/(1+np.exp(-np.asarray(z,dtype=float)))


def logit(p: float) -> float:
    return math.log(p/(1-p))


def reparameterization_experiment() -> dict[str,float]:
    target=0.8; p0=0.1; dt=0.01; steps=300
    time=np.arange(steps+1)*dt
    p_nat=np.empty(steps+1); p_nat[0]=p0
    z_nat=np.empty(steps+1); z_nat[0]=logit(p0)
    for t in range(steps):
        p_nat[t+1]=p_nat[t]-dt*(p_nat[t]-target)
        p=float(sigmoid(z_nat[t]))
        z_nat[t+1]=z_nat[t]-dt*(p-target)/(p*(1-p))
    pz_nat=sigmoid(z_nat)
    # Coordinate-dependent ordinary-gradient examples.
    p_gd=np.empty(steps+1); p_gd[0]=p0
    z_gd=np.empty(steps+1); z_gd[0]=logit(p0)
    for t in range(steps):
        p=np.clip(p_gd[t],1e-6,1-1e-6)
        grad=(p-target)/(p*(1-p))
        p_gd[t+1]=np.clip(p-0.002*grad,1e-6,1-1e-6)
        p=float(sigmoid(z_gd[t])); z_gd[t+1]=z_gd[t]-0.08*(p-target)
    pz_gd=sigmoid(z_gd)
    exact=target+(p0-target)*np.exp(-time)
    fig, ax = plt.subplots(figsize=(9.2,5.4))
    ax.plot(time,p_gd,label='ordinary GD in p')
    ax.plot(time,pz_gd,label='ordinary GD in logit')
    ax.plot(time,p_nat,'--',label='natural Euler in p')
    ax.plot(time,pz_nat,':',label='natural Euler in logit')
    ax.plot(time,exact,color='black',linewidth=1.2,label='exact natural flow')
    ax.set_xlabel('time / iteration scale'); ax.set_ylabel('Bernoulli p')
    ax.set_ylim(0,0.9); ax.set_title('Reparameterization: vector field vs finite steps')
    ax.legend(fontsize=9); fig.tight_layout(); fig.savefig(OUT/'M15_reparameterization.png',dpi=190); plt.close(fig)
    nat=float(np.max(np.abs(p_nat-pz_nat))); ordinary=float(np.max(np.abs(p_gd-pz_gd)))
    print('\nREPARAMETERIZATION')
    print(f'max ordinary trajectory difference={ordinary:.6f}')
    print(f'max natural Euler difference={nat:.6f}')
    return {'ordinary_difference':ordinary,'natural_difference':nat}


def empirical_fisher_experiment() -> dict[str,float]:
    q=0.8
    p=np.linspace(0.01,0.99,300)
    model=p*(1-p)
    data=q*(1-p)**2+(1-q)*p**2
    fig, ax=plt.subplots(figsize=(8.8,5.2))
    ax.plot(p,model,label=r'model Fisher $p(1-p)$')
    ax.plot(p,data,label='gradient second moment under data')
    ax.axvline(q,color='black',linewidth=0.8,alpha=0.5)
    ax.set_xlabel('model probability p'); ax.set_ylabel('matrix entry / second moment')
    ax.set_title('Model Fisher and observed-gradient second moment differ')
    ax.legend(); fig.tight_layout(); fig.savefig(OUT/'M15_empirical_fisher.png',dpi=190); plt.close(fig)
    print('\nMODEL VS EMPIRICAL')
    vals={}
    for p0 in (0.2,0.5,0.8):
        f=p0*(1-p0); e=q*(1-p0)**2+(1-q)*p0**2
        vals[p0]=(f,e); print(f'p={p0:.1f} model={f:.6f} data={e:.6f}')
    return {'model_p02':vals[0.2][0],'data_p02':vals[0.2][1]}


def categorical_kl(p: np.ndarray,q: np.ndarray)->float:
    return float(np.sum(p*np.log(p/q)))


def ppo_counterexample() -> dict[str,float]:
    old=np.array([0.50,0.49,0.01])
    q1=np.geomspace(0.49,1e-8,240)
    policies=np.column_stack([np.full_like(q1,0.50),q1,0.50-q1])
    kl=np.array([categorical_kl(old,q) for q in policies])
    ratio=policies[:,0]/old[0]
    surrogate=np.minimum(ratio,np.clip(ratio,0.8,1.2))
    x=-np.log10(q1)
    fig, ax=plt.subplots(figsize=(9.2,5.3))
    ax.plot(x,kl,label=r'$D_{KL}(\pi_{old}\|\pi_{new})$')
    ax.set_xlabel(r'$-\log_{10}\pi_{new}(a_1)$'); ax.set_ylabel('KL, nats')
    ax2=ax.twinx(); ax2.plot(x,surrogate,'--',color='black',label='sampled PPO surrogate')
    ax2.set_ylabel('sampled surrogate')
    lines,labels=ax.get_legend_handles_labels(); lines2,labels2=ax2.get_legend_handles_labels()
    ax.legend(lines+lines2,labels+labels2,loc='upper left',fontsize=9)
    ax.set_title('PPO clipping on sampled actions is not a global KL constraint')
    fig.tight_layout(); fig.savefig(OUT/'M15_ppo_counterexample.png',dpi=190); plt.close(fig)
    m=float(kl.max()); print('\nPPO COUNTEREXAMPLE'); print(f'max KL={m:.6f} nats; sampled ratio={ratio[0]:.1f}')
    return {'max_kl':m}


def flow_coupling_experiment() -> dict[str,float]:
    mu0,s0=-1.0,0.7; mu1,s1=2.0,1.4
    t=np.linspace(0,1,200)
    var_ot=((1-t)*s0+t*s1)**2
    var_ind=(1-t)**2*s0**2+t**2*s1**2
    action_ot=(mu1-mu0)**2+(s1-s0)**2
    action_ind=(mu1-mu0)**2+s0**2+s1**2
    fig, ax=plt.subplots(figsize=(8.8,5.2))
    ax.plot(t,var_ot,label='shared-noise / OT coupling')
    ax.plot(t,var_ind,label='independent coupling')
    ax.set_xlabel('t'); ax.set_ylabel(r'$\mathrm{Var}(X_t)$')
    ax.set_title('The coupling determines the Gaussian interpolation path')
    ax.legend(); fig.tight_layout(); fig.savefig(OUT/'M15_flow_couplings.png',dpi=190); plt.close(fig)
    print('\nFLOW COUPLINGS')
    print(f'OT action={action_ot:.6f}; independent action={action_ind:.6f}; gap={action_ind-action_ot:.6f}')
    return {'ot_action':action_ot,'ind_action':action_ind,'gap':action_ind-action_ot}


def main()->None:
    local_kl_experiment(); reparameterization_experiment(); empirical_fisher_experiment(); ppo_counterexample(); flow_coupling_experiment()
    print(f'\nSaved plots to {OUT}')

if __name__=='__main__':
    main()

Bonus. Probability-Flow ODE

For

dX_t=f(X_t,t)dt+g(t)dW_t,

derive the Fokker–Planck equation and show that the ODE

\frac{dX_t}{dt} =f(X_t,t) - \frac12g(t)^2\nabla\log p_t(X_t)

produces the same marginals $p_t$ . Why are the trajectories different?

Show answer and solution

For

dX_t=f(X_t,t)dt+g(t)dW_t,

with scalar state-independent $g(t)$ , the Fokker–Planck equation is

\partial_tp_t =-\nabla\cdot(fp_t) +\frac12g(t)^2\Delta p_t.

Let

v_t(x) =f(x,t)-\frac12g(t)^2\nabla\log p_t(x).

Then the ODE continuity equation gives

\begin{aligned} \partial_tp_t &=-\nabla\cdot(v_tp_t)\\ &=-\nabla\cdot(fp_t) +\frac12g(t)^2\nabla\cdot(p_t\nabla\log p_t)\\ &=-\nabla\cdot(fp_t) +\frac12g(t)^2\Delta p_t. \end{aligned}

The marginal-density equation matches, but paths differ: the SDE contains Brownian randomness, while the ODE is deterministic given its initial state.