CODE REPRESENTATION LEARNING AT SCALE
Authors: Dejiao Zhang
Wasi Uddin Ahmad
Ming Tan
Hantian Ding
Ramesh Nallapati
Dan Roth
Xiaofei Ma
Bing Xiang
Venue: ICLR 2024 poster (2024)
Keywords: code representation learning
customized token-level denoising objective for code
hard negative
hard positive
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: Recent studies have shown that code language model at scale demonstrate significant performance gains on downstream tasks, i.e., code generation. However, most of the existing works on code representation learning train models at a hundred million parameter scale using very limited pretraining corpora. In this work, we fuel code representation learning with a vast amount of code data via a two-stage pretraining scheme. We first train the encoders via a mix that leverages both randomness in masking language modeling and implicit structure and semantic aspects of programming language. We then enhance the representations via contrastive learning with hard negative and hard positive constructed in an unsupervised manner. We establish an off-the-shelf encoder model that persistently outperforms the existing models on a wide variety of downstream tasks by large margins. To comprehend the factors contributing to successful code representation learning, we conduct detailed ablations and share our findings on (i) a customized and effective token-level denoising scheme for source code; (ii) the importance of hard negatives and hard positives; (iii) how the proposed bimodal contrastive learning boost the cross-lingual semantic search performance; and (iv) how the pretraining schemes decide the downstream task performance scales with the model size.

Published as a conference paper at ICLR 2024
CODE REPRESENTATION LEARNING AT SCALE
Dejiao Zhang" & Wasi Uddin Ahmad {dejiade a wuahmad ama zon. com
Ming Tan & Hantian Ding
(mingtan, dhant aaaantiinn@amazorm.
Ramesh Nallapati & Dan Roth & Xiaofei Ma & Bing Xiang
(rnallapa, drot, xiaofeim, bxi bxiang) @amaammoo......... com
AWS AI Labs
ABSTRACT
Recent studies have shown that code language models scale demonstrate sig- nificant performance gains on downstream tasks, ie., code generation. However, most of the existing works on code representation learning train models hun- dred million parameter scale using very limited pretraining corpora. In this work, we fuel code representation learning with a vast amount of code data via a two- stage pretraining scheme. We first train the encoders via mix that leverages both randomness in masking language modeling and implicit structure and se- mantic aspects of programming language. We then enhance the representations via contrastive learning with hard negative and hard positive constructed in an unsupervised manner. We establish an off-the-shelf encoder model that persis- tently outperforms the existing models on a wide variety of downstream tasks. To comprehend the factors contributing to successful code representation learning, we conduct detailed ablations and share our findings on (i) a customized and ef- fective token-level denoising scheme for source code; (ii) the importance of hard negatives and hard positives; (iii) how the proposed bimodal contrastive learning boost the cross-lingual semantic search performance; and (iv) how the pretraining schemes decide the downstream task performance scales with the model size.
INTRODUCTION
Large language models (LLMs) pretrained on a massive amount of source code have reshaped the
landscape of code generation (Chen al. 202T Chowdhery al., 2022 Li et al. 2023 inter alia). As an example, the recent release of a 6TB dataset (Kocetkov et al. 2022) comprising source code under permissive licenses play pivotal roles in promoting the advancement of code language models in present times. Nonetheless, these large corpora are not fully utilized to develop general-purpose Programming Language (PL) embedding models. To date, most PL embedding models (Feng al.) 2020a Guo et al.) 2021 2022 inter alia) have no more than 125M parameters and are primarily trained on a few millions of training examples, e.g., CodeSearchNet (Husain et al., 2019).
Despite the undeniable significance of large-scale data, it's imperative to acknowledge the vital role of pretraining objectives. The prevailing approach for pretraining a bidirectional Transformer en- coder to learn representations is through the optimization of masked language modeling (MLM) objective, as proposed by Devlin et al. (2019b). The masking scheme in the standard MLM objective follows 80-10-10 practice However, development of suboptimal code embedding models. Since code snippets contain both natural lan- guage (NL.) statements (i.e., doestrings, comments) and pure code, hence replacing masked tokens with a random token following the 80-10-10 convention could result in replacing an NI. token with PL token, and vice versa (see statistics in Appendix (..)). We speculate such co-occurrence PL and NL together with the syntax nature of source code make it easier to disrupt both the semantics and structure of the masked code, resulting in sub-optimal learning of the language model.
have noticed that such masking scheme leads to the
"orresponding authors with equal Contribution.
"oode and models can found at http:////ooeeeepeeeettttee on- on- learn i ing g thub. o Under this scheme, 80% of the randomly selected tokens for prediction are replaced with the [MASK] token, 10% substituted with random tokens, and the remaining tokens remain unchanged.

---

Published as a conference paper at ICLR 2024
Ruby
DOBF implioit structure-sware masking
Hard Positive remove function signature and return statement
Pully p.os
Pr
Chhor
Cshorp
Typeseript
Customized MLM random masking without the 80-10-10 corruption convention
Hard Negative embedding distance- based hard negative
hypacipot
Invoeriipt
IavaSerpti
identification on-the-fly
Pythos
Python
Sequence-Level
Num Toiera [&I lio)
Token-level Denoising
Bimodal Contrastive Learning
Num Pairs (Million)
Figure 1: An overview of the key ingredients of CODESAGE for code representation learning.
While MLM pretraining yields contextual token representations, most downstream discriminative
tasks primarily function at the sequence level. When the objective is to enhance the representation
discrimination power for immediate application in sequence-level tasks, contrastive learning (CL) emerges as the go-to approach. Existing works have employed unimodal CL (using Code-Code pairs) (Guo et al. 2022 Jain et al. (02T) or bimodal CL (using Text-Code pairs) (Li et al.) (022) for representation learning. In unimodal CL, a popular choice is to utilize dropout augmentation Gao ctt al. (2021) to construct positive code pairs. However, we found that dropout augmentation suffers from supporting long training process, also reported by Zhou et al. (2022). In contrast, bimodal CL becomes an appealing choice, primarily because of the availability of naturally occurring pairs. Prior studies utilize functions and their corresponding docstrings establish the bimodal training pairs. Nonetheless, our preliminary experiments indicate that substantial overlap between docstrings and function signatures simplifies the contrastive learning process (see statistics in Appendix (A.6).
To this end, we present CODESAGE, a bidirectional encoder representation model for source code. We pretrain CODES using two-stage training scheme with a large amount of customized pre- training data (Kocetkov et al. 2022). We depict the key ingredients of CODES AGE in Figure] I We first train the bidirectional encoders via mix of two objectives complementing each other: identifier deobfuscation (DOBF) and MLM without the 80-10- 10- 10 practice. Similar to human programmer, finding meaningful names for obfuscated identifiers necessitates the model to acquire a profound comprehension of code semantics and structure. Meanwhile, as a more general objective, MLM covers other facets beyond identifiers of code this is important for enriching the training signals, especially for data examples with non-informative identifier names. In the second stage, we leverage the (texi, code) pairs for bimodal contrastive learning (CL). In contrast to existing approaches that primarily rely naturally occurring text and code pairs, propose strategy reduce the likeli- hood of the model learning shortcuts. Our approach involves exclusively utilizing the function body while disregarding the signature and return statements. We additionally harness CL based on hard negatives identified within the embedding space. We show that such a hard positive and negative construction strategy is simple, yet essential for effective bimodal contrastive learning.
We train three bidirectional encoder representation models, namely, CODESAGE-SMALL (130M), CODES (356M), and CODESAGE-LARGE (I.3B). We assess the effectiveness of our ap- proach over a wide variety of discriminative tasks, where CODES AGE substantially outperforms the previous of-e--of-theant models with similar model sizes on most tasks. To comprehend the factors contributing to successful code representation learning, we meticulously analyze the key compo- nents of our framework and present our findings for future research endeavors.
RELATED WORKS
Embedding for Programming Languages Recently, there has been a surge of interest in learn- ing general-purpose representations support a wide variety of downstream tasks in programming languages. Feng et al. 20200)); Kanade et al. (2020); Li et al. (2023) take the inspiration of the suc- cess in text and optimize the Masking Language Modeling (MLM) objective on the lincarized code data. Similar to text, they additionally optimize with replaced token detection objective (Clark et al.]
2020) the next sentence prediction objective (Devlin al. 2019b) for source code. Another line
2

---

