Rate-Distortion Theory

Table of Contents

• Introduction
• 1. Setup and Distortion Measures
  • 1.1 Source and Reproduction Alphabets
  • 1.2 Distortion Functions
  • 1.3 Canonical Examples
• 2. The Rate-Distortion Function
  • 2.1 Definition
  • 2.2 Properties
  • 2.3 Boundary Values
• 3. The Rate-Distortion Theorem
  • 3.1 Operational Definition of a Code
  • 3.2 Theorem Statement
  • 3.3 Achievability Proof Sketch
  • 3.4 Converse Proof Sketch
• 4. Gaussian Source with Squared Error
  • 4.1 Scalar Gaussian
  • 4.2 Derivation via the Test Channel
• 5. Parametric Form and Reverse Water-Filling
  • 5.1 The Gibbs Test Channel
  • 5.2 Parametric Equations
  • 5.3 Multivariate Gaussian and Reverse Water-Filling
• 6. The Blahut-Arimoto Algorithm
  • 6.1 Lagrangian Formulation
  • 6.2 Alternating Minimization
  • 6.3 Update Equations and Convergence
• 7. Connection to Learned Compression and Quantization
  • 7.1 Rate-Distortion Lagrangian and VAEs
  • 7.2 TurboQuant: Near-Optimal Vector Quantization
• References

Introduction 🗜️

Rate-distortion theory is the branch of information theory that answers a deceptively simple question: how many bits per symbol are needed to represent a source, if we are willing to tolerate some error?

Shannon’s lossless source coding theorem tells us that a source \(X \sim p(x)\) with entropy \(H(X)\) requires at least \(H(X)\) bits per symbol for exact reconstruction. But exact reconstruction is often unnecessary — a compressed image, a quantized audio signal, a low-precision neural network weight all involve controlled distortion. Once we permit error, the minimum achievable rate drops below \(H(X)\) and becomes a function of how much distortion we allow. Rate-distortion theory characterizes this trade-off precisely.

The core question

Suppose we encode a long block \(X^n\) into a codeword of rate \(R\) bits per symbol, then decode to a reconstruction \(\hat{X}^n\). We measure fidelity by an average distortion \(\frac{1}{n}\sum_i d(X_i, \hat{X}_i)\). The rate-distortion function \(R(D)\) is defined as the infimum over all rates \(R\) for which there exists a sequence of codes achieving expected distortion \(\leq D\):

\[R(D) = \min_{\substack{q(\hat{x}|x) \\ \mathbb{E}[d(X,\hat{X})] \leq D}} I(X; \hat{X}).\]

This is an information-theoretic quantity — a constrained mutual information minimization over test channels \(q(\hat{x}|x)\). The rate-distortion theorem (Section 3) establishes that \(R(D)\) is the exact operational threshold: rates above \(R(D)\) are achievable, rates below \(R(D)\) are not.

Why mutual information?

The mutual information \(I(X; \hat{X})\) emerges as the right quantity for the same reason it appears in channel capacity: it measures the essential information carried by the reconstruction about the source. The encoder-decoder system is, after all, a noisy channel in reverse — instead of asking how much information a channel can transmit, we ask how much information the encoder must transmit to achieve a given reconstruction quality. The two problems are Lagrangian duals of each other.

Main results to know

1. Setup and Distortion Measures 📐

1.1 Source and Reproduction Alphabets

Let \(\mathcal{X}\) denote a source alphabet and \(\hat{\mathcal{X}}\) a reproduction alphabet. In general these are not required to be equal, though in many examples they coincide. We are given a source random variable \(X \sim p(x)\) taking values in \(\mathcal{X}\), and we wish to represent \(X\) by a reconstruction \(\hat{X}\) taking values in \(\hat{\mathcal{X}}\).

Definition 1.1 (Block source code). An \((n, M)\) source code for the source \(X^n = (X_1, \ldots, X_n)\) consists of an encoding map \(f_n : \mathcal{X}^n \to \{1, \ldots, M\}\) and a decoding map \(g_n : \{1, \ldots, M\} \to \hat{\mathcal{X}}^n\). The rate of the code is \(R = \frac{1}{n}\log M\) bits per source symbol.

The source emits a sequence of i.i.d. symbols \((X_1, X_2, \ldots)\) drawn from the common marginal \(p(x)\). The decoder produces \(\hat{X}^n = g_n(f_n(X^n))\) as the reconstruction.

1.2 Distortion Functions

Definition 1.2 (Single-letter distortion). A distortion function (or distortion measure) is a map

\[d : \mathcal{X} \times \hat{\mathcal{X}} \to [0, \infty)\]

assigning a non-negative cost \(d(x, \hat{x})\) to reproducing source symbol \(x\) by reconstruction symbol \(\hat{x}\).

Definition 1.3 (Block distortion). For sequences \(x^n \in \mathcal{X}^n\) and \(\hat{x}^n \in \hat{\mathcal{X}}^n\), the block distortion is defined as the per-symbol average

\[d(x^n, \hat{x}^n) = \frac{1}{n} \sum_{i=1}^n d(x_i, \hat{x}_i).\]

This single-letter separability is a crucial structural assumption: it permits the problem to decompose across time steps and enables the typical-sequence arguments that underpin achievability.

Definition 1.4 (Expected distortion). Given a joint distribution \(p(x, \hat{x}) = p(x) \cdot q(\hat{x}|x)\) induced by a test channel \(q(\hat{x}|x)\), the expected distortion is

