Module 11. Comparing distributions: f-divergences, variational critics, and Wasserstein distance

Exercise 1. Generators and normalizations

Show that adding $a(t-1)$ to a generator does not change the f-divergence.
Derive the explicit forms of forward KL, reverse KL, TV, and standard squared Hellinger from the table.
Verify that the chosen JS generator yields
$\frac12D_{\mathrm{KL}}(P\|M) + \frac12D_{\mathrm{KL}}(Q\|M), \qquad M=\frac{P+Q}{2}.$
Explain why power divergence and Rényi divergence should not be treated as the same family without qualification.

Show answer and solution

Forward KL

For

f(t)=t\log t

we have $f(1)=0$ and $f''(t)=1/t>0$ . Hence

\begin{aligned} D_f(P\|Q) &=\int q\,\frac pq\log\frac pq\,dx\\ &=\int p\log\frac pq\,dx =D_{\mathrm{KL}}(P\|Q). \end{aligned}

Reverse KL

For $f(t)=-\log t$ ,

D_f(P\|Q) =\int q\log\frac qp\,dx =D_{\mathrm{KL}}(Q\|P).

Total variation

The function

f(t)=\frac12|t-1|

is convex, and

\begin{aligned} D_f(P\|Q) &=\frac12\int q\left|\frac pq-1\right|dx\\ &=\frac12\int|p-q|dx =\operatorname{TV}(P,Q). \end{aligned}

Squared Hellinger

For

f(t)=\frac12(\sqrt t-1)^2,

\begin{aligned} D_f(P\|Q) &=\frac12\int q \left(\sqrt{\frac pq}-1\right)^2dx\\ &=\frac12\int(\sqrt p-\sqrt q)^2dx\\ &=1-\int\sqrt{pq}\,dx. \end{aligned}

Pearson $\chi^2$

For $f(t)=(t-1)^2$ ,

D_f(P\|Q) =\int q\left(\frac pq-1\right)^2dx =\int\frac{(p-q)^2}{q}\,dx.

Jensen–Shannon

Take

f(t) =\frac12\left[ t\log t-(t+1)\log\frac{t+1}{2} \right].

Let $M=(P+Q)/2$ , with density $m=(p+q)/2$ . Substitution and regrouping give

\begin{aligned} D_f(P\|Q) &=\frac12\int p\log\frac p{m}\,dx +\frac12\int q\log\frac q{m}\,dx\\ &=D_{\mathrm{JS}}(P,Q). \end{aligned}

Why a linear term changes nothing

Let

\widetilde f(t)=f(t)+a(t-1).

Then

\begin{aligned} D_{\widetilde f}(P\|Q) &=D_f(P\|Q) +a\int q\left(\frac pq-1\right)dx\\ &=D_f(P\|Q)+a(1-1)\\ &=D_f(P\|Q). \end{aligned}

Thus one may affine-normalize a generator—for example, also impose $f'(1)=0$ —without changing the divergence.

Power divergence and Rényi divergence

A power divergence has the ordinary f-divergence form

D_{f_\alpha}(P\|Q) =\int q f_\alpha(p/q).

Rényi divergence applies the logarithm after integration:

D_\alpha^{\mathrm R}(P\|Q) = \frac1{\alpha-1} \log\int p^\alpha q^{1-\alpha}.

It is therefore a monotone transform of the corresponding power integral, but in general it is not itself of the form $\int qf(p/q)$ for one fixed $f$ .

Exercise 2. DPI through conditional Jensen

For $P\ll Q$ and a common channel $K(dy\mid x)$ , prove

\frac{dP_Y}{dQ_Y}(Y) = \mathbb E_Q \left[ \frac{dP}{dQ}(X) \middle| Y \right]

and derive

D_f(P_Y\|Q_Y) \le D_f(P\|Q).

State the meaning for a representation $Z=T(X)$ and a classifier that distinguishes two data groups.

Show answer and solution

Let

r(X)=\frac{dP}{dQ}(X).