Published as a conference paper at ICLR 2024
of work leverages the structure aspect of code to provide additional training signals. Among them, Guo et al. (022)) leverages the data flow to encode the relation of Shheer-the-value-comee from" between variables. Wang et al. (2021a) Diang et al. (202I) inject syntactical structure from the ab- stract syntax tree (AST) through variant auxiliary objectives. A more recent work (Guo et al. 2022) flattens the AST structure into sequence directly and encodes the syntax information via language modeling objectives. WWang et al. (2021b); (anne Lachaux et al. (202I) train sequence-to-sequenne language model to reconstruct the original code from an identifier-ofuuscated code where class, function, and variable names are replaced with special tokens. Deobfuscation implicitly encodes data flow and AST without involving auxiliary objectives or complex input with deep hierarchy, since the model needs to understand the dependency between variables as well as code structure so as to correctly predict the names for identifiers.
Contrastive Learning Ever since the early success attained by the Siamese (Hadsell et al. 2006) network, contrastive learning has been widely adopted in representation learning using deep neural networks. Song et al. (2016) extends the vanilla triplet loss by contrasting each positive example against all in-batch negatives, which has greatl improved the learning efficiency and is further popularized by SimCLR (Chen et al. 2020). However, different from the compute version domain where effective positives can be obtained by stochastic transformations of images in the input space, effective data augmentation has long been challenge in NLP due to the discrete nature of the input. Such challenge is further validated in Gao et al. (202T) which shows that dropout (Srivastava al.) 2014) as the minimum data augmentation is often more effective than those obtained by operating in the discrete input space, egg, word deletion and replacement.
Alternatively, various methods have been proposed to leverage naturally occurring pairs posi- tives. Zhou et al.
(2022) treat the consecutive utterances from dialogue data as positives, while Neelakantan ct al.
(2022) consider the neighboring texts mined from the internet. A very recent
work (Wang et al.
2022p leverages the question and answer or comment pairs from StackExchange
and Reddit. In similar vein for programming language, Guo et al. (222); Wang et al. (202a);; Neelakantan et al.] (2022) leverage (text, code) pairs with text mined from the docstrings. We take a
step further by focusing on hard positive and hard negative construction, which key ingredient for representation learning and allows us to attain off-the-shelf embedding models.
3 METHOD
3.1 MASK LANGUAGE MODELING AND DEOBFUSCATION PRE-TRAINING
Given an input sequence with tokens, i.e., x= objective (Devlin et al. 2019b) is formed as follows
the mask language modeling
[MLM(X) =- Σ log P (x,(x-M)
iEM
Here M denotes the mask applied on the given input x. Equation ① is essentially a denoising objective with the task to predict the original tokens given the masked sequence x
Deobfuscation We first consider identifier deobfuscation (DOBF) which pretrains the model to predict the masked-out names of the identifiers. Similar to human programmers, in order to deob- fuscate the code, the model needs to understand both the semantics and structure of the code. Also the NL tokens, i.e., docstring and comment, are excluded from code obfuscation. When the model trained predict the identifier names, can benefit from looking at and correlating with the NL tokens in comments or doestrings as those often carry rich semantics of code. Consequently, the model encouraged to learn improved shared representations between PI. and NI., as indicated by the better NL2Code search performance attained by DOBF than random masking in Table 3
DOBF is initially proposed for Seq2Seq models (anne Lachaux et al. 2021 Wang et al., 2021b). To the best of our knowledge, we are the first to apply it to the encoder-only models. The main challenge to adopting DOBF for encoder-only models is to construct the one-on-one mapping be- tween mask tokens (inputs to the LM) and identifier tokens (output labels) due to the differences in code tokenization (i.e., using vrtesr and model-specific tokenization (i.e., using a sentencepiece tokenizer). We briefly discuss the challenge in Appendix A.5
3

---

Published as a conference paper at ICLR 2024
Random Masking Additionally, we also involve the random token masking strategy in BERT Devlin et al. (2019b) for two main reasons. First, to promote better representations by promoting the model to learn beyond identifiers. Taking Python the code tokens associated with identifiers, hence better representations can be attained by encoding the information carried by the remaining 70% of tokens. Second, not every programmer follows the naming conventions, e.g., meaningless variable names like vl, v2, v3 can be used. Predicting such tokens unnecessarily hard and provides very limited training signal.
example, there are approximately 30% of
We do not follow the 80-10-10 masking convention proposed in the standard MLM for text (Devlin al. 2019b). Since source codes are composed of NL and PL tokens (i.e., identifiers, keywords, operators), random replacement of tokens could hurt both the structure and meaning of code and leads to deterioration in representation learning
Fl We show in Section 4.2.1 that the 80-10-10 con- vention consistently results in worse performance on downstream tasks. In this paper, we also set the random masking rate 15% which we find optimal through our ablation study in Appendix [A.4] For each training example, randomly pick DOBF random masking with equal probability.
3.2 BIMODAL CONTRASTIVE LEARNING WITH HARD NEGATIVE AND HARD POSITIVE
Let x, X,- denote a positive input pair and h, ha+ be the associated representations output by the last hidden layer of the encoder. Let B = (hi, h2+, h2, hg+,...,, hx, hy+) } denote the representa- tions of a randomly sampled batch with N pairs, we then minimize the following symmetrie loss,
La.. (hi, h(+) = -(1o5 exp(h, h, /T) EnEBI(UHt) 7 exp(h, o ha/T)
exp(h, o hTT
exp(h,+ o h1/T)
+ log exp(h,+ o hi/T))+ EkEB 7/- exp(h,+ o he/T))
Here, T is the temperature hyper-parameter which we set as 0.05 in this work. o denotes cosine similarity between two representation vectors, 7 is the weight parameter which we will detail next.
Hard Negative Without supervision, is tricky to identify hard negatives. We resort to a distance- based unsupervised approximation of hard negatives proposed in Zhang et al. (2021). For given anchor h, hard negatives refer to those semantically different examples but are mapped close to h; in the representation space. Thereby, the closer negative is to the anchor h, in the representation space, the larger 7 value is desired, which can be characterized as follows
exp(h; o ha/T)
* 71
(3)
exp(h; hg/r) + EjeBr .)) exp(h; h;/r)
That is. 7. approximates the relative importance of h& to the anchor h, among all 2. in in-batch negatives. Despite the semantic equivalence between training examples except the given positive pairs are not available in our case, the above approximation of hard negatives is still valid as each training batch is randomly sampled with much smaller size compared to that of the whole training data. Hence the presence of false negatives within each batch is negligible when the training data is large and diverse enough. We set the batch size to 8K in this paper, under which we observe monotonic increasing performance reported on the downstream tasks.
Hard Positive We consider naturally occurring (text, function) as positive pairs, where the text is mined from the function docstring (Husain et al.) 2019). The extracted text often summarizes the high-level semantics of the code. Therefore, contrastive learning with such bimodal data largely boosts the NL22Code semantic search performance in Section 4.2.2 Further, the extracted text of semantically equivalent code, no matter from the same different programming languages, is often less diverse compared to the code themselves. Thereby, semantically similar codes can be implicitly grouped together through the same very similar summary text. Our conjecture is validated by the large performance gain on both in-language and cross-language Code2Code search in Section 4.2.2
'For example, masking a couple of tokens randomly from tokenizer. convert.ids.tootooeee..... can yield tokenizer convert .is.ttiiditooooo wtt.ttosssssssssssss but random token replacement can result in tokeni zer. convert.jjt..... wtt.:jjttottttoeerrr..r.. representation leaming via the self-attention mechanism can therehy deteriorate. See Appendix A.3 for more.
Consequently, the code semanties are largely altered and
4

---

