Learning deep representations by mutual information estimation and maximization

1 INTRODUCTION . One core objective of deep learning is to discover useful 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 difficult 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 and show that , depending on the downstream task , maximizing MI between the complete input and the encoder output ( i.e. , global MI ) is often insufficient for learning useful representations . Rather , structure matters : maximizing the average MI between the representation and local regions of the input ( e.g . patches rather than the complete image ) can greatly improve the representation ’ s quality for , e.g. , classification 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 ; Hyvärinen & 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 local information , which we show can be used to tune the suitability of learned representations for classification or reconstruction-style tasks . • We use adversarial learning ( à la Makhzani et al. , 2015 ) to constrain the representation to have desired statistical characteristics specific 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 . 2 RELATED WORK . There are many popular methods for learning representations . Classic methods , such as independent component analysis ( ICA , Bell & Sejnowski , 1995 ) and self-organizing maps ( Kohonen , 1998 ) , generally lack the representational capacity of deep neural networks . More recent 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 ) , and self-supervised or co-learning ( Doersch & Zisserman , 2017 ; Dosovitskiy et al. , 2016 ; Sajjadi et al. , 2016 ) . Generative models are also commonly used 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 the quality of the representations they learn . 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 reconstruction error can be related to the MI as follows : Ie ( X , Y ) = He ( X ) −He ( X|Y ) ≥ He ( X ) −Re , d ( X|Y ) , ( 1 ) where X and Y denote the input and output of an encoder which is applied to inputs sampled from some source distribution . Re , d ( X|Y ) denotes the expected reconstruction error of X given the codes Y . He ( X ) andHe ( X|Y ) denote the marginal and conditional entropy ofX in the distribution formed by applying the encoder to inputs sampled from the source distribution . Thus , in typical settings , 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 or to minimize the reconstruction error ( Chen et al. , 2016 ) . Mutual-information estimation Methods based on mutual information have a long history in unsupervised feature learning . The infomax 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 , which can be nonlinear ( Hyvärinen & Pajunen , 1999 ; Almeida , 2003 ) but are often hard to adapt for use with deep networks . Mutual Information Neural Estimation ( MINE , Belghazi et al. , 2018 ) learns an 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 find that the generator is unnecessary . We also find it unnecessary to use the exact KL-based formulation of MI . For example , a simple alternative based on the Jensen-Shannon divergence ( JSD ) is more stable and provides better results . We will show that DIM can work with various MI estimators . Most significantly , DIM can leverage local structure in the input to improve the suitability of representations for classification . Leveraging known structure in the input when designing objectives based on MI maximization 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 ) . Our work investigates the suitability of representations learned across two different MI objectives that focus on local or global structure , a flexibility we believe is necessary for training representations intended for different applications . Proposed independently of DIM , Contrastive Predictive Coding ( CPC , Oord et al. , 2018 ) is a MIbased approach that , like DIM , maximizes MI between global and local representation pairs . CPC shares some motivations and computations with DIM , but there are important ways in which CPC and DIM differ . CPC processes local features sequentially to build partial “ summary features ” , which are used to make predictions about specific local features in the “ future ” of each summary feature . This equates to ordered autoregression over the local features , and requires training separate estimators for each temporal offset at which one would like to predict the future . In contrast , the basic version of DIM uses a single summary feature that is a function of all local features , and this “ global ” feature predicts all local features simultaneously in a single step using a single estimator . Note that , when using occlusions during training ( see Section 4.3 for details ) , DIM performs both “ self ” predictions and orderless autoregression . 3 DEEP INFOMAX Figure 1 : The base encoder model in the context of image data . An image ( in this case ) is encoded using a convnet until reaching a feature map of M × M feature vectors corresponding to M ×M input patches . 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 define 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 define 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.1 MUTUAL 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 classifier to distinguish between samples coming from the joint , J , and the Figure 3 : Maximizing mutual information between local features and global features . First we encode the image to a feature map that reflects 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 ) . 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 ; Y ) : = DKL ( J||M ) ≥ Î ( DV ) ω ( X ; Y ) : = EJ [ Tω ( x , y ) ] − logEM [ eTω ( x , y ) ] , ( 2 ) where Tω : X × Y → R is a discriminator function modeled by a neural network with parameters ω . At a high level , we optimize Eψ by simultaneously estimating and maximizing I ( X , Eψ ( X ) ) , ( ω̂ , ψ̂ ) G = arg max ω , ψ Îω ( X ; Eψ ( X ) ) , ( 3 ) 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ψ ) ,1 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 define a Jensen-Shannon MI estimator ( following the formulation of Nowozin et al. , 2016 ) , Î ( JSD ) ω , ψ ( X ; Eψ ( X ) ) : = EP [ −sp ( −Tψ , ω ( x , Eψ ( x ) ) ) ] − EP×P̃ [ sp ( Tψ , ω ( x ′ , Eψ ( x ) ) ) ] , ( 4 ) where x is an input sample , x′ is an input sampled from P̃ = P , and sp ( z ) = log ( 1+ez ) 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 . This is well-understood in terms of neural network optimization and we find works better in practice ( e.g. , is more stable ) than the DV-based objective ( e.g. , see App . A.3 ) . Intuitively , the Jensen-Shannon-based estimator should behave similarly to the DV-based estimator in Eq . 2 , since both act like classifiers 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 definition of mutual information . Noise-Contrastive Estimation ( NCE , Gutmann & Hyvärinen , 2010 ; 2012 ) was first used as a bound on MI in Oord et al . ( and called “ infoNCE ” , 2018 ) , and this loss can also be used with DIM by maximizing : Î ( infoNCE ) ω , ψ ( X ; Eψ ( X ) ) : = EP [ Tψ , ω ( x , Eψ ( x ) ) − EP̃ [ log ∑ x′ eTψ , ω ( x ′ , Eψ ( x ) ) ] ] . ( 5 ) For DIM , a key difference between the DV , JSD , and infoNCE 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 & Hyvärinen ( 2010 ) , which phrased unnormalized density estimation as binary classification 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 infoNCE 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. , 1Here we slightly abuse the notation and use ψ for both parts of Eψ . 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 infoNCE in our experiments , and find that using infoNCE often outperforms JSD on downstream tasks , though this effect diminishes with more challenging data . However , as we show in the App . ( A.3 ) , infoNCE 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 ×M2 ) 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 infoNCE as the number of negative samples becomes smaller .

Summary: This paper proposes Deep InfoMax (DIM), for learning representations by maximizing the mutual information between the input and a deep representation. By structuring the network and objectives to encode input locality or priors on the representation, DIM learns features that are useful for downstream tasks without relying on reconstruction or a generative model. DIM is evaluated on a number of standard image datasets and shown to learn features that outperform prior approaches based on autoencoders at classification.