\[\mathbb{E}[d(X, \hat{X})] = \sum_{x \in \mathcal{X}} \sum_{\hat{x} \in \hat{\mathcal{X}}} p(x) q(\hat{x}|x) d(x, \hat{x}).\]

1.3 Canonical Examples

🔑 The two most important distortion measures in practice are:

Definition 1.5 (Hamming distortion). For discrete alphabets \(\mathcal{X} = \hat{\mathcal{X}}\), the Hamming distortion is the indicator of disagreement:

\[d_H(x, \hat{x}) = \mathbf{1}[x \neq \hat{x}].\]

This measures the symbol error rate — the fraction of source symbols incorrectly reproduced.

Definition 1.6 (Squared error distortion). For \(\mathcal{X} = \hat{\mathcal{X}} = \mathbb{R}\), the squared error (or mean squared error, MSE) distortion is

\[d_{SE}(x, \hat{x}) = (x - \hat{x})^2.\]

For Gaussian sources this leads to a closed-form rate-distortion function. The expected block distortion is then the mean squared error \(\frac{1}{n}\sum_{i=1}^n (x_i - \hat{x}_i)^2\).

2. The Rate-Distortion Function 📊

2.1 Definition

The central object of the theory is the minimum achievable rate at a given distortion level.

Definition 2.1 (Information rate-distortion function). For a source \(X \sim p(x)\) and distortion measure \(d\), the information rate-distortion function is

\[R(D) = \min_{\substack{q(\hat{x}|x) \\ \mathbb{E}[d(X,\hat{X})] \leq D}} I(X; \hat{X}),\]

where the minimization is over all test channels \(q(\hat{x}|x)\), i.e., all conditional distributions on \(\hat{\mathcal{X}}\) given \(X\), subject to the average distortion constraint.

The mutual information is computed with respect to the joint \(p(x)q(\hat{x}|x)\):

\[I(X; \hat{X}) = \sum_{x, \hat{x}} p(x) q(\hat{x}|x) \log \frac{q(\hat{x}|x)}{q(\hat{x})},\]

where \(q(\hat{x}) = \sum_x p(x) q(\hat{x}|x)\) is the induced marginal on \(\hat{\mathcal{X}}\).

2.2 Properties

Proposition 2.2 (Properties of R(D)).

1. Non-increasing: \(R(D)\) is non-increasing in \(D\). Increasing the allowed distortion only relaxes the constraint set, so the minimum mutual information cannot increase.

2. Convex: \(R(D)\) is convex in \(D\). More bits are needed to achieve each successive unit of quality improvement (diminishing returns in the dual direction).

3. Continuous: \(R(D)\) is continuous on the interior of its domain.

Proof of convexity. Let \(D_1, D_2 \geq 0\) and \(\lambda \in [0,1]\). Let \(q_1, q_2\) be the optimal test channels achieving \(R(D_1), R(D_2)\) respectively. Define the mixture test channel \(q_\lambda(\hat{x}|x) = \lambda q_1(\hat{x}|x) + (1-\lambda)q_2(\hat{x}|x)\). By the convexity of mutual information in the conditional distribution (for fixed marginal \(p(x)\)):

\[R(\lambda D_1 + (1-\lambda)D_2) \leq I_{q_\lambda}(X;\hat{X}) \leq \lambda I_{q_1}(X;\hat{X}) + (1-\lambda) I_{q_2}(X;\hat{X}) = \lambda R(D_1) + (1-\lambda) R(D_2). \qquad \square\]

2.3 Boundary Values

Proposition 2.3 (Boundary behavior).

1. At \(D = 0\) (lossless): \(R(0) = H(X)\) for discrete sources (Shannon’s source coding theorem). For continuous sources with finite differential entropy the formula differs.

2. At maximum distortion \(D_{\max}\): \(R(D_{\max}) = 0\). The maximum distortion is achieved by the trivially uninformative encoder: \(q(\hat{x}|x) = q^*(\hat{x})\) for a deterministic constant \(\hat{x}^*\) that minimises \(\mathbb{E}[d(X, \hat{x}^*)]\). For this encoder \(I(X;\hat{X}) = 0\), which is feasible exactly when \(\mathbb{E}_{p(x)}[d(X, \hat{x}^*)] \leq D\).

3. \(R(D) = 0\) for all \(D \geq D_{\max}\) where \(D_{\max} = \min_{\hat{x} \in \hat{\mathcal{X}}} \mathbb{E}[d(X, \hat{x})]\).

💡 The entropy \(H(X)\) and mutual information \(I(X;\hat{X})\) that appear here are precisely the quantities developed in Entropy and Divergences (no note yet).

3. The Rate-Distortion Theorem 🔑

The rate-distortion function \(R(D)\) has an operational interpretation: it is the minimum rate required to reproduce the source within average distortion \(D\). This identification is the content of the rate-distortion theorem.

3.1 Operational Definition of a Code

Definition 3.1 (Achievable rate-distortion pair). A rate-distortion pair \((R, D)\) is achievable if there exists a sequence of \((n, \lceil 2^{nR} \rceil)\) source codes with encoder \(f_n\) and decoder \(g_n\) such that

\[\limsup_{n \to \infty} \mathbb{E}\!\left[d(X^n, g_n(f_n(X^n)))\right] \leq D.\]

