On Mutual Information Maximization for Representation Learning

Michael Tschannen; Josip Djolonga; Paul K. Rubenstein; Sylvain Gelly; Mario Lucic

Source-linked AI summary

On Mutual Information Maximization for Representation Learning

Michael Tschannen, Josip Djolonga, Paul K. Rubenstein, Sylvain Gelly, Mario Lucic

arXiv:1907.13625v2Jul 31, 2019cs.LGstat.ML

TL;DR

The paper examines why mutual-information maximization succeeds despite MI’s estimation difficulty and invariance to invertible transformations. It argues that encoder and estimator inductive biases, plus a connection to deep metric learning, better explain observed representation quality.

Problem
MI is difficult to estimate and invariant under invertible reparameterizations, so maximizing it alone may not ensure useful representations.
Method
The paper empirically studies how encoder architectures, critics, and MI estimators affect representations, and connects INCE to deep metric learning.
Results
The success of MI-based methods cannot be attributed to MI alone; looser bounds can produce better representations, while estimator and architecture choices strongly matter.
Takeaways & Limitations
Representation-learning theory and methods should account for representation geometry, function families, computational constraints, and negative-sampling strategy.
Takeaways & Limitations
Under linear evaluation, MI lower-bound maximization can yield bad representations, and MI may not correlate with downstream performance.

Abstract

from arXiv · show

Many recent methods for unsupervised or self-supervised representation learning train feature extractors by maximizing an estimate of the mutual information (MI) between different views of the data. This comes with several immediate problems: For example, MI is notoriously hard to estimate, and using it as an objective for representation learning may lead to highly entangled representations due to its invariance under arbitrary invertible transformations. Nevertheless, these methods have been repeatedly shown to excel in practice. In this paper we argue, and provide empirical evidence, that the success of these methods cannot be attributed to the properties of MI alone, and that they strongly depend on the inductive bias in both the choice of feature extractor architectures and the parametrization of the employed MI estimators. Finally, we establish a connection to deep metric learning and argue that this interpretation may be a plausible explanation for the success of the recently introduced methods.

1 INTRODUCTION

Unsupervised representation learning seeks useful lower-dimensional mappings from unlabeled data, but lacks a natural objective. MI-based methods face estimation and invariance problems, yet show strong empirical results whose success the paper argues depends on more than MI alone.

Unsupervised representation learning learns a usually lower-dimensional mapping that supports unknown supervised tasks using fewer labels from unlabeled data.
InfoMax-inspired methods choose representations by maximizing mutual information between inputs and representations, sometimes under structural constraints.MI measures information obtained about one random variable by observing another.
MI is invariant under smooth invertible reparametrizations, so its value does not determine a unique representation geometry.
High-dimensional MI estimation is notoriously difficult, leading practitioners to maximize tractable lower bounds instead.
Despite these challenges, MI-maximization methods have produced promising empirical results in representation learning.
The paper argues that these methods’ success depends on encoder architectures and MI-estimator parametrizations, and that tighter MI bounds can produce worse representations.

2 BACKGROUND AND RELATED WORK

Recent image representation methods maximize estimated MI between multiple views, using flexible encoder arrangements and lower-bound estimators such as InfoNCE and NWJ. These formulations reduce estimation dimensionality while introducing architectural and critic choices that shape the objective.

The paper focuses on unsupervised image methods that estimate MI between different, possibly overlapping views of the same image.Views may include image regions, augmentations, or modalities.
The multiview objective estimates MI between learned representations, with encoder function classes specifying structural constraints and the objective forming a lower bound on InfoMax.
Estimating MI between low-dimensional representations is advantageous because it avoids estimating MI directly in the typically higher-dimensional input space.
Multiview formulations support global-versus-local features, different augmentations, and multiple image modalities.
Because MI estimation is fundamentally limited, recent methods commonly optimize lower bounds based on distinguishing joint samples from samples drawn from the product of marginals.
InfoNCE uses a critic to identify which of K candidate views was jointly drawn with each input, assigning higher scores to matched pairs than mismatches.
The paper also considers NWJ, while common critic parametrizations include bilinear, separable, and concatenated forms.

3 BIASES IN APPROXIMATE INFORMATION MAXIMIZATION

The paper shows that approximate MI maximization favors representations through estimator, critic, and encoder biases rather than MI alone. Across experiments, downstream quality depends strongly on architecture and parametrization, while high or matched MI estimates do not guarantee useful representations.

