On Mutual Information Maximization for Representation Learning

1 INTRODUCTION . Unsupervised representation learning is a fundamental problem in machine learning . Intuitively , one aims to learn a function g which maps the data into some , usually lower-dimensional , space where one can solve some ( generally a priori unknown ) target supervised tasks more efficiently , i.e . with fewer labels . In contrast to supervised and semi-supervised learning , the learner has access only to unlabeled data . Even though the task seems ill-posed as there is no natural objective one should optimize , by leveraging domain knowledge this approach can be successfully applied to a variety of problem areas , including image ( Kolesnikov et al. , 2019 ; van den Oord et al. , 2018 ; Hénaff et al. , 2019 ; Tian et al. , 2019 ; Hjelm et al. , 2019 ; Bachman et al. , 2019 ) and video classification ( Wang and Gupta , 2015 ; Sun et al. , 2019 ) , and natural language understanding ( van den Oord et al. , 2018 ; Peters et al. , 2018 ; Devlin et al. , 2019 ) . Recently , there has been a revival of approaches inspired by the InfoMax principle ( Linsker , 1988 ) : Choose a representation g ( x ) maximizing the mutual information ( MI ) between the input and its representation , possibly subject to some structural constraints . MI measures the amount of information obtained about a random variable X by observing some other random variable Y 1 Formally , the MI between X and Y , with joint density p ( x , y ) and marginal densities p ( x ) and p ( y ) , is defined as the Kullback–Leibler ( KL ) divergence between the joint and the product of the marginals I ( X ; Y ) = DKL ( p ( x , y ) ‖ p ( x ) p ( y ) ) = Ep ( x , y ) [ log p ( x , y ) p ( x ) p ( y ) ] . ( 1 ) The fundamental properties of MI are well understood and have been extensively studied ( see e.g . Kraskov et al . ( 2004 ) ) . Firstly , MI is invariant under reparametrization of the variables — namely , if X ′ = f1 ( X ) and Y ′ = f2 ( Y ) are homeomorphisms ( i.e . smooth invertible maps ) , then I ( X ; Y ) = I ( X ′ ; Y ′ ) . Secondly , estimating MI in high-dimensional spaces is a notoriously difficult task , and in practice one often maximizes a tractable lower bound on this quantity ( Poole et al. , 2019 ) . ∗Equal contribution . Correspondence to Michael Tschannen ( tschannen @ google.com ) , Josip Djolonga ( josipd @ google.com ) , and Mario Lucic ( lucic @ google.com ) . †PhD student at University of Cambridge and the Max Planck Institute for Intelligent Systems , Tübingen . 1We denote random variables using upper-case letters ( e.g . X , Y ) , and their realizations by the corresponding lower-case letter ( e.g . x , y ) . Nonetheless , any distribution-free high-confidence lower bound on entropy requires a sample size exponential in the size of the bound ( McAllester and Statos , 2018 ) . Despite these fundamental challenges , several recent works have demonstrated promising empirical results in representation learning using MI maximization ( van den Oord et al. , 2018 ; Hénaff et al. , 2019 ; Tian et al. , 2019 ; Hjelm et al. , 2019 ; Bachman et al. , 2019 ; Sun et al. , 2019 ) . In this work we argue , and provide empirical evidence , that the success of these methods can not be attributed to the properties of MI alone . In fact , we show that maximizing tighter bounds on MI can result in worse representations . In addition , we establish a connection to deep metric learning and argue that this interpretation may be a plausible explanation of the success of the recently introduced methods.2 2 BACKGROUND AND RELATED WORK . Recent progress and the InfoMax principle While promising results in other domains have been presented in the literature , we will focus on unsupervised image representation learning techniques that have achieved state-of-the-art performance on image classification tasks ( Hénaff et al. , 2019 ; Tian et al. , 2019 ; Bachman et al. , 2019 ) . The usual problem setup dates back at least to Becker and Hinton ( 1992 ) and can conceptually be described as follows : For a given image X , let X ( 1 ) and X ( 2 ) be different , possibly overlapping views of X , for instance the top and bottom halves of the image . These are encoded using encoders g1 and g2 respectively , and the MI between the two representations g1 ( X ( 1 ) ) and g2 ( X ( 2 ) ) is maximized , max g1∈G1 , g2∈G2 IEST ( g1 ( X ( 1 ) ) ; g2 ( X ( 2 ) ) ) , ( 2 ) where IEST ( X ; Y ) is a sample-based estimator of the true MI I ( X ; Y ) and the function classes G1 and G2 can be used to specify structural constraints on the encoders . While not explicitly reflected in ( 2 ) , note that g1 and g2 can often share parameters . Furthermore , it can be shown that I ( g1 ( X ( 1 ) ) ; g2 ( X ( 2 ) ) ) ≤ I ( X ; g1 ( X ( 1 ) ) , g2 ( X ( 2 ) ) ) ,3 hence the objective in ( 2 ) can be seen as a lower bound on the InfoMax objective maxg∈G I ( X ; g ( X ) ) ( Linsker , 1988 ) . Practical advantages of multi-view formulations There are two main advantages in using ( 2 ) rather than the original InfoMax objective . First , the MI has to be estimated only between the learned representations of the two views , which typically lie on a much lower-dimensional space than the one where the original data X lives . Second , it gives us plenty of modeling flexibility , as the two views can be chosen to capture completely different aspects and modalities of the data , for example : 1 . In the basic form of DeepInfoMax ( Hjelm et al. , 2019 ) g1 extracts global features from the entire image X ( 1 ) and g2 local features from image patches X ( 2 ) , where g1 and g2 correspond to activations in different layers of the same convolutional network . Bachman et al . ( 2019 ) build on this and compute the two views from different augmentations of the same image . 2 . Contrastive multiview coding ( CMC ) ( Tian et al. , 2019 ) generalizes the objective in ( 2 ) to consider multiple views X ( i ) , where each X ( i ) corresponds to a different image modality ( e.g. , different color channels , or the image and its segmentation mask ) . 3 . Contrastive predictive coding ( CPC ) ( van den Oord et al. , 2018 ; Hénaff et al. , 2019 ) incorporates a sequential component of the data . Concretely , one extracts a sequence of patches from an image in some fixed order , maps each patch using an encoder , aggregates the resulting features of the first t patches into a context vector , and maximizes the MI between the context and features extracted from the patch at position t+ k. In ( 2 ) , X ( 1 ) would thus correspond to the first t patches and X ( 2 ) to the patch at location t+ k. Other approaches , such as those presented by Sermanet et al . ( 2018 ) , Hu et al . ( 2017 ) , and Ji et al . ( 2019 ) , can be similarly subsumed under the same objective . Lower bounds on MI As evident from ( 2 ) , another critical choice is the MI estimator IEST . Given the fundamental limitations of MI estimation ( McAllester and Statos , 2018 ) , recent work has focused on deriving lower bounds on MI ( Barber and Agakov , 2003 ; Belghazi et al. , 2018 ; Poole et al. , 2The code for running the experiments and visualizing the results is available at https : //github.com/googleresearch/google-research/tree/master/mutual_information_representation_learning . 3Follows from the data processing inequality ( see Prop . 1 in Appendix A ) . 2019 ) . Intuitively , these bounds are based on the following idea : If a classifier can accurately distinguish between samples drawn from the joint p ( x , y ) and those drawn from the product of marginals p ( x ) p ( y ) , then X and Y have a high MI . We will focus on two such estimators , which are most commonly used in the representation learning literature . The first of them , termed InfoNCE ( van den Oord et al. , 2018 ) , is defined as I ( X ; Y ) ≥ E [ 1 K K∑ i=1 log ef ( xi , yi ) 1 K ∑K j=1 e f ( xi , yj ) ] , INCE ( X ; Y ) , ( 3 ) where the expectation is over K independent samples { ( xi , yi ) } Ki=1 from the joint distribution p ( x , y ) ( Poole et al. , 2019 ) . In practice we estimate ( 3 ) using Monte Carlo estimation by averaging over multiple batches of samples . Intuitively , the critic function f tries to predict for each xi which of the K samples y1 , . . . , yk it was jointly drawn with , by assigning high values to the jointly drawn pair , and low values to all other pairs . The second estimator is based on the variational form of the KL divergence due to Nguyen , Wainwright , and Jordan ( NWJ ) ( Nguyen et al. , 2010 ) and takes the form I ( X ; Y ) ≥ Ep ( x , y ) [ f ( x , y ) ] − e−1Ep ( x ) [ Ep ( y ) ef ( x , y ) ] , INWJ ( X ; Y ) . ( 4 ) For detailed derivations we refer the reader to ( Ruderman et al. , 2012 ; Poole et al. , 2019 ) . Note that these bounds hold for any critic f and when used in ( 2 ) one in practice jointly maximizes over g1 , g2 and f . Furthermore , it can be shown that ( 3 ) is maximized by f∗ ( x , y ) = log p ( y|x ) and ( 4 ) by f∗ ( x , y ) = 1 + log p ( y|x ) ( Poole et al. , 2019 ) . Common choices for f include bilinear critics f ( x , y ) = x > Wy ( van den Oord et al. , 2018 ; Hénaff et al. , 2019 ; Tian et al. , 2019 ) , separable critics f ( x , y ) = φ1 ( x ) > φ2 ( y ) ( Bachman et al. , 2019 ) , and concatenated critics f ( x , y ) = φ ( [ x , y ] ) ( Hjelm et al. , 2019 ) ( here φ , φ1 , φ2 are typically shallow multi-layer perceptrons ( MLPs ) ) . When applying these estimators to solve ( 2 ) , the line between the critic and the encoders g1 , g2 can be blurry . For example , one can train with an inner product critic f ( x , y ) = x > y , but extract features from an intermediate layer of g1 , g2 , in which case the top layers of g1 , g2 form a separable critic . Nevertheless , this boundary is crucial for the interplay between MI estimation and the interpretation of the learned representations .

Summary: The paper addresses a question on whether mutual information (MI) based models for representation learning succeed primarily thanks to the MI maximization. The motivation of the work comes from the fact that although MI is known to be problematic in treatment, it has been successfully applied in a number of recent works in computer vision and natural language processing. The paper conducts a series of experiments that constitute a convincing evidence for a weak connection between the InfoMax principle and these practical successes by showing that maximizing established lower bounds on MI are not predictive of the downstream performance and that contrary to the theory higher capacity instantiations of the critics of MI may result in worse downstream performance of learned representations. The paper concludes that there is a considerable inductive bias in the architectural choices inside MI models that are beneficial for downstream tasks and note that at least one of the lower bounds on MI can be interpreted as a triplet loss connecting it with a metric learning approach.