For every bounded measurable function $\varphi(Y)$ ,

\begin{aligned} \mathbb E_{P_Y}[\varphi(Y)] &=\mathbb E_P[\varphi(Y)]\\ &=\mathbb E_Q[r(X)\varphi(Y)]\\ &=\mathbb E_Q[ \mathbb E_Q[r(X)\mid Y]\,\varphi(Y) ]. \end{aligned}

The Radon–Nikodym characterization therefore yields

\boxed{ \frac{dP_Y}{dQ_Y}(Y) = \mathbb E_Q[r(X)\mid Y] }.

Conditional Jensen now gives

\begin{aligned} D_f(P_Y\|Q_Y) &=\mathbb E_{Q_Y} f\!\left(\mathbb E_Q[r(X)\mid Y]\right)\\ &\le \mathbb E_Q[f(r(X))]\\ &=D_f(P\|Q). \end{aligned}

For a representation $Z=T(X)$ , fixed processing cannot make two population distributions more distinguishable than they were at the input. An encoder may make the distinction easier for a restricted classifier, but it cannot increase the population f-divergence.

Exercise 3. Optimal discriminator and JS

Assume real and model samples have equal class priors.

Minimize binary cross-entropy pointwise to obtain
$D^*(x)=\frac{p(x)}{p(x)+q(x)}.$
Show
$\operatorname{logit}D^*(x)=\log\frac{p(x)}{q(x)}.$
Substitute $D^*$ into the GAN value function and obtain
$-\log4+2D_{\mathrm{JS}}(P,Q).$
Explain why this does not imply that an alternating GAN follows the exact JS gradient at every step.

Show answer and solution

With equal class priors, for each $x$ we maximize

p(x)\log D(x)+q(x)\log(1-D(x)).

The stationarity condition is

\frac pD-\frac q{1-D}=0,

so

\boxed{ D^*(x)=\frac{p(x)}{p(x)+q(x)} }.

Its logit is

\operatorname{logit}D^*(x) =\log\frac{p(x)}{q(x)}.

Let $m=(p+q)/2$ . Then

\begin{aligned} V(D^*,G) &=\int p\log\frac p{p+q} +q\log\frac q{p+q}\,dx\\ &=D_{\mathrm{KL}}(P\|M) +D_{\mathrm{KL}}(Q\|M)-2\log2\\ &=2D_{\mathrm{JS}}(P,Q)-\log4. \end{aligned}

Thus

\boxed{ \max_DV(D,Q) =-\log4+2D_{\mathrm{JS}}(P,Q) }.

This exact statement assumes an optimal discriminator for the current $Q$ . Alternating training uses an imperfect discriminator, a restricted generator family, and often the non-saturating generator loss, so a practical update need not equal the exact JS gradient.

Exercise 4. From Fenchel KL to Donsker–Varadhan

Compute the conjugate of $f(t)=t\log t$ :
$f^*(s)=e^{s-1}.$
Derive the NWJ form
$D_{\mathrm{KL}}(P\|Q) = \sup_T \{\mathbb E_PT-\mathbb E_Qe^{T-1}\}.$
Set $T=G+c$ and maximize analytically over $c$ .
Recover the DV form and explain why merely shifting $T\mapsto T+1$ is insufficient.

Show answer and solution

For $f(t)=t\log t$ ,

f^*(s)=\sup_{t>0}\{st-t\log t\}.

Stationarity gives

t^*=e^{s-1},

hence

\boxed{f^*(s)=e^{s-1}}.

The Fenchel representation becomes

\boxed{ D_{\mathrm{KL}}(P\|Q) = \sup_T \left\{ \mathbb E_PT-\mathbb E_Qe^{T-1} \right\} }.

For a fixed function $G$ , put $T=G+c$ :

J(c) =\mathbb E_PG+c -e^{c-1}\mathbb E_Q[e^G].

The optimum is

c^*=1-\log\mathbb E_Q[e^G],

and therefore

J(c^*) =\mathbb E_PG- \log\mathbb E_Q[e^G].

