資訊瓶頸：為甚麼最佳表徵會忘掉大部分所見資訊

A neural network trained to classify images sees millions of pixels. A good representation keeps almost none of them. It keeps the shape of a cat, the texture of fur, the arrangement of ears, and throws away the exact shade of the background, the sensor noise, and the compression artefacts. This is not a failure of the model. It is the whole point of learning a representation.

The Information Bottleneck (IB), introduced by Tishby, Pereira, and Bialek in 1999, turns this intuition into a precise optimisation problem: find a representation that is as compressed as possible while still preserving everything relevant about a target variable. It gives us a language for asking what a model should remember, what it should forget, and how to trade one against the other.

The setup

Suppose we have an input variable X and a target variable Y, drawn from some joint distribution p(x, y). Think of X as an image and Y as its label. We want a representation Z, typically a stochastic function of X, that captures the information in X that matters for Y, and nothing else.

Two quantities matter:

• Compression: how much information Z retains about X, measured by the mutual information I(X; Z). Smaller is more compressed.
• Prediction: how much information Z retains about Y, measured by I(Z; Y). Larger is more predictive.

The Information Bottleneck formalises the trade-off as a constrained optimisation: minimise I(X; Z) subject to I(Z; Y) being at least some value, or equivalently, minimise the Lagrangian

L[p(z∣x)]=I(X;Z)−βI(Z;Y)

where β≥0 is a Lagrange multiplier that controls the trade-off. Small β means aggressive compression, even at the cost of predictive power. Large β means prediction dominates, and the representation is allowed to carry more of X.

This is deliberately different from ordinary supervised learning. A classifier trained with cross-entropy alone has no explicit pressure to forget anything. The bottleneck adds that pressure, and that pressure is where generalisation comes from: a representation that carries only task-relevant information cannot overfit to irrelevant details, because it never stored them.

Why mutual information, and not something simpler

It is tempting to ask why we need mutual information at all. Why not just penalise the dimensionality of Z, or the norm of the representation?

The answer is that mutual information is invariant to invertible transformations of Z. A representation can be rescaled, rotated, or reparametrised without changing what it knows. Dimensionality and norm penalties punish the encoding, not the information content. Mutual information punishes the information content itself, which is what we actually care about.

Mutual information is also symmetric and distribution-free in a useful sense. Recall its definition:

I(X;Z)=∑x,zp(x,z)log⁡p(x,z)p(x)p(z)

It measures how far the joint distribution is from independence, in nats or bits depending on the logarithm. When Z is independent of X, the ratio is 1, the logarithm is 0, and I(X; Z) = 0. When Z determines X exactly, I(X; Z) equals the full entropy H(X).

The catch, and it is a serious one, is that mutual information between high-dimensional continuous variables is notoriously hard to estimate. This is why the original IB remained mostly a theoretical tool for years, and why the practical versions of it look quite different from the clean Lagrangian above.

The self-consistent solution

For discrete variables, the IB problem has a beautiful self-consistent solution, analogous to the Blahut-Arimoto algorithm for rate-distortion theory. The optimal encoder takes the form

p(z∣x)=p(z)Z(x,β)exp⁡(−β DKL[p(y∣x) ∣ p(y∣z)])

where Z(x,β) is a normalising constant. Read this carefully, because it is the heart of the theory: an input x is assigned to a representation z with probability that decays exponentially in the KL divergence between what x predicts about y and what z predicts about y. Inputs that imply similar predictions about Y get mapped to similar representations. Inputs that imply different predictions get separated.

This is clustering, but with a very specific notion of similarity: two inputs are similar if they carry the same information about the target. The bottleneck does not cluster by how inputs look, but by what they mean for the task.

Varying β traces out a curve in the information plane, with I(X; Z) on the horizontal axis and I(Z; Y) on the vertical axis. At β=0 , the solution is trivial: Z is independent of everything, and both informations are 0. As β grows, the curve rises, eventually saturating at I(Z; Y) = I(X; Y), the total relevant information available. The curve is concave, and its slope at any point is 1/β . Points on this curve are optimal trade-offs. Points below it are achievable but wasteful. Points above it are impossible.

For jointly Gaussian X and Y, the solution is analytic: the optimal Z is a noisy linear projection of X, and the trade-off curve can be written in closed form in terms of the canonical correlations between X and Y. This Gaussian case is worth knowing because it connects the bottleneck to canonical correlation analysis and to linear dimensionality reduction, and it gives a sanity check for any approximate algorithm.

The information plane and deep learning

The IB became famous in deep learning through a 2017 paper by Shwartz-Ziv and Tishby, which claimed that stochastic gradient descent training of deep networks proceeds in two phases: a fast fitting phase, in which I(Z; Y) rises, followed by a slow compression phase, in which I(X; Z) falls while I(Z; Y) stays roughly constant. The claim was that this compression phase is what makes deep networks generalise, and that it is driven by the noise in SGD.

This claim sparked one of the liveliest debates in recent machine learning. Follow-up work showed that the compression phase depends heavily on the activation function (it appears clearly with saturating activations like tanh, and much less with ReLU), on the estimator used for mutual information, and on whether the network is deterministic or stochastic. Later studies argued that the observed compression was partly an artefact of the binning estimator used to measure mutual information in deterministic networks, where the true mutual information is actually infinite or vacuous.