Definition 3.2 (Operational rate-distortion function). The operational rate-distortion function is

\[R_{\mathrm{op}}(D) = \inf\{R : (R, D) \text{ is achievable}\}.\]

3.2 Theorem Statement

Theorem 3.3 (Rate-Distortion Theorem, Shannon 1948/1959). For a discrete memoryless source \(X\) with distribution \(p(x)\), a distortion measure \(d : \mathcal{X} \times \hat{\mathcal{X}} \to [0, \infty)\), and any \(D \geq 0\):

\[R_{\mathrm{op}}(D) = R(D) = \min_{\substack{q(\hat{x}|x) \\ \mathbb{E}[d(X,\hat{X})] \leq D}} I(X; \hat{X}).\]

Equivalently, for any \(\varepsilon > 0\) and sufficiently large \(n\): - (Achievability) There exists a \((n, 2^{n(R(D)+\varepsilon)})\) code with \(\mathbb{E}[d(X^n, \hat{X}^n)] \leq D + \varepsilon\). - (Converse) Any \((n, 2^{nR})\) code with \(\mathbb{E}[d(X^n, \hat{X}^n)] \leq D\) must have \(R \geq R(D) - \varepsilon\) for large enough \(n\).

The rate-distortion function is the tightest information-theoretic bound: no sequence of codes can do strictly better than \(R(D)\).

3.3 Achievability Proof Sketch

📐 We construct a random codebook and decode by joint typicality.

Step 1: Codebook generation. Fix the optimal test channel \(q^*(\hat{x}|x)\) achieving \(R(D)\), and let \(q^*(\hat{x}) = \sum_x p(x)q^*(\hat{x}|x)\) be the induced marginal. Generate \(2^{nR}\) codewords \(\hat{x}^n(1), \ldots, \hat{x}^n(2^{nR})\) independently, each drawn i.i.d. from \(\prod_{i=1}^n q^*(\hat{x}_i)\).

Step 2: Encoding. Given source sequence \(x^n\), the encoder finds an index \(m\) such that \((x^n, \hat{x}^n(m))\) is jointly \(\varepsilon\)-typical with respect to the joint distribution \(p(x)q^*(\hat{x}|x)\). If no such codeword exists, declare an error.

Step 3: Decoding. The decoder outputs \(\hat{x}^n(m)\).

Analysis of distortion. If encoding succeeds (no error), then \((x^n, \hat{x}^n(m))\) lies in the jointly typical set \(\mathcal{T}_\varepsilon^{(n)}(X, \hat{X})\). By the defining property of joint typicality, the per-symbol empirical average distortion satisfies

\[d(x^n, \hat{x}^n(m)) = \frac{1}{n} \sum_{i=1}^n d(x_i, \hat{x}_i(m)) \leq \mathbb{E}[d(X, \hat{X})] + \varepsilon \leq D + \varepsilon.\]

Probability of error. A codeword \(\hat{x}^n(m)\) is jointly typical with \(x^n\) with probability \(\approx 2^{-nI(X;\hat{X})}\) (the conditional typical set has size \(\approx 2^{nH(\hat{X}|X)}\) inside an ambient typical set of size \(\approx 2^{nH(\hat{X})}\)). The probability that none of the \(2^{nR}\) codewords is typical with \(x^n\) is approximately \((1 - 2^{-nI(X;\hat{X})})^{2^{nR}}\), which tends to zero when \(R > I(X;\hat{X})\). Since \(I(X;\hat{X}) = R(D)\) at the optimal test channel, any \(R > R(D)\) suffices. \(\square\)

3.4 Converse Proof Sketch

We must show that any \((n, 2^{nR})\) code achieving expected distortion \(\leq D\) must have \(R \geq R(D)\).

Let \(M = f_n(X^n) \in \{1, \ldots, 2^{nR}\}\) be the codeword index and \(\hat{X}^n = g_n(M)\) the reconstruction. The rate satisfies

\[nR \geq H(M) \geq I(X^n; M) = I(X^n; \hat{X}^n).\]

The second inequality uses \(H(M) \geq I(X^n; M)\), and the equality uses the fact that \(\hat{X}^n\) is a deterministic function of \(M\). Now expand:

\[I(X^n; \hat{X}^n) = \sum_{i=1}^n I(X_i; \hat{X}^n \mid X^{i-1}) \geq \sum_{i=1}^n I(X_i; \hat{X}_i),\]

where the inequality uses the fact that conditioning reduces entropy (the chain rule plus non-negativity of conditional mutual information) and the Markov chain \(X_i - \hat{X}^n - \hat{X}_i\). Since \(X_1, \ldots, X_n\) are i.i.d., each \(X_i \sim p(x)\). Define the induced marginal test channel \(q_i(\hat{x}|x) = p(\hat{X}_i = \hat{x} | X_i = x)\). Then

\[\frac{1}{n} \sum_{i=1}^n I(X_i; \hat{X}_i) \geq \frac{1}{n}\sum_{i=1}^n R\!\left(\mathbb{E}[d(X_i,\hat{X}_i)]\right) \geq R\!\left(\frac{1}{n}\sum_{i=1}^n \mathbb{E}[d(X_i,\hat{X}_i)]\right) \geq R(D),\]

where the second inequality uses the convexity of \(R(\cdot)\) (Jensen’s inequality) and the last uses the distortion constraint. Combining, \(R \geq R(D)\). \(\square\)