Thus

\boxed{ D_{\mathrm{KL}}(P\|Q) = \sup_G \left\{ \mathbb E_PG- \log\mathbb E_Q[e^G] \right\} }.

A fixed shift $T\mapsto T+1$ does not remove the exponential expectation. DV appears after optimizing the free constant, which depends on $G$ .

The optimal NWJ critic is

T^*(x)=1+\log\frac{p(x)}{q(x)},

whereas the optimal DV critic is

G^*(x)=\log\frac{p(x)}{q(x)}+c.

Exercise 5. InfoNCE: ceiling and interpretation

From $\mathcal L_{\mathrm{NCE}}\ge0$ , show that
$\log K-\mathcal L_{\mathrm{NCE}}$
cannot exceed $\log K$ .
For $K=32$ and loss $1.2$ nats, calculate the numerical lower bound.
Explain why a numerical MINE estimate need not have the same ceiling.
Give two mechanisms by which more negatives may hurt downstream representations despite raising the bound's ceiling.

Show answer and solution

Since cross-entropy loss is nonnegative,

\mathcal L_{\mathrm{NCE}}\ge0,

and hence

\boxed{ \log K-\mathcal L_{\mathrm{NCE}} \le\log K }.

For $K=32$ and loss $1.2$ nats,

\log32-1.2 \approx2.265736\ \text{nats}.

MINE uses the DV functional

\mathbb E[T]-\log\mathbb E[e^T],

which has no algebraic $\log K$ ceiling. A numerical MINE value can therefore exceed $\log K$ , although this does not make it a certified distribution-free lower bound after adaptive finite-sample training.

Increasing the number of negatives can hurt the downstream representation because, for example:

false negatives push semantically related objects apart;
under fixed compute, optimization may worsen or the number of independent updates may fall;
batch composition may fail to represent the intended marginal distribution;
a larger numerical bound need not correspond to more task-relevant information.

Any two mechanisms were sufficient.

Exercise 6. Python: critic, InfoNCE, and support geometry

Run M11_divergence_experiments.py.

It:

compares exact KL with optimized NWJ and DV objectives on a finite alphabet;
verifies that the optimal classifier logit recovers $\log(p/q)$ ;
checks the original GAN identity;
computes oracle InfoNCE bounds for correlated Gaussian variables;
compares JS, TV, and $W_1$ for shifted Dirac measures.

Answer:

why the NWJ critic differs from $\log(p/q)$ by exactly 1;
why the DV critic is defined only up to an additive constant;
why InfoNCE is far below true MI for $\rho=0.99$ at small $K$ ;
why $W_1$ changes with $|\theta|$ while JS and TV do not.

Show answer and solution

The script uses

P=(0.45,0.05,0.45,0.05),

Q=(0.05,0.45,0.05,0.45).

It obtains

D_{\mathrm{KL}}(P\|Q) \approx1.757779662\ \text{nats}.

Unrestricted NWJ and DV critics reproduce this value to numerical precision. The best linear critic reaches only

0.314646043\ \text{nats},

showing the variational gap caused by a restricted function class.

For $f(t)=t\log t$ ,

f'(t)=1+\log t,

so the NWJ optimum is

T^*=1+\log(p/q).

The DV functional is invariant under $G\mapsto G+c$ , hence

G^*=\log(p/q)+c.

The Bayes classifier recovered

\operatorname{logit}D^*=\log(p/q)

with zero maximum numerical error. The original-GAN identity also matched to about $10^{-16}$ :

V(D^*,G) =-\log4+2D_{\mathrm{JS}}(P,Q).

For a jointly Gaussian pair with $\rho=0.99$ ,

I(X;Y) =-\frac12\log(1-\rho^2) \approx1.958518\ \text{nats}.

The oracle InfoNCE estimates were approximately:

$K$	bound, nats
2	0.5705
4	1.0521
16	1.6600
64	1.8719
256	1.9281

For small $K$ , the bound is far from MI mainly because it cannot exceed $\log K$ . It tightens as $K$ grows, while remaining a Monte Carlo estimate of a particular multi-sample bound.

