Decomposing Mutual Information for Representation Learning

1 INTRODUCTION . The ability to extract actionable information from data in the absence of explicit supervision seems to be a core prerequisite for building systems that can , for instance , learn from few data points or quickly make analogies and transfer to other tasks . Approaches to this problem include generative models ( Hinton , 2012 ; Kingma & Welling , 2014 ) and self-supervised representation learning approaches , in which the objective is not to maximize likelihood , but to formulate a series of ( label-agnostic ) tasks that the model needs to solve through its representations ( Noroozi & Favaro , 2016 ; Devlin et al. , 2019 ; Gidaris et al. , 2018 ; Hjelm et al. , 2019 ) . Self-supervised learning includes successful models leveraging contrastive learning , which have recently attained comparable performance to their fully-supervised counterparts ( Bachman et al. , 2019 ; Chen et al. , 2020a ) . Many self-supervised learning methods train an encoder such that the representations of a pair of views x and y derived from the same input example are more similar to each other than to representations of views sampled from a contrastive negative sample distribution , which is usually the marginal distribution of the data . For images , different views can be built using random flipping , color jittering and cropping ( Bachman et al. , 2019 ; Chen et al. , 2020a ) . For sequential data such as conversational text , the views can be past and future utterances in a given dialogue . It can be shown that these methods maximize a lower bound on mutual information ( MI ) between the views , I ( x ; y ) , w.r.t . the encoder , i.e . the InfoNCE bound ( Oord et al. , 2018 ) . One significant shortcoming of this approach is the large number of contrastive samples required , which directly impacts the total amount of information which the bound can measure ( McAllester & Stratos , 2018 ; Poole et al. , 2019 ) . In this paper , we consider creating subviews of x by removing information from it in various ways , e.g . by masking some pixels . Then , we use representations from less informed subviews as a source of hard contrastive samples for representations from more informed subviews . For example , in Fig . 1 , one can mask a pixel region in x′ to obtain x′′ and ask ( the representation of ) x′′ to be closer to y than to random images of the corpus , and for x′ to be closer to y than to samples from p ( y|x′′ ) . This corresponds to decomposing the MI between x and y into I ( x ; y ) ≥ I ( x′′ ; y ) + I ( x′ ; y|x′′ ) . The conditional MI measures the information about y that the model has gained by looking at x′ beyond the information already contained in x′′ . In Fig . 1 ( left ) , standard contrastive approaches could focus on the overall “ shape ” of the object and would need many negative samples to capture other discriminative features . In our approach , the model is more directly encouraged to capture these additional features , e.g . the embossed detailing . In the context of predictive coding on sequential data such as dialogue , by setting x′′ to be the most recent utterance ( Fig . 1 , right ) , the encoder is directly encouraged to capture long-term dependencies that can not be explained by x′′ . We formally show that , by such decomposition , our representations can potentially capture more of the total information shared between the original views x and y . Maximizing MI between multiple views can be related to recent efforts in representation learning , amongst them AMDIM ( Bachman et al. , 2019 ) , CMC ( Tian et al. , 2019 ) and SwAV ( Caron et al. , 2020 ) . However , these models maximize the sum of MIs between views I ( { x′ , x′′ } ; y ) = I ( x′′ ; y ) + I ( x′ ; y ) . E.g. , in Bachman et al . ( 2019 ) , x′ and x′′ could be global and local representations of an image , and in Caron et al . ( 2020 ) , x′ and x′′ could be the views resulting from standard cropping and the aggressive multi-crop strategy . This equality is only valid when the views x′ and x′′ are statistically independent , which usually does not hold . Instead , we argue that a better decomposition is I ( { x′ , x′′ } ; y ) = I ( x′′ ; y ) + I ( x′ ; y|x′′ ) , which always holds . Most importantly , the conditional MI term encourages the encoder to capture more non-redundant information across views . To maximize our proposed decomposition , we present a novel lower bound on conditional MI in Section 3 . For the conditional MI maximization , we give a computationally tractable approximation that adds minimal overhead . In Section 4 , we first show in a synthetic setting that decomposing MI and using the proposed conditional MI bound leads to capturing more of the ground-truth MI . Finally , we present evidence of the effectiveness of the method in vision and in dialogue generation . 2 PROBLEM SETTING . The maximum MI predictive coding framework ( McAllester , 2018 ; Oord et al. , 2018 ; Hjelm et al. , 2019 ) prescribes learning representations of input data such that they maximize MI . Estimating MI is generally a hard problem that has received a lot of attention in the community ( Kraskov et al. , 2004 ; Barber & Agakov , 2003 ) . Let x and y be two random variables which can generally describe input data from various domains , e.g . text , images or sound . We can learn representations of x and y by maximizing the MI of the respective features produced by encoders f , g : X → Rd , which by the data processing inequality , is bounded by I ( x ; y ) : arg max f , g I ( f ( x ) ; g ( y ) ) ≤ I ( x ; y ) . ( 1 ) We assume that the encoders can be shared , i.e . f = g. The optimization in Eq . 1 is challenging but can be lower-bounded . Our starting point is the recently proposed InfoNCE lower bound on MI ( Oord et al. , 2018 ) and its application to self-supervised learning for visual representations ( Bachman et al. , 2019 ; Chen et al. , 2020a ) . In this setting , x and y are paired input images , or independentlyaugmented copies of the same image . These are encoded using a neural network encoder which is trained such that the representations of the two image copies are closer to each other in the embedding space than to other images drawn from the marginal distribution of the corpus . This can be viewed as a contrastive estimation of the MI ( Oord et al. , 2018 ) . We present the InfoNCE bound next . 2.1 INFONCE BOUND . InfoNCE ( Oord et al. , 2018 ) is a lower-bound on I ( x ; y ) obtained by comparing pairs sampled from the joint distribution x , y1 ∼ p ( x , y ) to a set of negative samples , y2 : K ∼ p ( y2 : K ) = ∏K k=2 p ( yk ) , also called contrastive , independently sampled from the marginal : INCE ( x ; y|E , K ) = Ep ( x , y1 ) p ( y2 : K ) [ log eE ( x , y1 ) 1 K ∑K k=1 e E ( x , yk ) ] ≤ I ( x , y ) , ( 2 ) where E is a critic assigning a real valued score to x , y pairs . We provide an exact derivation for this bound in the Appendix1 . For this bound , the optimal critic is the log-odds between the conditional distribution p ( y|x ) and the marginal distribution of y , E∗ ( x , y ) = log p ( y|x ) p ( y ) + c ( x ) ( Oord et al. , 2018 ; Poole et al. , 2019 ) . The InfoNCE bound is loose if the true mutual information I ( x ; y ) is larger than logK . In order to overcome this difficulty , recent methods either train with large batch sizes ( Chen et al. , 2020a ) or exploit an external memory of negative samples in order to reduce memory requirements ( Chen et al. , 2020b ; Tian et al. , 2020 ) . These methods rely on uniform sampling from the training set in order to form the contrastive sets . For further discussion of the limits of variational bounds of MI , see McAllester & Stratos ( 2018 ) . 3 DECOMPOSING MUTUAL INFORMATION . By the data processing inequality : I ( x ; y ) ≥ I ( { x1 , . . . , xN } ; y ) , where { x1 , . . . , xN } are different subviews of x – i.e. , views derived from x without adding any exogenous information . For example , { x1 , . . . , xN } can represent exchanges in a longer dialog x , sentences in a document x , or different augmentations of the same image x . Equality is obtained when the set of subviews retains all information about x , e.g . if x is in the set . Without loss of generality , we consider the case N = 2 , I ( x ; y ) ≥ I ( { x′ , x′′ } ; y ) , where { x′ , x′′ } indicates two subviews derived from the original x . We can apply the chain rule for MI : I ( x ; y ) ≥ I ( { x′ , x′′ } ; y ) = I ( x′′ ; y ) + I ( x′ ; y|x′′ ) , ( 3 ) where the equality is obtained if and only if I ( x ; y| { x′ , x′′ } ) = 0 , i.e . x doesn ’ t give any information about y in excess to { x′ , x′′ } 2 . This suggests that we can maximize I ( x ; y ) by maximizing each of the MI terms in the sum . The conditional MI term can be written as : I ( x′ ; y|x′′ ) = Ep ( x′ , x′′ , y ) [ log p ( y|x′ , x′′ ) p ( y|x′′ ) ] . ( 4 ) This conditional MI is different from the unconditional MI , I ( x′ ; y ) , insofar it measures the amount of information shared between x′ and y which can not be explained by x′′ . Note that the decomposition holds for arbitrary partitions of x′ , x′′ , e.g . I ( { x′ , x′′ } ; y ) = I ( x′ ; y ) + I ( x′′ ; y|x′ ) . When X is high-dimensional , the amount of mutual information between x and y will potentially be larger than the amount of MI that INCE can measure given computational constraints associated with large K and the poor log scaling properties of the bound . The idea that we put forward is to split the total MI into a sum of MI terms of smaller magnitude , thus for which INCE would have less bias for any given K , and estimate each of those terms in turn . The resulting decomposed bound can be written into a sum of unconditional and conditional MI terms : INCES ( x ; y ) = INCE ( x ′′ ; y ) + ICNCE ( x ′ ; y|x′′ ) ≤ I ( x ; y ) , ( 5 ) 1The derivation in Oord et al . ( 2018 ) presented an approximation and therefore was not properly a bound . An alternative , exact derivation of the bound can be found in Poole et al . ( 2019 ) . 2For a proof of this fact , it suffices to consider I ( { x , x′ , x′′ } ; y ) = I ( x ; y| { x′ , x′′ } ) + I ( { x′ , x′′ } ; y ) , given that I ( { x , x′ , x′′ } ; y ) = I ( x ; y ) , equality is obtained iff I ( x ; y| { x′ , x′′ } ) = 0. where ICNCE is a lower-bound on conditional MI and will be presented in the next section . Both conditional ( Eq . 6 ) and unconditional bounds on the MI ( Eq . 14 ) can capture at most logK nats of MI . Therefore , the bound that arises from the decomposition of the MI in Eq . 5 potentially allows to capture up to N log K nats of MI in total , where N is the number of subviews used to describe x . This shows that measuring mutual information by decomposing it in a sequence of estimation problems potentially allows to capture more nats of MI than with the standard INCE , which is bounded by log K .

Summary: This paper proposes a contrastive learning approach where one of the views, x, is converted into two subviews, x' and x'', and then separate InfoNCE style bounds constructed for each of I(x'';y) and I(x';y|x'') before being combined to form an overall training objective.  Critically, the second of these is based on the conditional MI, I(x';y|x''), distinguishing it from previous work using multiple views that just take the marginal I(x';y).  Estimating this conditional MI transpires to be somewhat trickier due to the additional intractability from p(y|x''), with approximations suggested to get around this.  Experiments are performed on both vision and NLP problems.