4. Gaussian Source with Squared Error 📐

4.1 Scalar Gaussian

Theorem 4.1 (Gaussian rate-distortion function). Let \(X \sim \mathcal{N}(0, \sigma^2)\) and \(d(x, \hat{x}) = (x - \hat{x})^2\). Then

\[R(D) = \begin{cases} \dfrac{1}{2}\log\dfrac{\sigma^2}{D} & 0 < D \leq \sigma^2, \\ 0 & D > \sigma^2. \end{cases}\]

The Gaussian source requires the most bits among all sources with the same variance, for any given MSE distortion level (Shannon lower bound).

4.2 Derivation via the Test Channel

We seek \(\min_{q(\hat{x}|x)} I(X; \hat{X})\) subject to \(\mathbb{E}[(X - \hat{X})^2] \leq D\), with \(X \sim \mathcal{N}(0, \sigma^2)\).

Claim: The optimal test channel is the additive Gaussian channel \(\hat{X} = X + Z\) where \(Z \sim \mathcal{N}(0, \sigma_Z^2)\) is independent of \(X\), with variance \(\sigma_Z^2\) chosen so \(\text{Var}(\hat{X} - X) = \sigma_Z^2 = D\).

Wait — the constraint is on the distortion \(\mathbb{E}[(X - \hat{X})^2]\), not the variance of noise in a forward channel. Let us think of it differently.

Step 1: Identify the distortion constraint. We need \(\mathbb{E}[(X - \hat{X})^2] \leq D\). With \(\hat{X} = \alpha X + Z\) for \(Z \sim \mathcal{N}(0, \sigma_Z^2)\) independent of \(X\):

\[\mathbb{E}[(X - \hat{X})^2] = \mathbb{E}[(X - \alpha X - Z)^2] = (1-\alpha)^2 \sigma^2 + \sigma_Z^2.\]

Step 2: Choose \(\alpha = 1\), \(\sigma_Z^2 = D\). The optimal choice turns out to be \(\alpha = 1\) (i.e., \(\hat{X} = X + Z\)), which is a valid test channel since the reconstruction is \(\hat{X}\) viewed as a noisy version of \(X\). The mutual information for this test channel is

\[I(X; \hat{X}) = h(\hat{X}) - h(\hat{X} | X) = h(X + Z) - h(Z).\]

Since \(\hat{X} = X + Z\) with \(X \sim \mathcal{N}(0, \sigma^2)\), \(Z \sim \mathcal{N}(0, D)\) independent, we have \(\hat{X} \sim \mathcal{N}(0, \sigma^2 + D)\). Thus

\[h(\hat{X}) = \frac{1}{2}\log(2\pi e(\sigma^2 + D)), \qquad h(Z) = \frac{1}{2}\log(2\pi eD).\]

Therefore \(I(X;\hat{X}) = \frac{1}{2}\log\frac{\sigma^2 + D}{D}\). But this is not the formula we want.

Step 3: Correct formulation — the backward channel. The rate-distortion function minimizes \(I(X;\hat{X})\) over test channels \(q(\hat{x}|x)\). For the Gaussian source, the backward optimal channel is: given \(\hat{X} \sim \mathcal{N}(0, \hat{\sigma}^2)\), set \(X | \hat{X} = \hat{x} \sim \mathcal{N}(\hat{x}, D)\) (the residual \(X - \hat{X}\) has variance \(D\)). With \(\hat{X} \sim \mathcal{N}(0, \sigma^2 - D)\) (where we enforce \(\text{Var}(X) = \sigma^2\)):

\[\mathbb{E}[(X - \hat{X})^2] = D, \qquad I(X; \hat{X}) = h(X) - h(X|\hat{X}) = \frac{1}{2}\log(2\pi e \sigma^2) - \frac{1}{2}\log(2\pi eD) = \frac{1}{2}\log\frac{\sigma^2}{D}.\]

Step 4: Optimality. The Gaussian test channel achieves the minimum by the Gaussian maximizes entropy lemma: among all distributions with fixed variance, the Gaussian has maximum differential entropy. Thus \(h(X|\hat{X}) = \frac{1}{2}\log(2\pi e D_\varepsilon)\) is maximized (hence \(I(X;\hat{X})\) minimized) precisely when the conditional \(X|\hat{X}\) is Gaussian, which is exactly the test channel above. \(\square\)

5. Parametric Form and Reverse Water-Filling 🌊

5.1 The Gibbs Test Channel

The minimisation defining \(R(D)\) is a convex optimisation. Introducing a Lagrange multiplier \(\beta \geq 0\) (the inverse temperature or slope parameter), we form the Lagrangian

\[\mathcal{L}(q, \beta) = I(X; \hat{X}) + \beta \left(\mathbb{E}[d(X, \hat{X})] - D\right).\]

Setting the functional derivative with respect to \(q(\hat{x}|x)\) to zero (using the calculus of variations with the normalisation constraint \(\sum_{\hat{x}} q(\hat{x}|x) = 1\) handled by a separate multiplier \(\mu(x)\)) yields the KKT stationarity condition:

\[\frac{\partial}{\partial q(\hat{x}|x)}\left[I(X;\hat{X}) + \beta \mathbb{E}[d(X,\hat{X})]\right] = 0.\]

Proposition 5.1 (Optimal test channel — Gibbs form). The optimal test channel satisfying the KKT conditions is