Published as a conference paper at ICLR 2024
It is also easy that function names and input variable names often share a significant similarity, especially in terms of the lexical overlap with the summary text (see Appendix AA.6 for statistics). We thereby form hard positives by removing both function signature and return statements We assess the effectiveness of such hard positive construction strategy in Section 4.2.2
4 EXPERIMENTS
Training Data and Model Architecture We train our models on The Stack dataset (Kocetkov al. 2022) over nine languages- - Python, Java, Javascript, Typescript, C#,, C. Ruby, Go, and PHP. As aforementioned, we train three embedding models with size 130M (COdeSAGE-SMALL) 356M CCODESAGE-BLL BASE), and 1.3B (CODESAGE-LARGE) AGE-LARGE) parameters. Please refer to Appendix A for training details at each stage and model hyper- parameerrssssss.
Evaluation Protocol We assess the performance of our models over two main categories of down- stream tasks, semantic search and classification. Our goal is to perform an evaluation of the encoder models for those practical scenarios where supervised fine-tuning data collection is costly. We thereby focus on zero-shot semantic search and only finetuning a linear classification layer on top of the frozen encoders for classification tasks (Peters et al. 2019 Chen et al. 2020; Wang et al. 2022) We report the fully finetuned classification results and finetuning hyper-paramett parameters in Appendix B.3
Baselines We compare models against four general purposs code representation learning en- (For convenience, we refer to it as OpemAI-Ada-002. and SereAl--eex-EEbbddiiige
coders and OpennAlEmbeedding For comparative analysis with OpenAI-CPT-Codee
please see Tables g & g in Appendix, where OpenAI-Ada-002 altains significantly betler or on par perfor- mance against the previous and latest models, respectively). Both CodeBERT (Feng al. (020b) and GraphCodeBERT (Guo et al., 2021l) are trained with standard MLM on six programming languages using CodeSearchNet (Husain et al. 2019 while the replaced token detection objective (Clark et al. (020) and data flow prediction objectives are adopted as auxiliary objectives, respec- tively. UnixCoder (Guo et al., 2022) is trained via three language modeling and two contrastive learning objectives using the same dataset. More recently, StarEncoder (Li et al. 2023) is trained with ML. and next sentence prediction (Deylin et al., 2019a) on 86 programming languages from The Stack (Kocetkov et al.) 2022p. We provide more details for each baseline model in Table 6 in Appendix. We also consider decoder-only baselines in Table & 9 in Appendix B
4.1 COMPARISON WITH THE BASELINES
We first compare CODESAGE against the aforementioned baselines the following tasks.
Code2Code semantic search is the task of retrieving relevant code fragments given a code fragment as query. In this work, we extend the Code2Code search evaluation set (Guo et al., (022) created from CodeNet to six more languages C, C#,, Javascript, Typescript, GO and PHP, for which we summarize the details in Appendix B.2] We report the in-language where query and candidate codes are in the same language, code2code search results in Table
NL2Code semantic search is the task of using natural language the query to retrieve the relevant
code. We consider three benchmarks in Table 2 CoSQA (Huang et al. 2021), AdvTest (Lu et al.) 2021. and CSN (Guo et al. 2021 Detailed data statisties can be found in Appendix B.2
Classification We consider three source code classification tasks. Code Defect detection is a bench- mark in C from CodeXGLUE (Lu et al., 2021). with binary label indicating whether code is insecure and may attack software systems. Code Complexity prediction (Teon et al. 2023) is Java benchmark that requires predicting the algorithmic complexity among 7 labels. The RunTime error prediction (Bieber et al.) (023) benchmark has 29 possible labels with highly imbalanced distribu- tion (see Table 10 in Appendix). For a more robust evaluation, we balance the dataset by aligning its total training examples of the "no_error" class with the cumulative count of the other 28 classes.
Overall Performance Summary On Code2Code search, Table shows that CODES (130M) persistently outperforms all the baseline models with known model size (i.e., exclude
Removal of function signature reduces the chance to learn shortcuts due to its similarity with the summary text. We remove the return statements to make a code look like generic code snippet.
'The dataset includes 2.3M functions paired with natural language documents.
5

---

Published as a conference paper at ICLR 2024
Model
Python
Java
JS
TS
C#
C
Ruby
PHP
GO
Avg.
CodeBERT
14.40
7.62
5.47
6.05
3.66
5.53
13.55
10.28
6.27
8.09
GraphCodeBERT
19.23
10.78
7.38
8.65
5.54
8,48
19.69
15.67
9.65
11.68
StarFncoder
19.17
11.65
9.0
10.52
5.69
9.72
21.57
16.98
10.81
12.79
UnixCoder
30.77
16.45
21.32
21.95
6.19
15.62
32.33
31.93
13.94
21.17
OpenAI-Ada-002 35.91
25.13
19.01
21.86
10.17
29.15
40.85
40.47
23.43
27.33
CODES AGE-SMAL 36.31
23.97
26.60
29.90
11.84
22.84
29.06
34.64
19.56
26.08
CODESAGE- 47.52
22.84
28.70
31.95
13.37
30.99
44.86
51.13
25.15
32.95
CODESAGE-LARGE 46.70
33.13
37.16
41.18
16.81
32.89
54.12
52.13
32.48
38.51
Table 1: MAP score (%) of the zero-shot code search task. The language names mentioned in the top row indicate the languages queries and candidates are written in.
NL2Code
Classification Complexity
Model
CoSQA
AdvTest
CSN
Defect
RunTime
CodeBERT
0.24
0.06
0.10
51.82..38
35.601.0
6.2g.00
GraphCodeBERT
16.20
5.58
11.26
55,260.28
55.54
10.63p.10
StarEncoder
10.78
0.93
2.69
53.2p.11
50.633.33
8.91
UnixCoder
42.11
27.32
46.39
60.28p.01
76.451.10
20.870.43
OpenAl-Ada-002
44.23
38.08
71.24
62.56p.11
79.82p.50
20.84p.36
CODESAGE-SMALL
49.92
41.28
63.86
57.52a.21
79.760.50
25.051.04
CODESAGE-BASE
48.50
49.08
68.72
57.749.09
85.32
24.70p.40
CODESAGE-LARGE 47.53
52.67
71.24
58.95o.13
90.322.10
24.420.28
Table Left. MRR score (%) of NL2Code search in zero-shot setting. For CSN, we report the average performance over six languages (see Table Appendix for the detailed results). Right. (macro) score of the source code classification tasks attained by only finetuning the classification head. We finetuned each model using three seeds and reported the mean and standard deviation (in subscript). The fully finetuned results can be found in Appendix B.3
OpenAI-Ada-002) on every language, with 23.19% relative (4.91% % absolute) improvement on the average performance when comparing with UnixCoder. With the increased model size, CODESAGE- outperform the best baseline model, i.e., Open (model
BASE and CODES AGF-LAGGFLAAREE size unknown), with 20.56% relative (5.62% absolute) and 40.91% % relative (1 111.18% absolute) im- provement on the average performance, respectively.
As shown in Table 2 CODES absolute) improvement over UnixCoder on NL2Code search. Compared to OpenAI-Ada-O00 CODESAGE-SMALL attains 12.86% relative (5.69% absolute) improvement on CosQA and an 8.4% relative (3. 12% absolute) improvement on AdvTest. On the other hand, OpenAl-Ada-O01 on CSN. However, we want to highlight
achieves 18.54% to 51.1% relative (7.81% to 13.96%
at-
tains the same average performance as CODESAGF-LE the performance gain attained by CODES on AdvTest which contains normalized Python func- tions (from CSN) with function and variable names replaced by dummy variables (see Figure 9 in Appendix). AdvTest constructed in this way better assesses the generalization performance as the model needs to understand what the obfuscated code does so as to identify the correct target code for given natural language query.
Compared to both UnixCoder and OpenAl-Ada-002, CODESAGE persistently performs better on code complexity and nuntime prediction with large margins in Table 2 We also notice that CODES AGE underperforms both models on code defect detection, whilst attaining better perfor- mance when finetuning the full models in Table 12 in Appendix.
4.2
ABLATION STUDY
4.2.1
MASKING STRATEGY
80-10-10 vs. Full Mask Given input sequence, standard MLM (Devlin al. 2019b) first ran- domly samples a subset of its tokens, of which 80% are replaced by a special token "[MASK]". 10% left unchanged, and the other 10% are replaced by random tokens from the vocabulary. We revisit

---

Published as a conference paper at ICLR 2024
def Sinary-search ( (arr, low, high x):
def meem_search (arr, low<MASK> high, x):
" "Returns index of, in arr if present, else -1. if high >= low:
1 RReurrss <MASSS
in
present, else -AAAS gelConfig
if sal >= lownone mid (high <MASS> low) <ASK
mid (high low) 11 2
arr return mid clif arr mid] > x:
if arr mid] = return mid elif <MASK> [id]] > X:
return binary-search arf low. mid 1,
M (
10
11
arr low. mid 1, <MAKK>
11
else:
return anaarysearcee
else:
13
arr mid + 1. high, x)
14
arr mid + 1. highsystemd x)
14
elso:
clse: return
return 15
(a) Sample code (left) and its corrupted version following the 80-10-10 10-10 rule (right).
30
Cross Languags Code M&F
20
50-10-10
300
Full
CodeSage BASE
Lanquage Code2
CusleSageeLAAGGE
15
20
HL.7Cole
-50
10
20
30
20-
5
(b) With fixed masking rate of 15%, assess the effectiveness of applying "Full Mask", Le.. replacing the sampled tokens with the [MASK] token only, and 80- 10-10-10 corruption strategy on downstream tasks.
Figure 2: 80- 10-10-10 vs. "Full Mask".
CODESAGE-SMALE
CODESAGE BASE
CODESAGE-LARGS
Model
R
D
S
P
R
D
S
P
R
D
S
P
NI,2Code
6.6
19.9
22.7
25.8
12.2
22.5
22.0
23.3
19.4
23.3
29.4
30.5
Code2Code (In)
16.8
14.6
17.9
19.7
28.2
23.7
25.3
29.2
30.7
28.2
30.2
33.9
Code2Code (Cross)
5.7
6.7
8.8
9.6
17.2
14.1
14.6
19.7
20.5
18.0
19.0
24.6
Classification 51.2
53.9
53.5
53.4
53.8
55.6
54.8
55.4
52.0
55.6
57.2
56.5
Table 3: We explore two options to leverage DOBF (D) and random masking (R) to complement each other. (1) Sequential (S): training the model with random masking first, then DOBF. (2) Parallel (P): randomly picking either DOBF random masking for training example our strategy.
the effectiveness of such convention, originally proposed for text, for code in Figure Surprisingly, compared to simply replacing all selected tokens with the [MASK] token, i.e., "Full Mask", the 80-10-10 masking scheme causes a large performance drop across different downstream tasks, as shown in Figure 2b A similar finding has been reported in Gao et al. (2022) for text. However, the degradation is more severe for source code. As Figure 2a indicates, when replacing with random tokens, both the semantics and structure of the masked code can be largely disrupted, which together with the presence of "[MASK]" tokens makes the learning too challenging (see Appendix AA.3 for more discussions). We hypothesize that excessive corruption may also account for the modest hancement observed in downstream tasks when scaling up the size a model trained with 80-10-10 in Figure 2b It would be intriguing explore whether this scaling trend would experience sudden expansion with a further increase in model size and training data, potentially identifying a phase transition point, provided that the computational resources permit such an investigation.
Deobfuscation Random Masking Complement Each Other We investigate DOBF and the random masking based MLM with "Full Mask" in Figure 3 DOBF persistently outperforms ran- dom masking on classification, which validates our motivation that the model is promoted to better capture (understand) the code structure better on NL2Code search than random masking. potential reason could be natural language in comments and docstrings often carry rich semantics of code while both being excluded from mask- ing in DOBF; hence when training the model to predict the identifier names, it will look at and correlate with the natural language and lead to better contextualized representations between natural
to predict the identifier names. DOBF also performs

---

Published as a conference paper at ICLR 2024
language and programming language. On the other hand, the random masking strategy (with "Full Mask") outperforms DOBF on both in-language and cross-language Code2Code search tasks. As examined in Appendix A.3 [
large portion of tokens in code snippets not identifiers. Therefore, the random masking strategy allows the model to learn beyond identifiers and enrich the semantics encoded in representations. In summary, Table 3 validates strategy of jointly optimizing DOBF and random masking so as to leverage their strengths to complement each other.
Baseline
Hard Only
Avg Code2Code MAP
30
Hard Nogative Only
20
Avg Code2code
25
Unimodal (Dropout)
Bimodal (Ours)
Her alions
CodeSage SMALL
CodeSage-BAS
(a) Effectiveness of hard negatives and hard positives.
(b) Unimodal bimodal contrastive learning.
Figure 3: (a) Hard negative and hard positive can independently boost performance over the baseline where neither is applied. Further improvement is attained when leveraging them simultaneously. (b) Unimodal contrastive learning with positives obtained via dropout requires longer training and hence cannot leverage vast amounts training data to further enhance the representations.
4.2.2 ON EFFECTIVENESS OF CONTRASTIVE LEARNING
Hard Positive and Hard Negative Effectively Boost Performance We first demonstrate the ef- fectiveness of the hard positive and hard negative construction strategy in Figure 3a| As it shows, both hard positive and hard negative can independently improve the performance by a large margin, while the combination of them persistently yields better performance across different model sizes. We also observe that large model size (i.e., CODES AGE- BASE) benefits more from the proposed hard negative construction strategy. This observation is unsurprising since larger models possess more capacity to leverage more challenging and effective learning objectives.
Unimodal Bimodal Contrastive Learning In Figure 3b
compare bimodal contrastive learning approach against the Dropout-bast unimodal contrastive learning where a positive pair is obtained by leveraging different dropout masks of the transformer in two forwarding passes of the same sequence (Gao et al.) 2022 1 Guo et al. 2022). For fair comparison, hard negative optimization applied both approaches. We can see that the dropout-basee -based unimodal contrastive learning suffers from supporting long training process and hence cannot effectively utilize a large amount of pretraining data to further improve the representations. A similar finding has been reported by (Zhou et al., (002).. Indeed, both Gao et al. (202T) nor [uoo et al.) (2022) demonstrate dropout as effective augmentation for text and code respectively, only use few million training examples that can be covered by the amount of training data in the first 500 iterations (with batch size 8K) in Figure 3h where the dropout- contrastive learning shows improvement over the baseline.
To gain a deeper understanding of the perfor- mance improvement achieved through contrastive learning during Stage of pretraining, delve into the analysis of semantic search performance. As Figure shows, contrastive learning persis- tently boosts the search performance with comparatively larger improvement on the cross-lingual scenarios, encompassing both NL2Code and cross lrass-language Code2Code search. We posit that the text extracted from docstring helps group semantically equivalent code together as the text often summarizes the high-level semantics of code and hence are likely less diverse than the code them- selves. particular, those parallel examples from different programming languages can share very similar or even the same summary. For NL2Code, the larger improvement can be credited to its alignment with the bimodal contrastive learning objective using (text, code) as positives. Such bi- modal objective also brings NL and PL closer in Figure4b). Compared to the model trained Stage-I only, contrastive learning pulls together NL and PL such that the relative similarity gap between par- allel NI.2Code pairs and cross-language Code2Code parallel examples largely decreased.
Larger Improvement on Cross-Lingual Search
8

---

Published as a conference paper at ICLR 2024
35
60
Staped Only
Avg Language coodccdo Moo
Stagel Saegell
1.3
30
Conc2Coo
-00
Arg NL2Codd MWA
0.8
Colire Sind by
20
5
10
0.0
Caso2Cate Languaget
0.2
Stagel
Stage: & Stage
(a) The performance of CODESAGE in semantic search, comparing sults between searches within the same language and across different languages. while varying model sizes training approaches.
(h) Cosine similarity between parallel examples randomly sampled pairs using CODESAGE representations.
Figure 4: Examining the effectiveness of contrastive learning (Stage-II) by comparing COD against those trained with the token-level denoising objective only (Stage-I). (a) Compared to the in-language Code2Code search, contrastive learning persistently leads to larger performance boost for cross-lingual search, including both NL2Code and cross-language Code2Code search. (b) Con- trastive learning leads to dispersed representation space with improved discrimination, in- dicated by the corresponding enlarged similarity gap between parallel and randomly sampled pairs, while simultancously bridging the relative similarity gap between NL2Code and Code2Code pairs.
4.3 ON OBJECTIVE AND DOWNSTREAM PERFORMANCE SCALING WITH MODEL SIZE
In Figure [ we study how the downstream task perfor- mance scales with the model size when pretrained with different schemes, iee, token-level objective only (Stage-
Avg In-lang Code MAP
D), contrastive learning only (Stage-II), and our proposed two-stage framework with Stage-I followed by Stage-II. We use zero-shot multilingual in-language code search performance (averaged over nine languages) for this ex- ploration. We can see that models pretrained from scratch with contrastive learning alone do not scale with the in- creased model size. NNeelakantan et al. (2022) report a similar finding that the contrastive objective its is not sufficient to learn useful representations. When train- ing from scratch with contrastive learning only, we find the training loss often converges at a large value, indicat- ing the model cannot well discriminate each positive pair from the other in-batch negatives. In other words, leveraging the token-level denoising objective to provide good embedding foundation is essential for contrastive learning be effective and further enhance the sequence-lewe presentations.
25
Stage-I Only
Stage-Il
15
Stage-I Stage
5.
SMALL
Figure 5: On the downstream task per- formance scaling with pretrained model size under different training schemes.
CONCLUSION
In this study, we unveiled CODESAGE, a cutting-edge encoder representation learning model for source code. We trained CODESAGE using an extensive dataset comprising 237 million code files and 75 million bimodal code and natural language pairs across nine languages. Our findings reveal that our model outperforms its predecessors significantly in tasks related to code search and code classification. We also delve into the essential factors contributing to enhanced code representation learning across various model sizes.
We hope our work will serve as an inspiration for future works not only in code representation learning by effectively utilizing publicly accessible extensive corpora for source code, but also in the broader field of iniver model training. This includes integrating language generation and embedding within single model, seen in the works of (Jain et al. 2023 Muennighoff et al. 2024]. Additionally, we aim to encourage advancements in cross-domain representation learning, such unifying text and code embedding within single model.

---

Published as a conference paper at ICLR 2024
REFERENCES
Marie anne Lachaux, Baptiste Roziere, Mare Szafraniec, and Guillaume Lample. DOBF: A de- obfuscation pre-training objective for programming languages. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. TRL
net/ wttp//iddmmddeeeee 9BSHTNT
David Bieher, Rishab Goel, Dan Zheng, Hugo Larochelle, and Daniel Tarlow. Static predic- tion of runtime errors by learning to execute programs with external resource descriptions. In The Eleventh International Conference on Learning Representations, 2023. URL https:
/openzeview
C5nTdJG
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. ArXiv preprint, abs/2107.03374, 2021. URL https:// arxiv.org/abe abs
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings URL http://pproceedings mmlr.
of Machine Learning Research, pp. 1597-1607, 2020.
pIess/vl19/ when20]cchen2]]].. .hmll
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022. URL
inttps://aran
org/ als 2204.222222 02311
Kevin Clark, Minh-Thang Thang Luong, Quoc V. Le. and Christopher D. Manning. Electra: Pre-training text encoders as discriminators rather than generators. International Conference on Learning Representations, 2020. URL httpssts:: httpttttss://pereteeett http:://opeeeeeeeeeeee net forum? ?id=r XMH1BtVB
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 4171-4186 Minneapolis, Minnesota, June 2019a. Association for Computational Linguistics. doi: 10. N0.18653/v1/N194 1423. URL https:
//aclanthology.oor///V 1423
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4171-4186, 2019b. doi: 18.18653/v1/
N19-1423. URL btttbtts::/// pettps://aallanthoooog.o wwwphhhwooooooo....111.o. org/ 9-119-9-433
Zhangyin Feng, Daya Guo, Duyu Tang. Nan Duan, Xiaocheng Feng, Ming Gong, Linjum Shou, Bing Qin. Ting Liu, Daxin Jiang, and Ming Zhou. CodeBERT: pre-trained model for programming and natural languages. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 1536-1547, 2020a. doi: 10. 18653/v1/2000..o.. hww.::/....nniinnnnnn...........
URL https:/
cclannho 1 ogy. org/2020, f f findiiii ngs eangg--mmmll emn p. 139
Zhangyin Feng, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, and Ming Zhou CodeBERT: A pre-trained model for pro- gramming and natural languages. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 1547, Online, November 2020b. Association for Computational Lin- org/
guistics. doi: 10.18653/v 1220/3//1//222.. 100.66652110000.mmmin--mmmpppppp 133. URL
2020. findingn -aweiggsseennl... 119
Jun Gao, Changlong Yu, Wei Wang, Huan Zhao, and Ruifeng Xu. Mask-then-fill: A flexible and ef- fective data augmentation framework for event extraction. In Findings of the Association for Com- putational Linguistics: EMNLP 2022, pp. 4537-4544, Abu Dhabi, United Arab Emirates, Decem- ber 2022. Association for Computational Linguistics. doi: 10. 10.18653
URL https: / antho hgw.llloooooogg///// nd htw.iinnnnnnn-eeeeee ngs 332
10