Finally, for

P=\delta_0, \qquad Q_\theta=\delta_\theta,

and $\theta\ne0$ ,

D_{\mathrm{JS}}=\log2, \qquad \operatorname{TV}=1,

independently of the distance between the points, whereas

W_1(P,Q_\theta)=|\theta|.

JS and TV depend only on the probability measures; Wasserstein additionally uses the ground metric.

Code solution · the same source file for RU/EN

PythonM11_divergence_experiments.pyDownload .py

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225

from __future__ import annotations

import json
import math
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import minimize
from scipy.special import logsumexp

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


def discrete_variational_experiment() -> dict[str, object]:
    """Check NWJ, DV, classifier-ratio, and original-GAN identities."""
    p = np.array([0.45, 0.05, 0.45, 0.05], dtype=np.float64)
    q = np.array([0.05, 0.45, 0.05, 0.45], dtype=np.float64)
    x = np.arange(len(p), dtype=np.float64)

    log_ratio = np.log(p / q)
    exact_kl = float(np.sum(p * log_ratio))

    # Exact population optima.
    t_nwj = 1.0 + log_ratio
    nwj_full = float(np.sum(p * t_nwj) - np.sum(q * np.exp(t_nwj - 1.0)))

    g_dv = log_ratio
    dv_full = float(np.sum(p * g_dv) - math.log(float(np.sum(q * np.exp(g_dv)))))

    # Restricted linear critics. NWJ uses a*x+b; DV fixes the irrelevant intercept.
    def nwj_negative(params: np.ndarray) -> float:
        a, b = params
        t = a * x + b
        value = np.sum(p * t) - np.sum(q * np.exp(t - 1.0))
        return -float(value)

    nwj_result = minimize(nwj_negative, x0=np.zeros(2), method="BFGS")
    if not nwj_result.success:
        raise RuntimeError(f"NWJ critic optimization failed: {nwj_result.message}")
    a_nwj, b_nwj = nwj_result.x
    nwj_linear = -float(nwj_result.fun)

    def dv_negative(a_arr: np.ndarray) -> float:
        a = float(a_arr[0])
        g = a * x
        value = np.sum(p * g) - math.log(float(np.sum(q * np.exp(g))))
        return -float(value)

    dv_result = minimize(dv_negative, x0=np.zeros(1), method="BFGS")
    if not dv_result.success:
        raise RuntimeError(f"DV critic optimization failed: {dv_result.message}")
    a_dv = float(dv_result.x[0])
    dv_linear = -float(dv_result.fun)

    # Bayes classifier recovers the log density ratio for equal class priors.
    d_star = p / (p + q)
    classifier_logit = np.log(d_star / (1.0 - d_star))
    classifier_max_abs_error = float(np.max(np.abs(classifier_logit - log_ratio)))

    # Original GAN identity.
    m = 0.5 * (p + q)
    js = 0.5 * float(np.sum(p * np.log(p / m))) + 0.5 * float(np.sum(q * np.log(q / m)))
    gan_value = float(np.sum(p * np.log(d_star)) + np.sum(q * np.log(1.0 - d_star)))
    gan_identity_rhs = -math.log(4.0) + 2.0 * js

    labels = ["exact KL", "NWJ full", "DV full", "NWJ linear", "DV linear"]
    values = [exact_kl, nwj_full, dv_full, nwj_linear, dv_linear]
    fig, ax = plt.subplots(figsize=(9.0, 5.1))
    ax.bar(np.arange(len(labels)), values)
    ax.set_xticks(np.arange(len(labels)), labels, rotation=15, ha="right")
    ax.set_ylabel("nats")
    ax.set_title("Restricted critics leave a variational gap")
    ax.axhline(exact_kl, linestyle="--", linewidth=1.0, label="exact KL")
    ax.legend()
    fig.tight_layout()
    fig.savefig(OUT / "M11_experiment_variational_bounds.png", dpi=190)
    plt.close(fig)

    return {
        "p": p.tolist(),
        "q": q.tolist(),
        "log_density_ratio": log_ratio.tolist(),
        "exact_kl_nats": exact_kl,
        "nwj_full_nats": nwj_full,
        "dv_full_nats": dv_full,
        "nwj_linear_nats": nwj_linear,
        "dv_linear_nats": dv_linear,
        "nwj_linear_a": float(a_nwj),
        "nwj_linear_b": float(b_nwj),
        "dv_linear_a": a_dv,
        "classifier_logit_max_abs_error": classifier_max_abs_error,
        "js_nats": js,
        "gan_value": gan_value,
        "gan_identity_rhs": gan_identity_rhs,
        "gan_identity_abs_error": abs(gan_value - gan_identity_rhs),
    }


