L EARNING DEEP REPRESENTATIONS BY MUTUAL INFORMATION ESTIMATION AND MAXIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 inﬂuence 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 performs favorably with fully-supervised learning on several classiﬁcation tasks in the speciﬁc settings studied.
DIM opens new avenues for unsupervised learning of representations and is an important step towards ﬂexible formulations of representation-learning objectives for speciﬁc end-goals.
I NTRODUCTION
One core objective of deep learning is to discover “good” representations, and the simple idea explored here is to train a representation-learning function (i.e., an encoder) to maximize the mutual information (MI) between its inputs and outputs.
MI is notoriously difﬁcult to compute, particularly in continuous and high-dimensional settings.
Fortunately, recent advances enable effective computation of MI between high dimensional input/output pairs of deep neural networks (Belghazi et al., 2018).
We leverage MI estimation for representation learning, but we will also show that, depending on the downstream task, maximizing MI between the complete input and the encoder output (i.e., global
MI) is often not sufﬁcient for learning useful representations.
Rather, structure matters : maximizing the average MI between the representation and local regions (e.g., representations of patches, rather than the complete image) of the input can greatly improve the representation’s quality for, e.g., classiﬁcation tasks, while global MI plays a stronger role in the ability to reconstruct the full input given the representation.
Usefulness of a representation is not just a matter of information content: representational characteristics like independence also play an important role (Gretton et al., 2012; Hyv ¨ arinen & Oja, 2000; Hinton, 2002; Schmidhuber, 1992; Bengio et al., 2013; Thomas et al., 2017).
We combine MI maximization with prior matching in a manner similar to adversarial autoencoders (AAE, Makhzani et al., 2015) to constrain representations according to desired statistical properties.
This approach is closely related to the infomax optimization principle (Linsker, 1988; Bell & Sejnowski, 1995), so we call our method Deep InfoMax (DIM).
Our main contributions are the following:
• We formalize Deep InfoMax (DIM), which simultaneously estimates and maximizes the mutual information between input data and learned high-level representations.
• Our mutual information maximization procedure can prioritize global or locally-consistent information, which we show can be used to tune the suitability of learned representations for classiﬁcation or reconstruction-style tasks.
• We use adversarial learning ( ` a la Makhzani et al., 2015) to constrain the representation to have desired statistical characteristics speciﬁc to a prior.
• We introduce two new measures of representation quality, one based on Mutual Information Neural Estimation (MINE, Belghazi et al., 2018) and a neural dependency measure (NDM) based on the work by Brakel & Bengio (2017), and we use these to bolster our comparison of DIM to different unsupervised methods.
R ELATED W ORK
There are many popular methods for learning representations.
Some of the older methods, such as independent component analysis (ICA, Bell & Sejnowski, 1995) and self-organizing maps (Kohonen, 1998) generally do not compete with deep neural networks in representational capacity.
Other approaches include deep volume-preserving maps (Dinh et al., 2014; 2016), deep clustering (Xie et al., 2016; Chang et al., 2017), noise as targets (NAT, Bojanowski & Joulin, 2017), self-supervised methods (Doersch & Zisserman, 2017), and co-learning (Dosovitskiy et al., 2016; Sajjadi et al., 2016).
Generative models are also commonly used as methods for building representations (Vincent et al., 2010;
Kingma et al., 2014; Salimans et al., 2016; Rezende et al., 2016; Donahue et al., 2016), and mutual information (MI) plays an important role in their representational quality.
In generative models that rely on reconstruction (e.g., denoising, variational, and adversarial autoencoders, Vincent et al., 2008; Rifai et al., 2012; Kingma & Welling, 2013; Makhzani et al., 2015), the negative reconstruction error can be related to the MI in the encoder as,
[Equation]
where X and Y are random variables corresponding to the input and an intermediate representation (e.g., the bottleneck), R e,d ( X ) is the reconstruction error of X , and H e ( Y ) is the marginal entropy of the encoder output.
Thus, under reasonable learning conditions 1 , models with reconstruction-type objectives provide some guarantees on the amount of information encoded in their intermediate representations.
Similar guarantees exist for bi-directional adversarial models (Dumoulin et al., 2016; Donahue et al., 2016), which adversarially train an encoder / decoder to match their respective joint distributions (which typically increases the marginal entropy in Equation 1) or to minimize the reconstruction error (Chen et al., 2016).
Mutual-information estimation Mutual information-based objectives have a long tradition in unsupervised learning of features.
The infomax optimization principle (Linsker, 1988; Bell &
Sejnowski, 1995), as prescribed for neural networks, advocates maximizing MI between the input and output.
This is the basis of numerous ICA algorithms, some of which are nonlinear (Hyv¨arinen & Pajunen, 1999; Almeida, 2003), but none of which are general enough to apply to deep networks.
Mutual information neural estimation (MINE, Belghazi et al., 2018) learns a neural estimate of the MI of continuous variables, is strongly consistent, and can be used to learn better implicit bi-directional generative models.
Deep InfoMax (DIM) follows MINE in this regard, though we ﬁnd that the generator is unnecessary.
In addition, it is not necessary to use the exact KL-based formulation of MI: for example, a simple classiﬁer based on the Jensen-Shannon divergence (JSD) is both more stable and provides better results.
This is good as the KL-based (continuous) MI is unbounded, while the
JSD is bounded.
DIM generally admits a variety of MI estimators (which we will show), but more importantly: our method can also leverage local structure in the input, which can be used to improve suitability of representations for classiﬁcation.
Leveraging known structure in the input in MI maximization-based objectives is nothing new (Becker, 1992; 1996; Wiskott & Sejnowski, 2002), and some very recent works also follow this intuition.
It has been shown in the case of discrete MI that data augmentations and other transformations can be used to avoid degenerate solutions (Hu et al., 2017).
Unsupervised clustering and segmentation is attainable by maximizing the MI between images associated by transforms or spatial proximity (Ji et al., 2018).
Mostly independent of our work, contrastive predictive coding (CPC, Oord et al., 2018) is a MI estimate-based approach that processes local features sequentially (ﬁxed-order autoregression) to build a partial “summary feature”, then makes separate predictions about several speciﬁc “future” local features that weren’t included in the “past” summary feature.
In contrast, the basic version of
DIM uses a single summary feature which is a function of all local features, and this “global” feature predicts all of those features simultaneously in a single step, rather than forming separate predictions for a few speciﬁc features as in CPC. However, we can also incorporate partial auto-regression, and we show this can provide performance gains.
Finally, our work further looks at the suitability of representations across two different MI maximization objectives (local vs. global), a ﬂexibility we believe is necessary for training representations intended for different end-goals. [Equation] [Figure]
These vectors are summarized into a single feature vector, Y . Our goal is to train this network such that useful information about the input is easily extracted from the high-level features.
Here we outline the general setting of training an encoder to maximize mutual information between its input and output.
Let X
and Y be the domain and range of a continuous and (almost everywhere) differentiable parametric function, E ψ : X → Y with parameters ψ
(e.g., a neural network).
These parameters deﬁne a family of encoders, E Φ = { E ψ } ψ ∈ Ψ over Ψ .
Assume that we are given a set of training examples on an input space, X : X :=
{ x ( i ) ∈ X} Ni =1 , with empirical probability distribution P .
We deﬁne U ψ, P to be the marginal distribution induced by pushing samples from P through E ψ .
I.e., U ψ, P is the distribution over encodings y ∈ Y produced by sampling observations x ∼ X and then sampling y ∼ E ψ ( x ) .
An example encoder for image data is given in Figure 1, which will be used in the following sections, but this approach can easily be adapted for temporal data.
Similar to the infomax optimization principle (Linsker, 1988), we assert our encoder should be trained according to the following criteria:
• Mutual information maximization: Find the set of parameters, ψ , such that the mutual information, I ( X ; E ψ ( X )) , is maximized.
Depending on the end-goal, this maximization can be done over the complete input, X , or some structured or “local” subset.
• Statistical constraints: Depending on the end-goal for the representation, the marginal U ψ, P should match a prior distribution, V .
Roughly speaking, this can be used to encourage the output of the encoder to have desired characteristics (e.g., independence).
The formulation of these two objectives covered below we call Deep InfoMax (DIM).
3. M UTUAL INFORMATION ESTIMATION AND MAXIMIZATION
Our basic mutual information maximization framework is presented in Figure 2.
The approach follows Mutual Information Neural Estimation (MINE, Belghazi et al., 2018), which estimates mutual information by training a classiﬁer to distinguish between samples coming from the joint, J , and the product of marginals, M , of random variables X and Y .
MINE uses a lower-bound to the MI based on the Donsker-Varadhan representation (DV, Donsker & Varadhan, 1983) of the KL-divergence, I ( X
[Equation]
where T ω : X × Y → R is a discriminator function modeled by a neural network with parameters ω .
First we encode the image to a feature map that reﬂects some structural aspect of the data, e.g. spatial locality, and we further summarize this feature map into a global feature vector (see Figure 1).
We then concatenate this feature vector with the lower-level feature map at every location .
A score is produced for each local-global pair through an additional function (see the Appendix A.2 for details).
At a high level, we optimize E ψ by simultaneously estimating and maximizing I ( X, E ψ ( X [Equation]
[Equation]
where the subscript G denotes “global” for reasons that will be clear later.
However, there are some important differences that distinguish our approach from MINE.
First, because the encoder and mutual information estimator are optimizing the same objective and require similar computations, we share layers between these functions, so that E ψ = f ψ ◦ C ψ and T ψ,ω
= D ω ◦ g ◦
( C ψ , E ψ ) 2 , where g is a function that combines the encoder output with the lower layer.
Second, as we are primarily interested in maximizing MI, and not concerned with its precise value, we can rely on non-KL divergences which may offer favourable trade-offs.
For example, one could deﬁne a Jensen-Shannon MI estimator (following the formulation of Nowozin et al., 2016),
[Equation]
where x is an input sample, x (cid:48) is an input sampled from ˜ P = P , and sp ( z )
= log(1+ e z ) is the softplus function.
A similar estimator appeared in Brakel & Bengio (2017) in the context of minimizing the total correlation, and it amounts to the familiar binary cross-entropy, which is well-understood in terms of neural network optimization and which we ﬁnd works better in practice (e.g., more stable) than the DV-based objective (e.g., see App. A.3).
Intuitively, the Jensen-Shannon-based estimator should behave similarly as the DV-based estimator in Eq. 2, since both act like classiﬁers whose objectives maximize the expected log -ratio of the joint over the product of marginals. We show in App.
A.1 the relationship between the JSD estimator and the formal deﬁnition of mutual information. Noise-Contrastive Estimation (NCE, Gutmann & Hyv ¨ arinen, 2010; 2012) was ﬁrst used as a bound on MI in Oord et al. (2018), and this loss can be used with DIM by maximizing:
[Equation]
For DIM, a key difference between the DV, JSD, and NCE-based formulations is whether an expectation over P / ˜ P appears inside or outside of a log .
In fact, the JSD-based objective mirrors the original NCE formulation in Gutmann & Hyv¨arinen (2010), which phrased unnormalized density estimation as binary classiﬁcation between the data distribution and a noise distribution.
DIM sets the noise distribution to the product of marginals over X/Y , and the data distribution to the true joint.
The NCE formulation in Eq.
5 follows a softmax-based version of NCE (Jozefowicz et al., 2016), similar to ones used in the language modeling community (Mnih & Kavukcuoglu, 2013;
Mikolov et al., 2013), and which has strong connections to the binary cross-entropy in the context of noise-contrastive learning (Ma & Collins, 2018).
In practice, implementations of these estimators appear quite similar and can reuse most of the same code.
We investigate JSD and NCE in our experiments, and ﬁnd that using NCE often outperforms JSD on downstream tasks, though this effect diminishes with more challenging data.
However, as we show in the App.
(A.3) , NCE and DV require a large number of negative samples (samples from ˜ P ) to be competitive.
We generate negative samples using all combinations of global and local features at all locations of the relevant feature map, across all images in a batch.
For a batch of size B , that gives O ( B × M 2 ) negative samples per positive example, which quickly becomes cumbersome with increasing batch size.
We found that DIM with the JSD loss is insensitive to the number of negative samples, and in fact outperforms NCE as the number of negative samples becomes smaller.
3. L OCAL MUTUAL INFORMATION MAXIMIZATION
The objective in Eq.
3 can be used to maximize MI between input and output, but ultimately this may be undesirable depending on the task.
For example, trivial pixel-level noise is useless for image classiﬁcation, so a representation may not beneﬁt from encoding this information (e.g., in zero-shot learning, transfer learning, etc.).
In order to obtain a representation more suitable for classiﬁcation, we can instead maximize the average MI between the high-level representation and local patches of the image.
Because the same representation is encouraged to have high MI with all the patches, this favours encoding aspects of the data that are shared across patches.
Suppose the feature vector is of limited capacity (number of units and range) and assume the encoder does not support inﬁnite output conﬁgurations.
For maximizing the MI between the whole input and the representation, the encoder can pick and choose what type of information in the input is passed through the encoder, such as noise speciﬁc to local patches or pixels.
However, if the encoder passes information speciﬁc to only some parts of the input, this does not increase the MI with any of the other patches that do not contain said noise.
This encourages the encoder to prefer information that is shared across the input, and this hypothesis is supported in our experiments below.
Our local DIM framework is presented in Figure 3.
First we encode the input to a feature map,
[Equation]
that reﬂects useful structure in the data (e.g., spatial locality), indexed in this case by i .
Next, we summarize this local feature map into a global feature, E ψ ( x ) =
f ψ ◦ C [Equation] ( x ) .
We then deﬁne our MI estimator on global/local pairs, maximizing the average estimated MI:
[Equation]
We found success optimizing this “local” objective with multiple easy-to-implement architectures, and further implementation details are provided in the App.
(A.2).
3. M ATCHING REPRESENTATIONS TO A PRIOR DISTRIBUTION
Absolute magnitude of information is only one desirable property of a representation; depending on the application, good representations can be compact (Gretton et al., 2012), independent (Hyv ¨ arinen & Oja, 2000; Hinton, 2002; Dinh et al., 2014; Brakel & Bengio, 2017), disentangled (Schmidhuber, 1992; Rifai et al., 2012; Bengio et al., 2013; Chen et al., 2018; Gonzalez-Garcia et al., 2018), or independently controllable (Thomas et al., 2017).
DIM imposes statistical constraints onto learned representations by implicitly training the encoder so that the push-forward distribution, U ψ, P , matches a prior, V .
This is done (see Figure 7 in the App. A.2) by training a discriminator, D φ : Y → R , to estimate the divergence, D ( V || U ψ, P ) , then training the encoder to minimize this estimate:
[Equation]
This approach is similar to what is done in adversarial autoencoders (AAE, Makhzani et al., 2015), but without a generator.
It is also similar to noise as targets (Bojanowski & Joulin, 2017), but trains the encoder to match the noise implicitly rather than using a priori noise samples as targets.
All three objectives – global and local MI maximization and prior matching – can be used together, and doing so we arrive at our complete objective for Deep InfoMax (DIM):
[Equation]
where ω 1 and ω 2 are the discriminator parameters for the global and local objectives, respectively, and α , β , and γ are hyperparameters.
We will show below that choices in these hyperparameters affect the learned representations in meaningful ways.
As an interesting aside, we also show in the
App. (A.8) that this prior matching can be used alone to train a generator of image data.
E XPERIMENTS
We test Deep InfoMax (DIM) on four imaging datasets to evaluate its representational properties:
• CIFAR10 and CIFAR100 (Krizhevsky & Hinton, 2009): two small-scale labeled datasets composed of 32 × 32 images with 10 and 100 classes respectively.
• Tiny ImageNet: A reduced version of ImageNet (Krizhevsky & Hinton, 2009) images scaled down to 64 × 64 with a total of 200 classes.
• STL-10 (Coates et al., 2011): a dataset derived from ImageNet composed of 96 × 96 images with a mixture of 100000 unlabeled training examples and 500 labeled examples per class.
We use data augmentation with this dataset, taking random 64 × 64 crops and ﬂipping horizontally during unsupervised learning.
• CelebA
(Yang et al., 2015, Appendix A.5 only): An image dataset composed of faces labeled with 40 binary attributes.
This dataset evaluates DIM’s ability to capture information that is more ﬁne-grained than the class label and coarser than individual pixels.
For our experiments, we compare DIM against various unsupervised methods: Variational AutoEncoders (VAE, Kingma & Welling, 2013), β -VAE (Higgins et al., 2016; Alemi et al., 2016), Adversarial AutoEncoders (AAE, Makhzani et al., 2015), BiGAN (a.k.a.
adversarially learned inference with a deterministic encoder:
Donahue et al., 2016; Dumoulin et al., 2016), Noise As Targets (NAT, Bojanowski & Joulin, 2017), and Contrastive Predictive Coding (CPC, Oord et al., 2018).
Note that we take CPC to mean ordered autoregressive predictor states that predict “future” representations, independent of the constrastive loss (JSD, NCE, or DV).
See the App.
(A.2) for details of the neural net architectures used in the experiments.
4. H OW DO WE EVALUATE THE QUALITY OF A REPRESENTATION ?
Evaluation of representations is case-driven and relies on various proxies.
Linear separability is commonly used as a proxy for disentanglement and mutual information (MI) between representations and class labels.
Unfortunately, this will not show whether the representation has high MI with the class labels when the representation is not disentangled.
Other works (Bojanowski & Joulin, 2017) have looked at transfer learning classiﬁcation tasks by freezing the weights of the encoder and training a small fully-connected neural network classiﬁer using the representation as input.
Others still have more directly measured the MI between the labels and the representation (Rifai et al., 2012;
Chen et al., 2018), which can also reveal the representation’s degree of entanglement.
Class labels have limited use in evaluating representations, as we are often interested in information encoded in the representation that is unknown to us.
Implicit models give us two new metrics for evaluating representation quality.
First, we can use mutual information neural estimation (MINE,
Belghazi et al., 2018) to more directly measure the MI between the input and output of the encoder.
Second, we can directly measure the independence of the representation using a discriminator.
Given a batch of representations, we generate a factor-wise independent distribution with the same per-factor marginals by randomly shufﬂing each factor along the batch dimension.
A similar trick has been used for learning maximally independent representations for sequential data (Brakel & Bengio, 2017).
We can train a discriminator to estimate the KL-divergence between the original representations (joint distribution of the factors) and the shufﬂed representations (product of the marginals) using the
Donsker-Varadhan representation as with MINE (see Figure 12).
The higher the KL divergence, the more dependent the factors.
We call this evaluation method Neural Dependency Measure (NDM) and show that it is sensible and empirically consistent in the App.
(A.6).
To summarize, we use the following metrics for evaluating representations.
For each of these, the encoder is held ﬁxed unless noted otherwise:
• Linear classiﬁcation using a support vector machine (SVM).
This is simultaneously a proxy for MI of the representation with linear separability.
• Non-linear classiﬁcation using a single hidden layer neural network ( 200 units) with dropout.
This is a proxy on MI of the representation with the labels separate from linear separability as measured with the SVM above.
[Table] • Semi-supervised learning (STL-10 here), that is, ﬁne-tuning the complete encoder by adding a small neural network on top of the last convolutional layer (matching architectures with a standard fully-supervised classiﬁer).
• MS-SSIM
(Wang et al., 2003), using a decoder trained on the L 2 reconstruction loss.
This is a proxy for the total MI between the input and the representation and can indicate the amount of encoded pixel-level information.
• Mutual information neural estimate (MINE) , (cid:98)
I ρ ( X, E ψ ( x )) , between the input, X , and the output representation, E ψ ( x ) , by training a discriminator with parameters ρ to maximize the DV estimator of the KL-divergence.
• Neural dependency measure (NDM) using a second discriminator that measures the KL between E ψ ( x ) and a batch-wise shufﬂed version of E ψ ( x ) .
For the neural network classiﬁcation evaluation above, we performed experiments on all datasets except CelebA, while for other measures we only looked at CIFAR10.
For all classiﬁcation tasks, we built separate classiﬁers on the high-level vector representation ( Y ), the output of the previous fully-connected layer (fc) and the last convolutional layer (conv).
Model selection for the classiﬁers was done by averaging the last 100 epochs of optimization, and the dropout rate and decaying learning rate schedule was set uniformly to alleviate over-ﬁtting on the test set across all models.
4. R EPRESENTATION LEARNING COMPARISON ACROSS MODELS
In the following experiments, DIM(G) refers to DIM with a global-only objective ( α = 1 , β = 0 , γ = 1 ) and DIM(L) refers to DIM with a local-only objective ( α = 0 , β = 1 , γ = 0 . 1 ), the latter chosen from the results of an ablation study presented in the App.
(A.5).
For the prior, we chose a compact uniform distribution on [0 , 1] 64 , which worked better in practice than other priors, such as Gaussian, unit ball, or unit sphere.
Our classiﬁcation results can be found in Tables 1, 2, and 3.
In general,
DIM with the local objective, DIM(L), outperformed all models presented here by a signiﬁcant margin on all datasets, regardless of which layer the representation was drawn from, with exception to CPC.
For the speciﬁc settings presented (architectures, no data augmentation for datasets except for STL10), DIM(L) performs as well as or outperforms a fully-supervised classiﬁer without ﬁne-tuning, which indicates that the representations are nearly as good as or better than the raw pixels given the model constraints in this setting. Note, however, that a fully supervised classiﬁer can perform much better on all of these benchmarks, especially when specialized architectures and carefully-chosen data augmentations are used. Competitive or better results on CIFAR10 also exist (albeit in different settings, e.g., Coates et al., 2011; Dosovitskiy et al., 2016), but to our knowledge our STL-10 results [Table]
[Table]
are SOTA for unsupervised learning.
The results in this setting supports the hypothesis that our local DIM objective is suitable for extracting class information.
Our results show that the NCE loss tends to perform best, but differences between NCE and JSD diminish with larger datasets.
DV can compete with JSD with smaller datasets, but DV performs much worse with larger datasets.
For CPC, we were only able to achieve marginally better performance than BiGAN with the settings above.
However, when we adopted the strided crop architecture found in Oord et al.
(2018), both CPC and DIM performance improved considerably. We chose a crop size of 25% of the image size in width and depth with a stride of 12 . 5% the image size (e.g., 8 × 8 crops with 4 × 4 strides for CIFAR10, 16 × 16 crops with 8 × 8 strides for STL-10), so that there were a total of 7 × 7 local features. For both
DIM(L) and CPC, we used the NCE loss as well as the same “encode-and-dot-product” architecture (can be thought of as a deep Bilinear model), rather than the shallow Bilinear model used in Oord et al.
For simplicity, we omitted the prior term, β , from DIM.
Without data augmentation on
{+CIFAR10, CPC performs worse than DIM(L) with a ResNet-
+}
(He et al., 2016) type architecture.
For experiments we ran on STL-10 with data augmentation (using the same encoder architecture as
Table 2), CPC and DIM were competitive, with CPC performing slightly better.
CPC uses multiple summary vectors when performing predictions, so in order to leverage more predictions with DIM, we incorporate partial autoregression by computing a set of 3 × 3 global representations through a single convolutional layer, each of which are optimized by maximizing the mutual information with all local representations.
The resulting model gave a signiﬁcant performance boost to DIM for
STL-10.
Surprisingly, this same architecture performed worse than using one global representation with CIFAR10.
Overall DIM only slightly outperforms CPC in this setting, so at most one can say that DIM and CPC are comparable, though the strictly ordered autoregression of CPC may be unnecessary for some tasks.
Extended comparisons Tables 4 shows results on linear separability, reconstruction (MS-SSIM), mutual information, and dependence (NDM) with the CIFAR10 dataset.
We did not compare to CPC due to the divergence of architectures.
For linear classiﬁer results (SVC), we trained ﬁve support vector machines with a simple hinge loss for each model, averaging the test accuracy.
For MINE, we used a decaying learning rate schedule, which helped reduce variance in estimates and provided faster convergence.
[Table]
DIM can improve the classiﬁcation accuracy, particularly with the highly-compressed global features.
[Table]
MS-SSIM correlated well with the MI estimate provided by MINE, indicating that these models encoded pixel-wise information well.
As our prior matching was done using a sigmoid function on the representation, Y , we measured NDM with and without (in parentheses) this nonlinearity.
Overall, all models showed much lower dependence than BiGAN, indicating the marginal of the encoder output is not matching to the generator’s spherical Gaussian input prior.
For MI, reconstruction-based models like VAE and AAE have high scores, and we found that combining local and global DIM objectives had very high scores ( α = 0 .
5 , β = 0 .
1 is presented here as DIM(L+G)).
For more in-depth analyses, please see the ablation studies and the nearest-neighbor analysis in the App.
(A.4, A.5).
4. A DDING COORDINATE INFORMATION AND OCCLUSIONS
Maximizing MI between global and local features is not the only way to leverage image structure.
We consider augmenting DIM by adding input occlusion when computing global features and by adding auxiliary tasks which maximize MI between local features and absolute or relative spatial coordinates given a global feature.
These additions improve classiﬁcation results (see Table 5).
For occlusion, we randomly occlude part of the input when computing the global features, but compute local features using the full input.
Maximizing MI between occluded global features and unoccluded local features aggressively encourages the global features to encode information which is shared across the entire image.
For coordinate prediction, we maximize the model’s ability to predict the coordinates ( i, j ) of a local feature c ( i,j ) =
C ( i,j ) ψ
( x ) after computing the global features
y =
E ψ ( x ) .
To accomplish this, we maximize E [log p θ (( i, j )
| y, c ( i,j ) )]
(i.e., minimize the crossentropy).
We can extend the task to maximize conditional MI given global features y between pairs of local features ( c ( i,j ) , c
( i (cid:48) ,j (cid:48) ) ) and their relative coordinates ( i − i (cid:48) , j − j (cid:48) ) .
This objective can be written as E [log p θ (( i − i (cid:48) , j − j (cid:48) )
| y, c ( i,j ) , c ( i (cid:48) ,j (cid:48) ) )] .
We use both these objectives in our results.
Additional implementation details can be found in the App.
(A.7).
Roughly speaking, our input occlusions and coordinate prediction tasks can be interpreted as generalizations of inpainting (Pathak et al., 2016) and context prediction (Doersch et al., 2015) tasks which have previously been proposed for self-supervised feature learning.
Augmenting DIM with these tasks helps move our method further towards learning representations which encode images (or other types of inputs) not just in terms of compressing their low-level (e.g. pixel) content, but in terms of distributions over relations among higher-level features extracted from their lower-level content.
C ONCLUSION
In this work we introduced Deep InfoMax (DIM), a new method for learning unsupervised representations by maximizing mutual information, allowing for representations that contain locally-consistent information across structural “locations” (e g., patches in an image).
This provides a straightforward and ﬂexible way to learn representations that perform well on a variety of tasks.
We believe that this is an important direction in learning higher-level representations
A A PPENDIX A. O N THE J ENSEN -S HANNON DIVERGENCE AND MUTUAL INFORMATION
Here we show the relationship between the Jensen-Shannon divergence (JSD) between the joint and the product of marginals and the pointwise mutual information (PMI).
Let p ( x ) and p ( y ) be two marginal densities, and deﬁne p ( y | x ) and p ( x, y )
=
p ( y | x ) p ( x ) as the conditional and joint distribution, respectively.
Construct a probability mixture density, m ( x, y )
=
12 ( p ( x ) p ( y )
+ p ( x, y )) .
It follows that m
( x ) = p ( x ) , m ( y ) =
p ( y ) , and m ( y | x ) = p ( y | x ) .
Note that:
[Equation]
Discarding some constants:
[Equation]
[Equation]
The quantity inside the expectation of Eqn.
10 is a concave, monotonically increasing function of the ratio p ( y | x ) p ( y ) , which is exactly e PMI ( x,y ) .
Note this relationship does not hold for the JSD of arbitrary distributions, as the the joint and product of marginals are intimately coupled.
We can verify our theoretical observation by plotting the JSD and KL divergences between the joint and the product of marginals, the latter of which is the formal deﬁnition of mutual information (MI).
As computing the continuous MI is difﬁcult, we assume a discrete input with uniform probability, p ( x ) (e.g., these could be one-hot variables indicating one of N i.i.d.
random samples), and a randomly initialized N × M joint distribution, p ( x, y ) , such that (cid:80)
Mj =1 p ( x i , y j ) = 1 ∀ i .
For this joint distribution, we sample from a uniform distribution, then apply dropout to encourage sparsity to simulate the situation when there is no bijective function between x and y , then apply a softmax.
As the distributions are discrete, we can compute the KL and JSD between p ( x, y ) and p ( x ) p ( y ) .
We ran these experiments with matched input / output dimensions of 8 , 16 , 32 , 64 , and 128 , randomly drawing 1000 joint distributions, and computed the KL and JSD divergences directly.
Our results (Figure A.1) indicate that the KL (traditional deﬁnition of mutual information) and the JSD have a strongly monotonic relationship.
Overall, the distributions with the highest mutual information also have the highest JSD.
A. E XPERIMENT AND ARCHITECTURE DETAILS
Here we provide architectural details for our experiments.
Example code for running Deep Infomax (DIM) can be found at REDACTED .
Encoder We used an encoder similar to a deep convolutional GAN (DCGAN, Radford et al., 2015) discriminator for CIFAR10 and CIFAR100, and for all other datasets we used an Alexnet (Krizhevsky et al., 2012) architecture similar to that found in Donahue et al. (2016).
ReLU activations and batch norm (Ioffe & Szegedy, 2015) were used on every hidden layer.
For the DCGAN architecture, a single hidden layer with 1024 units was used after the ﬁnal convolutional layer, and for the Alexnet architecture it was two hidden layers with 4096 .
For all experiments, the output of all encoders was a 64 dimensional vector. ( x )
( y with discrete inputs and a given randomized and sparse joint distribution, p ( x, y ) . 8 × 8 indicates a square joint distribution with 8 rows and 8 columns. Our experiments indicate a strong monotonic relationship between MI ( x ; y ) and JSD ( p ( x, y ) ||
p ( x ) p ( y ))
Overall, the distributions with the highest MI ( x ; y ) have the highest JSD ( p ( x, y ) ||
p ( x ) p ( y )) . Mutual information discriminators For the global mutual information objective, we ﬁrst encode the input into a feature map, C ψ ( x ) , which in this case is the output of the last convolutional layer.
We then encode this representation further using linear layers as detailed above to get E ψ ( x ) .
C ψ ( x ) is then ﬂattened, then concatenated with E ψ ( x ) .
We then pass this to a fully-connected network with two 512 -unit hidden layers (see Table 6).
[Table]
We tested two different architectures for the local objective.
The ﬁrst (Figure 5) concatenated the global feature vector with the feature map at every location, i.e., { [ C ( i ) ψ ( x ) , E ψ ( x )] } M × M i
= .
A 1 ×convolutional discriminator is then used to score the (feature map, feature vector) pair,
[Equation]
Fake samples are generated by combining global feature vectors with local feature maps coming from different images, x (cid:48) :
[Equation]
This architecture is featured in the results of Table 4, as well as the ablation and nearest-neighbor studies below.
We used a 1 × 1 convnet with two 512 -unit hidden layers as discriminator (Table 7).
[Table]
The other architecture we tested (Figure 6) is based on non-linearly embedding the global and local features in a (much) higher-dimensional space, and then computing pair-wise scores using dot products between their high-dimensional embeddings.
This enables efﬁcient evaluation of a large number of pair-wise scores, thus allowing us to use large numbers of positive/negative samples.
Given a sufﬁciently high-dimensional embedding space, this approach can represent (almost) arbitrary classes of pair-wise functions that are non-linear in the original, lower-dimensional features.
For more information, refer to Reproducing Kernel Hilbert Spaces.
We pass the global feature through a fully connected neural network to get the encoded global feature, S ω ( E ψ ( x )) .
In our experiments, we used a single hidden layer network with a linear shortcut (See Table 8)
[Table]
[Figure]
[Table]
We embed each local feature in the local feature map C ψ ( x ) using an architecture which matches the one for global feature embedding.
We apply it via 1 × 1 convolutions.
Details are in Table 9.
[Table]
Finally, the outputs of these two networks are combined by matrix multiplication, summing over the feature dimension ( 2048 in the example above).
As this is computed over a batch, this allows us to efﬁciently compute both positive and negative examples simultaneously.
This architecture is featured in our main classiﬁcation results in Tables 1, 2, and 5.
For the local objective, the feature map, C ψ ( x ) , can be taken from any level of the encoder, E ψ .
For the global objective, this is the last convolutional layer, and this objective was insensitive to which layer we used.
For the local objectives, we found that using the next-to-last layer worked best for
CIFAR10 and CIFAR100, while for the other larger datasets it was the previous layer.
This sensitivity is likely due to the relative size of the of the receptive ﬁelds, and further analysis is necessary to better understand this effect.
Note that all feature maps used for DIM included the ﬁnal batch normalization and ReLU activation.
Prior matching Figure 7 shows a high-level overview of the prior matching architecture.
The discriminator used to match the prior in DIM was a fully-connected network with two hidden layers of 1000 and 200 units (Table 10).
[Table]
Generative models For generative models, we used a similar setup as that found in Donahue et al.
(2016) for the generators / decoders, where we used a generator from DCGAN in all experiments.
All models were trained using Adam with a learning rate of 1 × 10 − 4 for 1000 epochs for CIFAR and CIFAR100 and for 200 epochs for all other datasets.
Contrastive Predictive Coding For Contrastive Predictive Coding (CPC, Oord et al., 2018), we used a simple a GRU-based PixelRNN (Oord et al., 2016) with the same number of hidden units as the feature map depth.
All experiments with CPC had the global state dimension matched with the size of these recurrent hidden units.
A. S AMPLING STRATEGIES
We found both the NCE- and DV-based estimators were sensitive to negative sampling strategies, while the JSD-based estimator was insensitive.
JSD worked better ( 1 − 2% accuracy improvement) by excluding positive samples from the product of marginals, so we exclude them in our implementation.
It is quite likely that this is because our batchwise sampling strategy overestimate the frequency of positive examples as measured across the complete dataset.
NCE was highly sensitive to the number of negative samples for estimating the log-expectation term (see Figure 9).
With high sample size, NCE outperformed JSD on many tasks, but performance drops quickly as we reduce the number of images used for this estimation.
This may become more problematic for larger datasets and networks where available memory is an issue.
DV was outperformed by JSD even with the maximum number of negative samples used in these experiments, and even worse was highly unstable as the number of negative samples dropped.
[Figure]
A. N EAREST - NEIGHBOR ANALYSIS
In order to better understand the metric structure of DIM’s representations we did a nearest-neighbor analysis randomly choosing a sample from each class in the test set ordering the test set in terms of
L 1 distance in the representation space (to reﬂect the uniform prior), then selecting the four with the lowest distance.
Our results in Figure 8 show that DIM with a local-only objective, DIM(L), learns a representation with a much more interpretable structure across the image.
However, our result potentially highlights an issue with using only consistent information across patches, as many of the nearest neighbors share patterns (colors, shapes, texture) but not class. [Figure] β γ =0 . 1: b
I ρ ( X,Y )
[Figure]
A. A BLATION STUDIES
To better understand the effects of hyperparameters α , β , and γ on the representational characteristics of the encoder, we performed several ablation studies.
These illuminate the relative importance of global verses local mutual information objectives as well as the role of the prior.
Local versus global mutual information maximization The results of our ablation study for
DIM on CIFAR10 are presented in Figure 10.
In general, good classiﬁcation performance is highly dependent on the local term, β , while good reconstruction is highly dependent on the global term, α .
However, a small amount of α helps in classiﬁcation accuracy and a small about of β improves reconstruction.
For mutual information, we found that having a combination of α and β yielded higher MINE estimates.
Finally, for CelebA (Figure 11), where the classiﬁcation task is more ﬁne-grained (is composed of potentially locally-speciﬁed labels, such as “lipstick” or “smiling”), the global objective plays a stronger role than with classiﬁcation on other datasets (e.g., CIFAR10). [Figure]
there a strong trend as we increase The effect of the prior We found including the prior term, γ , was absolutely necessary for ensuring low dependence between components of the high-level representation, E ψ ( x ) , as measured by NDM.
In addition a small amount of the prior term helps improve classiﬁcation results when used with the local term, β .
This may be because the additional constraints imposed on the representation help to encourage the local term to focus on consistent, rather than trivial, information.
A. E MPIRICAL CONSISTENCY OF N EURAL D EPENDENCY M EASURE (NDM)
Here we evaluate the Neural Dependency Measure (NDM) over a range of β -VAE (Alemi et al., 2016;
Higgins et al., 2016) models.
β -VAE encourages disentangled representations by increasing the role of the KL-divergence term in the ELBO objective.
We hypthesized that NDM would consistenly measure lower dependence (lower NDM)
as the β values increase, and our results in Figure A. conﬁrm this.
As we increase β , there is a strong downward trend in the NDM, though β = 0 .
5 and β = 1 . 0 give similar numbers.
In addition, the variance over estimates and models is relatively low, meaning the estimator is empirically consistent in this setting.
A. A DDITIONAL DETAILS ON OCCLUSION AND COORDINATE PREDICTION EXPERIMENTS
Here we present experimental details on the occlusion and coordinate prediction tasks.
[Figure]
Occlusions.
For the occlusion experiments, the sampling distribution for patches to occlude was ad-hoc.
Roughly, we randomly occlude the input image under the constraint that at least one 10 ×block of pixels remains visible and at least one 10 × 10 block of pixels is fully occluded.
We chose 10 × 10 based on the receptive ﬁelds of local features in our encoder, since it guarantees that occlusion leaves at least one local feature fully observed and at least one local feature fully unobserved.
Figure 14 shows the distribution of occlusions used in our tests.
Absolute coordinate prediction For absolute coordinate prediction, the global features y and local features c ( i,j ) are sampled by 1) feeding an image from the data distribution through the feature encoder, and 2) sampling a random spatial location ( i, j ) from which to take the local features c ( i,j ) .
Given y and c ( i,j ) , we treat the coordinates i and j as independent categorical variables and measure the required log probability using a sum of categorical cross-entropies.
In practice, we implement the prediction function p θ as an MLP with two hidden layers, each with 512 units, ReLU activations, and batchnorm.
We marginalize this objective over all local features associated with a given global feature when computing gradients.
Relative coordinate prediction For relative coordinate prediction, the global features y and local features c ( i,j ) /c
( i (cid:48) ,j (cid:48) ) are sampled by 1) feeding an image from the data distribution through the feature encoder, 2) sampling a random spatial location ( i, j ) from which to take source local features c ( i,j ) , and 3) sampling another random location ( i (cid:48) , j (cid:48) ) from which to take target
local features c ( i (cid:48) ,j (cid:48) ) .
In practice, our predictive model for this task uses the same architecture as for the task described previously.
For each global feature y we select one source feature c ( i,j ) and marginalize over all possible target features c ( i (cid:48) ,j (cid:48) ) when computing gradients. [Figure]
A. T RAINING A GENERATOR BY MATCHING TO A PRIOR IMPLICITLY
We show here and in our experiments below that we can use prior objective in DIM (Equation 7) to train a high-quality generator of images by training U ψ, P to map to a one-dimensional mixture of two Gaussians implicitly.
One component of this mixture will be a target for the push-forward distribution of P through the encoder while the other will be a target for the push-forward distribution of the generator, Q θ , through the same encoder.
Let G θ : Z → X be a generator function, where the input z ∈ Z is drawn from a simple prior, p ( z ) (such as a spherical Gaussian).
Let Q θ be the generated distribution and P be the empirical distribution of the training set.
Like in GANs, we will pass the samples of the generator or the training data through another function, E ψ , in order to get gradients to ﬁnd the parameters, θ .
However, unlike
GANs, we will not play the minimax game between the generator and this function.
Rather E ψ will be trained to generate a mixture of Gaussians conditioned on whether the input sample came from P or Q θ :
[Equation]
where N ( µ P , 1) and N ( µ Q , 1) are normal distributions with unit variances and means µ P and µ Q respectively.
In order to ﬁnd the parameters ψ , we introduce two discriminators, T Pφ P , T Qφ Q : Y → R , and use the lower bounds following deﬁned by the JSD f-GAN:
[Equation]
The generator is trained to move the ﬁrst-order moment of E U ψ, Q [ y ] =
E p ( z ) [ E ψ ( G θ ( z ))] to µ P :
[Equation]
Some intuition might help understand why this might work.
As discussed in Arjovsky & Bottou (2017), if P and Q θ have support on a low-dimensional manifolds on X , unless they are perfectly aligned, there exists a discriminator that will be able to perfectly distinguish between samples coming from P and Q θ , which means that U ψ, P and U ψ, Q must also be disjoint.
However, to train the generator, U ψ, P and U ψ, Q need to share support on Y in order to ensure stable and non-zero gradients for the generator.
Our own experiments by overtraining the discriminator (Figure 15) conﬁrm that lack of overlap between the two modes of the discriminator is symptomatic of poor training.
Suppose we start with the assumption that the encoder targets, V P and V Q , should overlap.
Unless P and Q θ are perfectly aligned (which according to Arjovsky & Bottou (2017) is almost guaranteed not to happen with natural images), then the discriminator can always accomplish this task by discarding information about P or Q θ .
This means that, by choosing the overlap, we ﬁx the strength of the encoder. [Table]
A. G ENERATION EXPERIMENTS AND RESULTS
For the generator and encoder, we use a ResNet architecture (He et al., 2016) identical to the one found in Gulrajani et al.
We used the contractive penalty (found in Mescheder et al.
but ﬁrst introduced in contractive autoencoders (Rifai et al., 2011)) on the encoder, gradient clipping on the discriminators, and no regularization on the generator.
Batch norm (Ioffe & Szegedy, 2015) was used on the generator, but not on the discriminator.
We trained on 64 × 64 dimensional LSUN (Yu et al., 2015), CelebA (Liu et al., 2015), and Tiny Imagenet dataset.
A. I MAGES G ENERATION
[Figure]
Here, we train a generator mapping to two Gaussian implicitly as described in Section A.8.
Our results (Figure 16) show highly realistic images qualitatively competitive to other methods (Gulrajani et al., 2017; Hjelm et al., 2018).
In order to quantitatively compare our method to GANs, we trained a non-saturating GAN with contractive penalty (NS-GAN-CP) and WGAN-GP (Gulrajani et al., 2017) with identical architectures and training procedures.
Our results (Table 11) show that, while our mehtod did not surpass NS-GAN-CP or WGAN-GP in our experiments, they came reasonably close.