Modern Generalization

Representation Learning Theory

What makes a good learned representation: the information bottleneck, contrastive learning, sufficient statistics, rate-distortion theory, and why representation learning is the central unsolved problem of deep learning.

AdvancedTier 2Current~60 min

Prerequisites

Information Theory Foundations Variational Autoencoders

Quiz (1)

Why This Matters

A neural network's intermediate layers compute a representation $Z = f(X)$ of the input $X$ . The quality of this representation determines transfer learning performance, robustness, and generalization. Two models with identical test accuracy can have wildly different representations, and the one with better representations will transfer better to new tasks, be more robust to distribution shift, and require less data for fine-tuning.

The theory of representation learning asks: given an input $X$ and a task $Y$ , what properties should $Z = f(X)$ have? Multiple frameworks give partial answers, but no single theory is complete.

What Makes a Good Representation

Three desirable properties appear across frameworks:

Sufficiency. $Z$ should contain all information in $X$ that is relevant to $Y$ . Formally, $X \to Z \to Y$ forms a Markov chain such that $I(Z; Y) = I(X; Y)$ , where $I$ denotes mutual information.

Minimality. $Z$ should discard information in $X$ that is irrelevant to $Y$ . This means $I(Z; X)$ should be as small as possible, subject to the sufficiency constraint.

Disentanglement. Different generative factors of the data should map to independent components of $Z$ . If $X$ is an image of a face, then pose, lighting, and identity should be encoded in separate dimensions of $Z$ .

Information Bottleneck

Definition

Information Bottleneck

The information bottleneck (Tishby et al., 2000) finds a representation $Z$ that solves:

$\min_{p(z|x)} I(X; Z) - \beta \, I(Z; Y)$

where $\beta > 0$ controls the trade-off between compression ( $I(X; Z)$ small) and prediction ( $I(Z; Y)$ large).

Theorem

Information Bottleneck Optimality

Statement

The optimal IB representation $Z^*$ satisfies:

$p(z|x) = \frac{p(z)}{Z(\beta, x)} \exp\left(\beta \, D_{\text{KL}}[p(y|x) \| p(y|z)]\right)$

where $Z(\beta, x)$ is a normalizing constant. The IB curve (plotting $I(Z;Y)$ vs $I(Z;X)$ as $\beta$ varies) is concave and traces the Pareto frontier of the compression-prediction trade-off.

Intuition

The optimal representation keeps information about $X$ only to the extent that it helps predict $Y$ . The KL divergence term measures how much predictive power about $Y$ is lost by compressing $X$ into $Z$ . The parameter $\beta$ controls how aggressively you compress.

Proof Sketch

This is a constrained optimization problem solvable via Lagrange multipliers. The Lagrangian is $I(X; Z) - \beta I(Z; Y)$ . Taking functional derivatives with respect to $p(z|x)$ and setting to zero yields the self-consistent equation above. The concavity of the IB curve follows from the concavity of mutual information in the conditional distribution.

Why It Matters

The IB framework formalizes the intuition that good representations compress irrelevant information while preserving task-relevant information. Tishby et al. (2017) conjectured that deep learning implicitly performs IB optimization during training (the "compression phase" hypothesis). This claim generated significant debate.

Failure Mode

The IB framework requires knowing $p(X, Y)$ , which is unavailable in practice. Variational approximations (VIB, Alemi et al., 2017) replace exact mutual information with tractable bounds, but the bounds can be loose. The "compression phase" hypothesis has been shown to depend on the activation function: networks with ReLU do not always exhibit compression (Saxe et al., 2018).

Contrastive Learning

Contrastive learning learns representations by pulling together representations of "positive pairs" (semantically similar inputs) and pushing apart "negative pairs" (dissimilar inputs).

Definition

InfoNCE Loss

Given a query $x$ , a positive example $x^+$ , and $K$ negative examples $\{x_1^-, \ldots, x_K^-\}$ , the InfoNCE loss is:

$\mathcal{L}_{\text{InfoNCE}} = -\log \frac{\exp(f(x)^T f(x^+) / \tau)}{\exp(f(x)^T f(x^+) / \tau) + \sum_{k=1}^{K} \exp(f(x)^T f(x_k^-) / \tau)}$

where $f$ is the encoder and $\tau$ is a temperature parameter.

Theorem

InfoNCE Bounds Mutual Information

Statement

The InfoNCE loss provides a lower bound on the mutual information between $X$ and $X^+$ :

$I(X; X^+) \geq \log(K+1) - \mathcal{L}_{\text{InfoNCE}}$

The bound becomes tighter as $K \to \infty$ , but is capped at $\log(K+1)$ regardless of the true mutual information.

Intuition