def gaussian_log_ratio(x: np.ndarray, y: np.ndarray, rho: float) -> np.ndarray:
    """Return log p(y|x)/p(y) for a standard jointly Gaussian pair."""
    var = 1.0 - rho * rho
    return -0.5 * math.log(var) - 0.5 * (y - rho * x) ** 2 / var + 0.5 * y**2


def estimate_infonce_bound(
    rho: float,
    k: int,
    n_trials: int,
    rng: np.random.Generator,
    chunk_size: int = 500,
) -> tuple[float, float, float]:
    """Monte Carlo InfoNCE bound using the Bayes-optimal score."""
    losses: list[np.ndarray] = []
    remaining = n_trials
    sigma = math.sqrt(1.0 - rho * rho)
    while remaining > 0:
        n = min(chunk_size, remaining)
        x = rng.standard_normal((n, 1))
        y_pos = rho * x + sigma * rng.standard_normal((n, 1))
        y_neg = rng.standard_normal((n, k - 1)) if k > 1 else np.empty((n, 0))
        y_all = np.concatenate([y_pos, y_neg], axis=1)
        scores = gaussian_log_ratio(x, y_all, rho)
        loss = logsumexp(scores, axis=1) - scores[:, 0]
        losses.append(loss)
        remaining -= n
    all_losses = np.concatenate(losses)
    mean_loss = float(np.mean(all_losses))
    stderr = float(np.std(all_losses, ddof=1) / math.sqrt(len(all_losses)))
    return math.log(k) - mean_loss, mean_loss, stderr


def infonce_experiment() -> dict[str, object]:
    rng = np.random.default_rng(20260713)
    rhos = [0.5, 0.9, 0.99, 0.9999]
    ks = [2, 4, 16, 64, 256]
    n_trials = 20_000

    rows: list[dict[str, float | int]] = []
    fig, ax = plt.subplots(figsize=(9.2, 5.6))
    for rho in rhos:
        exact = -0.5 * math.log(1.0 - rho * rho)
        bounds = []
        for k in ks:
            bound, loss, stderr = estimate_infonce_bound(rho, k, n_trials, rng)
            bounds.append(bound)
            rows.append({
                "rho": rho,
                "K": k,
                "exact_mi_nats": exact,
                "infonce_loss_nats": loss,
                "infonce_bound_nats": bound,
                "loss_stderr": stderr,
                "log_K": math.log(k),
            })
        ax.plot(ks, bounds, marker="o", label=fr"$\rho={rho}$ bound")
        ax.axhline(exact, linestyle="--", linewidth=0.9, alpha=0.55)
    ax.set_xscale("log", base=2)
    ax.set_xlabel("number of candidates K")
    ax.set_ylabel("nats")
    ax.set_title("InfoNCE approaches MI but is capped by log K")
    ax.legend(fontsize=8, ncol=2)
    fig.tight_layout()
    fig.savefig(OUT / "M11_experiment_infonce_saturation.png", dpi=190)
    plt.close(fig)

    return {"n_trials": n_trials, "rows": rows}


