DIM: Learning Deep Representations by Mutual Information Estimation and Maximization
By: Sherwin Chen
Tags: Machine Learning, Artificial Intelligence, Unsupervised Learning, Featured, Deep Learning

Introduction

This is our second article of the series about mutual information. In the previous articles, we have seen how to maximizes the mutual information between two variables via the MINE estimator and some practical applications of maximizing mutual information. In this article, we focus on representation learning with mutual information maximization. Specifically, we will discuss an adversarial architecture for representation learning and two other objectives of mutual information maximization that has been experimentally shown to outperform MINE estimator for downstream tasks.

This article is organized into four parts. First, we briefly review MINE and introduce two additional methods for mutual information maximization that have been shown to experimentally outperform MINE. Next, we address the deficiency of traditional pixel-level representation learning methods, disclosing the reason why we maximize mutual information for representation learning. We then discuss the components of Deep InfoMax(DIM) along with the corresponding objectives. We will prove the validity of the JSD-based objective in the end for completeness.

Methods for Mutual Information Maximization

MINE defines the greatest lower-bound on mutual information based on the Donsker-Varadhan representation(DV) of the KL divergence:

DV objective

As we are primarily interested in maximizing mutual information, not its precise value, the authors suggest to replace the KL divergence with the Jensen-Shannon Divergence(JSD) and thereby have the objective as below

JSD-based objective. We will prove these in the end

Another objective, InfoNCE proposed by Aaron van den Oord et al. [2], could also be used as a bound on mutual information:

InfoNCE objective. This is just Noise Contrastive Estimation(NCE). For its connection with mutual information, please refer to Appendix A.1 of Representation Learning with Contrastive Predictive Coding

Comparison between DV, JSD, and InfoNCE

Classification accuracy with representation learned using different objectives. Source: Learning Deep Representations by Mutual Information Estimation and Maximization

MINE(or the DV objective) is an estimator for mutual information, but in practice, it does not work well because it is potentially unbounded and its direct gradient is biased(though methods to alleviate these issues have been proposed in the original MINE paper, see our previous post for more details).

The InfoNCE objective experimentally performs best on downstream tasks among the three with large negative samples, but the performance drops quickly as we reduce the number of samples used for the estimation.

The JSD-based objective is insensitive to the number of negative samples and performs well in practice. Unlike InfoNCD, which includes the positive sample in the second term, excluding positive samples from the product of the marginals improves the downstream classification accuracy.

Deficiency of Traditional Pixel-Level Representation Learning

Pixel-level representation-learning algorithms could be detrimental when only a small fraction of the bits of the signal actually matter at a semantic level. Such issues become more apparent when it comes to the representation learned by generative models: as the dimension of the data becomes large, the objectives of generative models become dominated by the pressure to represent all of the pixels well, whereas it may not leave sufficient dimensions to capture the important characteristics of the distribution. In particular, generative models must represent all statistical variability in the data to generate (or even discriminate) properly, and much of this variability could be unimportant and even detrimental to the suitability of the representation for downstream tasks.

R Devon Hjelm et al. [1] argue that representation learning should be more directly in terms of their information content and statistical or structural constraints. To address the first quality, the authors propose learning unsupervised representations by maximizing mutual information. To address the second, they consider adversarially matching it prior to imposing structural constraints on the learned representation.

Deep InfoMax

Deep InfoMax(DIM) defines an encoder network that distills a high-level representation and multiple discriminators for different objectives. In this section, we will walk through all these networks, along with their corresponding objectives.

Encoder Network

encoder network

The encoder network, E_ψ=f_ψ∘ C_ψ, first encodes the input X​ into an M×M ​ feature map, then summarizes this feature map into a high-level feature vector, ​Y=E_ψ(X).

Global Mutual Information Maximization

The global structure for representation learning with mutual information maximization

The global structure first flattens the M​×M feature map and concatenate it with ​Y. We then pass this to the discriminator D_ω to maximizes the mutual information between the​ feature map and Y​ using the following objective

​I could be any objective function discussed in the first section.

Notice that we could generate negative samples by combining Y​ with different images in the same batch to save additional computational overhead.

Local Mutual Information Maximization

The global mutual information maximization may introduce some dependence irrelevant to the task. For example, trivial noise local to pixels is useless for image classification, so a representation may not be benefit from this information if the end-goal is to classify. Moreover, since the capacity of the high-level representation is fixed, such irrelevant information may squeeze out some valuable information.

To eliminate the unwanted local noise, the authors suggest maximizing the average mutual information between the high-level representation Y​ and local patches of image (feature vectors in the M×M​ feature map). Because now the high-level representation Y​ is encouraged to have high mutual information with all patches, this favors encoding aspects of the data that share across patches. Experiments also verify the intuition that this encourages the encoder to prefer information shared across the input and significantly boosts classification performance.

Concat-and-convolve architecture for local mutual information maximization

The authors propose two different architectures for the local objective. The first, depicted in the above figure, concatenate Y​ with the M×M feature map at every location. A 1×1​ convolutional discriminator is then used to score every feature vector in the location-wise concatenated feature map.

Encode-and-dot-product architecture for local mutual information maximization

The other architecture, shown in the above figure, further encodes the ​ M×M feature map and Y​ through a​ convolutional network and a fully-connected network, respectively. We then take the dot product between the feature at each location of the feature map encoding and the encoded global vector Y’ for scores.

The objective used for local mutual information maximization is defined as

I^i could be any objective function discussed in the first section.

To reduce computational overhead, the fake samples are generated by shuffling the feature map at each location across a batch independently.

Match Representation to a Prior Distribution

prior matching

To satisfy the structural constraints (i.e., match the representation to a prior distribution ​P(Z)), we define a new discriminator which distinguishes if the input distribution is from the output of the encoder or from the prior distribution. The objective is simply a cross-entropy

The reason why we impose the structural constraints on Y​ is that, during the course of approximating to the prior ​Z, we expect E_ψ​ to learn a representation with some desirable properties such as independence and disentanglement if the prior Z​ inherently owns such properties.

Put All Together

The objective for the whole DIM is

where ​α, β​ and γ​ are hyperparameters. The following figure shows some ablation studies done by the authors, where y​-axis is β​. In particular, we can see good classification performance is highly dependent on the local term, β​, while good reconstruction highly depends on the global term ​α. This aligns with our discussion in this section. However, a small amount of α​ also helps in classification accuracy and a small amount of ​β improves reconstruction.