Sharp Learning Bounds for Contrastive Unsupervised Representation Learning

1 INTRODUCTION . The contrastive loss ( Chopra et al. , 2005 ) is one of the popular loss functions in metric learning ( Kulis , 2012 ) and representation learning ( Bengio et al. , 2013 ) . The contrastive loss forces data representation of semantically similar pairs closer in some metric space than multiple random samples , called negative samples . Many state-of-the-art representation learning algorithms use a type of contrastive losses in natural language processing ( Mikolov et al. , 2013 ; Logeswaran & Lee , 2018 ) , vision ( Chopra et al. , 2005 ; He et al. , 2019 ; Chen et al. , 2020 ) , and graph ( Lirong et al. , 2021 ) domains . A simple model built on top of the learned representation can achieve almost the same accuracy as supervised learning does . See surveys ( Le-Khac et al. , 2020 ; Schmarje et al. , 2021 ) and references therein for recent empirical progress . Compared to its empirical success , we still do not know much about the theoretical perspective of the contrastive loss function . Since the first theoretical analysis of contrastive unsupervised representation learning ( CURL ) ( Arora et al. , 2019 ) , a number of follow-up researches ( Nozawa et al. , 2020 ; Chuang et al. , 2020 ; Nozawa & Sato , 2021 ; Ash et al. , 2021 ) improved theoretical understanding of CURL . They provided upper bounds of a downstream classification loss by the contrastive loss . While they partially explain why CURL provides good representation in terms of downstream classification , a controversy has been held on the impact of the negative sample size , denoted byK in this work . From the empirical viewpoint , it has been widely shown that downstream classification accuracy tends to be better with the largerK ( He et al. , 2020 ; Chen et al. , 2020 ) .1 From the theoretical viewpoint , although Ash et al . ( 2021 ) and Nozawa & Sato ( 2021 ) addressed the issue of the earlier analysis ( Arora et al. , 2019 ) —the bound grows exponentially in K—their improved bounds are still so huge that they may not capture the actual behavior of CURL related to K , as we will see in Figure 1 . In addition , Ash et al . ( 2021 ) claims that the optimal K exists in smaller K , which does not agree well with the empirical observations . Henceforth , the relationship between the performance of downstream classification and K needs to be investigated carefully . This study aims to clarify the relationship between the performance of downstream classification and the negative sample size K by deriving an upper bound of the downstream classification loss 1The empirical evidence has also not reached a consensus yet . Chen et al . ( 2021 ) reported that SimCLR with the smaller K could be competitive in linear evaluation than K used by the original SimCLR ( Chen et al. , 2020 ) . ( Theorem 1 ) . By deriving the lower bound ( Theorem 2 ) , we found that the intercept of our upper bound is tight in K. Notably , such a lower bound has yet to be known in CURL so far . Our bounds support that the larger K reduces estimation gap of the downstream classification loss ( Section 3.2 ) , while the smaller K could perform well . This finding is evidence for the empirical success of the larger K. In addition , our bounds are compared with the existing bounds to demonstrate that there is no optimal point in K and how accurately the derived bounds capture the downstream loss ( Section 4 ) . Finally , we empirically verify our theory by experiments ( Section 6 ) on a synthetic dataset , CIFAR-10/100 ( Krizhevsky , 2009 ) datasets , and Wiki-3029 dataset ( Arora et al. , 2019 ) . 2 FORMULATION OF CONTRASTIVE LEARNING . First , this section briefly summarizes the problem setup and formulation of CURL . Notation . The C-dimensional vector whose elements are all one is denoted by 1C : = [ 1 1 . . . 1 ] > . When it is clear from context , the subscript is abbreviated . For a vector a ∈ Rp , a ( i ) denotes the i-th largest element of a , namely , a ( 1 ) ≥ a ( 2 ) ≥ · · · ≥ a ( p ) . Likewise , a ( −i ) denotes the i-th smallest element of the vector a . The indicator function is denoted by 1 { A } for a predicate A . Let 4C : = { p ∈ [ 0 , 1 ] C | p > 1 = 1 } be the C-dimensional probability simplex . For p ∈ 4C , the Shannon entropy is denoted by H ( p ) : = − ∑ c∈ [ C ] pc ln pc . Let Hn be the n-th harmonic number . Supervised classification . One of the goals in machine learning is supervised classification , while we consider the setup where supervision is unavailable . Here , we first formulate the multiclass classification problem . Assume that there exists an unknown number of latent classes denoted by C ∈ N. Let X be d-dimensional feature space and Y : = [ C ] be the supervised label set.2 In the supervised setup , we are interested in the following risk quantity , the supervised loss , for a multi-class classifier g : X → RC : Rsupv ( g ) : = E x , y∼P [ − ln exp ( gy ( x ) ) ∑ c∈Y exp ( gc ( x ) ) ] , ( 1 ) which is specialized for the softmax cross-entropy loss . The expectation is taken over the unknown underlying joint distribution P. The classifier is eventually used for prediction by argmaxy∈Y gy ( · ) . Contrastive unsupervised representation learning ( CURL ) . In the CURL framework ( Arora et al. , 2019 ) , we target to learn meaningful data representation in an unsupervised manner . By doing so , a similarity model is trained to make the representation of positive pairs more similar than randomly drawn K negative samples using a contrastive loss . The class-conditional distribution is denoted by Dc : = P ( X | Y = c ) for each c ∈ Y and the class-prior distribution by π : = [ P ( Y = c ) ] c∈Y ∈ 4C . The data generating process is described as follows : ( i ) draw latent positive/negative classes : c+ , 2Strictly speaking , latent classes should be distinguished from the supervised labels as they may contain concepts that human annotators are not interested in . In the CURL framework ( Arora et al. , 2019 ) , latent classes [ C ] and the supervised label set Y are allowed to be different , i.e. , Y ⊆ [ C ] . To focus on the relationship between the supervised loss and contrastive loss , we analyze under the assumption Y = [ C ] for simplicity . { c−k } Kk=1 ∼ πK+1 ( ii ) draw an anchor sample x ∼ Dc+ ( iii ) draw a positive sample x+ ∼ Dc+3 ( iv ) draw K negative samples x−k ∼ Dc−k ( for each k ∈ [ K ] ) . Now we specify our model . A multi-class classifier g : X → RC consists of representation f : X → Rh and linear parameters W ∈ RC×h as g ( · ) : = Wf ( · ) , where h ∈ N denotes the dimensionality of the representation given in advance . In CURL , the representation is learned through minimization of the following contrastive loss Rcont ( f ) : = E c+ , { c−k } K k=1 ∼πK+1 E x , x+∼D2 c+ x−k ∼Dc− k [ − ln exp ( f ( x ) > f ( x+ ) ) exp ( f ( x ) > f ( x+ ) ) + ∑ k∈ [ K ] exp ( f ( x ) > f ( x−k ) ) ] . ( 2 ) In the downstream task , the representation f shall be frozen , and the parameters W are to be learned via minimizing the supervised loss Rsupv . Evaluation of representations . For the sake of evaluation , a specific linear classifier called mean classifier is introduced . Given representation f , the mean classifier Wµ is defined as Wµ : = [ µ1 · · · , µC ] > , where µc : = Ex∼Dc [ f ( x ) ] . This will later be used for evaluating the representation f combined with the supervised loss , which is denoted by Rµ-supv ( f ) : = Rsupv ( Wµf ) . We call it the mean supervised loss . If the mean supervised loss is successfully bounded from above , we end up a bound on the supervised loss through infW∈RC×h Rsupv ( Wf ) ≤ Rµ-supv ( f ) . For this reason , an upper bound on Rµ-supv is an intermediate milestone that we seek throughout this paper . 3 LEARNING BOUNDS FOR CONTRASTIVE LEARNING . In this section , our main theoretical results are provided . We aim at showing that the contrastive loss Rcont ( f ) —the computable functional in CURL—serves as a good estimator of the mean supervised loss Rµ-supv ( f ) —the inaccessible target in CURL—for any f . We show this by establishing upper and lower bounds of Rµ-supv ( f ) by Rcont ( f ) . Eventually , the minimization of Rcont ( f ) may lead to a good minimizer of Rµ-supv ( f ) . All proofs are provided in Appendix B . 3.1 MAIN RESULTS . First , we show a sharp upper bound of the mean supervised loss . Unlike any of the existing learning bounds of CURL , the upper bound obtained here has a constant coefficient in the contrastive loss and is applicable for all C and K ( see discussions in Section 4 ) .4 Theorem 1 . For all f such that ‖f ( x ) ‖2 ≤ L ( ∀x ∈ X ) , the following inequality holds . Rµ-supv ( f ) ≤ Rcont ( f ) + ∆U , ( 3 ) where ∆U : = 2 ln ( ( CKπ ( −1 ) + 1 ) cosh ( L2 ) ) − 2 ln ( 1 +K ) − lnKπ ( −1 ) . Next , the first lower bound of the mean supervised loss is provided . While the existing theoretical analyses only provided upper bounds , they often entail a huge coefficient in the contrastive loss . The lower bound provided below has the same constant coefficient and intercept ( ∆U and ∆L ) rate as our upper bound , ensuring the tightness of our analysis . Theorem 2 . For all f such that ‖f ( x ) ‖2 ≤ L ( ∀x ∈ X ) , the following inequality holds . Rµ-supv ( f ) ≥ Rcont ( f ) + ∆L , ( 4 ) where ∆L : = H ( π ) + ln K ( K+1 ) 2 − 2 ln cosh ( L 2 ) . 3In self-supervised learning , data augmentation ( DA ) is commonly used to obtain a positive sample . While Nozawa & Sato ( 2021 ) handled DA in this formulation , we omit it for simplicity to focus on the effect of K. 4Even if Y = [ C ] is assumed in Section 2 , Theorem 1 can be extended to the case Y ⊆ [ C ] ( subset ) and Y is a coarse-grained set of [ C ] . In contrast , Theorem 2 can only be extended to the coarse-grained Y . The details are discussed in their proofs . Our proofs leverage that the contrastive loss and mean supervised loss share the similar log-sum-exp functional form , which leads to the direct relationships between the two losses . This is in contrast to the existing works including Arora et al . ( 2019 ) , which approximate the mean supervised loss with the contrastive loss by taking the expectation over latent classes , leading to an exponentially large coefficient . In Theorems 1 and 2 , the size of the representation ‖f ( x ) ‖2 is assumed to be bounded . This assumption is reasonable from the experimental perspective since it is common to normalize representation to employ the cosine similarity as the similarity metric . Several works reported that the normalized embeddings improve the performance ( Chen et al. , 2020 ; Wang & Isola , 2020 ) . Unlike the existing analyses ( reviewed in Section 4 ) , we take advantage of this assumption to derive the sharp bounds . As we see in Section 3.2 , ∆U and ∆L are the same order in K under the uniform class prior assumption . By applying either the high-probability bound ( Arora et al. , 2019 ) or PAC-Bayesian analysis ( Nozawa et al. , 2020 ) , Theorem 1 ( Theorem 2 as well ) can be naturally extended to the form Rµ-supv ( f̂ ) ≤ Rcont ( f ) + ∆U + χ with a complexity term χ , where f̂ is the empirical minimizer of the contrastive loss . Since this is a routine , we omit the high-probability bounds . 3.2 DISCUSSION Subsequently , we discuss implications of our main results on the relationship between the mean supervised loss and the negative sample size K. For the sake of simplicity , we assume πc = 1/C for all c ∈ [ C ] ( the uniform class prior ) in this section . Gap between learning upper and lower bounds . Both of our upper ( Theorem 1 ) and lower ( Theorem 2 ) bounds draw the linear relationship between the mean supervised loss Rµ-supv and the contrastive loss Rcont , with the additional intercept terms ∆U and ∆L . Under the uniform class prior assumption , the intercepts are in the same order : ∆U = ln C K + 2 ln cosh ( L2 ) = O ( ln 1 K ) , ∆L = ln CK ( K + 1 ) 2 − 2 ln cosh ( L2 ) = O ( ln 1 K ) , ( 5 ) and the gap between two bounds is ∆U−∆L = 4 ln cosh ( L2 ) + 2 ln ( 1 + 1K ) , meaning that the gap shrinks to 4 ln cosh ( L2 ) as K increases . Hence , our bounds have the tight intercepts , and the larger K is beneficial for CURL from the viewpoint of the estimation gap of the mean supervised loss . Learning bounds and feasible region . Next , we consider the ( Rcont , Rµ-supv ) -plot , in which a point indicates ( Rcont ( f ) , Rµ-supv ( f ) ) for some f ( see Figure 2 ) . Here , let us focus on the feasible region in the ( Rcont , Rµ-supv ) -plot by assuming ‖f‖2 ≤ L for any f ( same as Theorems 1 and 2 ) . Then , the mean supervised loss and contrastive loss are essentially lower-bounded by the constants5 R∗µ-supv : = ln ( 1 + ( C − 1 ) exp ( −2L2 ) ) , ( 6 ) R∗cont : = K∑ m=0 ( K m ) ( 1 C ) m ( 1− 1 C ) K−m ln { 1 +m+ ( K −m ) exp ( −2L2 ) } , ( 7 ) respectively . Hence , the feasible region is Rµ-supv ( f ) ≤ Rcont ( f ) + ∆U , Rµ-supv ( f ) ≥ Rcont ( f ) + ∆L , ( 8a ) Rµ-supv ( f ) ≥ R∗µ-supv , Rcont ( f ) ≥ R∗cont , ( 8b ) as illustrated in Figure 2 . The first two bounds ( 8a ) restrict the mean supervised loss by the contrastive loss . We specifically refer to these bounds as learning bounds . The remaining two bounds ( 8b ) 5The derivations of R∗µ-supv and R∗cont are detailed in Appendix C. represent the achievable limits for each loss separately . One of the important questions is how the smallest possible value of Rµ-supv in the feasible region ( 8 ) changes as K and C change . In other words , we are interested in whether the optimal point ( R∗cont , R ∗ µ-supv ) ( F in Figure 2 ) is always achievable regardless of the values of K and C. To investigate it , we check the following two points . • The feasible region at Rcont ( f ) = R∗cont ( Figure 3a ) : We plot the value R ∗ cont + ∆U ( solid line ; the Rµ-supv-value of the point • in Figure 2 ) and the minimum possible Rµ-supv ( R∗µ-supv ; dotted line ) numerically . These two curves do not cross for allK , which meansRµ-supv ( f ) = R∗µ-supv is attainable no matter the values K and C. In addition , the bound becomes sharper as K increases , but the gap between the upper bound and R∗µ-supv does remain even at the limit K ↗∞ . • The feasible region at Rµ-supv ( f ) = R∗µ-supv ( Figure 3b ) : When Rµ-supv ( f ) = R ∗ µ-supv , the con- trastive loss Rcont ( f ) is upper-bounded by R∗µ-supv − ∆L ( the Rcont-value of the point in Figure 2 ) . The curve of this value does not cross R∗cont , which tells us that the lower bound does not exclude the optimal point ( R∗cont , R ∗ µ-supv ) from the feasible region ( 8 ) at any K. Note that the gap between R∗cont and R ∗ µ-supv −∆L gradually increases , meaning that it becomes much easier to attain R∗µ-supv as K increases . Hence , the optimal point F stays in the feasible region ( 8 ) no matter the value K. From this viewpoint , smaller K is not necessarily disadvantageous because the optimal point F remains in the feasible region . Note again that the estimation of Rµ-supv may become harder with the smaller K because of the bound gap ∆U −∆L = O ( ln 1/K ) , even if the optimal solution is unaffected by K. Summary . We draw a connection between the mean supervised loss and the negative sample size K by the following claim : if we regard CURL as an estimation process of the mean supervised loss by the contrastive loss , the small K does not necessarily have the estimation bias , but the larger K helps reduce the estimation gap .

Summary: === After Rebuttal ===  The author(s) have made many clarifications and improvements to the paper per reviewers' requests. I am happy to bump up my score by 1 to reflect the changes & promises made.   =================  This paper provides some new theoretical insights on how contrastive representation help to improve downstream classification performance. In particular, it fixes the mismatch between prior theoretical predictions and empirical observations, that a larger negative sample size leads to better downstream performance.