def support_geometry_experiment() -> dict[str, object]:
    theta = np.linspace(-3.0, 3.0, 601)
    nonzero = np.abs(theta) > 1e-12
    js = np.where(nonzero, math.log(2.0), 0.0)
    tv = np.where(nonzero, 1.0, 0.0)
    w1 = np.abs(theta)

    fig, ax = plt.subplots(figsize=(8.8, 5.1))
    ax.plot(theta, w1, label=r"$W_1(\delta_0,\delta_\theta)=|\theta|$")
    ax.plot(theta, js, label=r"$D_{JS}$")
    ax.plot(theta, tv, label="TV")
    ax.axvline(0.0, linewidth=0.8)
    ax.set_xlabel(r"shift $\theta$")
    ax.set_ylabel("discrepancy")
    ax.set_title("Ground geometry distinguishes disjoint supports")
    ax.legend()
    fig.tight_layout()
    fig.savefig(OUT / "M11_experiment_support_geometry.png", dpi=190)
    plt.close(fig)

    return {
        "theta_examples": [-2.0, -0.1, 0.0, 0.1, 2.0],
        "js_for_nonzero_theta": math.log(2.0),
        "tv_for_nonzero_theta": 1.0,
        "kl_for_nonzero_theta": "infinity",
        "w1_formula": "abs(theta)",
    }


def main() -> None:
    result = {
        "discrete_variational": discrete_variational_experiment(),
        "infonce": infonce_experiment(),
        "support_geometry": support_geometry_experiment(),
    }
    path = OUT / "M11_experiment_results.json"
    path.write_text(json.dumps(result, indent=2, ensure_ascii=False), encoding="utf-8")

    d = result["discrete_variational"]
    print("Exact KL:", f"{d['exact_kl_nats']:.9f}", "nats")
    print("NWJ full:", f"{d['nwj_full_nats']:.9f}", "nats")
    print("DV full:", f"{d['dv_full_nats']:.9f}", "nats")
    print("NWJ linear:", f"{d['nwj_linear_nats']:.9f}", "nats")
    print("DV linear:", f"{d['dv_linear_nats']:.9f}", "nats")
    print("Classifier logit max error:", f"{d['classifier_logit_max_abs_error']:.3e}")
    print("GAN identity error:", f"{d['gan_identity_abs_error']:.3e}")
    print("Saved results to", path)


if __name__ == "__main__":
    main()

Bonus. Total variation as both an f-divergence and an IPM

Show both identities:

\operatorname{TV}(P,Q) = \frac12\int|p-q|,

and

\operatorname{TV}(P,Q) = \sup_{0\le h\le1} \left| \mathbb E_Ph- \mathbb E_Qh \right|.

Find an optimal $h$ using

A=\{x:p(x)\ge q(x)\}.

Show answer and solution

As an f-divergence, TV uses

f(t)=\frac12|t-1|:

D_f(P\|Q) =\frac12\int|p-q| =\operatorname{TV}(P,Q).

Let

A=\{x:p(x)\ge q(x)\}.

For every $0\le h\le1$ ,

\mathbb E_Ph-\mathbb E_Qh =\int h(p-q) \le\int_A(p-q).

Equality is attained by $h^*=\mathbf1_A$ , and

\int_A(p-q) =\frac12\int|p-q| =\operatorname{TV}(P,Q).

Taking the complement handles the opposite sign, so

\boxed{ \operatorname{TV}(P,Q) = \sup_{0\le h\le1} \left| \mathbb E_Ph-\mathbb E_Qh \right| }.

Exercise 1. Generators and normalizations

Forward KL

Reverse KL

Total variation

Squared Hellinger

Pearson χ2\chi^2χ2

Jensen–Shannon

Why a linear term changes nothing

Power divergence and Rényi divergence

Exercise 2. DPI through conditional Jensen

Exercise 3. Optimal discriminator and JS

Exercise 4. From Fenchel KL to Donsker–Varadhan

Exercise 5. InfoNCE: ceiling and interpretation

Exercise 6. Python: critic, InfoNCE, and support geometry

Bonus. Total variation as both an f-divergence and an IPM

Pearson $\chi^2$