\[q^*(\hat{x}|x) = \frac{q^*(\hat{x})\exp(-\beta d(x, \hat{x}))}{\sum_{\hat{x}'} q^*(\hat{x}')\exp(-\beta d(x, \hat{x}'))},\]

where \(q^*(\hat{x}) = \sum_x p(x) q^*(\hat{x}|x)\) is the self-consistent marginal.

Sketch. Write \(I(X;\hat{X}) = \sum_{x,\hat{x}} p(x)q(\hat{x}|x)\log \frac{q(\hat{x}|x)}{q(\hat{x})}\). Taking the derivative with respect to \(q(\hat{x}|x)\) and setting to zero (with the normalisation constraint absorbed via a multiplier \(\mu(x)\)):

\[p(x)\!\left[\log q(\hat{x}|x) - \log q(\hat{x}) + 1 + \beta d(x,\hat{x})\right] = p(x)\mu(x).\]

Solving for \(q(\hat{x}|x)\) and normalizing gives the Gibbs form above, with \(\beta\) and \(q^*(\hat{x})\) implicitly coupled through the self-consistency equation. \(\square\)

5.2 Parametric Equations

As \(\beta\) ranges over \([0, \infty)\), the pair \((D(\beta), R(\beta))\) traces the rate-distortion curve:

\[D(\beta) = \mathbb{E}_{q^*}[d(X, \hat{X})], \qquad R(\beta) = I_{q^*}(X; \hat{X}).\]

One shows that \(-\beta = R'(D) = dR/dD\) — the slope of the rate-distortion function at the operating point. Thus:

• \(\beta = 0\): no penalty for distortion; optimal \(q^*\) is the independent channel \(q^*(\hat{x}|x) = q^*(\hat{x})\), giving \(D = D_{\max}\), \(R = 0\).
• \(\beta \to \infty\): distortion is strongly penalized; \(D \to 0\), \(R \to H(X)\) (or \(h(X)\) for continuous sources).

The trade-off slope \(\beta\) is the Lagrange multiplier, and sweeping \(\beta\) over \([0,\infty)\) traces the full rate-distortion boundary.

5.3 Multivariate Gaussian and Reverse Water-Filling

💡 Consider now a vector Gaussian source \(X \sim \mathcal{N}(0, \Sigma)\) with \(\Sigma = \text{diag}(\sigma_1^2, \ldots, \sigma_k^2)\) (diagonal covariance — independent components). The distortion measure is \(d(x, \hat{x}) = \|x - \hat{x}\|^2 = \sum_{j=1}^k (x_j - \hat{x}_j)^2\).

Theorem 5.2 (Reverse water-filling). The rate-distortion function for independent parallel Gaussian sources is

\[R(D) = \sum_{j=1}^k \frac{1}{2}\left[\log\frac{\sigma_j^2}{\lambda}\right]^+,\]

where \([\cdot]^+ = \max(\cdot, 0)\) and the water level \(\lambda\) is determined by the distortion constraint:

\[D = \sum_{j=1}^k \min(\lambda, \sigma_j^2).\]

The optimal per-component distortion allocation is \(D_j^* = \min(\lambda, \sigma_j^2)\).

Sketch. By the product structure of the source and distortion, the problem separates: minimize \(\sum_j R_j\) subject to \(\sum_j D_j \leq D\). Each component \(j\) has rate-distortion function \(R_j(D_j) = \frac{1}{2}[\log(\sigma_j^2/D_j)]^+\). By convex duality (introducing a common Lagrange multiplier \(\lambda\)), the optimal \(D_j\) satisfies the KKT conditions

\[\frac{dR_j}{dD_j} = -\frac{1}{2D_j} = -\frac{1}{2\lambda}, \quad \text{if } D_j < \sigma_j^2.\]

This gives \(D_j = \lambda\) for all active components (those with \(\sigma_j^2 > \lambda\)), and \(D_j = \sigma_j^2\), \(R_j = 0\) for inactive ones. \(\square\)

The name reverse water-filling reflects the geometric picture: one “pours water downward” to fill the basins of low-variance components, in contrast to forward water-filling (which allocates power to high-SNR channels in capacity problems).

flowchart LR
    A["Source variances
σ₁², ..., σₖ²"]
    B["Choose water level λ"]
    C["Active components
σⱼ² > λ"]
    D["Inactive components
σⱼ² ≤ λ"]
    E["Rate Rⱼ = ½ log(σⱼ²/λ)"]
    F["Rate Rⱼ = 0"]
    G["Total R(D) = Σ Rⱼ"]

    A --> B
    B --> C
    B --> D
    C --> E
    D --> F
    E --> G
    F --> G

6. The Blahut-Arimoto Algorithm ⚙️

6.1 Lagrangian Formulation

For a fixed Lagrange multiplier \(\beta > 0\), define the free energy (Lagrangian):

\[F_\beta(q, r) = \sum_{x, \hat{x}} p(x) q(\hat{x}|x) \left[\log \frac{q(\hat{x}|x)}{r(\hat{x})} + \beta d(x, \hat{x})\right],\]

where \(r(\hat{x})\) is an auxiliary marginal distribution on \(\hat{\mathcal{X}}\). Note that \(F_\beta(q, q(\cdot))\) (where \(q(\hat{x}) = \sum_x p(x)q(\hat{x}|x)\)) equals \(I(X;\hat{X}) + \beta \mathbb{E}[d(X,\hat{X})]\), the Lagrangian to be minimized.

Lemma 6.1 (Saddle point). The following saddle-point identity holds:

\[\min_{q(\hat{x}|x)} \left[I(X;\hat{X}) + \beta \mathbb{E}[d]\right] = \min_{q(\hat{x}|x)} \min_{r(\hat{x})} F_\beta(q, r).\]

Proof sketch. For any fixed \(q\), minimizing \(F_\beta(q, r)\) over \(r\) gives \(r^*({\hat{x}}) = \sum_x p(x)q(\hat{x}|x) = q(\hat{x})\) (by differentiation, the minimizer is exactly the induced marginal). Substituting back yields \(F_\beta(q, r^*) = I(X;\hat{X}) + \beta\mathbb{E}[d]\). \(\square\)

The key insight is that \(F_\beta\) is jointly convex in \((q, r)\): it is linear in \(q\) for fixed \(r\) (and conversely linear in \(r\) for fixed \(q\)). This licenses alternating minimization.

6.2 Alternating Minimization

The Blahut-Arimoto algorithm is alternating minimization on \(F_\beta\):

1. M-step: Fix \(r^{(t)}\), minimize over \(q\): \(q^{(t+1)} = \arg\min_{q} F_\beta(q, r^{(t)})\).
2. E-step: Fix \(q^{(t+1)}\), minimize over \(r\): \(r^{(t+1)} = \arg\min_r F_\beta(q^{(t+1)}, r)\).

The minimum over \(r\) in step 2 is always the induced marginal. The minimum over \(q\) in step 1 gives the Gibbs form derived in Proposition 5.1.

Proposition 6.2 (Monotone convergence). Each alternating step does not increase \(F_\beta\), so \(F_\beta(q^{(t)}, r^{(t)})\) is non-increasing. The algorithm converges to the global minimum of \(F_\beta\).

Proof sketch. Non-increase per step is immediate. Global convergence follows from the fact that \(F_\beta\) is strictly convex in \(r\) (for fixed \(q\)) and the feasible set is compact, so the alternating minimization converges to a stationary point. Uniqueness of the fixed point follows from the strict convexity of \(F_\beta\) in the joint argument, as established by Csiszár and Tusnády (1984). \(\square\)

6.3 Update Equations and Convergence

Starting from an arbitrary initial \(r^{(0)}(\hat{x}) > 0\) on \(\hat{\mathcal{X}}\), the update equations are:

Step 1 (Update the test channel):

\[q^{(t+1)}(\hat{x} | x) = \frac{r^{(t)}(\hat{x}) \exp\!\left(-\beta\, d(x, \hat{x})\right)}{\displaystyle\sum_{\hat{x}'} r^{(t)}(\hat{x}') \exp\!\left(-\beta\, d(x, \hat{x}')\right)}.\]

Step 2 (Update the marginal):

\[r^{(t+1)}(\hat{x}) = \sum_{x \in \mathcal{X}} p(x)\, q^{(t+1)}(\hat{x} | x).\]

Convergence criterion: Terminate when \(\left|F_\beta(q^{(t+1)}, r^{(t+1)}) - F_\beta(q^{(t)}, r^{(t)})\right| < \varepsilon_{\text{tol}}\).

The value \(R(\beta) = F_\beta(q^*, r^*)\) at convergence, combined with \(D(\beta) = \mathbb{E}_{q^*}[d]\), gives one point on the rate-distortion curve. By sweeping \(\beta\) over a grid of values, the full \(R(D)\) curve is traced.

import numpy as np

def blahut_arimoto(p_x, d_matrix, beta, n_iter=1000, tol=1e-9):
    """
    Blahut-Arimoto algorithm for computing R(D) at inverse temperature beta.

    Args:
        p_x:      np.ndarray of shape (|X|,), source distribution
        d_matrix: np.ndarray of shape (|X|, |X_hat|), distortion d(x, x_hat)
        beta:     float, Lagrange multiplier (slope parameter, >= 0)
        n_iter:   int, maximum iterations
        tol:      float, convergence tolerance

    Returns:
        R: float, mutual information I(X; X_hat) at convergence
        D: float, expected distortion at convergence
        q_cond: np.ndarray of shape (|X|, |X_hat|), optimal test channel
    """
    n_x, n_xhat = d_matrix.shape

    # Initialize r uniformly
    r = np.ones(n_xhat) / n_xhat

    F_prev = np.inf
    q_cond = np.zeros((n_x, n_xhat))

    for _ in range(n_iter):
        # Step 1: update test channel (Gibbs softmax)
        log_num = np.log(r + 1e-300)[None, :] - beta * d_matrix  # (|X|, |X_hat|)
        log_num -= log_num.max(axis=1, keepdims=True)             # numerical stability
        q_cond = np.exp(log_num)
        q_cond /= q_cond.sum(axis=1, keepdims=True)              # normalize rows

        # Step 2: update marginal
        r = p_x @ q_cond  # shape (|X_hat|,)

        # Compute free energy F_beta = I(X; X_hat) + beta * E[d]
        joint = p_x[:, None] * q_cond                            # p(x) q(x_hat|x)
        safe_log = np.where(q_cond > 0,
                            np.log(q_cond / (r[None, :] + 1e-300) + 1e-300),
                            0.0)
        I_XY  = np.sum(joint * safe_log)
        E_d   = np.sum(joint * d_matrix)
        F_cur = I_XY + beta * E_d

        if abs(F_cur - F_prev) < tol:
            break
        F_prev = F_cur

    R = np.sum(p_x[:, None] * q_cond * np.log(
        q_cond / (r[None, :] + 1e-300) + 1e-300
    ))
    D = np.sum(p_x[:, None] * q_cond * d_matrix)
    return float(R), float(D), q_cond


def trace_rd_curve(p_x, d_matrix, beta_grid):
    """Sweep over beta values to trace the full R(D) curve."""
    points = []
    for beta in beta_grid:
        R, D, _ = blahut_arimoto(p_x, d_matrix, beta)
        points.append((D, R))
    return sorted(points)   # sort by distortion

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

# Discretize N(0,1) on a grid
N = 200
x_grid = np.linspace(-4, 4, N)
dx = x_grid[1] - x_grid[0]
p_x = np.exp(-0.5 * x_grid**2) / np.sqrt(2 * np.pi) * dx
p_x /= p_x.sum()   # normalize

# Squared error distortion matrix
d_matrix = (x_grid[:, None] - x_grid[None, :])**2

# Sweep over beta
beta_grid = np.logspace(-1, 2, 20)
rd_points = []
for beta in beta_grid:
    R, D, _ = blahut_arimoto(p_x, d_matrix, beta)
    rd_points.append((D, R))

D_vals, R_vals = zip(*sorted(rd_points))

# Theoretical curve
D_theory = np.linspace(0.01, 1.0, 200)
R_theory = 0.5 * np.log(1.0 / D_theory)

# Maximum absolute error (interpolate computed onto theory grid)
D_interp = np.array(D_vals)
R_interp = np.array(R_vals)
# (Interpolation and error computation omitted for brevity)

fig, ax = plt.subplots()
ax.plot(D_theory, R_theory, label='Theoretical R(D)')
ax.scatter(D_vals, R_vals, label='Blahut-Arimoto')
ax.set_xlabel('D'); ax.set_ylabel('R (nats)')
ax.legend()
fig.savefig('/tmp/rd_curve.png')

import numpy as np

# Binary symmetric source
p_x = np.array([0.5, 0.5])
d_matrix = np.array([[0.0, 1.0], [1.0, 0.0]])
beta = 2.0

r = np.ones(2) / 2.0
F_values = []
for _ in range(50):
    # Step 1
    log_num = np.log(r + 1e-300)[None, :] - beta * d_matrix
    log_num -= log_num.max(axis=1, keepdims=True)
    q = np.exp(log_num); q /= q.sum(axis=1, keepdims=True)
    # Step 2
    r = p_x @ q
    # Free energy
    joint = p_x[:, None] * q
    I_XY = np.sum(joint * np.log(q / (r[None,:] + 1e-300) + 1e-300))
    E_d  = np.sum(joint * d_matrix)
    F_values.append(I_XY + beta * E_d)

F_star = F_values[-1]
errors = np.abs(np.array(F_values[:-1]) - F_star)
# Fit log(errors) vs t to estimate linear rate
t = np.arange(len(errors))
slope, _ = np.polyfit(t[5:], np.log(errors[5:] + 1e-30), 1)
print(f"Estimated convergence rate: exp({slope:.4f}) per iteration")

7. Connection to Learned Compression and Quantization 🤖

7.1 Rate-Distortion Lagrangian and VAEs

Modern neural data compression — as surveyed by Yang, Mandt & Theis (2023) (no note yet) — directly optimizes a finite-block-length version of the rate-distortion Lagrangian:

\[\mathcal{L}_\lambda(\theta, \phi) = \mathbb{E}_{x \sim p_{\text{data}}}[d(x, g_\theta(z))] + \lambda \cdot \mathbb{E}[-\log q_\phi(z|x)],\]

where \(z = f_\phi(x)\) is the latent code, \(g_\theta\) is the decoder, and \(q_\phi(z|x)\) is the encoder distribution. The rate term \(-\log q_\phi(z|x)\) approximates the bit cost of transmitting \(z\) under a learned entropy model.

Proposition 7.1 (VAE-RD equivalence). When the distortion \(d(x, \hat{x}) = -\log p_\theta(x|\hat{x})\) is the negative log-likelihood of a decoder model \(p_\theta(x|z)\), the Lagrangian \(\mathcal{L}_\lambda\) is exactly the variational autoencoder (VAE) evidence lower bound (ELBO):

\[\mathcal{L}_\text{ELBO} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_{\text{KL}}(q_\phi(z|x) \| p(z)).\]

Sketch. Identify \(\lambda = 1\), distortion \(= -\log p_\theta(x|z)\), and rate \(\approx D_{\text{KL}}(q_\phi(z|x) \| p(z))\). The KL divergence bounds the code length: \(-\log q_\phi(z|x) \geq -\log p(z) - \log \frac{q_\phi(z|x)}{p(z)}\). The ELBO is the exact bound. \(\square\)

The operational rate of a practical codec additionally requires a fixed-precision entropy coder: the bits used to transmit \(z\) under the entropy model \(q_\phi(z|x)\) approach \(-\log_2 q_\phi(z|x)\) for long messages by Shannon’s source coding theorem. The gap between \(R(D)\) and the ELBO-minimizing operating point measures the redundancy of the learned codec.

7.2 TurboQuant: Near-Optimal Vector Quantization

A striking application of rate-distortion theory to vector quantization (VQ) for weight/activation compression is the TurboQuant algorithm of Zandieh et al. (2025).

Setup. Quantize a unit-norm vector \(x \in \mathbb{R}^d\) to \(b\) bits per coordinate using a randomized scheme. The goal is to minimize the expected MSE distortion.

Information-theoretic lower bound. Any randomized quantizer using \(b\) bits per coordinate must satisfy

\[D_{\text{mse}} \geq \frac{1}{4^b}.\]

This follows from the Gaussian rate-distortion bound: treating each coordinate as approximately \(\mathcal{N}(0, 1/d)\) (by high-dimensional concentration), the bound \(D \geq \sigma^2 / 2^{2R} = (1/d)/4^b\) per coordinate sums to \(1/4^b\) across \(d\) coordinates.

TurboQuant upper bound. The algorithm applies a random rotation \(R \in \mathcal{O}(d)\) (drawn from the Haar measure) to \(x\), then applies the Panter-Dite optimal scalar quantizer to each coordinate of \(Rx\):

\[\hat{x} = R^\top \hat{z}, \quad \hat{z} = Q_b(Rx),\]

where \(Q_b\) is optimal scalar quantization to \(b\) bits. After rotation, each coordinate of \(Rx\) follows a Beta distribution that concentrates near \(\mathcal{N}(0, 1/d)\) as \(d \to \infty\). The MSE of TurboQuant satisfies

\[D_{\text{mse}} \leq \frac{\sqrt{3}\pi}{2} \cdot \frac{1}{4^b},\]

where the constant \(\frac{\sqrt{3}\pi}{2} \approx 2.72\) arises from the Panter-Dite formula for the MSE of an optimal \(b\)-bit scalar quantizer applied to a Gaussian source:

\[C(f_X, b) = \frac{1}{12}\left(\int f_X(x)^{1/3}\,dx\right)^3 \cdot \frac{1}{4^b}.\]

Derivation sketch. After random rotation, by high-dimensional concentration each coordinate \(z_j = (Rx)_j\) has variance \(1/d\) and follows approximately \(\mathcal{N}(0, 1/d)\). The Panter-Dite formula gives the distortion per coordinate as \(C(f_{z_j}, b) = \frac{1}{12} \cdot (2\pi e \cdot (1/d))^{1/2} \cdot \frac{1}{4^b} \cdot d = \frac{\sqrt{2\pi e}}{12} \cdot \frac{1}{4^b}\). Summing over coordinates and evaluating the constant yields the stated bound. \(\square\)

The multiplicative gap between TurboQuant and the information-theoretic lower bound is a universal constant \(\approx 2.72\), independent of \(d\) and \(b\). This is a remarkable result: a simple two-step randomized scheme achieves within a constant factor of the optimal rate-distortion bound for all vectors simultaneously.

import numpy as np

variances = np.array([9.0, 4.0, 1.0, 0.25])

def reverse_water_fill(variances, D_target, tol=1e-9):
    """Compute R(D) for parallel Gaussian source via binary search on lambda."""
    lo, hi = 0.0, variances.max()
    # Increase hi until sum of min(lambda, sigma_j^2) >= D_target
    while np.sum(np.minimum(hi, variances)) < D_target:
        hi *= 2
    for _ in range(100):
        lam = (lo + hi) / 2
        D_cur = np.sum(np.minimum(lam, variances))
        if abs(D_cur - D_target) < tol:
            break
        if D_cur < D_target:
            lo = lam
        else:
            hi = lam
    R = 0.5 * np.sum(np.maximum(0, np.log(variances / lam)))
    return R

D_max = variances.sum()
D_grid = np.linspace(0.05, D_max - 0.01, 100)
R_grid = np.array([reverse_water_fill(variances, D) for D in D_grid])

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
ax.plot(D_grid, R_grid)
ax.set_xlabel('D (MSE)'); ax.set_ylabel('R (nats)')
ax.set_title('Reverse Water-Filling R(D) for 4-Component Gaussian')
fig.savefig('/tmp/rw_rd.png')

import numpy as np
from scipy.stats import ortho_group

d = 64
b = 3
n_vecs = 1000

# Generate unit-norm vectors
X = np.random.randn(n_vecs, d)
X /= np.linalg.norm(X, axis=1, keepdims=True)

# Haar-random rotation (reused across vectors for efficiency)
R = ortho_group.rvs(d)

# Quantization: uniform b-bit in [-max_val, max_val]
Z = X @ R.T                         # rotate: (n_vecs, d)
max_val = np.abs(Z).max()
levels = 2**b
step = 2 * max_val / levels
Z_q = np.round(Z / step) * step     # quantize
Z_q = np.clip(Z_q, -max_val, max_val)

# Invert rotation
X_hat = Z_q @ R                     # (n_vecs, d)

mse = np.mean(np.sum((X - X_hat)**2, axis=1))
lb = 1 / (4**b)
ub = (np.sqrt(3) * np.pi / 2) / (4**b)
print(f"Empirical MSE: {mse:.5f}")
print(f"Lower bound:   {lb:.5f}")
print(f"Upper bound:   {ub:.5f}")
print(f"Within bounds: {lb <= mse <= ub + 0.01}")  # small slack for finite d

References