---

Published as a conference paper at ICLR 2024
Tianyu Gao, Xingcheng Yao, and Danqi Chen. SimCSE: Simple contrastive learning of sentence embeddings. In Proceedings of the 2021 Conference on Empirical Methods in Natural Lan- guage Processing, pp. 6894-6910, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational I. Linguistics. doi: 10. 11055//112221..mmpmaai.... htttss://aclanthoo logy. org 2021 emnle-mae
URL
552
Daya Guo, Shuo Ren, Shuai Lu, Zhangyin Feng, Duyu Tang, Shujie LIU, Long Zhou, Nan Duan, Alexey Svyatkovskiy, Shengyu Fu, Michele Tufano, Shao Kun Deng, Colin Clement, Dawn Drain, Neel Sundaresan, Jian Yin, Daxin Jiang, and Ming Zhou. Graphcode (bert): Pre-training code representations with data flow, In International Conference Learning Representations, 2021.
net/ forum?id JLCA4eZ43P
Daya Guo, Shuai Lu, Nan Duan, Yanlin Wang, Ming Zhou, and Jian Yin. UniXcoder: Unified cross-modal pre-training for code representation. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 7212-7225, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 1011118661122222 499. URL ht tptt:p::0. http:::///aaaatthooogoeo org. wttp///gg22222 2022 acw-long..cn.c.... 499
Raia Hadsell, Sumit Chopra, and Yann LeCun. Dimensionality reduction by learning invariant mapping. In 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recogni- fion (CVPR 06), volume 2. pp. 1735-1742. IEEE, 2006.
Junjie Huang, Duyu Tang, Linjun Shou, Ming Gong, Ke Xu, Daxin Jiang, Ming Zhou, and Nan Duan. CoSQA: 20,000+ web queries for code search and question answering. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th In- ternational Joint Conference Natural Language Processing (Volume 1: Long Papers), pp. 5690-5700, Online, August 2021, Association for Computational Linguistics. doi: 10.18653/v1/ 2021 aclloooooog
442
Hamel Husain, Ho-Hsiang Wu, Tiferet Gazit, Miltiadis Allamanis, and Marc Brockschmidt. CodeSearchNet challenge: Evaluating the state of semantic code search. arXiv preprint
arXiv:1909.09453. 094909990999,..... 2019. httl/httts///axtttt www.aawwww.wiaaaaaaaaaaaoreeeeea abaa 190.9.99996666
Nihal Jain, Dejiao Zhang, Wasi Uddin Ahmad, Zijian Wang, Feng Nan, Xiaopeng Li, Ming Tan, Ramesh Nallapati, Baishakhi Ray, Parminder Bhatia, Xiaofei Ma, and Bing Xiang. Contra- CLM: Contrastive learning for causal language model. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.). Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pp. 6436-6459, Toronto, Canada, July URL.
2023. Association for Computational Linguistics. doi: 10.116633/1/22233acll/
http//aclaann http:///aaaatthhogg... wtw.:o/....o22222..111ooooo1 acl -ong. 355
Paras Jain, Ajay Jain, Tianjun Zhang, Pieter Abbeel, Joseph Gonzalez, and Ion Stoica. Contrastive code representation learning- Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 5954-5971, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/1120011mmmm.. main. 482. UtL.htt.p:://accaamthooooo
482.
Mingi Jeon, Seung yeop Baik, Joonghyuk Hahn, Yo-Sub Han, and Sang- Ki Ko. Deep learning-base based source code complexity prediction. 2023. URL http:::: net forum?? id
9irBKvxsw9
Xue Jiang, Zhuoran Zheng, Chen Lyu, Liang Li, and Lei Lyu. Treebert: A brsed pre-trained model for programming language. In Uncertainty in Artificial Intelligence, pp. 54-63. PMLR, 2021.
Aditya Kanade, Petros Maniatis, Gogul Balakrishnan, and Kensen Shi. Learning and evaluating contextual embedding of source code. In International conference on machine learning, pp- 5110-5121. PMLR, 2020.
Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferran- dis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, et al. The stack: 3 tb of permissively licensed source code. arXiv preprint arXiv:2211.15533 2022.
11