The debate is instructive rather than embarrassing. It forced the field to be precise about what it means for a deterministic network to compress information at all, and it pushed the practical work toward stochastic networks, where mutual information is well-defined. Whether or not every deep network literally traverses the information plane, the IB remains the cleanest framework we have for stating what a good representation should do.

The Variational Information Bottleneck

The practical breakthrough came with the Variational Information Bottleneck (VIB), introduced by Alemi and co-authors in 2017. The idea is to replace the intractable mutual information terms with variational bounds that can be optimised with stochastic gradient descent.

The compression term I(X; Z) is upper-bounded using a variational marginal r(z):

I(X;Z)≤Ex[DKL(p(z∣x) ∣ r(z))]

The prediction term I(Z; Y) is lower-bounded using a variational decoder q(y|z):

I(Z;Y)≥Ex,yEz∼p(z∣x)[log⁡q(y∣z)]+H(Y)

Dropping the constant H(Y), the VIB objective becomes

LVIB=Ex,yEz∼p(z∣x)[−log⁡q(y∣z)]+β Ex[DKL(p(z∣x) ∣ r(z))]

This is immediately recognisable: the first term is a cross-entropy loss, and the second is a KL regulariser pulling the encoder toward a prior. If you have trained a variational autoencoder, you have optimised something very close to this. The difference is the target: a VAE reconstructs X, while the VIB predicts Y.

Here is a minimal PyTorch sketch of the idea, with a Gaussian encoder and a standard normal prior:

import torch
import torch.nn as nn

class VIBEncoder(nn.Module):
    def __init__(self, input_dim, z_dim):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(input_dim, 256), nn.ReLU())
        self.mu = nn.Linear(256, z_dim)
        self.logvar = nn.Linear(256, z_dim)

    def forward(self, x):
        h = self.net(x)
        mu, logvar = self.mu(h), self.logvar(h)
        std = torch.exp(0.5 * logvar)
        z = mu + std * torch.randn_like(std)  # reparametrisation trick
        kl = -0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp(), dim=1)
        return z, kl

# Training step: cross-entropy on the prediction plus beta times the KL term
z, kl = encoder(x)
logits = classifier(z)
loss = nn.functional.cross_entropy(logits, y) + beta * kl.mean()
loss.backward()

Two practical details matter here. First, the encoder must be stochastic: without the sampling step, the KL term is meaningless and the bottleneck does nothing. Second, β is a genuine hyperparameter that needs tuning, not a value to set and forget. Too small, and the KL term is ignored. Too large, and the representation collapses to the prior and predicts nothing. A common practice is to anneal β from a small value upward during training.

What the bottleneck buys you

The VIB and its relatives have delivered real, measurable benefits in several settings:

• Robustness. Representations trained with a bottleneck are less sensitive to input perturbations, including adversarial ones, because irrelevant details were never encoded in the first place. An adversary cannot exploit information the model does not have.
• Out-of-distribution detection. The KL term gives a natural signal: inputs far from the training distribution tend to produce encodings far from the prior, which can be used as an uncertainty or anomaly score.
• Disentanglement and interpretability. By limiting the information budget, the bottleneck forces the representation to spend its capacity on the most predictive factors, which often correspond to interpretable structure.
• Generalisation bounds. The IB connects directly to generalisation theory: bounds on the generalisation gap can be expressed in terms of I(X; Z), so compressing the representation provably limits overfitting in a precise, distribution-dependent sense.

None of this is magic, and none of it is free. The bottleneck trades accuracy on the training distribution for robustness elsewhere, and the trade is only worth making when the irrelevant information in X is genuinely harmful rather than merely unused.

Limitations and common pitfalls

• Mutual information estimation is hard. The variational bounds are bounds, not equalities, and a loose bound means you are optimising a different objective than the one you wrote down. If your variational marginal r(z) is too rigid, the compression penalty is mismeasured.
• The bottleneck can discard information you later need. A representation optimised for one target Y may be useless for a related target Y'. This is the transfer learning trap: aggressive compression is task-specific, and a representation that forgot everything except what one task needed cannot be reused for another.
• Deterministic networks and the IB do not mix cleanly. For a deterministic function Z = f(X), the mutual information I(X; Z) is infinite for continuous X (or equal to H(X) for discrete X), so the Lagrangian is vacuous. Claims about compression in deterministic networks need a different formalism, such as the recent work on stochastic relaxations or on geometric notions of compression.
• Beta is not a free lunch. There is no universal value of β that works across tasks. It must be tuned, and the information plane for your problem should be inspected, not assumed.
• Do not confuse the bottleneck with ordinary regularisation. Weight decay and dropout also limit capacity, but they do so in parameter space, not information space. The bottleneck is a statement about what the representation knows, which is a stronger and more interpretable claim.

The takeaway

The Information Bottleneck asks the right question about representation learning: not "how much can this model remember?" but "how little does it need to remember to do the job?" The answer is a curve, not a number, and the interesting work is choosing where on that curve to sit.

When you train your next model, ask what it is carrying that it does not need. The chances are that a large fraction of its representation is memorised noise, and that the model would be simpler, more robust, and easier to understand if it forgot most of what it saw.

Connect with me on LinkedIn: linkedin.com/in/umeshtharukaofficial