Lossy compression: rate–distortion, the information bottleneck and the β-VAE

Before this lesson: lesson 05 (lossless coding, rates in bits per symbol), lesson 07 (mutual information, the data processing inequality), lesson 09 (the ELBO and the VAE). "Conditioning reduces entropy" is from lesson 02; the maximum entropy of the Gaussian is from lesson 06. Reverse water-filling uses PCA (linear algebra lesson 09). Time: about 75 minutes, plus the notebook. This lesson is dense; take a break after section 3.4.

---

Why this matters #

• Learned compression. Neural image and video codecs (Ballé et al.) are trained with the loss rate $+\ \lambda \cdot$ distortion. JPEG, MP3 and video codecs are hand-made points on the same plane.
• VAEs and latent diffusion. The two terms of the VAE loss are a distortion and a rate. The β-VAE moves along the curve, and posterior collapse is the point "rate = 0". The autoencoder of a latent diffusion model chooses a point with a high rate, so that details survive.
• Quantization of models. "How many bits per weight for a given error?" is a rate–distortion question. The rule "6 dB per bit" comes from the Gaussian $R(D)$.
• Representation learning. The information bottleneck and its variational form are used as regularizers, and as one (debated) theory of what deep networks learn.
• ML for science. Climate, weather and simulation archives are compressed with error bounds, and compressors are compared by bits per value against error.

As in the whole track, $\log$ is the natural log (nats) and $\log_2$ gives bits. Rates are in bits unless a formula says nats. $H_b(x) = -x\log_2 x - (1-x)\log_2(1-x)$ is the binary entropy in bits.

1. Predict first #

2. The picture #

Encoder, bits, decoder #

An encoder sees a block of $n$ source symbols $x^n = (x_1, \dots, x_n)$ and sends one of $2^{nR}$ messages, that is, $nR$ bits. A decoder turns the message into a reconstruction $\hat x^n$. A distortion measure $d(x, \hat x) \ge 0$ scores each symbol:

• Hamming distortion $d(x, \hat x) = \mathbf{1}[x \ne \hat x]$ for symbols: the average is the error rate;
• squared error $d(x, \hat x) = (x - \hat x)^2$ for real numbers: the average is the mean squared error (MSE).

The rate $R$ is in bits per symbol. The question: what is the smallest rate for an average distortion of at most $D$?

A fair coin: store or guess #

Take fair coin flips and Hamming distortion. A simple scheme: store a fraction $\alpha$ of the flips exactly (1 bit each), and output a fixed guess for the rest. A fixed guess is wrong half of the time. So

This is a straight line from $(D, R) = (0, 1)$ to $(\frac12, 0)$. To allow 11% errors, it needs $R = 1 - 0.22 = 0.78$ bits per flip.

Shannon's answer is the curve $R(D) = 1 - H_b(D)$. At $D = 0.11$, $H_b(0.11) = 0.50$, so only 0.50 bits per flip are needed. The curve lies below the line everywhere between the end points.

Interactive figure: loading…

Why block coding wins (a counting picture). Code $n$ flips at once, for a large $n$. Around each reconstruction sequence, the sequences within Hamming distance $nD$ form a "ball". This ball contains about $2^{nH_b(D)}$ sequences (for $D \le \frac12$). There are $2^n$ possible source sequences, so you need at least

reconstruction sequences to come within distance $nD$ of every one. That is $n(1 - H_b(D))$ bits, a rate of $1 - H_b(D)$. Shannon showed that a random set of slightly more sequences is enough. Store-or-guess treats each flip alone and wastes this geometry. For $n = 100$ and $D = 0.11$: about $2^{50}$ reconstruction sequences for $2^{100}$ source sequences.

The test channel #

Think of encoder and decoder together as one noisy channel from $X$ to $\hat X$, with a conditional distribution $q(\hat x \mid x)$: the test channel. The message has $nR$ bits, so $\hat X^n$ can carry at most $nR$ bits of information about $X^n$. Rate–distortion asks for the channel that carries the least information about $X$ while it still meets the distortion budget.

