Module 9. Rate–distortion and the Information Bottleneck: managing information loss — exercises and worked solutions · Information Theory for ML

Exercise 1. Bernoulli rate–distortion

Let $X\sim\operatorname{Bernoulli}(p)$ , $p\le1/2$ , with

d(x,\widehat x)=\mathbf1\{x\ne\widehat x\}.

Prove the lower bound
$I(X;\widehat X)\ge h_2(p)-h_2(D)$
for $0\le D\le p$ .
Construct a test channel that attains equality.
Explain why $R(D)=0$ for $D\ge p$ .

Show answer and solution

Let

P_e=P(X\ne\widehat X)\le D.

Because $X$ is binary, $\widehat X$ together with the error indicator $E=\mathbf1\{X\ne\widehat X\}$ determines $X$ . Hence

H(X\mid\widehat X) \le H(E\mid\widehat X) \le H(E) =h_2(P_e).

For $0\le D\le p\le1/2$ , $h_2$ is increasing, so

H(X\mid\widehat X) \le h_2(P_e) \le h_2(D).

Therefore

\boxed{ I(X;\widehat X) =H(X)-H(X\mid\widehat X) \ge h_2(p)-h_2(D) }.

To attain equality, it is convenient to specify the backward test channel. Let

\widehat X\sim\operatorname{Bernoulli}(q), \qquad Z\sim\operatorname{Bernoulli}(D)

be independent, and define

X=\widehat X\oplus Z.

Then $P(X\ne\widehat X)=D$ , while

P(X=1)=D+q(1-2D).

Choose

q=\frac{p-D}{1-2D},

which is valid for $0\le D\le p\le1/2$ (with the boundary case $p=D=1/2$ handled separately). The marginal of $X$ is then $\operatorname{Bernoulli}(p)$ , and

H(X\mid\widehat X)=H(Z)=h_2(D).

Thus

I(X;\widehat X)=h_2(p)-h_2(D).

The forward channel $P(\widehat x\mid x)$ follows from Bayes' rule applied to this joint distribution.

For $D\ge p$ , the decoder can always emit $\widehat X=0$ . The rate is then zero and the average distortion is $P(X=1)=p\le D$ . Therefore

R(D)=0, \qquad D\ge p.

Exercise 2. Predictive IB distortion

For the Markov chain $Y\to X\to T$ , prove

I(X;Y)-I(T;Y) = \mathbb E_{X,T} D_{\mathrm{KL}} \bigl(P_{Y\mid X}\|P_{Y\mid T}\bigr).

Show that minimizing $I(X;T)-\beta I(T;Y)$ is equivalent to a rate–distortion Lagrangian with this distortion.

Show answer and solution

The chain rule gives

I(X,T;Y)=I(T;Y)+I(X;Y\mid T)

and also

I(X,T;Y)=I(X;Y)+I(T;Y\mid X).

For the Markov chain $Y\to X\to T$ ,

I(T;Y\mid X)=0.

Therefore

I(X;Y)-I(T;Y)=I(X;Y\mid T).

Conditional mutual information is an average KL divergence:

I(X;Y\mid T) = \mathbb E_{X,T} D_{\mathrm{KL}} \bigl(P_{Y\mid X,T}\|P_{Y\mid T}\bigr).

Markovity gives $P_{Y\mid X,T}=P_{Y\mid X}$ , hence

\boxed{ I(X;Y)-I(T;Y) = \mathbb E_{X,T} D_{\mathrm{KL}} \bigl(P_{Y\mid X}\|P_{Y\mid T}\bigr) }.

Calling the KL term $d_{\mathrm{IB}}(X,T)$ ,

I(X;T)-\beta I(T;Y) = I(X;T)+\beta\mathbb E[d_{\mathrm{IB}}(X,T)]-\beta I(X;Y).

The last term does not depend on $P_{T\mid X}$ . The minimizers therefore coincide with those of the rate–distortion Lagrangian

I(X;T)+\beta\mathbb E[d_{\mathrm{IB}}(X,T)].

