C ONTRA G EN : E FFECTIVE C ONTRASTIVE L EARNING F OR C AUSAL L ANGUAGE M ODEL
Anonymous authors
Paper under double-blind review
A BSTRACT
Despite exciting progress in large-scale causal language models, the expressiveness of its representations is largely limited by the anisotropy issue where the hidden representations are distributed into a narrow cone in the vector space.
To resolve this problem, we present C ONTRA G EN , a novel contrastive learning framework at both token-level and sequence-level.
We assess C ONTRA G EN on a wide range of downstream tasks and show that C ONTRA G EN can effectively enhance both isotropy and discrimination of the representations. This leads to the desired improvement on various language understanding tasks, which helps bridge the gap with the encoder-only models and makes causal language models more suited for tasks beyond language generation.
Specifically, we attain 44% relative improvement on the Semantic Textual Similarity tasks and 34% on Code-to-Code
Search tasks.
Furthermore, by improving the expressiveness of the representations, C ONTRA G EN also boosts the source code generation capability with 9% relative improvement on execution accuracy on the HumanEval benchmark.
I NTRODUCTION
Causal Language Models (CLM) have seen remarkable success in language generation, both in natural language (Radford et al., 2018; 2019;
Brown et al., 2020) and programming language (Chen et al., 2021; Nijkamp et al., 2022).
However, one limitation at its core is the anisotropy issue, where the hidden representations, especially those output from the final layers of transformers, are squeezed into a tiny cone in the representation space.
As observed by Ethayarajh (2019), the average cosine similarity of any two words from two randomly sampled sequences is almost at one when evaluating at the last hidden layer of GPT-2 (Radford et al., 2019).
Such anisotropic often causes large performance gap (see Appendix D.1) with the encoder-based models on the discriminative tasks and hence limits the wide usage of CLM beyond language generation.
Two main reasons have been posited for causing the degenerated representations.
Dong et al.
(2021) theoretically prove that the key ingredient of language models – self-attention (Vaswani et al., 2017) possesses a strong inductive bias towards the representation collapse.
Another line of work argues that the autoregressive next token prediction objective is the main cause for the non-isotropic representations.
We thereby face a conundrum.
On the one hand, both self-attention and next token prediction are the driving force behind the widespread successes of CLM on language generation (Lewis et al., 2020;
Radford et al., 2019; Chen et al., 2021; Chowdhery et al., 2022).
On the other hand, as shown in Section 4 and Appendix D.1, the resulting anisotropic representations cause inferior performance on language understanding tasks, e.g., Semantic Textual Similarity (Agirre et al., 2012; 2013; 2014; Marelli et al., 2014; Agirre et al., 2015; 2016; Cer et al., 2017) and Code Search (Lu et al., 2021;
Huang et al., 2021; Guo et al., 2022).
We tackle this challenge by leveraging contrastive learning (Chen et al., 2020;
He et al., 2020; Gao et al., 2021).
Intuitively, by separating each instance apart from each other, contrastive learning can produce more uniformly distributed representations.
Previous studies (Wang & Isola, 2020; Wang & Liu, 2021) have shown that uniformity is important for yielding discriminative representations.
We thereby present C ONTRA G EN , a novel contrastive learning framework at both token-level (C ONTRA G EN -T OK ) and sequence-level (C ONTRA G EN -S EQ ).
Our goals are two-fold.
First, an
R e p r e s e n t a t i o n D
i s c
r
i m i n a t i o n
CLM SimCTG
ContraGen-Tok ContraGen-Seq
ContraGen
Model Discrimination
Generation
STS-AVG
Pass@1 Rank@
CLM 27.
SimCTG 36.
C ONTRA G EN -S EQ 37.
C ONTRA G EN -T OK 42.
C ONTRA G EN 45.
ideal CLM should be able to better leverage the representation space by dispersing apart semantically different tokens or sequences.
Second, we aim to attain separable features by mapping tokens or sequences under similar contexts to distinct but comparatively closer locations in the vector space.
We conduct extensive experiments and analyses to assess the effectiveness of our proposed framework with particular focus to address the following – (1) Does
C ONTRA G EN improve the discrimination ability of representations learned by CLM? (2) Do the enhanced representations lead to better performance on language generation tasks? (3) Is the joint contrastive learning at both token-level and sequence-level necessary and how do they benefit from each other?
(4) How does the impact of C ONTRA G EN vary across language domains?
R ELATED W ORK
Anisotropic Representation of Language Models Despite the remarkable success achieved by large scale language models (Devlin et al., 2019; Radford et al., 2019;
Yang et al., 2019; Raffel et al., 2020; Lewis et al., 2020), they suffer from the anisotropy issue where the representations are distributed into a tiny cone in the vector space (Gao et al., 2019; Ethayarajh, 2019; Li et al., 2020; Wang et al., 2020).
In particular, Ethayarajh (2019) shows that the degeneration is much severer on CLM, where the average cosine similarity between two words sampled from randomly selected sequences is almost at one when evaluating at the output from the last hidden layer of GPT- (Radford et al., 2019).
Inspired by the previous findings (Arora et al., 2017; Mu & Viswanath, 2018), several efforts have focused on regularizing the representations to isotropic distribution either via postprocessing (Su et al., 2021; Li et al., 2020) or directly optimizing variant regularization terms during training (Gao et al., 2019; Wang et al., 2020).
Although promising improvement has been observed, performance is still inadequate.
Contrastive Learning
An alternative comes to light as the popular contrastive learning techniques
Chen et al.
; He et al.
(2020) have seen remarkable successes in Natural Language Processing (NLP).
A large amount of research has focused on sentence representation learning for the encoderonly model, with the main differences lying in how the augmentations are generated (Fang & Xie, 2020; Giorgi et al., 2021;
Wu et al., 2020;
Meng et al., 2021; Yan et al., 2021;
Kim et al., 2021; Gao et al., 2021; Zhang et al., 2022).
Recently there is emerging interest in developing effective contrastive learning approach for text generation models.
However, most existing work mainly focus on the encoder-decoder structure Dong et al.
Raffel et al.
; Lewis et al.
by contrasting the suboptimal model generation, obtained via diverse sampling (An et al., 2022) or adding perturbations on the embedding space (Lee et al., 2021), against the ground truth.
On the other hand, it is not intuitive and even elusive to develop effective contrastive learning strategy for the decoderonly models.
A recent work (Su et al., 2022) proposes, SimCTG, a token-level contrastive learning approach that aims to separate each token apart from others within the same sequence by a prefixed distance.
However, we find that our temperature-based token-level contrastive learning approach, C ONTRA G EN -T OK , consistently outperforms SimCTG across different tasks.
We conjecture that the fixed margin based contrastive learning allows less flexibility for the token-level representation separation within each sequence, especially considering that the semantic relevance among tokens can vary across contexts (sequences).
Code Generation and Beyond Language modeling for source code is an actively growing area of research.
Various model architectures have been explored recently, including the encoder-only (Feng et al., 2020;
Guo et al., 2021), encoder-decoder (Ahmad et al., 2021; Wang et al., 2021;
Li et al., 2022), and decoder-only models (Chen et al., 2021; Nijkamp et al., 2022; Chowdhery et al., 2022).
Among them, the decoder-only models found effective on the code generation front.
However, as shown in Section 4.3.2 and Appendix D.1, they suffer from the less satisfied performance on the discriminative tasks (Lu et al., 2021;
Huang et al., 2021; Guo et al., 2022).
It motivates us to improve the decoder-only models on the discriminative tasks so as to extend its main usage beyond language generation.
Furthermore, code is fundamentally different from natural language in that it is more structured, which helps validate the generalizability of our approach beyond plain text.
C ONTRASTIVE L EARNING FOR L ANGUAGE G ENERATION
Before we present the main model, we first call out the following notations that we will use throughout this section.
Let x =
[ x 1 , x 2 , · · · , x | x | ] denote a sequence with variable length | x | , e.g., a text document or a code snippet, and
h =
[ h 1 , h 2 , · · · , h | x | ] be its representation output by the last hidden layer of the decoder.
For a randomly sampled batch B = (cid:8) x j (cid:9) N j =1 with N sequences, we use x ji and h ji to denote the i th token and its representations in the j th sequence, respectively.
Let h j , h j + denote the representation pair of sequence x j and h ji , h j + i correspond to the representations of the i -th token.
Such representation pairs are referred as the positive pairs in contrastive learning, which are often obtained via data augmentation.
3. L ANGUAGE G ENERATION M ODELING
Language generation modeling is usually formulated as sequence distribution estimation over a set of examples, B = { x j } Nj =1 .
For tractable estimation, common practice is to factorize the joint distribution of each sequence into the product of conditional token prediction probabilities.
The language generation model is then trained via maximum likelihood estimation as L CLM = − 1 N (cid:80) Nj =1 (cid:80) | x j | i =1 log p ( x ji | x j<i ) where x j<i =
[ x j 1 , · · · , x ji − 1 ] denotes the subsequence before x ji and | x j | refers to the sequence length.
In this paper, we mainly focus on the decoder-only Causal
Language Model (CLM), though we believe our strategy can be applied to most language models.
Despite promising progress in language generation, CLM suffers from the anisotropy issue, which limits the wide usage of it beyond language generation.
Ideally, to better leverage the representation space, we aim to disperse apart the representations of tokens or sequences from different contexts while simultaneously pulling close those from the similar contexts for better discrimination.
To this end, we present a novel contrastive learning framework at both token-level and sequence-level.
3. T OKEN -L EVEL C ONTRASTIVE L EARNING
As aforementioned, h ji , h j + i are a pair of representations for x ji , the i -th token in the j -th sequence.
Let I
j = { 1 , 2 , . .
, | x j |} denote the indices of tokens in x j .
{+Further let τ denote the temperature hyper-parameter and ⋄ denotes the cosine similarity, i.e., a ⋄ b = a T b / ∥ a ∥ 
+}
∥ b ∥ 2 .
Then for the token-level contrastive learning, we minimize the following,
[Equation]
[-Here τ is the temperature hyper-parameter and ⋄ denotes the cosine similarity, i.e., a ⋄ b = a T b / ∥ a ∥ 
-][Table] log
[Table]
  . The above objective tries to separate the positive representation pairs of each token at distinct locations apart from the others within the same sequence.
