Self-supervised Representation Learning with Relative Predictive Coding

1 INTRODUCTION . Unsupervised learning has drawn tremendous attention recently because it can extract rich representations without label supervision . Self-supervised learning , a subset of unsupervised learning , learns representations by allowing the data to provide supervision ( Devlin et al. , 2018 ) . Among its mainstream strategies , self-supervised contrastive learning has been successful in visual object recognition ( He et al. , 2020 ; Tian et al. , 2019 ; Chen et al. , 2020c ) , speech recognition ( Oord et al. , 2018 ; Rivière et al. , 2020 ) , language modeling ( Kong et al. , 2019 ) , graph representation learning ( Velickovic et al. , 2019 ) and reinforcement learning ( Kipf et al. , 2019 ) . The idea of self-supervised contrastive learning is to learn latent representations such that related instances ( e.g. , patches from the same image ; defined as positive pairs ) will have representations within close distance , while unrelated instances ( e.g. , patches from two different images ; defined as negative pairs ) will have distant representations ( Arora et al. , 2019 ) . Prior work has formulated the contrastive learning objectives as maximizing the divergence between the distribution of related and unrelated instances . In this regard , different divergence measurement often leads to different loss function design . For example , variational mutual information ( MI ) estimation ( Poole et al. , 2019 ) inspires Contrastive Predictive Coding ( CPC ) ( Oord et al. , 2018 ) . Note that MI is also the KL-divergence between the distributions of related and unrelated instances ( Cover & Thomas , 2012 ) . While the choices of the contrastive learning objectives are abundant ( Hjelm et al. , 2018 ; Poole et al. , 2019 ; Ozair et al. , 2019 ) , we point out that there are three challenges faced by existing methods . The first challenge is the training stability , where an unstable training process with high variance may be problematic . For example , Hjelm et al . ( 2018 ) ; Tschannen et al . ( 2019 ) ; Tsai et al . ( 2020b ) show that the contrastive objectives with large variance cause numerical issues and have a poor downstream performance with their learned representations . The second challenge is the sensitivity to minibatch size , where the objectives requiring a huge minibatch size may restrict their practical usage . For instance , SimCLRv2 ( Chen et al. , 2020c ) utilizes CPC as its contrastive objective and reaches state-of-the-art performances on multiple self-supervised and semi-supervised benchmarks . Nonetheless , the objective is trained with a minibatch size of 8 , 192 , and this scale of training requires enormous computational power . The third challenge is the downstream task performance , which is the one that we would like to emphasize the most . For this reason , in most cases , CPC 1Project page : https : //github.com/martinmamql/relative_predictive_coding is the objective that we would adopt for contrastive representation learning , due to its favorable performance in downstream tasks ( Tschannen et al. , 2019 ; Baevski et al. , 2020 ) . This paper presents a new contrastive representation learning objective : the Relative Predictive Coding ( RPC ) , which attempts to achieve a good balance among these three challenges : training stability , sensitivity to minibatch size , and downstream task performance . At the core of RPC is the relative parameters , which are used to regularize RPC for its boundedness and low variance . From a modeling perspective , the relative parameters act as a ` 2 regularization for RPC . From a statistical perspective , the relative parameters prevent RPC from growing to extreme values , as well as upper bound its variance . In addition to the relative parameters , RPC contains no logarithm and exponential , which are the main cause of the training instability for prior contrastive learning objectives ( Song & Ermon , 2019 ) . To empirically verify the effectiveness of RPC , we consider benchmark self-supervised representation learning tasks , including visual object classification on CIFAR-10/-100 ( Krizhevsky et al. , 2009 ) , STL-10 ( Coates et al. , 2011 ) , and ImageNet ( Russakovsky et al. , 2015 ) and speech recognition on LibriSpeech ( Panayotov et al. , 2015 ) . Comparing RPC to prior contrastive learning objectives , we observe a lower variance during training , a lower minibatch size sensitivity , and consistent performance improvement . Lastly , we also relate RPC with MI estimation , empirically showing that RPC can estimate MI with low variance . 2 PROPOSED METHOD . This paper presents a new contrastive representation learning objective - the Relative Predictive Coding ( RPC ) . At a high level , RPC 1 ) introduces the relative parameters to regularize the objective for boundedness and low variance ; and 2 ) achieves a good balance among the three challenges in the contrastive representation learning objectives : training stability , sensitivity to minibatch size , and downstream task performance . We begin by describing prior contrastive objectives along with their limitations on the three challenges in Section 2.1 . Then , we detail our presented objective and its modeling benefits in Section 2.2 . An overview of different contrastive learning objectives is provided in Table 1 . We defer all the proofs in Appendix . Notation We use an uppercase letter to denote a random variable ( e.g. , X ) , a lower case letter to denote the outcome of this random variable ( e.g. , x ) , and a calligraphy letter to denote the sample space of this random variable ( e.g. , X ) . Next , if the samples ( x , y ) are related ( or positively-paired ) , we refer ( x , y ) ∼ PXY with PXY being the joint distribution of X × Y . If the samples ( x , y ) are unrelated ( negatively-paired ) , we refer ( x , y ) ∼ PXPY with PXPY being the product of marginal distributions overX×Y . Last , we define f ∈ F for F being any class of functions f : X ×Y → R . 2.1 PRELIMINARY . Contrastive representation learning encourages the contrastiveness between the positive and the negative pairs of the representations from the related data X and Y . Specifically , when sampling a pair of representations ( x , y ) from their joint distribution ( ( x , y ) ∼ PXY ) , this pair is defined as a positive pair ; when sampling from the product of marginals ( ( x , y ) ∼ PXPY ) , this pair is defined as a negative pair . Then , Tsai et al . ( 2020b ) formalizes this idea such that the contrastiveness of the representations can be measured by the divergence between PXY and PXPY , where higher divergence suggests better contrastiveness . To better understand prior contrastive learning objectives , we categorize them in terms of different divergence measurements between PXY and PXPY , with their detailed objectives presented in Table 1 . We instantiate the discussion using Contrastive Predictive Coding ( Oord et al. , 2018 , JCPC ) , which is a lower bound of DKL ( PXY ‖PXPY ) with DKL referring to the KL-divergence : JCPC ( X , Y ) : = sup f∈F E ( x , y1 ) ∼PXY , { yj } Nj=2∼PY [ log ef ( x , y1 ) 1 N ∑N j=1 e f ( x , yj ) ] . ( 1 ) Then , Oord et al . ( 2018 ) presents to maximize JCPC ( X , Y ) , so that the learned representations X and Y have high contrastiveness . We note that JCPC has been commonly used in many recent self-supervised representation learning frameworks ( He et al. , 2020 ; Chen et al. , 2020b ) , where they constrain the function to be f ( x , y ) = cosine ( x , y ) with cosine ( · ) being cosine similarity . Under this function design , maximizing JCPC leads the representations of related pairs to be close and representations of unrelated pairs to be distant . The category of modeling DKL ( PXY ‖PXPY ) also includes the Donsker-Varadhan objective ( JDV ( Donsker & Varadhan , 1975 ; Belghazi et al. , 2018 ) ) and the Nguyen-Wainright-Jordan objective ( JNWJ ( Nguyen et al. , 2010 ; Belghazi et al. , 2018 ) ) , where Belghazi et al . ( 2018 ) ; Tsai et al . ( 2020b ) show that JDV ( X , Y ) = JNWJ ( X , Y ) = DKL ( PXY ‖PXPY ) . The other divergence measurements considered in prior work are DJS ( PXY ‖PXPY ) ( with DJS referring to the Jenson-Shannon divergence ) and DWass ( PXY ‖PXPY ) ( with DWass referring to the Wassersteindivergence ) . The instance of modeling DJS ( PXY ‖PXPY ) is the Jensen-Shannon f-GAN objective ( JJS ( Nowozin et al. , 2016 ; Hjelm et al. , 2018 ) ) , where JJS ( X , Y ) = 2 ( DJS ( PXY ‖PXPY ) − log 2 ) .2 The instance of modeling DWass ( PXY ‖PXPY ) is the Wasserstein Predictive Coding ( JWPC ( Ozair et al. , 2019 ) ) , where JWPC ( X , Y ) modifies JCPC ( X , Y ) objective ( equation 1 ) by searching the function from F to FL . FL denotes any class of 1-Lipschitz continuous functions from ( X × Y ) to R , and thus FL ⊂ F . Ozair et al . ( 2019 ) shows that JWPC ( X , Y ) is the lower bound of bothDKL ( PXY ‖PXPY ) andDWass ( PXY ‖PXPY ) . See Table 1 for all the equations . To conclude , the contrastive representation learning objectives are unsupervised representation learning methods that maximize the distribution divergence between PXY and PXPY . The learned representations cause high contrastiveness , and recent work ( Arora et al. , 2019 ; Tsai et al. , 2020a ) theoretically show that highly-contrastive representations could improve the performance on downstream tasks . After discussing prior contrastive representation learning objectives , we point out three challenges in their practical deployments : training stability , sensitivity to minibatch training size , and downstream task performance . In particular , the three challenges can hardly be handled well at the same time , where we highlight the conclusions in Table 1 . Training Stability : The training stability highly relates to the variance of the objectives , where Song & Ermon ( 2019 ) shows that JDV and JNWJ exhibit inevitable high variance due to their inclusion of exponential function . As pointed out by Tsai et al . ( 2020b ) , JCPC , JWPC , and JJS have better training stability because JCPC and JWPC can be realized as a multi-class classification task and JJS can be realized as a binary classification task . The cross-entropy loss adopted in JCPC , JWPC , and JJS is highly-optimized and stable in existing optimization package ( Abadi et al. , 2016 ; Paszke et al. , 2019 ) . Sensitivity to minibatch training size : Among all the prior contrastive representation learning methods , JCPC is known to be sensitive to the minibatch training size ( Ozair et al. , 2019 ) . Taking a closer look at equation 1 , JCPC deploys an instance selection such that y1 should be selected from { y1 , y2 , · · · , yN } , with ( x , y1 ) ∼ PXY , ( x , yj > 1 ) ∼ PXPY with N being the minibatch size . Previous work ( Poole et al. , 2019 ; Song & Ermon , 2019 ; Chen et al. , 2020b ; Caron et al. , 2020 ) showed that a large N results in a more challenging instance selection and forces JCPC to have a better contrastiveness of y1 ( related instance for x ) against { yj } Nj=2 ( unrelated instance for x ) . JDV , JNWJ , and JJS do not consider 2JJS ( X , Y ) achieves its supreme value when f∗ ( x , y ) = log ( p ( x , y ) /p ( x ) p ( y ) ) ( Tsai et al. , 2020b ) . Plugin f∗ ( x , y ) into JJS ( X , Y ) , we can conclude JJS ( X , Y ) = 2 ( DJS ( PXY ‖PXPY ) − log 2 ) . the instance selection , and JWPC reduces the minibatch training size sensitivity by enforcing 1- Lipschitz constraint . Downstream Task Performance : The downstream task performance is what we care the most among all the three challenges . JCPC has been the most popular objective as it manifests superior performance over the other alternatives ( Tschannen et al. , 2019 ; Tsai et al. , 2020b ; a ) . We note that although JWPC shows better performance on Omniglot ( Lake et al. , 2015 ) and CelebA ( Liu et al. , 2015 ) datasets , we empirically find it not generalizing well to CIFAR-10/100 ( Krizhevsky et al. , 2009 ) and ImageNet ( Russakovsky et al. , 2015 ) .

Summary: This paper presents a new contrastive representation objective that has good training stability, minibatch size sensitivity, and downstream task performance. This objective is a generalization of Chi-square divergence, the optimal solution is the density ratio of joint distribution and product of marginal distributions, the estimation is consistent and its variance goes to 0 as sample size goes to infinite, so the paper is theoretical sound. The authors conduct comprehensive experiments to show that the training based on this objective is stable, not sensitive to batch size and leads to good downstream task performance  in vision and phoneme and speaker classification.