InfoNCE is a $(K+1)$ -way classification problem: identify the positive from $K$ negatives. A perfect classifier achieves loss $0$ , giving the bound $I \geq \log(K+1)$ , which is the capacity of the classification task. You cannot estimate mutual information higher than $\log(K+1)$ with only $K$ negatives, no matter how good the encoder is.

Proof Sketch

The InfoNCE loss is the cross-entropy for the $(K+1)$ -way classification problem. The Bayes-optimal classifier uses the density ratio $p(x^+|x) / p(x^+)$ . Substituting the optimal classifier into the loss gives $\log(K+1) - I(X; X^+)$ when $I(X;X^+) \leq \log(K+1)$ , yielding the bound after rearranging.

Why It Matters

This result explains the empirical observation that contrastive learning needs large batch sizes (large $K$ ). SimCLR uses batches of 4096-8192. With small $K$ , the bound is loose and the encoder cannot distinguish fine-grained features. With large $K$ , the bound tightens and the encoder must learn more detailed representations to discriminate positives from negatives.

Failure Mode

The bound saturates at $\log(K+1)$ . If the true mutual information is much larger (e.g., for high-resolution images), you need exponentially many negatives to get a tight bound. Also, the bound says nothing about which information is captured. The encoder might learn features that distinguish positives from negatives without capturing task-relevant information.

Sufficient Statistics View

Definition

Sufficient Representation

A representation $Z = f(X)$ is sufficient for task $Y$ if $I(Z; Y) = I(X; Y)$ , equivalently if $Y \perp X \mid Z$ . A sufficient representation captures all task-relevant information and discards nothing useful.

A sufficient representation is the ideal. In practice, you want a representation that is approximately sufficient for a family of downstream tasks, not just one. This connects to the notion of a universal representation: one that is sufficient for many tasks simultaneously.

Rate-distortion connection. Rate-distortion theory (Shannon, 1959) asks: what is the minimum number of bits needed to describe $X$ such that the reconstruction error is at most $D$ ? The information bottleneck extends this by replacing reconstruction error with prediction error. The IB curve is a generalization of the rate-distortion curve from compression theory to supervised learning.

Common Confusions

Watch Out

Disentanglement is not well-defined without inductive bias

Locatello et al. (2019) proved that unsupervised disentanglement is impossible without inductive biases: for any disentangled representation, there exists a generative model consistent with the data that maps to an entangled representation. You need assumptions about the data-generating process (e.g., independence of factors) or supervision to achieve disentanglement.

Watch Out

Maximizing mutual information is not always good

A representation that maximizes $I(Z; X)$ is just the identity function. You want to maximize $I(Z; Y)$ while keeping $I(Z; X)$ small. Representations that capture too much about $X$ encode noise and spurious correlations that hurt generalization.

Watch Out

Contrastive learning does not require labels

Contrastive learning constructs positive pairs from data augmentation (two views of the same image) rather than from labels. This is self-supervised. However, the quality of the representation depends heavily on the augmentation strategy: the augmentations implicitly define what information the representation should be invariant to.

Exercises

ExerciseCore

Problem

In the InfoNCE loss with $K = 255$ negative samples, what is the maximum mutual information you can estimate? If your dataset has images with true mutual information of 20 nats between two augmented views, approximately how many negatives do you need for a tight bound?

ExerciseAdvanced

Problem

Prove that a minimal sufficient statistic $T(X)$ for $Y$ achieves the endpoint of the information bottleneck curve (maximum $I(T; Y)$ with minimum $I(T; X)$ ). Why does this imply that the IB objective at $\beta \to \infty$ recovers the minimal sufficient statistic?

References

Canonical:

Tishby, Pereira, Bialek, "The Information Bottleneck Method" (2000)
Oord, Li, Vinyals, "Representation Learning with Contrastive Predictive Coding" (2018)

Current:

Locatello et al., "Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations" (ICML 2019)
Alemi et al., "Deep Variational Information Bottleneck" (ICLR 2017)

Last reviewed: April 2026

Prerequisites

Foundations this topic depends on.

Information Theory FoundationsLayer 0B
Variational AutoencodersLayer 3
AutoencodersLayer 2
Feedforward Networks and BackpropagationLayer 2
Differentiation in RnLayer 0A
Sets, Functions, and RelationsLayer 0A
Basic Logic and Proof TechniquesLayer 0A
Matrix CalculusLayer 1
The Jacobian MatrixLayer 0A
The Hessian MatrixLayer 0A
Activation FunctionsLayer 1
Convex Optimization BasicsLayer 1
Matrix Operations and PropertiesLayer 0A
Vectors, Matrices, and Linear MapsLayer 0A
Maximum Likelihood EstimationLayer 0B
Common Probability DistributionsLayer 0A