Intuitively, pushing representations of different tokens away from each other should indeed cause them to be more uniformly distributed.
On the other hand, as suggested in Section 4, the token-level contrastive learning leads to the grouping of semantically similar sequences, which we conjecture is due to implicit clustering effect of contrastive learning (Wang & Liu, 2021; Zhang et al., 2021) on tokens and so as the sequences consist of semantically similar tokens.
3. S EQUENCE -L EVEL C ONTRASTIVE L EARNING
Let I B = { 1 , 2 , . . .
, N } ∪ { 1 + , 2 + , . . .
, N + } denote indices of all 2 N sequence-level representations for batch B .
The sequence-level contrastive loss then follows as
[Equation]
The sequence-level contrastive loss aims to separate each sequence apart by separating its positive representation pairs from those of randomly sampled sequences (in-batch negatives).
The tokenlevel and sequence-level can complement each other for two main reasons.
First, the sequence-level discrimination can help cross-sequence token-level separation.
Second, as semantically irrelevant sequences can still share some similar tokens, better token-level separation can help the sequencelevel discrimination be more driven by the semantically irrelevant tokens.
3. C ONTRA G EN
In summary, C ONTRA G EN optimizes both the token-level and sequence-level contrastive learning objectives in addition to the standard causal language modeling objective as follows
[Equation]
[Table]
Unless otherwise specified, we weight each loss in C ONTRA G EN equally and set the temperature τ = 0 .
for both contrastive losses.
Although better performance can be achieved by hyperparameter optimization, we mainly investigate how C ONTRA G EN improves the representation quality and boost the zero-shot transfer learning performance on both discrimination and generation tasks.
We therefore leave hyperparameter optimization in a supervised setting for future work.
Positive pair of representations We consider the simple yet effective dropout-based augmentation (Gao et al., 2021), where the representation pairs are obtained by forwarding each sequence twice.
Unlike the existing findings that the dropout-based augmentation can boost the contrastive learning performance when (continually) training a language model, we find that the trends can be different when evaluating on the discrimination tasks and the generation tasks.
We discuss with detailed ablation study in both Section 4.4 and Appendix D.2.
E XPERIMENTS
We design experiments in both natural language and programming language tasks to address the four main questions listed in the Introduction section.
4. D ATA AND M ODELS
Data & Models For text
, we continue train GPT-
(Radford et al., 2019) on WikiText-103, a collection of over 100 million tokens extracted from the set of verified Good and Featured articles on
Wikipedia (Merity et al., 2017).
For code, we continue train
CodeGen 350M monolingual
(Nijkamp et al., 2022) on collected permissively licensed Python code from GitHub.
Please refer to Appendix
A.1 for the training details.
We consider the following objectives for continual training both models.
• CLM .
The standard left to right autoregression objective for training causal language models, which is also the objective used for pretraining both GPT-2 and CodeGen.
• SimCTG (Su et al., 2022).
A margin-based token-level contrastive learning framework aims to separate tokens at each distinct location within a sequence apart from each other based on a predefined margin value.• C ONTRA G EN -T OK & C ONTRA G EN -S EQ .
In addition to our full model C ONTRA G EN , we also evaluate its components C ONTRA G EN -T OK and C ONTRA G EN -S EQ defined in Section 3.4.
4. E VALUATE ON N ATURAL L ANGUAGE
We first evaluate our model on discrimination and generation tasks in natural language.
[Table]
4.2. S EMANTIC T EXTUAL S IMILARITY
We first evaluate semantic textual similarity (STS), the most commonly used benchmark for evaluating the semantic discrimination capability of representations.
STS consists of seven tasks, namely
STS 2012-2016 (Agirre et al., 2012; 2013; 2014; 2015; 2016), the STS Benchmark (Cer et al., 2017), and the SICK-Relatedness (Marelli et al., 2014).
Human annotators provide a fine-grained similarity score ranging from 0 to 5 for each sequence pair in STS.
Same as Reimers & Gurevych (2019), for the sequence pairs in each dataset, we report the overall Spearman’s correlation between the cosine similarities of representations and the human-provided similarity scores in Table 1.
Contrastive Learning Leads to Discriminative Representations Table 1 shows both GPT-2 and the one continually trained with CLM perform poorly on STS, which is a consequence of the poorly discriminative representations where the cosine similarities between both semantically similar and dissimilar pairs are almost at one (see Figure 4 in Appendix B.1).
Also, continuing to train GPTwith CLM on WikiText-103 worsens performance.
We hypothesize this is because WikiText-103 is much smaller than the original dataset used for GPT-2 pretraining and hence degenerates the transfer learning performance on STS.
In contrast, both C ONTRA G EN and SimCTG largely outperforms GPT-2, yet still, C ONTRA G EN attains 25% relative improvement over SimCTG.
Moreover, with the token-level contrastive learning only, C ONTRA G EN -T OK still outperforms SimCTG on almost all
STS benchmarks and the trend remains the same even without the dropout-based augmentation (see Appendix D.3).
Therefore, we posit our temperature-based contrastive learning allows more flexibility in encouraging the semantic-dependent separations among tokens, while requiring a prefixed separation margin between tokens within the same sequence (context) as SimCTG does, is not ideal.
Sequence-level vs. Token-level Contrastive Learning Table 1 also indicates that C ONTRA G ENT OK and C ONTRA G EN -S EQ complement each other, as C ONTRA G EN consistently performs better than them on all STS tasks.
Note that, C ONTRA G EN -S EQ performs worse than C ONTRA G EN -T OK .
It is surprising, especially since STS mainly assesses the sequence-level representation quality.
We investigate this by dividing the sequence pairs into two groups – semantically similar pairs with human-annotated similarity scores no less than 0.7 and dissimilar pairs with human scores no larger than 0.3.
3 We plot the rank of the model inferred similarity scores against the human similarity scores in Figure 2 (left).
As we can see, C ONTRA G EN -S EQ struggles in ranking semantically dissimilar sequence pairs higher and similar pairs lower.
This suggests that the token-level contrastive loss is essential for making the sequence-level representations robust to spurious patterns of tokens
[Figure]
[Figure] PPL
[Table]
MAUVE [Table] GPT [Table]
[Figure]
[Figure]
CLM 22.48 0.
[Figure]
[Figure]
[Figure]
SimCTG 22.51 0.
[Figure]
[Figure]
[Figure]
[Table]
or phrases, e.g., ranking semantically similar sequences with different synonyms low and ranking dissimilar sequences high even in presence of the same phrase (see Figure 2 (right)).
4.2. T EXT G ENERATION
Next, we assess the open-ended language generation capability, where each model is required to generate text continuations given the prefixes from the WikiText-103 test set.
Following Su et al.
(2022), we set the lengths of prefix and continuation to 32 and 128, respectively.
We use nucleus sampling (Holtzman et al., 2020) with top- p = 0 .
We report the model perplexity (PPL) in Table 2.
PPL evaluates the model prediction confidence on the ground truth continuations in our case.
The lower the PPL value, the higher the prediction probabilities assigned to the human-written text.
Additionally, we propose the following for evaluating representation quality.
• Intra-Similarity (ISim) is the average cosine similarity between the representations of a sequence and its constituent tokens.
In this paper, the sequence representation is obtained via mean pooling.
• Inter-Similarity (ESim) is the average cosine similarity between the representations of a sequence and each token of another sequence.
Regarding contextual representations, it is desired that tokens under S imilar contexts (sequences) are pulled closer and hence attaining higher ESim(S) score, while tokens from Different contexts are separated apart with lower ESim(D) score.
• Discrimination (Disc) is defined as Disc = 1 - ESim / ISim where Disc(S) and Disc(D) associate with ESim(S) and ESim(D) defined above, respectively.
Thereby, a lower discrimination score (Disc(S)) between each ground truth and generation is desired for better semantic coherence, while higher discrimination scores (Disc(D)) are desired for generations of randomly sampled prompts for more diverse generations and better isotropic representations.
Contrastive Learning Trades Off between Model Perplexity and Representation Quality As reported in Table 2, despite a slight increase in PPL, contrastive learning improves MAUVE (Pillutla et al., 2021) and the representations with better isotropy and discrimination. Furthermore, C ONTRAG EN attains better semantic coherence between the ground truth and generation as indicated by the
[Table]
lower Disc(S) score and more diverse generations for randomly sampled prompts as indicated by the higher Disc(D) value.
We stress that the zero valued Disc(S) score attained by GPT-2 and CLM is indeed a natural consequence of the inferior representations, since the discrimination score between semantically irrelevant sequences is also zero. Also, the slight increase on PPL is probably expected considering PPL is better aligned with the standard CLM objective.
Thereby, contrastive learning can be interpreted as regularization that trades-off between PPL and the desired representation properties. We hypothesize C ONTRA G EN can improve PPL with new decoding strategy that better leverages the enhanced representations.
4. E VALUATE ON P ROGRAMMING L ANGUAGE
In this section, we study the effectiveness of our proposed contrastive learning framework on programming language applications – code search, code completion, and code re-ranking.
Different from Section 4.2, CodeGen is pretrained without using dropout from the very beginning.
We follow the same setting for all models in this subsection so as to study the effectiveness of C ONTRA G EN without the dropout-based augmentation.
We also further investigate how dropout would affect the decoder-only models when evaluated on the downstream tasks in Section 4.4 and Appendix D.2.
4.3. C ODE S EARCH
Code search is the task of retrieving relevant code fragments given a code fragment as a query .
We perform in-language (query and relevant code are in the same language) and cross-language (query and relevant code are in different languages) code search.
We provide an example in Figure(Appendix C.1).
In this study, we experiment in the zero-shot setting - we use the models described in Section 4.1 to generate dense representations of code and perform nearest neighbor search to retrieve relevant code fragments.
We use publicly available implementations of Guo et al.
Contrastive Learning Yields Discriminative Code Representations For code-to-code search task, Guo et al.
(2022) used problem solutions in Ruby/Python/
Java languages from CodeNet (Puri et al., 2021).
They propose to use each program as a query and retrieve all programs that solve the same problem.
We present a detailed statistics of the dataset in Table 5 (Appendix C.1).
We set the maximum sequence length as 512 5 and use cosine similarity between two mean vectors of the last hidden states as relevance scores.
We then sort the candidates by their scores to calculate the Mean
Average Precision (MAP) score.
We present the results for the code search tasks in Table 3.
We observe C ONTRA G EN -T OK and C ONTRA
G EN frameworks improve CodeGen trained with CLM by 33.5% (absolute 2.12) and 32.6% (absolute 2.06) on average, respectively.
We also point out that the performance gap between C ONTRA G EN -T OK and SimCTG are apples-to-apples comparison since the dropout-based augmentation is not used in either cases. As aforementioned, the consistent better performance of C ONTRA G EN -T OK suggests the superiority of our temperature based contrastive learning objective.
On the other hand, C ONTRA G EN -S EQ improves over the [Figure] CLM baseline by 10.4% only.
Code search results indicate that C ONTRA G EN -S EQ performs poorly compared to C ONTRA G EN -T OK , the performance gap is larger than what we observed in natural language evaluation.
We conjecture that C ONTRA G EN -T OK generates better discriminative representations for code sequences since the finer-grained understanding of the code tokens is crucial to understanding the code sequences’ functionality (semantics).
To verify, we check if non-semantic factors impact model performances in the following section.
Token-level Contrastive Learning is Effective for Code Understanding We break down the code search performance based on edit similarities and length differences between query code and their relevant code fragments.
While edit similarity indicates how much queries and their relevant code overlap, the length difference indicates whether models effectively capture relevance between two code fragments if they are similar in length or differ significantly.
We present the results for
Python language in Figure 3 (for all the languages, see Figures 7 & 8 in Appendix C.3).
The results show that C ONTRA G EN -T OK outperforms CLM, SimCTG, and C ONTRA G EN -S EQ irrespective of edit similarities and length differences.
Therefore, we can conclude that sequence overlap or length are not the reasons for improvements in C ONTRA G EN -T
OK .
Presumably, a finer-grained understanding of code tokens makes C ONTRA G EN -T
OK more effective for code representations.
4.3. C ODE C OMPLETION AND R E -R ANKING
Given a sequence of tokens composed of natural language, function signature, and input-output examples (as a whole, we call them prompt), the goal of the code completion task is to complete the function.
To evaluate the functional correctness of a complete code, we use existing benchmarks that include unit tests.
If the generated code successfully passes the unit tests, we refer to this as successful execution.
We compute pass@ k for k ≤ n following Chen et al.
In addition, we compare the models in the code re-ranking task - given n sampled code using a code completion model, the goal of code re-ranking is to find an ordering of the code.
We use the mean log probability of the sampled code to order them (Chen et al., 2021).
For code re-ranking evaluation, we report ranked pass@ k
(Inala et al., 2022).
We present an example in Figure 6 (Appendix C.1) to illustrate the code completion and re-ranking tasks.
We detail the evaluation metrics in Appendix C.2.
Contrastive Learning Improves Source Code Generation Chen et al.
introduced HumanEval, a collection of 164 handwritten programming problems and their respective unit tests.
Each problem in this dataset is presented using a function signature and supporting docstring, and the task is to complete the body of the function, such that the complete function can pass all unit tests.
In all our experiments, we use nucleus sampling (Holtzman et al., 2020) with top p = 0 .
We sample n = 10 completions per problem with sampling temperature 0 .
Table 4 presents the evaluation results on the HumanEval benchmark.
[Table]
The numbers in the subscript indicates the difference between ranked pass@k and pass@k.
While C ONTRA G ENT OK and C ONTRA G EN -S EQ perform competitively to the baselines,
C ONTRA G EN significantly outperforms them.
We see that C ONTRA G EN -T OK
and
C ONTRA G EN -S EQ perform comparably to CLM and SimCTG, while C ONTRA G EN outperforms them significantly (by 9% and 10.3% in terms of pass@1 accuracy).
Similarly, in the code re-ranking task, C ONTRA G EN outperforms CLM and SimCTG baselines by 11% and 12% in terms of ranked pass@ accuracy, respectively.
While C ONTRA G EN -S EQ underperforms in code completion, it boosts code re-ranking significantly.
We hypothesize the improvement is due to the contrastive learning objective aligns with the mean log probability-based re-ranking choice.
4. D ISCUSSION
Impact of Dropout Dropout-based augmentation (Gao et al., 2021) has proven to be effective for contrastive learning on language models, which often leads to significant improvement when evaluating on discriminative tasks.
We observe the same trend on both GPT-2 and CodeGen (see Tables 8b & 8a in Appendix D.2).
However, we observed the opposite when evaluating on language generation, no matter when training with CLM only or together with contrastive learning (see Appendix D.2).
Dropout has been one of the key ingredients for training large models, further investigation on the proper ways to use and evaluate it is indeed required.
On the bright side, even without dropout, Section 4.3 shows C ONTRA G EN still consistently yield considerable improvement.
Bridge the Gap between Causal and Bidirectional Attention Models on Discriminative Tasks
In comparison with the causal (left-to-right) attention mechanism of the decoder-only models, the bidirectional attention mechanism allows to better leverage the context of the sequence and hence yield better representations for discrimination tasks.
Take the encoder-only model as an example, as
Table 6 in Appendix shows, both BERT-Base (Devlin et al., 2019) and RoBERTa-Base (Liu et al., 2019) outperform GPT-2 by at least 60% relative performance on STS.
Despite the performance gap between CodeGen and the encoder-only or encoder-decoder models decreases in Table 7, it is still significant considering that both the model size and pretraining data size used by CodeGen are much larger.
Such large performance gap severely limits the decoder-only models being used in many discriminative tasks.
On the bright side, contrastive learning shows the promise to bridge the gap.
As seen in Table 6, on STS, C ONTRA G EN reduce the relative performance gap from 67.24% (absolute 21.12%) to 16.17% (absolute 7.33%) regarding BERT-Base, and from 84.62% (absolute 26.64%) to 28.24% (absolute 12.8%) regarding RoBERTa-Base.
Similarly, Table 7 shows that C ONTRA G EN outperforms encoder-decoder models and perform comparably to encoder-only model GraphCodeBERT.
Please refer to Appendix D.1 for more detailed discussions.
C ONCLUSION
In this paper, we present C ONTRA G EN , an effective contrastive learning framework to resolve the representation degeneration issue of causal language models trained with the standard autoregression objective.
We assess the effectiveness of the proposed strategy on a variety of downstream tasks in both natural language domain and programming language domain, where we attain significant improvement on both discrimination tasks and language generation tasks.
Moreover, we conduct in-depth analyses on our proposed token-level and sequence-level contrastive losses.
Although we only explored on the decoder-only causal language models, our proposed contrastive learning framework can serve as a drop-in term for encoder-decoder, encoder-only, or prefixLM.
We leave the explorations as future work.
A T RAINING D ETAILS A. S ETUP D ETAILS
Training Data For text, we use WikiText-103, a collection of over 100 million tokens extracted from the set of verified Good and Featured articles on Wikipedia (Merity et al., 2017).
For code, we collected permissively licensed Python code from GitHub.
Following Chen et al.
; Nijkamp et al.
(2022), we perform filtering and deduplication and further remove data that contains a significant use of non-English languages or is not parsable, resulting in a dataset of 101GB code.
Model We use GPT-
(Radford et al., 2019) and CodeGen 350M monolingual (Nijkamp et al., 2022) for all experiments on natural language (text) and programming language (code), respectively.
We set the batch size to 512 and continue to train GPT-2 on WikiText-103 and CodeGen on the
GitHub data for 12 and 2 epochs, respectively.
We trained both models using a max sequence length of 512 tokens and 1024 for WikiText-103 and Code data, respectively.
We set the learning rate to 2e-5, warm-up steps as 500 with linear annealing after peak learning rate, weight decay of 0.1, the temperature of 0.05 (when using contrastive losses), and gradient clipping of 1.0.
We use AdamW optimizer (Loshchilov & Hutter, 2019) with β 1 = 0 .
9 , β 2 = 0 .
999 , and ϵ =
10 −following (Nijkamp et al., 2022).
Our training pipeline is based on PyTorch Lightning 7 and we use
DeepSpeed (Rasley et al., 2020) for training optimization.
Processing Code Training Data
Our preprocessing strategy for code datasets used for training is designed to ensure that we optimize for data utilization while retaining the syntactic structure of programming language sequences.
We also eliminate duplicate sequences since this benefits training large language models (Lee et al., 2022).
Specifically, we break long sequences into chunked sequences of smaller lengths to retain most parts of the original program.
Further, we maintain syntactic structure in the chunks by ensuring that each chunk ends with a ‘ \ n ’ character.
Each chunk obtained this way contains at most max chars per seq characters where max chars per seq = max tokens per seq * chars per tok .
In our experiments, we fix chars per tok = 3 . 2 and max tokens per seq = 1024 .
We also perform deduplication using character-based exact matches between chunked sequences over the entire dataset.
This step helps eliminate exact-duplicates that might be present after the chunking stage.
B M ORE ON N ATURAL L ANGUAGE E VALUATION B. R EPRESENTATION Q UALITY E VALUATION ON STS
In this section, we dive deep into the model performance on the STS benchmarks.
For each sequence pair in STS, a fine-grained similarity score ranging from 0 to 5 is provided, with a high similarity score indicating semantically similar pairs and low similarity scores suggesting semantically dissimilar or irrelevant pairs.
For better illustration, we scale the human-annotated similarity scores to [0 , 1] to align with the model-predicted cosine similarity scores.
This does not affect the evaluation as the spearman correlation reported in Section 4.2 is a rank-based correlation metric.
CLM yields poorly discriminative representations We report the model predicted similarity scores of sequence pairs in Figure 4a.
A good model is expected to yield representations that attain higher similar scores between similar sequence pairs and lower similarity values for dissimilar sequences.
A large gap between the similarity scores of similar and dissimilar pairs is desired.
However, as seen in the left column of Figure 4, the similarity scores attained by the model trained with the standard CLM only objective are almost at one for both similar sequence pairs and dissimilar pairs.
This again suggests that the representations yielded by CLM are squeezed into a tiny cone in the representation space rather than being scattered apart to leverage the vector space’s capacity better.
Despite the resulting similarity ranks are not entirely flattened, as shown in the right column in Figure 4b, CLM struggles in ranking similar sequences lower and dissimilar sequences higher as a consequence of its poorly discriminative representations such that low similarity scores can be assigned to the semantically similar and high similarity scores to dissimilar sequence pairs.
[Figure]
r i
[Figure]
y
[Figure]
r
[Figure] [Figure]
r
[Figure]
t y
[Figure]
that the resulting ranking results are as close as the ranks provided by human as possible.
C ONTRA G EN leads to representations with better discrimination and contextulization In comparison, Figure 4a also validates that contrastive learning effectively yields more discriminative representations with a comparatively larger similarity gap between the similar pairs and dissimilar pairs.
Thereby, the similarity ranking results of the sequence pairs are more aligned with those obtained according to similarity scores provided by a human.
Figure 4b shows that C ONTRA G EN consistently outperforms the other models, with the ranking results better matching the ground truth.
C M ORE ON P ROGRAMMING L ANGUAGE E VALUATION C. T ASKS E XAMPLES AND S TATISTICS
In Figure 5, we present an example of query code fragment in Python and relevant code fragments in Python and Java, respectively.
While in-language code-to-code search refers to retrieving relevant code fragments in the same language, cross-language code-to-code search refers to retrieving code fragments in a different language.
We present the statistics of the code search dataset in Table 5.
To demonstrate the code completion task, we illustrate an example in Figure 6.
C. E VALUATION M ETRICS
Mean Average Precision (MAP)
For a set of queries, it indicates the mean of the average precision scores for each query.
[Equation]
Q where Q is the number of queries.
[Table]
Pass@k Given a problem (code prompt as as shown in Figure 6), pass@k indicates functional correctness of model generated code samples.
A problem is considered solved if any sample passes the unit tests.
Following Chen et al.
(2021), we generate n ≥ k samples per problem (in this paper, we use n = 10 and k ∈ { 1 , 5 } ), count the number of correct samples c ≤ n that pass unit tests, and calculate the unbiased estimator of pass@ k as: pass@ k :=
E Problems (cid:34) 1 − (cid:0)
[Equation]
(cid:1) (cid:0) nk (cid:1) (cid:35) .
Ranked Pass@k Unlike Pass@k, where we randomly chose k out of n samples, in ranked pass@k, we chose the top- k samples based on model-provided scores and then compute pass@k.
C. D ETAILED C ODE S EARCH R ESULTS
We provide comparison between encoder-only (Feng et al., 2020;
Guo et al., 2021), encoder-decoder (Ahmad et al., 2021; Wang et al., 2021), and decoder-only models (main focus of this work) on the zero-shot code-to-code search task in Table 7.
We see that
C ONTRA G EN -T OK and C ONTRA
G EN outperforms the encoder-only model CodeBERT and both the encoder-decoder models.
It is important to note that the comparison across these models is not apple-to-apple as these models differ in size, the scale of pretraining, and language settings.
This comparison’s purpose is to show the promise of decoder-only models being used in discriminative tasks like code search.
We further break down the code search performances based on edit similarities and length differences between query code and their relevant code fragments.
We present the results in Figureand 8.
We observe a similar performance trend in all three languages, although cross-lingual search performance still needs to improve.
Nonetheless, the objective of this performance analysis is to show that sequence overlap or length are not the reasons for improvements in C ONTRA G EN -T OK .
Instead, a finer-grained understanding of code tokens due to the token-level contrastive learning makes C ONTRA G EN -T OK more effective for code representations.
D M ORE A NALYSIS AND D ISCUSSIONS D. C ONTRASTIVE L EARNING B RIDGE THE G AP B ETWEEN C AUSAL AND B IDIRECTIONAL A TTENTION M ODELS ON D ISCRIMINATIVE T ASKS
Compared to the causal (left-to-right) attention mechanism of the decoder-only models, the bidirectional attention mechanism in both encoder-only and encoder-decoder models allows to better leverage the context of the sequence.
It hence leads to better representations when evaluated on discrimination tasks.
Taking the encoder-only models as an example.
As shown in Table 6, on average, BERT-Base (Devlin et al., 2019) and Roberta-Base (Liu et al., 2019) outperform GPT-2 with 67.24% (absolute 21.12%) and 84.62% (absolute 26.64%) relative improvement on STS, respectively.
Although the performance gap between CodeGen and the BERT models trained on programming languages (CodeBERT (Feng et al., 2020), GraphCodeBERT (Guo et al., 2021)) decreases or even diminishes when evaluated on the Code-Search tasks, the performance gap is still significant because both the model size and pretraining data in CodeGen are much larger than those by the encoder-only models in Table 7.
Similar trends were observed regarding the performance gap between the decoder-only and encoder-decoder models on both natural language (Lewis et al., 2020; Raffel et al., 2020) and programming language (Ahmad et al., 2021; Wang et al., 2021).
[Table] [Table]
The large performance gap severely limits the decoder-only models used in many discriminative tasks.
To this end, contrastive learning shows the promise to largely bridge the gap.
As seen in Table 6, on STS, C ONTRA G EN reduce the relative performance gap from 67.24% (absolute 21.12%) to 16.17% (absolute 7.33%) regarding BERT-Base
[Table]
and from 84.62% (absolute 26.64%) to 28.24% (absolute 12.8%).
Similarly, Table 7 shows that C ONTRA G EN outperfors encoder-decoder models and perform comparably to encoder-only model GraphCodeBERT.
[Table]
Dropout (0.1) to all the layers of the models.
D. D ROPOUT FOR C ONTRASTIVE L EARNING WITH D ECODER -O NLY M ODEL
Gao et al.
(2021) showed that dropout-based augmentation is an effective strategy for unsupervised contrastive learning and followup works (Chuang et al., 2022; Wu et al., 2022) endorse the effectiveness.
This motivates us to study dropout-based augmentation in our proposed contrastive learning framework.
We present the results on discriminative and generation tasks in Table 8 and 9, respectively.
From results, it is evident that the adoption of dropout-based augmentation improves the discrimination task performances, which corroborates the findings of Gao et al.
[Table]
On the other hand, dropout-based augmentation hurts the generation task performances.
While for code completion
[Table]
we anticipated that dropout-based augmentation would hurt the performances since we use CodeGen (Nijkamp et al., 2022) in this work which does not use dropout during pretraining.
However, to our surprise
[Table]
we observe a drop in perplexity due to the dropout-based augmentation that does not go with our anticipation (unlike CodeGen, GPT-2 is pretrained with dropout).
Therefore
[Table]
we leave this as a future exploration to dive deeper into the reasoning behind this finding.
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Figure] [Figure]
[Figure] [Figure]
D. C ONTRA G EN CONSISTENTLY OUTPERFORM S IM CTG
To better understand the performance gap between C ONTRA G EN and SimCTG (Su et al., 2022)
[Table]
we run the following ablations on GPT-2 and report the evaluations on STS.
In Table
[Table]
we report the results of (1) running C ONTRA G EN w/o dropout-based data augmentation and compare it with the original SimCTG model; and (2) augmenting SimCTG with both the sequence-level contrastive loss and dropout-based augmentation and compare it with our proposed
C ONTRA G EN model.
As we can see, C ONTRA G EN consistently outperforms SimCTG.
Figuretogether with our results reported in Section 4.3 where we disabled the dropout-based augmentation for C ONTRA G EN and its variations
[Table]
but still observed consistent better performance than SimCTG on both discrimination and generation tasks, conclude that C ONTRA G EN is better than SimCTG across domains and settings.
[Table]
[Figure]
Relevant Program in Python
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
Relevant Program in Java
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
Completion 1 (passed; mean logp: -0.1146)
[Table]
for
[Table]
in
[Table]
for
[Table]
numbers)
[Table] [Table]
if
[Table]
Completion 2 (failed; mean logp: -0.1211) for i in range(len(numbers) -
1):for j in range(i + 1, len(numbers)):if abs(numbers[i] - numbers[j]) > threshold:return Truereturn False
Completion 3 (failed; mean logp: -0.1598) for i in range(len(numbers) - 1):if abs(numbers[i] - numbers[i + 1]) >
threshold:return Truereturn False
Completion 4 (failed; mean logp: -0.1607) for i in range(len(numbers) - 1):if abs(numbers[i] - numbers[i + 1]) >
threshold:return Falsereturn True
Completion 5 (failed; mean logp: -0.3904) return abs(sum(numbers) - sum(numbers) / len(numbers))
< threshold
[Figure] [Figure]