Exercise 3. Variational VIB bounds

Prove
$I(T;Y) \ge H(Y)+\mathbb E[\log r_\psi(Y\mid T)].$
Prove
$\mathbb E_XD_{\mathrm{KL}}(q(T\mid X)\|r(T)) = I(X;T)+D_{\mathrm{KL}}(q(T)\|r(T)).$
Explain when both bounds are tight simultaneously.

Show answer and solution

1. Lower bound on relevance

For any decoder $r_\psi(y\mid t)$ ,

\begin{aligned} \mathbb E[-\log r_\psi(Y\mid T)] &=H(Y\mid T) + \mathbb E_TD_{\mathrm{KL}} \bigl(P_{Y\mid T}\|r_\psi(\cdot\mid T)\bigr)\\ &\ge H(Y\mid T). \end{aligned}

Therefore

\begin{aligned} I(T;Y) &=H(Y)-H(Y\mid T)\\ &\ge H(Y)+\mathbb E[\log r_\psi(Y\mid T)]. \end{aligned}

The bound is tight exactly when

r_\psi(y\mid t)=P(y\mid t)

almost surely under $P_T$ .

2. Upper bound on rate

Let

q(t)=\int q(t\mid x)p(x)\,dx.

Then

\begin{aligned} \mathbb E_XD_{\mathrm{KL}}(q(T\mid X)\|r(T)) &=\mathbb E_{X,T}\log\frac{q(t\mid x)}{r(t)}\\ &=\mathbb E_{X,T}\log\frac{q(t\mid x)}{q(t)} + \mathbb E_T\log\frac{q(t)}{r(t)}\\ &=I_q(X;T)+D_{\mathrm{KL}}(q(T)\|r(T)). \end{aligned}

Thus KL to the chosen prior upper-bounds $I(X;T)$ .

3. Simultaneous tightness

Both bounds are tight when

r_\psi(y\mid t)=P(y\mid t)

and

r(t)=q(t).

The first requires an exact variational decoder; the second requires the prior to equal the aggregate posterior. A restricted parametric model may not satisfy either condition.

Exercise 4. VAE: the likelihood defines distortion

Consider

p_\theta(x\mid z) = \mathcal N(\mu_\theta(z),\sigma^2I).

Derive the negative log-likelihood and its connection to MSE.
Explain why changing $\sigma^2$ changes the relative weighting of distortion and rate even when $\beta=1$ .
Derive
$\mathcal R = I(X;Z)+D_{\mathrm{KL}}(q(Z)\|p(Z)).$

Show answer and solution

For a $d$ -dimensional observation,

p_\theta(x\mid z) =(2\pi\sigma^2)^{-d/2} \exp\left[-\frac{\|x-\mu_\theta(z)\|^2}{2\sigma^2}\right].

Hence

\boxed{ -\log p_\theta(x\mid z) = \frac d2\log(2\pi\sigma^2) + \frac{\|x-\mu_\theta(z)\|^2}{2\sigma^2} }.

At fixed $\sigma^2$ , minimizing NLL is equivalent to minimizing MSE. The coefficient of MSE is $1/(2\sigma^2)$ , however, so changing $\sigma^2$ changes the relative strength of reconstruction and KL-rate even at formal $\beta=1$ .

For the rate decomposition, add and subtract $\log q(z)$ :

\begin{aligned} \mathcal R &=\mathbb E_{X,Z}\log\frac{q(z\mid x)}{p(z)}\\ &=\mathbb E_{X,Z}\log\frac{q(z\mid x)}{q(z)} + \mathbb E_Z\log\frac{q(z)}{p(z)}\\ &=I_q(X;Z)+D_{\mathrm{KL}}(q(Z)\|p(Z)). \end{aligned}

Exercise 5. Quantization: which rate–distortion problem is being posed?

For weight-only quantization, list:

the source alphabet;
everything included in the actual rate;
three possible distortion measures;
why the Gaussian MSE bound cannot be converted directly into a bound on downstream accuracy.

