Towards Robust and Efficient Contrastive Textual Representation Learning

1 INTRODUCTION . Representation learning is one of the pivotal topics in natural language processing ( NLP ) , in both supervised and unsupervised settings . It has been widely recognized that some forms of “ general representation ” exist beyond specific applications ( Oord et al. , 2018 ) . To extract such generalizable features , unsupervised representation models are generally pretrained on large-scale text corpora ( e.g. , BERT ( Devlin et al. , 2018 ; Liu et al. , 2019 ; Clark et al. , 2020 ; Lagler et al. , 2013 ) ) to avoid data bias . In supervised learning , models are typically built on top of these pre-trained representations and further fine-tuned on downstream tasks . Representation learning greatly expedites model deployment and meanwhile yields performance gains . There has been growing interest in exploiting contrastive learning ( CL ) techniques to refine context representations in NLP ( Mikolov et al. , 2013a ; b ) . These techniques aim to avoid representation collapse for downstream tasks , i.e. , getting similar output sentences with different input in conditional generation tasks ( Dai & Lin , 2017 ) . Intuitively , these methods carefully engineer features from crafted ( “ negative ” ) examples to contrast against the features from real ( “ positive ” ) examples . A feature encoder can then enhance its representation power by characterizing input texts at a finer granularity . Efforts have been made to empirically investigate and theoretically understand the effectiveness of CL in NLP , including noise contrastive estimation ( NCE ) of word embeddings ( Mikolov et al. , 2013b ) and probabilistic machine translation ( Vaswani et al. , 2013 ) with theoretical developments ( Gutmann & Hyvärinen , 2010 ) . More recently , InfoNCE ( Oord et al. , 2018 ) further links the CL to the optimization of mutual information , which inspired a series of practical followup works ( Tian et al. , 2020 ; Hjelm et al. , 2019a ; He et al. , 2020 ; Chen et al. , 2020 ) . Despite the significant empirical success of CL , there are still many open challenges in its application to NLP , including ( i ) the propagation of stable contrastive signals . An unregularized critic function in CL can suffer from unstable training and gradient vanishing issues , especially in NLP tasks due to the discrete nature of text . The inherent differences between positive and negative textual features make those examples easily distinguished , resulting in a weak learning signal in contrastive schemes ( Arora et al. , 2019 ) . ( ii ) Empirical evidence ( Wu et al. , 2017 ) shows that it is crucial to compare each positive example with adequate negative examples . However , recent works suggest using abundant negative examples , which are not akin to the positive examples , which can result in sub-optimal results and unstable training with additional computational overhead ( Ozair et al. , 2019 ; McAllester & Stratos , 2020 ) . In this paper , we propose two methods to mitigate the above issues . In order to stabilize the training and enhance the model ’ s generalization ability , we propose to use the Wasserstein dependency measure ( Ozair et al. , 2019 ) as a substitute for the Kullback-Leibler ( KL ) measure in the vanilla CL objective . We further actively select K high-quality negative samples to contrast with each positive sample under the current learned representations . These supply the training procedure with necessarily large and non-trivial contrastive samples , encouraging the representation network to generate more distinguishable features . Notably , our approach also significantly alleviates the computational burden of massive features compared with previous works ( Tian et al. , 2020 ; Hjelm et al. , 2019b ) . Contributions : ( i ) We propose a Wasserstein-regularized critic to stabilize training in a generic CL framework for learning better textual representations . ( ii ) We further employ an active negativesample selection method to find high-quality contrastive samples , thus reducing the gradient noise and mitigating the computation concerns . ( iii ) We empirically verify the effectiveness of our approach under various NLP tasks , including variational text generation ( Bowman et al. , 2016 ) , natural language understanding tasks on GLUE with supervised and semi-supervised setups ( Wang et al. , 2018 ) , and image-text retrieval ( Lee et al. , 2018 ) . 2 BACKGROUND . 2.1 NOISE CONTRASTIVE ESTIMATION . Our formulation is inspired by Noise Contrastive Estimation ( NCE ) ( Gutmann & Hyvärinen , 2010 ) , which was originally introduced for unnormalized density estimation , where the partition functions is intractable . To estimate a parametric distribution p , which we refer to as our target distribution , NCE leverages not only the observed samples A = ( a1 , a2 , ... , an1 ) ( positive samples ) , but also the samples drawn from a reference distribution q , denoted as B = ( b1 , b2 , ... , bn2 ) ( negative samples ) . Instead of estimating p directly , the density ratio p/q is estimated by training a critic between samples from A and B . Specifically , let Z = ( z1 , ... , zn1+n2 ) denote the union of A and B . A binary class label Ct is assigned to each zt , where Ct = 1 if zt ∈ A and Ct = 0 otherwise . The label probability is therefore P ( C = 1|z ) = p ( z ) p ( z ) + γq ( z ) , P ( C = 0|z ) = γq ( z ) p ( z ) + γq ( z ) , ( 1 ) where γ = n2n1 is a balancing hyperparameter accounting for the difference in number of samples between A and B . In practice , we do not know the analytic form of p ; therefore , a classifier g : z 7→ [ 0 , 1 ] to estimate p ( C = 1|z ) is trained . To get an estimation of the critic function g , NCE maximizes the log likelihood of the data for a binary classification task : L ( A , B ) = n1∑ t=1 log [ g ( at ) ] + n2∑ t=1 log [ 1− g ( bt ) ] . ( 2 ) 2.2 CONTRASTIVE TEXTUAL REPRESENTATION LEARNING AND ITS CHALLENGES . Let { wi } ni=1 be the observed text instances . We are interested in finding a vector representation u of the text w , i.e. , via an encoder u = Enc ( w ) , that can be repurposed for downstream tasks . A positive pair refers to paired instances ai = ( ui , vi ) associated with wi , where we are mostly interested in learning u ; v is a feature at a different representation level . In unsupervised scenarios , vi can be the feature representation at the layer next to the input text wi . In supervised scenarios , vi can be either the feature representation layer immediately after wi or immediately before the label yi that corresponds to the input wi . We will use π ( u , v ) to denote the joint distribution of the positive pairs , with πu ( u ) and πv ( v ) for the respective marginals . Contrastive learning follows the principle of “ learning by comparison. ” Specifically , one designs a negative sample distribution τ ( u′ , v′ ) , and attempts to distinguish samples from π ( u , v ) and τ ( u′ , v′ ) with a critic function g ( u , v ) . The heuristic is that , using samples from τ as references ( i.e. , to contrast against ) , the learner is advantaged to capture important properties that could have been otherwise missed ( Hjelm et al. , 2019a ; Oord et al. , 2018 ) . A popular choice of τ is the product of marginals , i.e. , τ ← π0 ( u′ , v′ ) = πu ( u′ ) πv ( v′ ) where ( u , v ) are independent of each other , so that bi = ( u ′ i , v ′ i ) ∼ π0 . Inputting the new ai and bi to ( 2 ) , we obtain the new CL loss : LNCE = Eu , v∼π [ log g ( u , v ) ] + γEu′ , v′∼π0 [ log ( 1− g ( u′ , v′ ) ) ] . ( 3 ) Note that when g is trained to optimality g∗ ( u , v ) = p ( C = 1|u , v ) under π0 , it establishes a lower bound of mutual information ( MI ) between u and v for the positive distribution ( Tian et al. , 2020 ; Neyshabur et al. , 2018 ) : MI ( πu , πv ) = KL ( π ( u , v ) ||πu ( u ) πv ( v ) ) ≥ Eu , v∼π [ log g∗ ( u , v ) ] + log γ . ( 4 ) However , there are three concerns regarding why directly applying equation 3 might not be good in practice for learning the contrastive representation of the input text w. • Robustness . The first issue concerns the MI ’ s strong sensitivity to small differences in data samples ( Ozair et al. , 2019 ; Tschannen et al. , 2020 ) . By definition in equation 4 , mutual information is a KL divergence . It is well known that the KL divergence is not a metric-aware divergence measure , which implies a minor difference in representation can induce drastic changes in the mutual information , as a special case of KL . Consequently , the learned g could be numerically unstable ( Ozair et al. , 2019 ) , which makes the learned representations less robust and does not generalize well to downstream tasks , especially when features come from text ( Chen et al. , 2018 ) . • Weak/vanishing contrastive signal . With a poor initialization or a poor choice of negative samples , the MI will vanish as the π ( u , v ) and πu ( u ) πv ( v ) become far apart , delivering a faint and nonsmooth gradient for training . In an extreme case , the support for π ( u , v ) and πu ( u ) πv ( v ) do not overlap , and the MI and the gradient will vanish to zero ( Arjovsky et al. , 2017 ) . • Negative-sample selection strategy . Learning MI is generally considered sample inefficient . This point can be corroborated from several perspectives , ranging from theoretical arguments to practical considerations . To confidently estimate a lower bound to the MI , one would need a sample size exponential to the mutual information ( i.e. , N ≥ exp ( Iπ ( u , v ) ) ) ( Ozair et al. , 2019 ; McAllester & Stratos , 2018 ) . Also , both theoretical prediction and empirical evidence suggest a large ratio γ is needed for good performance ( Tian et al. , 2020 ; Hjelm et al. , 2019a ) , imposing potential computational concerns for large training datasets . On the other hand , some studies report a large γ can instead deteriorate model performance ( Tschannen et al. , 2020 ; Arora et al. , 2019 ) . Such a large γ is also believed to be problematic especially when a strong association is expected between u and v. In that case , the majority of negative samples are so different from positive samples that the comparisons do not lend effective learning signals , but instead randomly drift the training ( Gutmann & Hyvärinen , 2010 ) .

Summary: This paper proposes an approach to improve (supervised and unsupervised) representation learning for text using constrastive learning. The proposed approach augments standard contrastive learning with: (1) Spectral-norm regularization of the critic to estimate the Wasserstein distance instead of the KL (as in the Wasserstein GAN-style approach), (2) Active negative sampling to select hard negative examples, and (3) momentum-based updates of intermediate features. The resulting contrastive learning objective can be combined with standard supervised and unsupervised objectives to improve downstream tasks.