3 BIASES IN APPROXIMATE INFORMATION MAXIMIZATION: The paper therefore argues that practical success is only loosely connected to MI and may instead reflect a connection to triplet losses in deep metric learning.The proposed alternative explanation concerns how practical estimator and architecture choices shape representation geometry.
3 BIASES IN APPROXIMATE INFORMATION MAXIMIZATION: Approximate estimators provide gradient feedback that selects representations useful for linear classification even when exact MI gradients would vanish.For invertible encoders, the true MI is constant across parameters, but INCE and INWJ increase alongside downstream accuracy during training.
3 BIASES IN APPROXIMATE INFORMATION MAXIMIZATION: MI maximization can improve or worsen downstream representations among invertible encoders, although all preserve the same true MI.Some invertible encoders improve linear classification, while others perform worse than raw pixels despite globally maximal MI.
3 BIASES IN APPROXIMATE INFORMATION MAXIMIZATION: When encoders can be non-invertible, increasing estimated MI improves accuracy but makes the encoder increasingly ill-conditioned and harder to invert.The estimator favors attenuating noise over preserving it through a well-conditioned mapping such as the identity.
3 BIASES IN APPROXIMATE INFORMATION MAXIMIZATION: Critic architecture changes representation quality: bilinear and separable critics outperform MLP critics despite reaching lower estimated MI values.Thus, tighter estimated MI bounds from higher-capacity critics do not necessarily yield better downstream representations.
3 BIASES IN APPROXIMATE INFORMATION MAXIMIZATION: Encoder architecture can matter more than the estimator: ConvNet encoders achieve higher classification accuracy than MLP encoders while matching the same MI lower bound.This comparison uses a bilinear critic and targets the same estimated objective value for both architectures.

4 CONNECTION TO DEEP METRIC LEARNING AND TRIPLET LOSSES

The paper shows that InfoNCE is closely connected to triplet-based metric learning, making negative-sample selection and critic structure central design choices rather than merely MI-estimation details.

The metric learning view: Metric learning seeks embeddings where each anchor is closer to its positive instance than to its negative instance.The paper describes this using ℓ2 distances between represented triplets.
Connection to InfoNCE: InfoNCE with a symmetric separable critic and shared encoder is equivalent to metric learning based on a multi-class K-pair loss.With asymmetric critics or separate view encoders, the corresponding objective becomes an asymmetric variant.
The metric learning view: The inner-product critic acts as a simple metric on the embedding space, although its associated MI lower bound can be weak.This reframes the practical role of the critic from approximating MI to shaping embedding geometry.
InfoNCE and the importance of negative sampling: InfoNCE is an MI lower bound only when negative samples are drawn from the true marginal distribution.Dependent negative sampling can invalidate that interpretation.
InfoNCE and the importance of negative sampling: Increasing negative samples often helps, but theory and experiments also show that it can fail to help or even deteriorate performance.The resulting negative-sampling intricacies remain a key research challenge.

5 CONCLUSION

The conclusion argues that approximate MI maximization alone does not ensure useful representations and that estimator, critic, encoder, sampling, and evaluation choices jointly matter. It proposes metric learning as a plausible lens for designing improved unsupervised methods.

Conclusion: Maximizing MI lower bounds can produce bad representations under the common linear evaluation protocol.The paper reports that looser bounds can sometimes yield better representations.
Alternative measures of information: MI is hard to estimate, invariant to bijections, and insufficiently aligned with downstream representation quality.The paper therefore calls for information measures that also account for representation geometry and observer constraints.
A holistic view: Downstream performance depends on the balance among critics, encoders, and evaluation protocols, which the paper argues should be co-designed.The authors describe this holistic view as under-explored because theory and systematic studies remain limited.
Going beyond the widely used linear evaluation protocol: Highly nonlinear evaluation may improve downstream performance but undermines the goal of efficiently transferable representations.The paper notes that linear evaluation is also connected to reduced downstream sample complexity.
Systematic investigations into design decisions that matter: A metric-learning perspective could motivate methods that prioritize influential design choices such as negative sampling over explicitly estimating MI.The paper cites a multi-view video method as an example of a similar approach without independently sampled negatives or an MI-based motivation.

A RELATION BETWEEN (2) AND THE INFOMAX OBJECTIVE

For two arbitrary functions of X, the mutual information between their outputs is bounded above by the information that X contains about the pair of outputs. The proof applies the data processing inequality twice.

Proposition 1: For X1 = g1(X) and X2 = g2(X), I(X1; X2) ≤ I(X; (X1, X2)).This is the stated proposition for arbitrary functions g1 and g2.
Proof: The proof relies on two applications of the data processing inequality along appropriate Markov chains.That inequality states that processing cannot increase mutual information in a Markov relation X → Y → Z.
Proof: Because X1 and X2 are conditionally independent given X, the data processing inequality gives I(X1; X2) ≤ I(X; X1).The Markov relation is X1 ← X → X2, equivalently X1 → X → X2.
Proof: Since X → (X1, X2) → X1, the data processing inequality gives I(X; X1) ≤ I(X; (X1, X2)).Combining this with the first inequality proves the proposition.