---

Published as a conference paper at ICLR 2024
Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Mare Marone, Christopher Akiki, Jia Li. Jenny Chim, et al. Starcoder: may the source be with you! arXiv preprint wrXiv22305006660 2023.
Xiaonan Li, Yeyun Gong, Yelong Shen, Xipeng Qiu, Hang Zhang, Bolun Yao, Weizhen Qi, Daxin Jiang. Weizhu Chen, and Nan Duan. CodeRetriever: A large scale contrastive pre-training method for code search. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 2898-29188 Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguisties. doi: 10. (0..1863/1120222mmmmm/ main. URL wrtp www.pemmnnpppm.mmmm.mm
Shuai Lu, Daya Guo, Shuo Ren, Junjie Huang, Alexey Svyatkovskiy, Ambrosio Blanco, Colin Clement, Dawn Drain, Daxin Jiang. Duyu Tang, Ge Li, Lidong Zhou, Linjun Shou, Long Zhou, Michele Tufano, MING GONG, Ming Zhou, Nan Duan, Neel Sundaresan, Shao Kun Deng, Shengyu Fu, and Shujie LIU. CodeXGLUE: A machine learning benchmark dataset for code understanding and generation. In Thirty-fitth Conference on Neural Information Processing Sys- tems Datasets and Benchmarks Track (Round 1), 2021. htthttps: htw.tttps///ppeereeeeeee net.
Niklas Muennighoff, Hongjin Su, Liang Wang, Nan Yang, Furu Wei, Tao Yu, Amanpreet Singh, and Douwe Kiela. Generative representational instruction tuning, 2024.
Arvind Neelakantan, Tao Xu, Raul Puri, Alec Radford, Jesse Michael Han, Jerry Tworek, Qim- ing Yuan, Nikolas Tezak, Jong Wook Kim, Chris Hallacy, et al. Text and code embeddings by contrastive pre-training. arXiv preprint arXiv:2201.10005 2022.
Matthew E. Peters, Sebastian Ruder, and Noah A. Smith. To tune or not to tune? adapting pre- trained representations to diverse tasks. In Proceedings of the 4th Workshop on Representation Learning for NLP (RepLANLP-2019), pp. 7-14, Florence, Italy, August 2019. Association for Computational Lingusstic. doi: 10. 10.1653//1/W19-44000 URL
org/W19 4302
Hyun Oh Song. Yu Xiang, Stefanie Jegelka, and Silvio Savarese. Deep metric learning via lifted structured feature embedding. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pp. 4012, 2016. doi: 2..1009CCVP.2201..42 URL httthhttp:///d www.oioo.oooooooooooooooooooooe oo.orrg/100 1109//1PPPR 2016. 434
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, lly Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15()):1299-198, 2014.
Liang Wang, Nan Yang, Xiaolong Huang, Binxing Jiao, Linjun Yang, Daxin Jiang, Rangan Ma- jumder, and Furu Wei. Text embeddings by weakly-supervised contrastive pre-training. arXiv preprint arXiv:2212.0355. 2022.
Xin Wang. Yasheng Wang. Fei Mi, Pingyi Zhou, Yao Wan, Xiao Liu, Li Li, Hao Wu, Jin Liu, and Xin Jiang. Syncobert: Syntax- guided mudti-modal contrastive pre-training for code representation. arXiv preprini 20ti::2208004555. 2021a.
Yue Wang, Weishi Wang, Shafiq Joty, and Steven C.H. Hoi. CodeT5: Identifier- unified pre- pre-trained encoder-decoder models for code understanding and generation. Proceedings of the 2021 Conference Empirical Methods in Natural Language Processing, pp. 8696 8708, 2021b. doi: 10. 1o.118633/1/221111111 www.emmmmi.n...a 2021 emnlp-main. 685
URL htap::::/aallanhhooogy..
Dejiao Zhang, Shang-Wen Li, Wei Xiao, Henghui Zhu, Ramesh Nallapati, Andrew O. Arnold, and Bing Xiang. Pairwise supervised contrastive learning of sentence representations. In Proceed- ings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 5786- 5798, 2021. doi: 122116653/112221.omml
haalp-mmin.461 URL https: www.aaattntooooo.......oo.....a.
2021 emn 1p pp-mai wemmain... 467
12

---

Published as a conference paper at ICLR 2024
Zhihan Zhou, Dejiao Zhang, Wei Xiao, Nicholas Dingwall, Xiaofei Ma, Andrew Arnold, and Bing Xiang. Learning dialogue representations from consecutive utterances. In Proceed- ings of the 2022 Conference of the North American Chapter of the Association for Computa- tional Linguistics: Human Language Technologies, pp. 754-768, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 2atzg//222.aaac..... naacl mammainn 55
waw.0222.aac-mmmin...m URL
13

---

Published as a conference paper at ICLR 2024
A DATA, MODEL, AND HYPER-PARAMETERS DETAILS
A.1 PRETRAINING DATA
Masked Language Modeling (MLM) and Identifier Deobsfucation (DOBF) For both MLM and DOBE, we use the Stack dataset (Kocetkov al., 2022). We set the maximum sequence length to 1024 with concatenation and block attention.
Contrastive Learning (CL) In CL, we focus on bimodal data, i.e., code and natural language pairs, denoted as (lexi, function). Text is extracted as the first sentence from the docstring of a function (Husain al. 2019). For better interpretation, refer such text "summary" in this section as it often summarizes the high-level semantics of function. We filter or modify summaries based on the following practices.
1. Filter summary if it is not in English.
2. Filter summary the number of tokens summary is <3 or >256.
3. Remove URLs, HTML tags, and doctags from the summaries.
4. Fix bad Unicode text in the summaries.
5. Filter functions with no more than one line of code in the function body.
We summarize the statistics of our pretraining data at each stage in Figure and Table
Language
Total files
#Functions
#Func. w/ docstring
#Func. w/ summary
Python
24,214,270
67,264,716
24,321,126
18, 18,14,.327
Java
42,429,211
84,828,833
17,613,636
13,118,303
Javascript
40.112,121
35,469,803
7,450,153
4.796,101
C#
21,702,269
37,284,300
9,325,665
7,350,191
21,383,832
16,253,435
4,392,973
2,958,699
Ruby
1,205,146
5,475,047
1,217,764
1,049, ,356
GO
11,653,185
31,067,259
11,247,051
9,739,86 861
PHP
34,851,418
42,373,766
22,191,329
13,416,574
Typescript 19,589,267
16,612,988
2,637,244
1,863,436
Total 237,961,548
367,905,026
105,760,862
75,389,347
Table 4: Statistics of the data used in pre-training via Masked Language Modeling (MLM) and Identifier Deobsfucation (DOBF), followed by contrastive learning (CL). The data is collected from The Stack (Kocetkov et al. 2022).
A.2 MODEL AND TRAINING POPPER-PRRMMTERS
We pretrain three sizes of model architecture which we refer to as CODESAGE-SMALL CODESAGE- BASE, and CODESAGE-LARGE, We summarize the model hyper-parameters in Table
A.3 ON TOKEN DISTRIBUTION AND STAGE-I PRETRAINING OBJECTIVE
In our preliminary study, we perform data analysis where we examine the ratio of natural language (NL) and programming language (PL) tokens. In source code, tokens are broadly categorized into five groups: (1) identifiers, (2) keywords, (3) operators, (4) delimiters, and (5) literals. We tag the String literals (i.e., docstring, comments) as NL tokens, while all other tokens are considered PL tokens. We use free-sitier parse source code and extract the five categories of code tokens. Then we tokenize them using Starcoder tokenizer (Li et al. (023). From Stack- Python corpora, we compute the following statistics using Starcoder tokenized tokens.
Approximate PL Tokens: 57.8% tokens belong to fidentifiers, keywords, delimiters, operators). Among them, 53.8% tokens belong to identifiers and 46.2% other tokens. 2. Approximate NL Tokens: 42.2% tokens Literals [Boolean, Numeric, String)- Among them, 92.9% tokens belong to String literals and 7.1 % tokens belong others.
As we can tell from the above numbers, the approximate NL tokens account for roughly 40% of the overall tokens for particular programming language. Therefore, when replacing masked to-
14