Show answer and solution

A careful answer may use the following choices.

Source alphabet

One must decide whether the source consists of individual weights, weight blocks, matrix rows or columns, or channel groups together with activation statistics. Classical theory also requires a probabilistic ensemble, not merely one fixed matrix without a source model.

Actual rate

It may include quantized values or indices, scales and zero-points, codebooks, group boundaries, sparsity metadata, padding and alignment, and container headers. Nominal bit-width is therefore not always the actual number of bits per parameter.

Possible distortions

Examples are:

$\|W-\widehat W\|_2^2$ ;
$\mathbb E\|WX-\widehat W X\|_2^2$ on calibration activations;
an increase in NLL or KL between the original and quantized predictive distributions;
a change in a downstream metric.

Why the Gaussian bound does not bound accuracy

$D\ge\sigma^2 2^{-2R}$ assumes an i.i.d. Gaussian source, squared error, and an optimal long-block code. Downstream accuracy is a nonlinear function of the entire network, data distribution, and decision boundaries. Equal weight MSE can have very different functional effects in different parameter directions.

Exercise 6. Python: discrete IB and the VAE-rate decomposition

Run M9_ib_rate_distortion_experiments.py. The script:

plots theoretical rate–distortion curves;
solves a toy Information Bottleneck problem with self-consistent iterations;
verifies the exact decomposition of VAE rate into mutual information and prior mismatch.

Answer:

why the IB representation is trivial at small $\beta$ ;
why $I(T;Y)$ saturates before $I(X;T)$ ;
how a poor prior increases KL-rate without adding information about the input.

Show answer and solution

The script performs three independent checks.

Rate–distortion curves

The Bernoulli curve decreases from $h_2(p)$ at $D=0$ to zero at $D=p$ . The Gaussian curve diverges as $D\downarrow0$ and becomes zero for $D\ge\sigma^2$ .

Toy Information Bottleneck

At small $\beta$ , the penalty on $I(X;T)$ dominates. The optimal representation is nearly unused:

I(X;T)\approx0, \qquad I(T;Y)\approx0.

As $\beta$ grows, the representation separates inputs with different $P(Y\mid x)$ . Relevance saturates near $I(X;Y)$ because $X$ contains no more information about $Y$ . Meanwhile $I(X;T)$ can continue to grow if the code preserves input details that add no predictive relevance.

VAE-rate decomposition

Numerically one should find

\mathbb E_XD_{\mathrm{KL}}(q(Z\mid X)\|p(Z)) = I(X;Z)+D_{\mathrm{KL}}(q(Z)\|p(Z))

to floating-point accuracy.

Replacing a good prior with a poor one leaves $I(X;Z)$ unchanged because the encoder and joint distribution are unchanged. It increases $D_{\mathrm{KL}}(q(Z)\|p(Z))$ and hence the KL-rate. This is an extra mismatch cost, not extra information about the data.

Code solution · the same source file for RU/EN

PythonM9_ib_rate_distortion_experiments.pyDownload .py

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276

from __future__ import annotations

"""Reproducible numerical experiments for Module 9.

The script:
1. plots Bernoulli and Gaussian rate-distortion curves;
2. solves a finite Information Bottleneck problem by self-consistent updates;
3. verifies the exact VAE-rate decomposition on a discrete toy encoder.

All logarithms used for information quantities are base 2 unless stated otherwise.
"""

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)
EPS = 1e-15


def normalize(a: np.ndarray, axis: int = -1) -> np.ndarray:
    total = a.sum(axis=axis, keepdims=True)
    if np.any(total <= 0):
        raise ValueError("Cannot normalize non-positive mass")
    return a / total


def xlog2x(x: np.ndarray) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    return np.where(x > 0, x * np.log2(x), 0.0)


def binary_entropy(p: np.ndarray | float) -> np.ndarray:
    p = np.asarray(p, dtype=float)
    p = np.clip(p, EPS, 1 - EPS)
    return -(p * np.log2(p) + (1 - p) * np.log2(1 - p))