B EXPERIMENT DETAILS: ADVERSARIALLY TRAINED ENCODER (SECTION 3.1)

The experiment adversarially trains an invertible RealNVP encoder using a linear classifier attached to its logits, with encoder and classifier optimized through opposing objectives.

Training setup: The encoder g1 uses the same RealNVP architecture as in the first experiment, while g2 is not modeled.A linear layer on g1(X(1)) produces 10 logits.
Training setup: The linear layer minimizes cross-entropy against the true label of the source image for X(1).The passage states that encoder parameters are trained adversarially, though the remainder is truncated.

C CONNECTION BETWEEN METRIC LEARNING AND INFONCE

This section rewrites INCE, presenting an alternative form of the estimator.

INCE is rewritten in an alternative form.
The rewrite concerns the mathematical expression for INCE.
The section introduces the rewritten form before subsequent analysis.

D INFONCE UNDER NON-I.I.D. SAMPLING

The standard InfoNCE lower-bound proof depends critically on negative samples being drawn independently from the true marginal distribution. With dependent negatives, the resulting estimator is neither a lower nor an upper bound on true MI.

The InfoNCE proof starts from the NWJ lower bound on KL divergence.The bound holds for any critic function.
A specific critic choice and averaging over K samples yield the familiar InfoNCE lower bound.The averaging step cancels the constant term under the identical-distribution assumption.
The proof uses i.i.d. paired samples to equate I(X_i; Y_i) with I(X_1:K; Y_i).This equality enables multiple samples to estimate MI between two variables.
Dependent negative samples instead give I(X_i; Y_i) ≤ I(X_1:K; Y_i).The remaining proof steps still hold, but the key equality does not.
Under dependent negative sampling, InfoNCE is neither a lower nor an upper bound on true MI.

E EXPERIMENT DETAILS: NON-I.I.D. SAMPLING (SECTION 4)

The experiment uses coupled Gaussian batches with a shared latent variable and independently sampled noise, compares them with i.i.d. pairs, and evaluates InfoNCE with a neural critic.

The experiment models each pair as a shared Gaussian latent variable plus Gaussian noise.The true MI is analytically calculable because the pair is jointly Gaussian with known covariance.
Within each non-i.i.d. batch, the noise is sampled independently while the shared latent variable is sampled once.
The study compares non-i.i.d. batches with identically distributed i.i.d. sampled pairs.
The critic is a five-hidden-layer MLP with ten units per layer, ReLU activations, and a final linear layer.InfoNCE is maximized using non-i.i.d. samples with batch size 128.
The batch size is chosen so I(X; Y) ≤ log K does not trivially lower-bound the true MI.The InfoNCE bound always satisfies INCE ≤ log K for batch size K.

F ADDITIONAL FIGURES

The additional figures examine Jacobian conditioning, INCE values across critic architectures, and estimator behavior under target-matched training losses.

Figure 5 plots Jacobian condition-number percentiles for inputs sampled from the data distribution.The empirical distribution aggregates condition numbers from 128 sampled inputs across runs.
Figure 5 indicates that g1 becomes increasingly ill-conditioned as its Jacobian condition number deteriorates.
Figure 6 tests INCE values for MLP encoders and different critic architectures.
Figure 7 tests loss across encoder architectures and MI estimators with a bilinear critic.The critic is trained to match a target estimator value t using L_t(g1, g2) = |I_EST(g1(X^(1)); g1(X^(2))) − t|.

G RESULTS FOR THE EXPERIMENTS FROM SEC. 3.2 AND 3.3 ON CIFAR10

On CIFAR10, the experiments reproduce the qualitative MNIST behavior: critic and encoder architectures affect downstream accuracy and testing IEST, with ConvNet encoders outperforming MLP encoders.

Experimental setup: The CIFAR10 experiments use adapted encoders for flattened 32 × 14 × 3 image halves, a learning rate of 10^-5, and three times as many training iterations.The encoder and critic architectures otherwise remain the same as in the earlier experiments.
Overall conclusion: The qualitative behavior of encoder and critic architectures on CIFAR10 is very similar to that observed on MNIST.The paper states that the MNIST conclusions therefore carry over to CIFAR10.
Critic architectures: Bilinear and separable critics achieve higher downstream accuracy than MLP critics while reaching lower INWJ.INWJ also exhibits high variance in this experiment.
Encoder architectures: For a fixed estimator and target IEST, ConvNet encoders clearly outperform MLP encoders in downstream testing accuracy.The comparison uses a bilinear critic trained to match the target IEST.