---

Published as a conference paper at ICLR 2024
CODESAGE-SMALL CODESAGE BASE CODESAGE-LARGE
#layers
24
24
#heads
16
Model dim
1024
1024
2048
Vocab size
49,152
49,152
49,152
Max sequence length
1024
1024
1024
Total parameters
130M
356M
1.3B
Stagel: Masked Language Modeling
Dropout
0.1
0.1
0.1
Max steps
250,000
250,000
250,000
Warmup steps
5000
5000
5000
Batch size
2048
2048
2048
Base learning rate
3c-4
3e-4
3e-4
Stage2: Contrastive Learning
Dropout
0.1
0.1
0.1
Max steps
20,000
20,000
20,000
Warmup steps
500
500
500
Batch size
8192
8192
8192
Bae learning rate
5e-06
5e-06
5c-06
Table 5: Model architecture and pre-training related Syper-parameete
kens with random token could result in replacing a PL token with an NL token, and vice However, there are often no clear boundaries between PI. and NL tokens in many scenarios, as PI. tokens, e.g., those identifier-related tokens, often expected to carry clear semanties that the code snippets are interpretable by humans. Therefore, given masked input tokens following the 80-10-10 10 convention, can be a non-trivial task for the model decide which tokens are from corruption. This together with the structure nature of PL makes it possible for those random tokens largely disrupt both the semantics and structure of code and make the representation learning too challenging to be effective.
Take the example in Figure 6 (right) for illustration, the function name "binary search" is being corrupted with random tokens at all three places it appears, which has the potential to alter the semantic meaning of the code. Although we may expect the model to still be able correctly recover "binary _search" from the corrupted code, it is challenging task given (I) the syntax of the code has been disrupted by another random token "getConfig": (2) the presence of MASK > tokens; and (3) the bidirectional self-attention mechanism can drive the model to leverage those random tokens to form prediction of the masked tokens.
def canary-scarcee
high x): Returns index of in arr if present. else -1. high >= low:
def chem search (arr, low<MASK> high x):
Returrs <MASK> of x in arr if
present. else -CMAS getConfig if sal >= lownone mid = (high MASKS low) <MAK 2
mid (high low)
if arr midd] ==
arr [mid] == x: return mid elif MASK>-[mid] return BOBT <MASK-MAKK
return mid
elif arr [mdd] > x:
return ginary-searee
10
arr, low, mid 1, x)
arr low. mid 1, <ASKK>
else:
return binary
else:
return instead meserc_seearch arr mid + 1, highsystemd x)
arr. mid 1. high. x)
14
else:
else: return
15
return -1
Figure 6: A code snippet(on the left) and its corresponding masked version, were created using the 80-10-10 10 practice with 15% masking rate (on the right).
A.4 MASKING RATE
With "Full Mask", fee. MLM without the 80-10-10 corruption strategy, we investigate the optimal masking rate in Figure 7 We consider three constant masking rates, 7.5%, 15%, and 30%, well
15

---

Published as conference paper at ICLR 2024
dynamic masking strategy with the masking rate being randomly selected from the range [10%, 50%] for each training example. We find 15% remains the optimal masking rate among the four variants we investigate.
Cross-Language Code2Code MAP
20
In-Language Code2Code MAP
Full Mask
30
15
20
10
10%6-50%
Figure 7: Maksing rate and zero-shot Code2Code search performance investigated on CODESAGE- BASE. We consider three constant masking rates, 7.5%, 15%, and 30%, as well as a dynamic masking strategy with the masking rate being randomly selected from the range [10%, 50%] for each training example.
A.5 IDENTIFIER OBFUSCATION
In this research, we employ an identifier deobfuscation (DOBF) objective to train bidirectional encoder representation models. While our inspiration for this approach comes from the DOBF method introduced by anne Lachaux et al. (202I), our adoption strategy differs from theirs. In their work, Janne Lachaux al. (2021) trained sequence-to-seqqeeee language model to recon- struet the original code from an obfuscated version where class, function, and variable names were replaced with special tokens. In contrast, approach applies this technique to encoder- only models. This adaptation involves a non-trivial effort to establish a 1-1 mapping between mask tokens and identifier tokens (will be masked and encoders will predict them) due to dis- parities in code tokenization (i.e., using free-sitter) and model-specific tokenization (i.e., utilizing sentencepiece tokenizer). To illustrate, let's consider the tokenization process. Tree-sitter kenizes "def function _nmme():") as {def, function name, (),;}, whereas a model-specific tok- enizer might tokenize it {def, forf.uuncee function. -, name(, (, ).:). Consequently, we encounter challenge to construct the mapping from masked tokens to prediction tokens: {[mask], [mask], [mask]] {J |functionn namne). -,aiii,,,,,,,,,,, by skipping "(" token that part of the identifier token "name". perform obfuscation and construct the mask map, we developed obfuscation (OBF) tool.
OBF Tool We developed tool that takes entire source code file function input and outputs an identifier obfuscated code along with a token map. We provide an example in Figure 8
We used Iree-siteer parse code snippet and extract all identifiers and their ancestor node types. Based on the node types, we identify class names, function names, function arguments, and function calls. Then we replace them with special tokens (c,, fs, v2 for class names, function names, and variable names, respectively). Then we include the special tokens into the model tokenizer (i.e., Starcoder tokenizer) and tokenize the obfuscated code such that special tokens are retained in the output tokens. Finally, we use the model tokenizer to tokenize the identifiers individually and replace the special tokens (C,, fi, v;) with the identifier tokens.
A.6 ON LEXICAL OVERLAP AND HARD POSITIVE DESIGN
As detailed in Appendix A.1 mary text. We then examine the lexical overlap between docstring (summary) and function signature versus that between docstring (summary) and function body. In Stack-Python corpora, found
extract the first sentence from the function docstring the sum-
1. 22.3% of tokens function signature and 23.1% of tokens in function body overlap with docstring tokens.
2. 12.3% of tokens in function signature and 11.6% of tokens in function body overlap with summary tokens.
16

---