Channel capacity (the other half of Shannon's theory) is the opposite problem: there the channel is fixed, and you maximize $I(X; Y)$ over the input distribution.

3. Definitions and main results #

3.1 The rate–distortion function #

Notation. The source $X \sim p(x)$ is i.i.d. The reconstruction $\hat X$ takes values in a set $\hat{\mathcal X}$. $d(x, \hat x) \ge 0$ is the distortion and $q(\hat x \mid x)$ is a test channel.

Definition (rate–distortion function).

where the expectation and $I$ are under the joint distribution $p(x)\,q(\hat x \mid x)$. Let $D_{\max} = \min_{\hat x}\mathbb{E}[d(X, \hat x)]$: the distortion of the best constant output. For $D \ge D_{\max}$, a constant output meets the budget and carries no information, so $R(D) = 0$.

Theorem (Shannon's rate–distortion theorem; stated without proof). For an i.i.d. source and a bounded distortion measure:

• for every rate $R > R(D)$ and every $\varepsilon > 0$, there are block codes of rate $R$ with expected distortion at most $D + \varepsilon$, for $n$ large enough;
• no code of rate $R < R(D)$ has expected distortion at most $D$.

(Cover and Thomas, Elements of Information Theory, ch. 10.) So the mathematical minimum in the definition is also the operational limit.

Properties.

• $R(D)$ is non-increasing: a larger $D$ allows more channels, so the minimum can only go down.
• For a discrete source with Hamming distortion, $R(0) = H(X)$: $D = 0$ forces $\hat X = X$.
• $R(D)$ is convex.

Proof of convexity.

1. Let $q_1$ and $q_2$ be optimal channels for $D_1$ and $D_2$, and $\lambda \in [0, 1]$. Mix them: $q_\lambda = \lambda q_1 + (1 - \lambda) q_2$.
2. The expected distortion is linear in the channel, so $q_\lambda$ has distortion at most $\lambda D_1 + (1 - \lambda) D_2$.
3. For a fixed $p(x)$, $I(X; \hat X)$ is a convex function of $q(\hat x \mid x)$ (Cover and Thomas, Theorem 2.7.4). So $R(\lambda D_1 + (1 - \lambda) D_2) \le I(q_\lambda) \le \lambda R(D_1) + (1 - \lambda) R(D_2)$. $\blacksquare$

Convexity explains the coin picture: store-or-guess mixes the two end points $(0, H(X))$ and $(D_{\max}, 0)$, so it is the straight line between them, and a convex curve lies below that line.

3.2 A binary source with Hamming distortion #

Theorem. Let $X \sim$ Bernoulli($p$) with $p \le \frac12$, and $d(x, \hat x) = \mathbf{1}[x \ne \hat x]$. Then

(For $p > \frac12$ swap the roles of 0 and 1: the formula holds with $\min(p, 1-p)$ in place of $p$.)

Proof. First a lower bound, for any channel with $\mathbb{E}[d] \le D \le p$.

1. Let $E = X \oplus \hat X$ (XOR). $E = 1$ exactly when the reconstruction is wrong, so $\Pr(E = 1) = \mathbb{E}[d] \le D$.
2. $I(X; \hat X) = H(X) - H(X \mid \hat X) = H_b(p) - H(E \mid \hat X)$. Once $\hat X$ is known, knowing $X$ is the same as knowing $E$.
3. $H(E \mid \hat X) \le H(E) = H_b(\Pr(E = 1)) \le H_b(D)$. The first step is "conditioning reduces entropy" (lesson 02). The second holds because $H_b$ is increasing on $[0, \frac12]$.
4. So $I(X; \hat X) \ge H_b(p) - H_b(D)$.

Then a channel that reaches the bound.

5. Build the joint distribution backwards: $\hat X \sim$ Bernoulli($s$) and $X = \hat X \oplus Z$, with $Z \sim$ Bernoulli($D$) independent of $\hat X$. Choose $s = \frac{p - D}{1 - 2D}$, so that $\Pr(X = 1) = s(1 - D) + (1 - s)D = p$.
6. Then $\mathbb{E}[d] = \Pr(Z = 1) = D$ and $I(X; \hat X) = H(X) - H(X \mid \hat X) = H_b(p) - H_b(D)$, because given $\hat X$, the only uncertainty in $X$ is $Z$. The forward channel $q(\hat x \mid x)$ follows from Bayes' rule (section 4). $\blacksquare$

Example. $p = 0.3$, $D = 0.1$: $R = H_b(0.3) - H_b(0.1) = 0.881 - 0.469 = 0.412$ bits per symbol. Lossless coding needs 0.881 bits; allowing 10% errors saves more than half.

3.3 A Gaussian source with squared error #

Theorem. Let $X \sim \mathcal N(0, \sigma^2)$ and $d(x, \hat x) = (x - \hat x)^2$. Then

Equivalently, the smallest MSE at rate $R$ is $D(R) = \sigma^2\, 2^{-2R}$.

Why this is true (intuition). The reconstruction must "explain" all of the variance except $D$. The best reconstruction is itself Gaussian with variance $\sigma^2 - D$, and the leftover error is Gaussian noise with variance $D$. The information is the log of the ratio of the two "volumes", $\frac12\log_2(\sigma^2/D)$.

Proof. Lower bound, for any channel with $\mathbb{E}[(X - \hat X)^2] \le D$. Here $h$ is the differential entropy in bits (lesson 06).

1. $I(X; \hat X) = h(X) - h(X \mid \hat X) = h(X) - h(X - \hat X \mid \hat X) \ge h(X) - h(X - \hat X)$. Given $\hat X$, subtracting it is a shift, which does not change $h$; then conditioning reduces entropy.
2. Among all distributions with a given variance, the Gaussian has the largest differential entropy, and the variance is at most the second moment. So $h(X - \hat X) \le \frac12\log_2\big(2\pi e\,\mathbb{E}[(X - \hat X)^2]\big) \le \frac12\log_2(2\pi e D)$.
3. With $h(X) = \frac12\log_2(2\pi e\sigma^2)$: $I(X; \hat X) \ge \frac12\log_2(\sigma^2/D)$.

Achievability:

4. Backwards again: $\hat X \sim \mathcal N(0, \sigma^2 - D)$ and $X = \hat X + Z$ with $Z \sim \mathcal N(0, D)$ independent of $\hat X$. Then $X \sim \mathcal N(0, \sigma^2)$, $\mathbb{E}[(X - \hat X)^2] = D$, and $I(X; \hat X) = h(X) - h(Z) = \frac12\log_2(\sigma^2/D)$. $\blacksquare$

Each extra bit divides $D$ by 4, that is, $10\log_{10} 4 = 6.02$ dB per bit (it halves the standard deviation of the error). The Gaussian is also the hardest source for squared error: every source with variance $\sigma^2$ has $R(D) \le \frac12\log_2(\sigma^2/D)$.

Interactive figure: loading…

3.4 Many Gaussian components: reverse water-filling #

Now let $X = (X_1, \dots, X_k)$ have independent Gaussian components with variances $\lambda_1, \dots, \lambda_k$ (for example, the principal components of data), and let the distortion be the total squared error. How should you share the bits?

Theorem (reverse water-filling). Choose a "water level" $\theta > 0$. Component $i$ gets the distortion $D_i = \min(\theta, \lambda_i)$ and the rate $R_i = \frac12\log_2(\lambda_i / D_i)$. Then

where $\theta$ is chosen so that the $D_i$ add up to $D$. Components with $\lambda_i \le \theta$ get no bits at all.

Interactive figure: loading…

Why water-filling is optimal (sketch)

1. For independent components, $I(X; \hat X) \ge \sum_i I(X_i; \hat X_i)$ (Cover and Thomas, §10.3.3), and each term is at least $R_i(D_i) = \max\big(0, \frac12\log_2(\lambda_i / D_i)\big)$. So $R(D) \ge \min \sum_i R_i(D_i)$ over $D_1 + \dots + D_k = D$. Coding each component with its own optimal channel reaches this sum, so equality holds.
2. Minimize $\sum_i \frac12\log_2(\lambda_i / D_i)$ subject to $\sum_i D_i = D$ and $0 < D_i \le \lambda_i$. With a Lagrange multiplier, the derivative $-\frac{1}{2 D_i \ln 2}$ must be the same for every component that gets bits: all such $D_i$ are equal, to $\theta$.
3. If $\lambda_i < \theta$, the best choice is $D_i = \lambda_i$ (send nothing): reducing its error below $\lambda_i$ would cost more than spending the bits elsewhere. $\blacksquare$

This is the theory behind transform coding: rotate the data to its principal components (PCA, the DCT in JPEG), then spend bits only on components with large variance. Dropping the small components is exactly what truncated PCA does.

3.5 The Blahut–Arimoto algorithm #

For a general discrete source there is no closed form. The Blahut–Arimoto (BA) algorithm computes points of the curve.

1. Lagrangian. For $\beta > 0$, minimize $I(X; \hat X) + \beta\,\mathbb{E}[d(X, \hat X)]$ over $q(\hat x \mid x)$. The minimizer is the point of the curve where the slope is $dR/dD = -\beta$ (with $R$ in nats). Large $\beta$: low distortion and high rate. Small $\beta$: low rate.
2. A variational form of $I$. For any distribution $r(\hat x)$ on $\hat{\mathcal X}$,

where $\bar q(\hat x) = \sum_x p(x)\,q(\hat x \mid x)$ is the output marginal. Equality holds if and only if $r = \bar q$. So BA minimizes, over both $q$ and $r$,

3. Alternate two exact minimizations (see the box).
4. Convergence. No step increases $F$, and $F$ is jointly convex in $(q, r)$, so the iteration converges to the minimum (Csiszár and Tusnády, 1984). At the end, report $R = I(X; \hat X)$ and $D = \mathbb{E}[d]$. A sweep over $\beta$ traces the whole curve.

Interactive figure: loading…

For the Bernoulli source, the slope of $H_b(p) - H_b(D)$ in nats is $-\log\frac{1 - D}{D}$. Setting it to $-\beta$ gives $D = 1/(1 + e^{\beta})$, which does not depend on $p$ (as long as $D \le \min(p, 1 - p)$).

3.6 The information bottleneck #

Now the goal is not to reconstruct $X$ but to keep what matters for a target $Y$. Input $X$, target $Y$, and a representation $T$ computed from $X$ alone, with a distribution $p(t \mid x)$. So $Y \to X \to T$ is a Markov chain. The information bottleneck (IB; Tishby, Pereira and Bialek, 1999) is, for $\beta > 0$:

Compress $X$ (small $I(X; T)$) while keeping the information about $Y$ (large $I(T; Y)$). By the data processing inequality (lesson 07), $I(T; Y) \le I(X; Y)$: $T$ cannot know more about $Y$ than $X$ does.

IB is a rate–distortion problem. Take the rate $I(X; T)$ and the distortion $d(x, t) = \mathrm{KL}(p(y \mid x)\,\|\,p(y \mid t))$: how much the prediction of $Y$ gets worse when you use $t$ instead of $x$. The Markov chain gives $p(y \mid x, t) = p(y \mid x)$, so

$I(X; Y)$ does not depend on $p(t \mid x)$. So minimizing $I(X; T) + \beta\,\mathbb{E}[d]$ is the IB objective. Its solution has the BA form $p(t \mid x) \propto p(t)\exp\big(-\beta\,\mathrm{KL}(p(y \mid x)\,\|\,p(y \mid t))\big)$; the iteration also updates $p(y \mid t)$, because here the distortion itself depends on the channel.

Variational IB (Alemi et al., 2017) makes this trainable with neural networks. For any $r(t)$, $I(X; T) \le \mathbb{E}_x\,\mathrm{KL}(p(t \mid x)\,\|\,r(t))$. For any decoder $q(y \mid t)$, $I(T; Y) \ge H(Y) + \mathbb{E}[\log q(y \mid t)]$. The resulting loss is

a β-VAE whose decoder predicts $y$ instead of $x$.

3.7 The β-VAE as a rate–distortion trade-off #

A VAE has an encoder $q_\phi(z \mid x)$, a decoder $p_\theta(x \mid z)$ and a prior $p(z)$ (lesson 09). For one data point $x$:

• Distortion is the expected reconstruction loss in nats. For a Gaussian decoder with fixed variance $\sigma_x^2$, $-\log p_\theta(x \mid z) = \|x - \mu_\theta(z)\|^2 / (2\sigma_x^2) + \text{const}$: a scaled squared error.
• Rate is the KL term. Averaged over the data it bounds the information that $z$ carries about $x$.

Theorem (the rate is an upper bound on $I(X; Z)$). Average over the data $p_{\text{data}}(x)$ and let $q(z) = \mathbb{E}_{x}\,q_\phi(z \mid x)$ be the aggregate posterior. Then

where $I_q$ is the mutual information under $p_{\text{data}}(x)\,q_\phi(z \mid x)$.

Proof.

1. Split the log: $\log\frac{q_\phi(z\mid x)}{p(z)} = \log\frac{q_\phi(z\mid x)}{q(z)} + \log\frac{q(z)}{p(z)}$.
2. Take the expectation over $p_{\text{data}}(x)\,q_\phi(z \mid x)$. The first term gives $I_q(X; Z)$ by the definition of mutual information.
3. In the second term only $z$ appears, and its marginal is $q(z)$. So it gives $\mathrm{KL}(q(z)\,\|\,p(z)) \ge 0$. $\blacksquare$

The bound is tight when the aggregate posterior equals the prior. On the other side, $I_q(X; Z) \ge H - D$, where $H$ is the entropy of the data, because $H(X \mid Z) \le D$ (a cross-entropy is at least the entropy). So $H - D \le I_q(X; Z) \le R$ (Alemi et al., 2018).

The β-VAE (Higgins et al., 2017) minimizes $D + \beta R$.

• $\beta = 1$ is the negative ELBO.
• $\beta > 1$ buys a lower rate with more distortion: more compressed codes, often blurrier reconstructions.
• $\beta < 1$ does the opposite: more information in $z$, sharper reconstructions.
• Same ELBO, different models. $-\mathrm{ELBO} = D + R$, so every point on a line $D + R = \text{const}$ has the same ELBO. A model with $R \approx 0$ ignores $z$: this is posterior collapse. A powerful decoder can reach a good ELBO this way. So report $R$ and $D$ separately, not only the ELBO.

A linear-Gaussian example. Let $X \sim \mathcal N(0, 1)$, the encoder $z = a x + s\,\varepsilon$ with $\varepsilon \sim \mathcal N(0, 1)$, the prior $p(z) = \mathcal N(0, 1)$, and the best decoder $\hat x = \mathbb{E}[X \mid z] = \frac{a}{a^2 + s^2} z$. With $v = a^2 + s^2$, all quantities have closed forms (in nats):

So $I(X; Z) = \frac12\ln(1/\text{MSE})$: measured by mutual information, every such encoder sits exactly on the Gaussian curve. The rate equals $I$ only when $v = 1$, that is, when $q(z) = \mathcal N(0, v)$ is the prior. With a decoder variance of $\frac12$ the distortion is the MSE (plus a constant), and the β-VAE optimum is $\text{MSE} = \beta/2$ on the curve for $\beta < 2$, and the collapsed point $R = 0$, $\text{MSE} = 1$ for $\beta \ge 2$.

Interactive figure: loading…

Learned image compression (Ballé et al., 2017) minimizes $R + \lambda D$, with the rate given by a learned entropy model of a quantized latent (coded with arithmetic coding, lesson 05). It is the β-VAE objective with $\beta = 1/\lambda$.

Answers to section 1. (1) $1 - H_b(0.11) = 0.50$ bits per flip, not 0.78: coding long blocks beats store-or-guess. (2) By a factor of 4: $D(R) = \sigma^2 2^{-2R}$, that is 6.02 dB per bit. (3) The KL term (rate) goes down and the reconstruction error goes up; for large enough $\beta$ the rate can reach 0 (posterior collapse).

4. Worked example #

A Bernoulli(0.3) source with Hamming distortion, at $D = 0.1$.

1. Rate. $H_b(0.3) = 0.8813$ and $H_b(0.1) = 0.4690$, so $R(0.1) = 0.4123$ bits per symbol. $D_{\max} = 0.3$ (always output 0).
2. Backward channel. $s = \frac{0.3 - 0.1}{1 - 0.2} = 0.25$: $\hat X \sim$ Bernoulli(0.25), and $X$ is $\hat X$ flipped with probability 0.1. Check: $\Pr(X = 1) = 0.25 \cdot 0.9 + 0.75 \cdot 0.1 = 0.225 + 0.075 = 0.3$ ✓.
3. Forward channel, by Bayes' rule. $q(\hat x{=}1 \mid x{=}1) = \frac{\Pr(\hat X = 1)\Pr(X = 1 \mid \hat X = 1)}{\Pr(X = 1)} = \frac{0.25 \cdot 0.9}{0.3} = 0.75$, and $q(\hat x{=}1 \mid x{=}0) = \frac{0.25 \cdot 0.1}{0.7} = \frac1{28} = 0.0357$. The encoder keeps almost all 0s and turns a quarter of the 1s into 0s: it errs toward the more likely symbol. Distortion: $0.3 \cdot 0.25 + 0.7 \cdot 0.0357 = 0.075 + 0.025 = 0.1$ ✓.
4. Slope. $dR/dD = -\ln\frac{0.9}{0.1} = -\ln 9 = -2.197$ nats per unit of distortion. So BA with $\beta = \ln 9$ ends here: $D = 1/(1 + 9) = 0.1$.
5. This is a fixed point of BA. The output marginal is $r = (0.75, 0.25)$ (vectors list $\hat x = 0$, then $\hat x = 1$), and $e^{-\beta} = \frac19$. For $x = 1$: $q \propto (0.75 \cdot \frac19,\ 0.25 \cdot 1) = (0.0833, 0.25)$, which normalizes to $(0.25, 0.75)$. For $x = 0$: $q \propto (0.75 \cdot 1,\ 0.25 \cdot \frac19) = (0.75, 0.0278)$, which normalizes to $(0.9643, 0.0357)$. Both match step 3.
6. Gaussian. $\sigma^2 = 1$ and $D = 0.25$ give $R = \frac12\log_2 4 = 1$ bit per sample. A uniform scalar quantizer with ideal entropy coding needs about $\frac12\log_2\frac{\pi e}{6} = 0.254$ bits more per sample at high rate (1.53 dB). Coding long blocks closes this gap.

5. In code #

The code checks the Bernoulli formula by brute force over all binary test channels, runs Blahut–Arimoto at $\beta = \ln 9$, and evaluates the linear-Gaussian VAE in closed form. SciPy computes $I(X; \hat X)$ as $\mathrm{KL}(p(x, \hat x)\,\|\,p(x)\,p(\hat x))$.

import numpy as np
from scipy.stats import entropy

# 1. R(D) of a Bernoulli(0.3) source by brute force over all binary test channels.
#    a = q(x̂=1 | x=0), b = q(x̂=0 | x=1) on a grid; I(X; X̂) = KL(joint || product of marginals).
p1 = 0.3
a, b = np.meshgrid(np.linspace(0, 1, 1001), np.linspace(0, 1, 1001))
joint = np.stack([(1 - p1) * (1 - a), (1 - p1) * a, p1 * b, p1 * (1 - b)])   # (x, x̂) = 00, 01, 10, 11
out1 = joint[1] + joint[3]                                                   # P(x̂ = 1)
prod = np.stack([(1 - p1) * (1 - out1), (1 - p1) * out1, p1 * (1 - out1), p1 * out1])
info_bits = entropy(joint, prod, axis=0, base=2)
dist = joint[1] + joint[2]                                                   # P(x̂ != x)
Hb = lambda x: entropy([x, 1 - x], base=2)
for D in [0.05, 0.1, 0.2]:
    print(D, info_bits[dist <= D].min().round(3), round(Hb(p1) - Hb(D), 3))
# 0.05 0.595 0.595
# 0.1 0.412 0.412
# 0.2 0.159 0.159

# 2. Blahut-Arimoto for the same source at beta = ln 9.
px, d = np.array([0.7, 0.3]), 1.0 - np.eye(2)
beta, q = np.log(9), np.full((2, 2), 0.5)
for _ in range(500):
    r = px @ q                                   # output marginal r(x̂)
    q = r * np.exp(-beta * d)                    # Gibbs form, then normalize each row
    q /= q.sum(axis=1, keepdims=True)
r = px @ q
I = np.sum(px[:, None] * q * np.log2(q / r))
print(q.round(4).tolist(), round(I, 4), round(np.sum(px[:, None] * q * d), 4))
# [[0.9643, 0.0357], [0.25, 0.75]] 0.4123 0.1

# 3. Linear-Gaussian "VAE" on X ~ N(0, 1): encoder z = a x + s eps, prior N(0, 1).
for a_, s_ in [(0.8, 0.6), (0.8, 0.3)]:
    v = a_**2 + s_**2
    rate = 0.5 * (v - 1 - np.log(s_**2))         # E_x KL(N(a x, s^2) || N(0, 1))
    mi = 0.5 * np.log(v / s_**2)                 # I(X; Z) = h(Z) - h(Z | X)
    mse = s_**2 / v                              # best linear decoder x̂ = (a / v) z
    print(f"rate {rate:.3f}  I {mi:.3f}  1/2 ln(1/MSE) {0.5 * np.log(1 / mse):.3f}  nats")
# rate 0.511  I 0.511  1/2 ln(1/MSE) 0.511  nats    (a^2 + s^2 = 1: on the curve)
# rate 1.069  I 1.047  1/2 ln(1/MSE) 1.047  nats    (rate = I + KL(q(z) || p(z)))

Part 3: measured by $I(X; Z)$, a linear-Gaussian encoder with the best decoder sits exactly on the Gaussian curve $\frac12\ln(\sigma^2/D)$. The VAE rate adds $\mathrm{KL}(q(z)\,\|\,p(z))$, which is zero only when $a^2 + s^2 = 1$.

6. Where it shows up #

• Learned compression. End-to-end image codecs (Ballé et al., 2017; with a hyperprior, 2018) minimize $R + \lambda D$, with the rate given by a learned entropy model and arithmetic coding. Neural codecs for video, audio (neural audio codecs with residual vector quantization) and scientific fields follow the same recipe.
• VAEs and latent diffusion. β-VAEs move along the curve; free bits set a floor on the rate per latent dimension to avoid collapse. The first-stage autoencoder of a latent diffusion model uses a very small KL weight: a high rate, so details survive, and the diffusion model then learns the distribution of the latents.
• Quantization. A uniform quantizer with step $\Delta$ has MSE about $\Delta^2/12$ at high resolution. Each extra bit halves $\Delta$ (for a fixed range) and divides the error by 4. Going from 8 to 4 bits per weight at a fixed range multiplies the error by $4^4 = 256$. Post-training quantization decides where to spend bits for the smallest output error: a rate–distortion problem, often solved layer by layer.
• Transform coding. JPEG uses the DCT, an approximation of PCA for images, then spends more bits on the large low-frequency coefficients: reverse water-filling in practice.
• ML for science. Lossy compressors for climate, weather and seismic archives (error-bounded SZ and ZFP, or neural codecs) are compared by bits per value against error. Reverse water-filling explains why coding in principal components (EOFs of a climate field) works.
• Representations. Variational IB is used as a regularizer for classifiers and for robustness. Rate and distortion curves are a useful way to compare representation learners.

7. Common mistakes #

• Mixing up $R(D)$ and capacity. $R(D)$ minimizes $I$ over channels for a fixed source. Capacity maximizes $I$ over inputs for a fixed channel.
• Using $H_b(p) - H_b(D)$ beyond $D = \min(p, 1 - p)$. There $R(D) = 0$; the formula would give a wrong, often negative, value. The same holds for the Gaussian formula when $D > \sigma^2$.
• Units. BA's slope $-\beta$ is in nats per unit of distortion when $I$ is in nats. Convert $R$ to bits at the end.
• Mixing up the β's. In BA and the IB, a large $\beta$ means a high rate; in the β-VAE, a large $\beta$ means a low rate.
• Reading the β-VAE rate as $I(X; Z)$. It is an upper bound; the gap is $\mathrm{KL}(q(z)\,\|\,p(z))$.
• Comparing VAEs by the ELBO alone. Equal ELBOs can hide very different $(R, D)$ pairs, including $R = 0$.
• Expecting a scalar quantizer to reach $R(D)$. For a Gaussian at high rate, even with ideal entropy coding, it stays 0.254 bits per sample (1.53 dB) above the curve. The bound needs coding over long blocks (vector quantization).
• IB with a deterministic network. For a continuous $X$ and an injective deterministic $T = g(X)$, $I(X; T)$ is infinite. A finite estimate comes from the estimator, not from the network.

8. Check yourself #

1. A fair coin, Hamming distortion, 0.25 bits per flip. What is the smallest error rate? What does store-or-guess reach?
   Answer

   Solve $1 - H_b(D) = 0.25$, so $H_b(D) = 0.75$ and $D = 0.2145$. Store-or-guess: $1 - 2D = 0.25$ gives $D = 0.375$.
2. A Gaussian source with $\sigma^2 = 4$ and squared error. What MSE can you reach with 3 bits per sample?
   Answer

   $D = \sigma^2 2^{-2R} = 4 \cdot 2^{-6} = \frac1{16} = 0.0625$.
3. Why is $R(D)$ convex?
   Answer

   Mix the optimal channels for $D_1$ and $D_2$ with weights $\lambda$ and $1 - \lambda$. The distortion is linear in the channel, so the mixture meets $\lambda D_1 + (1-\lambda)D_2$. $I(X; \hat X)$ is convex in the channel for a fixed source, so $R(\lambda D_1 + (1 - \lambda)D_2) \le \lambda R(D_1) + (1 - \lambda)R(D_2)$.
4. Three independent Gaussian components have variances $4, 1, 0.25$. With the water level $\theta = 0.5$, give each $D_i$, each rate and the totals.
   Answer

   $D_i = \min(0.5, \lambda_i) = 0.5, 0.5, 0.25$. Rates $\frac12\log_2(4/0.5) = 1.5$, $\frac12\log_2(1/0.5) = 0.5$, and 0 bits. Total $R = 2$ bits for $D = 1.25$. The third component is not sent.
5. Prove $\mathbb{E}_x\,\mathrm{KL}(q(z \mid x)\,\|\,p(z)) = I_q(X; Z) + \mathrm{KL}(q(z)\,\|\,p(z))$ with $q(z) = \mathbb{E}_x\,q(z \mid x)$.
   Answer

   Write $\log\frac{q(z\mid x)}{p(z)} = \log\frac{q(z\mid x)}{q(z)} + \log\frac{q(z)}{p(z)}$ and take the expectation over $p_{\text{data}}(x)\,q(z \mid x)$. The first term gives $I_q(X; Z)$. In the second term $z$ has marginal $q(z)$, so it gives $\mathrm{KL}(q(z)\,\|\,p(z))$.
6. A β-VAE with $\beta \ge 1$ and a powerful autoregressive decoder. What can happen, and why?
   Answer

   Posterior collapse: the decoder can model $p(x)$ without $z$, so the optimum can set $q(z \mid x) = p(z)$, which gives $R = 0$ with no loss in $D$. The code then carries no information. $\beta < 1$, free bits or a weaker decoder push the rate up.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Define the rate–distortion function $R(D)$.

   Show answer

2. What is $R(D)$ for a Bernoulli($p$) source with Hamming distortion?

   Show answer

3. What is $R(D)$ for a Gaussian $\mathcal N(0, \sigma^2)$ source with squared error?

   Show answer

4. By what factor does one extra bit per sample reduce the smallest MSE of a Gaussian source?

   Show answer

5. Compute $R(D)$ for a fair coin at $D = 0.11$.

   Show answer

6. How does reverse water-filling share bits among independent Gaussian components?

   Show answer

7. What are the two steps of the Blahut–Arimoto algorithm?

   Show answer

8. Write the information bottleneck objective.

   Show answer

9. Why is the β-VAE rate an upper bound on $I(X;Z)$?

   Show answer

10. In a β-VAE, what does a larger $\beta$ do?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "What is the rate–distortion function? How is it different from channel capacity?" The minimum of $I(X; \hat X)$ over channels that meet the distortion budget; capacity maximizes $I$ over inputs for a fixed channel.
• "How many bits per sample does a Gaussian signal need at an SNR of 30 dB?" $\sigma^2/D = 1000$, so $R = \frac12\log_2 1000 = 4.98$ bits.
• "Why does a β-VAE with $\beta > 1$ give blurrier reconstructions?" It pays more for rate, so it moves to a point with lower rate and higher distortion.
• "What is posterior collapse in rate–distortion language?" The point $R = 0$: the code carries no information, and the decoder alone models the data.
• "What is the information bottleneck?" $\min I(X;T) - \beta I(T;Y)$: rate–distortion with the distortion $\mathrm{KL}(p(y \mid x)\,\|\,p(y \mid t))$. Mention that the "compression phase" claim is disputed.
• "Why do codecs transform the data before quantizing?" In decorrelated components, reverse water-filling gives the bits to high-variance components and drops the rest.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. Brute-force $R(D)$ of a Bernoulli source over binary test channels on a grid, compared with $H_b(p) - H_b(D)$.
2. blahut_arimoto(p_x, distortion, beta), then trace the whole curve with a sweep over $\beta$.
3. Gaussian $R(D)$ and reverse water-filling on the PCA spectrum of the load_digits data set.
4. A linear-Gaussian "VAE": rate, $I(X; Z)$ and MSE in closed form for several encoders, placed on the $R(D)$ plane.
5. Uniform scalar quantization of a Gaussian: entropy of the quantized values against MSE, compared with $R(D)$ (the gap of about 0.25 bits at high rate).

11. Watch and read #

Lecture 9: Rate Distortion Theory and Information Bottleneck.Gregory Falkovich · 59 min

Rate–distortion as minimal mutual information, the Gaussian case, and the information bottleneck.

Digital Image Processing I - Lecture 43 - Rate-Distortion TheoryCharles Bouman · 52 min

The Gaussian R(D) and reverse water-filling, from the point of view of image coding.

Variational AutoencodersArxiv Insights · 15 min

A visual introduction to VAEs, including the β-VAE and its trade-off between reconstruction and the KL term.

Stanford Seminar - Information Theory of Deep Learning, Naftali TishbyStanford Online · 85 min

The information bottleneck view of deep learning, by its author; present it as one view, not a settled fact (see section 3.6).

• Stanford EE274 notes: rate–distortion theory (the Bernoulli and Gaussian examples) and quantization.
• Kevin Murphy, Probabilistic Machine Learning: Advanced Topics (free PDF): §5.4.2 (lossy compression and the rate–distortion trade-off), §5.6 (the information bottleneck and variational IB), §21.3.1 (β-VAE).
• Yury Polyanskiy and Yihong Wu, Information Theory: From Coding to Learning (free draft): the chapters on rate–distortion theory.
• Tishby, Pereira and Bialek, The information bottleneck method (1999); Alemi et al., Deep Variational Information Bottleneck (ICLR 2017).
• Alemi et al., Fixing a Broken ELBO (ICML 2018), §2: rate, distortion and the β-VAE; Higgins et al., β-VAE (ICLR 2017).
• Ballé, Laparra and Simoncelli, End-to-end optimized image compression (ICLR 2017); Saxe et al., On the information bottleneck theory of deep learning (ICLR 2018).