def mutual_information(joint: np.ndarray) -> float:
    joint = np.asarray(joint, dtype=float)
    joint = joint / joint.sum()
    px = joint.sum(axis=1, keepdims=True)
    py = joint.sum(axis=0, keepdims=True)
    ratio = np.divide(joint, px @ py, out=np.ones_like(joint), where=joint > 0)
    return float(np.sum(np.where(joint > 0, joint * np.log2(ratio), 0.0)))


def kl_rows(p: np.ndarray, q: np.ndarray) -> np.ndarray:
    """Return all pairwise KL divergences D(p_i || q_j) in bits."""
    p = np.asarray(p, dtype=float)
    q = np.asarray(q, dtype=float)
    p_safe = np.clip(p, EPS, 1)
    q_safe = np.clip(q, EPS, 1)
    return np.sum(p[:, None, :] * (np.log2(p_safe[:, None, :]) - np.log2(q_safe[None, :, :])), axis=-1)


def plot_rate_distortion() -> None:
    p = 0.25
    d_b = np.linspace(0, p, 250)
    r_b = binary_entropy(p) - binary_entropy(d_b)
    r_b[-1] = 0.0

    sigma2 = 1.0
    d_g = np.geomspace(0.01, sigma2, 250)
    r_g = 0.5 * np.log2(sigma2 / d_g)

    fig, ax = plt.subplots(figsize=(8.6, 5.0))
    ax.plot(d_b, r_b, label=r"Bernoulli($p=0.25$), Hamming")
    ax.plot(d_g, r_g, label=r"Gaussian($\sigma^2=1$), MSE")
    ax.set_xlabel("Allowed distortion D")
    ax.set_ylabel("Rate R(D), bits/source symbol")
    ax.set_title("Two rate-distortion benchmarks")
    ax.legend()
    ax.grid(alpha=0.25)
    fig.tight_layout()
    fig.savefig(OUT / "M9_experiment_rate_distortion.png", dpi=180)
    plt.close(fig)


def ib_metrics(px: np.ndarray, py_given_x: np.ndarray, qt_given_x: np.ndarray) -> tuple[float, float]:
    pxt = px[:, None] * qt_given_x
    ixt = mutual_information(pxt)
    pty = np.einsum("x,xt,xy->ty", px, qt_given_x, py_given_x)
    ity = mutual_information(pty)
    return ixt, ity


def solve_ib(
    px: np.ndarray,
    py_given_x: np.ndarray,
    n_t: int,
    beta: float,
    rng: np.random.Generator,
    restarts: int = 24,
    max_iter: int = 5000,
    tol: float = 1e-12,
) -> tuple[np.ndarray, float, float, float]:
    """Solve a finite IB problem by self-consistent updates.

    Returns q(t|x), I(X;T), I(T;Y), and the IB objective I(X;T)-beta I(T;Y).
    Multiple random restarts are used because the fixed-point equations may have
    several stationary solutions.
    """
    best: tuple[np.ndarray, float, float, float] | None = None
    n_x = px.size

    for restart in range(restarts):
        if restart == 0:
            qtx = np.full((n_x, n_t), 1.0 / n_t)
            qtx += 1e-3 * rng.random((n_x, n_t))
        else:
            qtx = rng.dirichlet(np.ones(n_t), size=n_x)
        qtx = normalize(qtx, axis=1)

        for _ in range(max_iter):
            pt = px @ qtx
            pt = np.clip(pt, EPS, None)
            pyt = np.einsum("x,xt,xy->ty", px, qtx, py_given_x)
            pyt = normalize(pyt, axis=1)
            distortion = kl_rows(py_given_x, pyt)

            # Exponents use natural exp. Convert KL bits to nats by ln(2).
            log_new = np.log(pt)[None, :] - beta * math.log(2.0) * distortion
            log_new -= log_new.max(axis=1, keepdims=True)
            new = normalize(np.exp(log_new), axis=1)
            if np.max(np.abs(new - qtx)) < tol:
                qtx = new
                break
            qtx = new

        ixt, ity = ib_metrics(px, py_given_x, qtx)
        objective = ixt - beta * ity
        candidate = (qtx.copy(), ixt, ity, objective)
        if best is None or objective < best[3] - 1e-11:
            best = candidate

    assert best is not None
    return best