Published as a conference paper at ICLR 2024
class Node:
class c.0:
2 def --init-- (self, v):
def f_0 v.l):
3 self. data V
3.
v_0, v.2 vv1
self left None
v_0 v_3 None
self. right = None
v_0. v_4 None
7 # Function to print postorder traversal
7 Function to print postorder traversal
8 def printPostorder (node). if node == None: return
8 def f. (v.5):
if v_5 = None
return
12 First recur on the left subtree
12
First recur on the left subtree
node.... left)
(v.(v..vv.3
Then the right subtree print P Postorder (node... ri tight)
# Then
the right subtree (v-5. v_4)
16
17
17
Now deal with the node print node. data end=
# Now deal with the node
print (v.5.v.2 end=
Figure 8: An example of a Python code (at the left) and its corresponding obfuscated version (at the right) generated by our developed obfuscation tool. The class names, function names, and variables are replaced by special tokens. Given the code on the left, our developed OBF tool produces the obfuscated code and the identifier map: {co, vo, U1, U2, V3, U4, US, So, fi (Node, self. data, left, right, node, -init prott prstorrottrorrrr
This validates our intuition that the docstring summary of function often has large lexical overlap with the function signature. Thereby, when contrasting doestring or summary with the entire function, the model tends learn shorteut by leveraging such overlap, and hence fail to capture the semantic equivalence concept in the representations. Consequently, poor generalization is attained.
B EVALUATION OF DOWNSTREAM TASKS
B.
BASELINE MODELS
We summarize the baseline model size and output representation dimension in Table
Embedding Dimension
Max Sequence
Training
Model
Model Size
Length
Data Source
CodeBERT
125M
768
512
CodeSearchNet
GrRR
125M
768
512
CodeSearchNet
StarEncoder
125M
768
1024
The Stack
UnixCoder
125M
768
1024
CodeSearchNet
OpenAl-Embeedding-Adaae
Unknown
1536
8191
Unknown
Table 6: Model size and dimension of the embeddings. The GitHub Code dataset is available
https://huggingface co/ ht////aaatts/ceeeoe cctaasets/ccoeppaooo//ttbbbbbbb/
B.2
CODE SEARCH
We summarize the data statisties of NI.2Code and Code2Code benchmarks in Table 7 Below, we provide more context each dataset.
Code2Code search is the task of retrieving relevant code fragments given a code fragment as a query. In this work, we extend the code2code search dataset (Guo et al. 2022 created from CodeNet to six more languages- C. C#. Javascript, Typescript, GO, and PHP. The original dataset includes 2 to 10 solutions for each problem in Java, Python, and Ruby. At first, collect the problem identifiers and aggregate solutions in those six languages from CodeNet. Also, CodeNet provides cluster identifiers for each solution where solutions within cluster near duplicates of each other. We collect one
17

---

Published as a conference paper at ICLR 2024
Code2Code Semantic Search Data Statistics in Each Language
Python
Java
JS
TS
C#
C
Ruby
PHP
GO
Num Queries
15,594
23,530
6,866
3,385
11,952
11,260
11,744
6,782
9,720
Num Candidates
15,594
23,530
6,866
3,385
11,952
1 11,260
1 11,744
6,782
9,720
NL2Code Semantic Search Data Statistics in Benchmark and Language
CoSQA
AdvTest
CSN
Python
Python
Python
Java
JS
PhP
Go
Ruby
Num Queries
500
19,120
14,918
10,955
3,291
14,014
8,122
1,261
Num Candidates
6,268
19,120
43,827
40,347
13,981
52,660
28,120
4,360
Table 7: Evaluation data statistics of both NL2Code and Code2Code search.
Model
Python
Java
JS
TS
C#
C
Ruby
РHP
GO
Avg
(07)Gen2.5(7BB)
16.45
10.18
7
8.46
4.24
7.99
17.28
15.56
9.38
10.73
Maarcoder(15.5B)
7.09
3,91
3.19
4.36
1.72
2.35
6.82
6.06
3.26
4.31
CodeT5+(16B) Encoder
18.24
9.85
5.84
6,86
4.19
8.16
16.47
13.88
8.01
10.17
CodeBERT
14.40
7.62
5.47
6.05
3.66
5.53
13.55
10.28
6.27
8.09
GraphCodeBERT
19.23
10.78
7.38
8.65
5.54
8.48
19.69
15.67
9.65
11.68
StarEncoder
19.17
11.65
9.0
10.52
5.69
9.72
21.57
16.98
10.81
12.79
UnixCoder
30.77
16.45
21.32
21.95
6.19
15.62
32.33
31.93
13.94
21.17
OpenAl-CPT-Coad--001
21.92
8.90
4.90
5.70
3.15
11.58
26.25
16.60
9.40
12.04
OpenAT-Ada-102
35.91
25.13
19.01
21.86
10.17
29.15
40.85
40.47
23.43
27.33
OpenAl-Text-3--S5
25.18
12.61
8.00
9.44
5.46
15.86
30.70
23.33
11.2
15.75
OpenAl-Text-3-Laan 3-age
40.57
25.33
20.09
22.00
11.84
31.9
42.54
41,84
21.75
28.65
36.31
23.97
26.60
29.90
11.84
22.84
29,06
34.64
19.56
26.08
CODESAGE-BASE
47.52
22.84
28.70
31.95
13.37
30.99
44.86
51.13
25.15
32.95
CODESAGE-LARGE
46.70
33.13
37.16
41.18
16.81
32.89
54.12
52.13
32.48
38.51
Table 8: MAP score (%) of the zero-shot code search task. The language names mentioned in the top row indicate the languages queries and candidates are written in.
solution from each cluster and randomly pick 2 to 10 solutions per problem. We summarize in- language (query and candidates are in the same language) code2code search results in Table
NL2Code search refers to the task of using natural language as the query to retrieve the relevant code. We consider three benchmarks in this paper. CoSQA where the NL queries (in NL) are from the search logs of the Bing search engine and the candidate functions Net (in Python). Total queries 500 and number of candidates 6,268. CSN is constructed from the CodeSearchNet dataset of six programming languages, including Python, Java, JavaScript, PHP, Go, and Ruby. AdyTest which normalizes Python functions (from CSN) and variable names to test the understanding and generalization capabilities of models (an example shown in Figure 9
from CodeSearch-
Additional Baselines are considered in Tables and 9 We constantly find decoder-only models yield poor performance on semantic search. Finetuning prompt engineering may help improve the performance of decoder-only models, which we leave as future work.
We also compare against OpenAI-CPT-Code-ee code-001 model, along with OpenAl--ext--3Large and OpenAl-Text- S-mall. 3-aalll For the OpenAI- Text-3 models, we utilized their default embedding sizes of 3072 and 1536 for the large and small models, respectively. When submitting our paper in September 2023, based on OpenAl's recom- mendation, we evaluated against the ApenAl-Ada-001 model. Our findings, presented Tables & g indeed demonstrate that OpenAl-Ada-O0 surpasses OpenAI-CPT-Code- and NL2Codde search tasks. Furthermore, while OpenAI-Text-33 L-rge shows superior performance OpenAl-Ada-005 OpenAl-Text- 3 S-aall falls short. Nevertheless, model, CodeSage, sistently outperforms -een-l--extt--Laar in Code2Code search tasks and shows slightly inferior results in NI NI.2Code search.
specifically choosing the code-scarch-babbet babbage-
in both 2ode2Code
18

---

Published as a conference paper at ICLR 2024
CoSQA
AdvTest
CSN
Model
Python
Python
Python
Java
JS
PhP
Go
Ruby
CodeGin2255 (7B)
0.02
0.01
0.06
0.02
0.05
0.18
6.03
2.04
Starcoder (15.5B)
0.02
0.06
0.03
0.01
0.05
0.59
0.06
0.05
CodcT5+ (16B) Encoder
22.96
20.36
19.93
14.05
12.26
26.08
20.37
13.05
CodeBERT
0.24
0.06
0.05
0.03
0.04
0.02
0.14
0.34
GraphCodeBERT
16,20
5.58
10.37
8.59
7.29
8.07
12.47
20.79
StarEncoder
10.78
0.93
2.81
2.51
1.87
0.74
2.65
5.54
UnixCoder
42.11
27.32
42.17
43.92
40.46
35.21
61.39
55.22
OpenAL-CPT-Code-000
52.20
36.03
63.13
67.85
62.30
57.47
85.22
69.28
OpenAl-Tict-AAd--002
44.23
38.08
68.02
71.49
67.50
60.62
85.63
74.20
OpenAl--rett-3-Sman
52.48
34.10
62.62
65.87
60.28
54.85
81.96
67.57
OpenAL-Text-3-Lagge
55.21
46.83
70.81
72.89
68.12
59.58
87.6
75,22
CODESAGE-SMALL
49.92
41.28
64.38
63.19
60.01
54.71
77.66
63.20
CoDESAGE-BASE
48.50
49.08
67.99
68.02
66.95
58.15
83.21
68.00
CODESAGE-LARGE 47.53
52.67
70.77
70.21
69.50
61.33
83.71
71.92
Table 9: MRR score (%) of NL2Code search in zero-shot setting.
NL query: Try loading the given cache file.
# Original Python function
def from ffll from_file (cls, file, sargs, s*kwargs):
try:
cache shelve open file)
5
return cls (file, cache, *args. **kwargs)
except OSError as e: logger debug ("Loading {0) failed". (filee(file)) raise
# AdvTest Python function
2 def Fune ( arg .0, arg-l. earg-2, **arg-3):
try:
arg-4 shelve.. open arg.l)
return (arg-1, arg_4, #arg_2, **arg_3)
except OSError as e: logger debug (" ("Loading {0} failed".formatt format 1))
raise
Figure 9: An example of natural language query and the associated ground truth function from the AdvTest dataset. The function names and variables in the original function the top) are replaced by special tokens (at the bottom) to obfuscate the code.
Target Class
Train #
Valid #
Test #
Target Class
Train
Valid #
Test
No error
1,20,503
13,049
13,745
Importbrror
259
37
22
ZcroDivisionError
25,087
3,087
2,828
TabError
74
OSError
21540
2,427
2,422
re.error
62
6
11
LnbluodLocalErroo
21,414
2.641
2,603
AtrrbuteError
47
decimal
10,026
509
1,674
Stoplteration
24
5
ValucError
8,585
991
833
OverflowErroor
19
AssertionError
7,816
1.072
691
Timeout
18
8
FileeoodounndError
7,676
727
797
IndexError
10
O
12
IndentationError
7,645
841
ModuueNooFFonndErro
Error
7,505
965
733
RecursionError
0
NrmeError
1.876
186
110
EOFError
0
numpy.AxisError
437
47
125
SyntaxError
0
Earrnorain
362
39
22
RuntimeError
2
Table 10: Distribution of target classes in the Python Runtime Errors dataset.
B.3 CODE CLASSIFICATION
We present the label distribution for the RunTime error prediction dataset in Table 10 We present hyper-parane parameersss that we used while finetuning models for code classification tasks Table
19

---

Published as a conference paper at ICLR 2024
Ft. linear classification head only
Ft. full model end-to-end
Defect Complexity AdamW
Runtime
Defect
Complexity
Runtime
Optimizer
Adam
5e-5 (baselines)
1e-5 (CODES
Learning rate (LR)
le-3
5e (E-GCODSSAAG-LAGGE)
LR schedule
Linear
Linear
Batch size
32
32
# Epoch
10
10
2
5
5
2
Table 11: Hyperparameters for fine-tuning baseline models and CODESAGE on code classification tasks. Across all models, we used mean pooling to form sequence representations from contextual- ized token representations.
Finetuning models end-to-end classification tasks In the main body of this paper, pre- sented the evaluation results (in Table 2 of finetuning a linear classification head on top of the frozen code representation learning encoders. Furthermore, we finetune the code encoder models end-to-end on the classification tasks and present the results in Table 12 It's evident from these results that CODESAGE outperforms the bascline models.
Classification Complexity
Model
Defect
RunTime
CodeBERT
64.370.37
85.81 2.53
42.082,49
GraphCodeBERT
65.361.00
87.982.45
44.290.97
StarEncoder
65. 65.20o.11
92.87
38.063.86
CodeT5+ Embedding
64. 64.72p.65
90.631.47
38.362.54
UnixCoder
65. 35.741.00
93.75%.67
47.142.71
CODESAGE-SMALL
66. 14.67
94.740.29
44.461.50
CODES BBSE
66.52.. 48
95,90%.43
46.402.90
CODESAGE-LARGE
66.38
96.20
49.25
Table 12: F1 (macro) score of the code classification tasks in the full finetuning setup. We finetuned using three seeds and reported the mean and standard deviation (in subscript).
20


Reviews:
[{"review_id": "p6Hc21Hw9s", "reviewer": "Reviewer_zUhF", "summary": "This paper proposes CodeSage, a new series of models for learning representations of source code, ranging from 130M to 1.3B parameters. CodeSage is pretrained on the Stack dataset, with masked language modeling, identifier deobfuscation, and contrastive learning objectives, and all three objectives are carefully tailored specifically for code representation learning. Experiments are conducted on code search and classification tasks, against strong baselines. Results show that CodeSage not only outperforms all previous open-source code embedding models of similar sizes (i.e., ~125M), scaling it up to 1.3B further improves the performance and it also outperforms the `text-embedding-ada-002` embedding model from OpenAI.", "questions": "Q1: Do you plan to release the model and/or the code?  \nQ2: About not using random token replacement, this makes a lot of sense to me. But we would probably have the same problem for natural language, if we consider multiple languages, right? Is the same practice common in training multi-lingual text embeddings?  \nQ3: What is the pretraining hardware and compute?   \nQ4: Can you comment on the possibility to train even larger code embedding models using this method? It seems that most of the benchmarks are still improving by going from CodeSage-base to CodeSage-large in Table 1 and Table 2.   \nQ5: For Table 1, are the baseline models also trained on the 9 programming languages being evaluated here?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "S1: All the pretraining objectives are specifically designed for code, the discussions on why such objectives need to be changed (e.g., no random token replacement) for code are very useful for future research in this domain;\nS2: To the best of my knowledge, CodeSage 1.3B is the largest (and best performing) code embedding model to date. Scaling up encoder-based models, especially for code representation learning, is an underexplored area;  \nS3: The experiments are very comprehensive. CodeSage is compared with strong baselines and it was able to beat all of them on the same model size, it also shows that the performance is significantly improved when scaled up to 1.3B. The ablations studies are rather thorough in showcasing the design choices;  \nS4: The writing of this paper is also great, making it very easy to follow, and conclusions from the figures and tables are very clear. For example, I found Figure 2(a) very helpful in showing why random token replacement makes less sense for code.", "weaknesses": "I do not have any significant concerns about this work. But I do have a couple of questions, which are listed in the \"Questions\" section. \n\nOne suggestion: it would be great if we could have an ablation on how each of the training objectives affects the performance (e.g., remove one at a time and see how the performance changes). But I also understand that such ablation is quite expensive.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "d11MX4keDl", "reviewer": "Reviewer_iHpZ", "summary": "The authors propose pre-training of encoder style for the domain of programming languages at a larger scale (pre-training data and model parameters) than most encoder style models. These models are then tested on tasks in the code understanding domain like text to code search (or semantic code search) and classification tasks on code. The authors also study the impact of different pre-training choices on the model's performance through their ablations to back their methodology.", "questions": "- Why are decoder only models not studied as baselines for the code understanding tasks?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Scaling the pre-training stage of an encoder only model is a very interesting idea, and hasn't received much attention compared to decoder only models that support generation tasks.\n\n- Exhaustive ablation studies that justify the need for different objectives proposed in pre-training.\n\n- Impressive performance in zero-shot settings on search tasks compared to previous encoder style models", "weaknesses": "- Missing comparison with decoder only models: Proposing to scale encoder style models is interesting, but the current results are not convincing as to why this is needed when large scale decoder style pre-trained models have been proposed. While pre-trained decoder only models like CodeGen2.5, CodeLlama, Llama are not specifically proposed for code search or classification, they could be adapted for all the downstream tasks studied by the authors in Sec 4.1. For instance, the last token's encoding from a decoder can be used for search related tasks. Fine-tuning for classification can be performed by framing the classification as a seq2seq problem or by training a classification head on top of the last token's representation. Zero-shot or few shot prompting performance of these decoder style models can be reported for comparison.\n\n- Common Code Understanding benchmarks ignored: The choice of benchmarks is consistent with contemporary literature for the text to code search task, but not for other tasks in the classification settings. This makes it hard for a reader to assess the quality of the pre-trained CodeSage model and compare it to other pre-training strategies from previous works. There are many code understanding tasks (eg. ones from the CodeXGlue benchmark) which have been used in previous works (CodeBERT, GraphCodeBERT, UniXCoder, CodeT5) that are not a part of the evaluation suite here. See Table 1 of https://arxiv.org/abs/2203.03850 - Clone Detection datasets.\n\n- Missing fine-tuned results on CSN, AdvTest and CoSQA: While the benchmark is consistent with previous work for text to code search, the setting chosen is zero shot whereas many previous works show fine-tuned results on the sizeable training datasets provided by these benchmarks (CodeSearchNet, CoSQA, AdvTest). Given that the size of these models is reasonable (130M - 1.3B params) compared to the models studied today (>6-10B params), fine-tuning is an important aspect to be studied with these models.\n\nFor instance, MRR scores in Table 2 on the CodeSearchNet benchmark are reported only in the zero-shot setting.  Fine-tuning improves many of these baselines substantially, and beats CodeSage's zero shot performance by significant margin. Post fine-tuning CodeT5+ is at 77 MRR overall, while CodeT5 is at 71, UniXCoder is at 74.4. See Table 6 of https://arxiv.org/abs/2305.07922. Same is the case with CoSQA and AdvTest results. Also, large scale pre-trained decoder only models CodeGen & Llama2 are missing when evaluating in the zero shot setting. While these are not encoder only models, their embeddings can be leveraged for search related tasks.\n\nAlso, it's strange why text-embedding-ada-002 is chosen as a baseline here. OpenAI's cpt-code (https://openai.com/blog/introducing-text-and-code-embeddings) would have been a stronger baseline to compare with instead of text-embedding-ada-002 model for code related search tasks.\n\nWriting:\nAbsract: unclear what downstream task is, is it gen?\nIntro: No discussion of next token pred objective from decoder only or encoder-decoder models\n\nTypos:\n\nWriting:\nAbsract: unclear what downstream task is, is it gen?\nstructure --> structural aspect\npretraining schemes decide --> dictate/influence/impact\n\nIntro:\npara 1: Don't agree that only encoder models are embedding models. Even Llama, CodeLlama, CodeGen can serve as embedding models.\npara 2: make it --> makes it\npara 3: representation's discriminative power\n\nCodeT5 appears twice in the bib", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XBD9elRVgp", "reviewer": "Reviewer_RZCJ", "summary": "The paper considers the problem of learning better representations of code. The authors propose a model called CodeSage which is based on the Transformer architecture trained on a large corpus of code (The Stack) with various techniques such as masked language modeling, deobfuscation pre-training, and bimodal contrastive learning with hard negative and hard positive examples. \n\nThe authors showed that the learned representations can be used for various downstream tasks such as zero-shot code-to-code, natural language-to-code search and code classification. They compared the performance of the proposed model with other prior models such as CodeBERT, GraphCodeBERT, StarEncoder, CodeT5, UniXcoder, and OpenAI's text-embedding-ada-002. The results show that the proposed model outperforms all the other models on nearly all the tasks except for the code classification on code defect detection benchmark.", "questions": "How would you explain why the performance of smaller models (e.g., CodeSage-Small and CodeSage-Base) is sometimes better than larger model? Does this mean that the proposed idea is not scalable to larger models or more effective for smaller models?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "- The paper is well written and the proposed model is well motivated. \n- The experiments are thorough and the results are convincing. \n- The paper is a good contribution to the field of code representation learning.", "weaknesses": "- The paper is not very clear about the differences between the proposed model and the prior models. \n- For instance, it seems that the proposed model is a combination of the prior models (CodeBERT and DOBF) with some additional techniques, such as bimodal contrastive learning. It would be good to clarify the differences between the proposed model and the prior models.\n- page 3, line 27: remove the last comma in the following square brackets: [x_1, x_2, \\ldots, x_N ,]", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hAyTXmty46", "reviewer": "Reviewer_n8qP", "summary": "This paper introduces a novel two-step pretraining methodology for encoder-only code language models (LMs). The approach starts with masked language modeling (MLM) and follows up with contrastive learning to foster robust code representations that can be effectively employed in downstream tasks such as retrieval or classification. The resulting model consistently outshines existing baselines across various downstream tasks, including notable ones from the widely recognized CodeXGLUE benchmark. Furthermore, comprehensive ablation studies are conducted to dissect and understand the individual contributions of each component within the proposed framework.", "questions": "1. Is there a reason why some tasks from the CodeXGLUE benchmark, like POJ104, were excluded from the evaluation?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Code representation learning is a significant and often overlooked aspect of this domain.\n\n2. The evaluation methodology used in this study is solid, showing the proposed method's efficacy through extensive experiments on various downstream tasks from the well-established CodeXGLUE benchmark. Notably, the method exhibits remarkable performance, substantially outperforming established baselines such as CodeT5 and StarEncoder.\n\n3. The inclusion of detailed ablation studies is a strong point of the paper, shedding light on the impact and effectiveness of each component within the proposed method. The insights gleaned from these analyses, especially concerning the methodologies used for contrastive learning, are indeed meaningful.", "weaknesses": "1. The naming convention for the model configurations is somewhat confusing. For instance, \"CodeSage-Small\" (6 layers, 1024 dim, and 130M params) is roughly equivalent to BERT-Base (12 layers, 768 dim, 110M params) in terms of its parameters, while \"CodeSage-**Base**\" is on par with BERT-**Large** (24 layers, 1024 dim, 330M params). This inconsistency could potentially lead to confusion when making comparisons.\n\n2. The LaTeX typesetting requires significant improvement, as there are inconsistencies in font styles and sizes throughout the paper, notably in Section 4, Figures 2(a) and 6, and Appendix A.5. Additionally, the small text in Figure 4(b) is difficult to read without zooming in, and there are redundant references for CodeBERT and CodeXGLUE.\n\n3. The paper's title, \"Code Representation Learning at Scale,\" is rather broad and lacks specificity.\n\n4. The evaluation could be more comprehensive by including additional tasks, such as POJ104 from the CodeXGLUE benchmark, which is relevant but absent from the current evaluation.\n\n5. The baselines utilized for comparison are generally smaller than the proposed model. It would be beneficial to include larger baselines, such as CodeT5+ (16B) and CodeLLaMa (34B), for a more balanced comparison.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)