Mutual Information and Representations

A transformer's hidden layer produces a dense, several-thousand-dimensional vector for every token. How do you know whether that vector actually encodes something real, say, the sentiment of the sentence, versus encoding nothing useful about sentiment at all, with a powerful enough downstream probe finding spurious signal anyway? Mutual information (MI) is the quantity built specifically to answer this kind of question, and it underlies a surprising amount of representation-learning and interpretability practice once you know to look for it.

Definition and two equivalent readings

I(X; Y) = H(X) - H(X | Y) = KL( P(X,Y) || P(X)P(Y) )

Two readings of the same quantity. First: how much observing Y reduces your uncertainty about X, in bits, the drop from X's unconditional entropy to its entropy once you already know Y. Second: how far the true joint distribution of X and Y sits from the hypothetical world where they are independent, measured in KL divergence (see cross-entropy-and-kl). Unlike KL itself, mutual information is symmetric: I(X;Y) = I(Y;X). It is exactly zero if and only if X and Y are independent, and it can never be negative.

MI as a probe for what a representation actually holds

A useful consequence of the data processing inequality: no function computed from a representation Z can recover more information about a target Y than Z itself carries about Y. That means I(Z; Y) is a hard upper bound on how well any classifier, however clever, can ever do at predicting Y from Z. This reframes a whole line of interpretability work, training a small "probing classifier" on frozen hidden states to check whether some property is linearly or nonlinearly recoverable, as approximately estimating a lower bound on true mutual information: a high-accuracy probe proves the information is at least present (or the probe would not find it), but a low-accuracy probe only proves the probe failed to extract it, not that the information is absent, because probes are themselves imperfect estimators.

InfoNCE: the estimator that made this tractable

Computing MI directly in high dimensions requires the true joint density, which is exactly the thing you never have. Contrastive learning objectives sidestep this. van den Oord, Li, and Vinyals, 2018 (arXiv:1807.03748) introduced InfoNCE as part of Contrastive Predictive Coding and proved it is a tractable lower bound on the true mutual information between two related views of data (in their case, a context and the correct continuation that follows it). Training a contrastive objective, pulling representations of matching pairs together while pushing apart random negatives, is, in expectation, directly maximising this lower bound. This is the actual theoretical justification behind most modern contrastive representation learning, not an ad hoc heuristic: "make positive pairs similar and negative pairs dissimilar" is optimising a proxy for mutual information because true MI cannot be computed.

The information bottleneck framing

Tishby, Pereira, and Bialek's information bottleneck principle frames representation learning itself as a tradeoff: a good representation Z of input X for predicting target Y should maximise I(Z; Y) (keep everything predictive) while minimising I(Z; X) (discard everything else, i.e. compress). This casts representation learning as a compression-versus-relevance tradeoff, and connects directly back to the theme in language-modelling-as-compression: good modelling and good compression describe the same underlying act, viewed through different formalisms.

When it falls down

• Neural MI estimators are badly biased with finite samples. Several popular contrastive-objective-based MI estimators, InfoNCE included, have been shown to be provably upper-bounded by log(batch_size), meaning reported "mutual information" from a small-batch neural estimator can sit far below the true value simply as an artefact of batch size, not a fact about the representation.
• MI measures presence of information, not ease of access. MI is invariant to any invertible reparameterisation of either variable, so "the representation has high MI with the label" does not imply a simple probe can extract it; the extracting function might need to be as complex as the original task itself.
• The "compression phase" story is contested, not settled. Claims that deep network training dynamics show a distinct information-bottleneck compression phase that explains generalisation have been disputed by follow-up work, which found the effect depends heavily on activation function choice rather than being a universal property of deep learning. Treat it as an open research question, not an established mechanism.

Further reading

• van den Oord, Li, Vinyals, 2018, Representation Learning with Contrastive Predictive Coding, arXiv:1807.03748 - InfoNCE as a tractable mutual-information lower bound.
• Tishby, Pereira, Bialek, The Information Bottleneck Method, arXiv:physics/0004057 - the compression-versus-relevance framing of representation learning.