def run_ib_experiment() -> list[tuple[float, float, float]]:
    px = np.full(6, 1 / 6)
    # Three progressively more positive regions. Neighboring states have similar
    # predictive distributions, making the bottleneck geometry visible.
    p1 = np.array([0.03, 0.08, 0.30, 0.70, 0.92, 0.97])
    py_given_x = np.column_stack([1 - p1, p1])
    rng = np.random.default_rng(20260713)
    betas = np.array([0.05, 0.10, 0.20, 0.40, 0.75, 1.25, 2.0, 3.5, 6.0, 10.0, 18.0])
    rows: list[tuple[float, float, float]] = []

    for beta in betas:
        _, ixt, ity, _ = solve_ib(px, py_given_x, n_t=6, beta=float(beta), rng=rng)
        rows.append((float(beta), ixt, ity))

    ixy = mutual_information(px[:, None] * py_given_x)

    fig, ax = plt.subplots(figsize=(7.2, 5.4))
    ixt = np.array([r[1] for r in rows])
    ity = np.array([r[2] for r in rows])
    ax.plot(ixt, ity, marker="o", label="IB fixed points")
    for beta, x, y in rows[::2]:
        ax.annotate(f"β={beta:g}", (x, y), xytext=(5, 5), textcoords="offset points", fontsize=8)
    ax.axhline(ixy, linestyle="--", linewidth=1.2, label=f"I(X;Y)={ixy:.3f} bits")
    ax.set_xlabel("I(X;T), bits")
    ax.set_ylabel("I(T;Y), bits")
    ax.set_title("Toy Information Bottleneck frontier")
    ax.legend()
    ax.grid(alpha=0.25)
    fig.tight_layout()
    fig.savefig(OUT / "M9_experiment_ib_frontier.png", dpi=180)
    plt.close(fig)
    return rows


def expected_conditional_kl(px: np.ndarray, qz_given_x: np.ndarray, prior: np.ndarray) -> float:
    """Average KL in bits for normalized probability arrays, without smoothing."""
    joint = px[:, None] * qz_given_x
    positive = joint > 0
    prior_grid = np.broadcast_to(prior, qz_given_x.shape)
    # A prior must cover every latent with positive joint mass. Rows with
    # px[x] = 0 and terms with q(z|x) = 0 contribute nothing, including 0 log 0.
    if np.any(prior_grid[positive] == 0):
        return math.inf
    return float(np.sum(joint[positive] * (
        np.log2(qz_given_x[positive]) - np.log2(prior_grid[positive])
    )))


