[Paper Review] Learning deep representations by mutual information estimation and maximization

R Devon Hjelm, Alex Fedorov+5 more|arXiv (Cornell University)|Aug 20, 2018

Adversarial Robustness in Machine LearningComputer Science1,393 citations

TL;DR

DIM learns unsupervised representations by maximizing mutual information between inputs and encoder outputs, with global and local MI objectives and adversarial prior matching to shape representations.

ABSTRACT

In this work, we perform unsupervised learning of representations by maximizing mutual information between an input and the output of a deep neural network encoder. Importantly, we show that structure matters: incorporating knowledge about locality of the input to the objective can greatly influence a representation's suitability for downstream tasks. We further control characteristics of the representation by matching to a prior distribution adversarially. Our method, which we call Deep InfoMax (DIM), outperforms a number of popular unsupervised learning methods and competes with fully-supervised learning on several classification tasks. DIM opens new avenues for unsupervised learning of representations and is an important step towards flexible formulations of representation-learning objectives for specific end-goals.

Motivation & Objective

• Motivate unsupervised learning of representations by maximizing mutual information between inputs and encoder outputs.
• Show that leveraging local input structure (patch-level MI) yields representations better for classification than global MI alone.
• Introduce prior matching to constrain representations to desirable statistical properties via adversarial learning.
• Demonstrate that combining global/local MI objectives with prior matching yields competitive or superior results on several datasets.

Proposed method

• Estimate and maximize mutual information between input X and encoder output Y using MI estimators (MINE-inspired) with a DV/JSD/infoNCE flavor.
• Optionally maximize MI between Y and local input patches C^{(i)}(X) to emphasize local structure (local DIM).
• Train a discriminator to match the encoder’s output distribution U_{ψ,P} to a prior V, enforcing desired statistics (adversarial prior matching).
• Combine global MI, local MI, and prior matching into the DIM objective with tunable weights α, β, γ as in the joint objective (Eq. 8).
• Compare MI estimators (DV, JSD, infoNCE) and analyze robustness to negative samples and architectural choices.
• Explore occlusion and coordinate-predictive augmentations to further exploit structure and improve classification performance.

Experimental results

Research questions

• RQ1 Can mutual information between inputs and encoder outputs yield useful representations without supervision?
• RQ2 Does emphasizing local structure via local MI maximization improve downstream classification performance compared to global MI?
• RQ3 Does aligning the encoder’s output distribution to a prior via adversarial matching improve representation properties such as independence or disentanglement?
• RQ4 How do different MI estimators (DV, JSD, infoNCE) affect learning stability and downstream task performance?
• RQ5 What is the impact of augmentations like occlusion or coordinate prediction on representation quality?

Key findings

• DIM with local MI (DIM(L)) substantially outperforms other unsupervised methods on several datasets and can approach or surpass fully supervised performance in some settings.
• DIM with global MI (DIM(G)) is competitive with some models but generally lags behind DIM(L) and generative baselines on larger classification tasks.
• Using adversarial prior matching to constrain the representation improves its statistical properties and aligns it with desired priors, as shown in MI/NDM analyses.
• infoNCE often yields strong downstream performance and can be more robust to negative sample counts than DV-based formulations, while JSD remains competitive in certain regimes.
• Occlusion and coordinate-prediction augmentations further boost classification accuracy for DIM, indicating that leveraging spatial structure benefits representation quality.