def run_rate_decomposition(bad_prior: np.ndarray | None = None) -> dict[str, float]:
    px = np.array([0.40, 0.35, 0.25])
    qz_given_x = np.array([
        [0.80, 0.15, 0.05],
        [0.20, 0.65, 0.15],
        [0.05, 0.20, 0.75],
    ])
    qz = px @ qz_given_x
    good_prior = qz.copy()
    if bad_prior is None:
        bad_prior = np.array([0.78, 0.12, 0.10])

    joint = px[:, None] * qz_given_x
    ixz = mutual_information(joint)
    # Ordinary KL is the same calculation with a single source state.
    mismatch_good = expected_conditional_kl(np.ones(1), qz[None, :], good_prior)
    mismatch_bad = expected_conditional_kl(np.ones(1), qz[None, :], bad_prior)
    rate_good = expected_conditional_kl(px, qz_given_x, good_prior)
    rate_bad = expected_conditional_kl(px, qz_given_x, bad_prior)

    if not np.isclose(rate_good, ixz + mismatch_good, atol=1e-12):
        raise AssertionError("Good-prior decomposition failed")
    if not np.isclose(rate_bad, ixz + mismatch_bad, atol=1e-12):
        raise AssertionError("Bad-prior decomposition failed")

    labels = ["I(X;Z)", "prior mismatch", "total KL-rate"]
    good = [ixz, mismatch_good, rate_good]
    bad = [ixz, mismatch_bad, rate_bad]
    x = np.arange(len(labels))
    width = 0.35
    fig, ax = plt.subplots(figsize=(7.7, 5.0))
    ax.bar(x - width / 2, good, width, label="prior = aggregate posterior")
    finite = np.isfinite(bad)
    ax.bar(x[finite] + width / 2, np.asarray(bad)[finite], width,
           label="mismatched prior")
    for position in x[~finite] + width / 2:
        ax.text(position, 0.78, "∞", transform=ax.get_xaxis_transform(),
                ha="center", color="C1", fontsize=22)
    ax.set_xticks(x, labels)
    ax.set_ylabel("Bits")
    ax.set_title("VAE-rate decomposition: information plus prior mismatch")
    ax.legend()
    ax.grid(axis="y", alpha=0.25)
    if not np.all(finite):
        # Text annotations do not participate in Matplotlib's autoscaling.
        ax.set_xlim(-0.6, len(labels) - 0.4)
        fig.text(0.5, 0.02, "∞: prior assigns zero mass to a possible latent",
                 ha="center", fontsize=10)
    fig.tight_layout(rect=(0, 0.06 if not np.all(finite) else 0, 1, 1))
    fig.savefig(OUT / "M9_experiment_rate_decomposition.png", dpi=180)
    plt.close(fig)

    return {
        "I(X;Z)": ixz,
        "good_mismatch": mismatch_good,
        "good_rate": rate_good,
        "bad_mismatch": mismatch_bad,
        "bad_rate": rate_bad,
    }


def main() -> None:
    plot_rate_distortion()
    rows = run_ib_experiment()
    decomposition = run_rate_decomposition()

    print("Toy IB results (bits):")
    print("beta\tI(X;T)\tI(T;Y)")
    for beta, ixt, ity in rows:
        print(f"{beta:>5g}\t{ixt:.6f}\t{ity:.6f}")

    print("\nVAE-rate decomposition (bits):")
    for key, value in decomposition.items():
        print(f"{key}: {value:.9f}")
    print("\nGenerated figures:")
    for path in sorted(OUT.glob("M9_experiment_*.png")):
        print(path.name)


if __name__ == "__main__":
    main()

Bonus. I–MMSE for a Gaussian input

For $X\sim\mathcal N(0,1)$ , derive the posterior mean, $\operatorname{mmse}(\gamma)$ , and $I(X;Y_\gamma)$ . Verify the identity and explain why the integral to infinity diverges even though differential entropy is finite.

Show answer and solution

Let

Y_\gamma=\sqrt\gamma X+N, \qquad X,N\sim\mathcal N(0,1)

independently. Gaussian conditioning gives

\mathbb E[X\mid Y_\gamma] = \frac{\sqrt\gamma}{1+\gamma}Y_\gamma

and

\operatorname{Var}(X\mid Y_\gamma)=\frac1{1+\gamma}.

Therefore

\boxed{ \operatorname{mmse}(\gamma)=\frac1{1+\gamma} }.

The Gaussian mutual information is

\boxed{ I(X;Y_\gamma)=\frac12\ln(1+\gamma) }.

Indeed,

\frac d{d\gamma}I(X;Y_\gamma) = \frac1{2(1+\gamma)} = \frac12\operatorname{mmse}(\gamma).

But

\frac12\int_0^\infty\frac{d\gamma}{1+\gamma}=\infty.

This does not contradict the finiteness of

h(X)=\frac12\ln(2\pi e).

I–MMSE integrates the mutual information between a continuous variable and an observation of ever increasing precision, which diverges. A relation to differential entropy requires subtraction of a divergent reference term; the unrenormalized equality is false.