Paper: CODE REPRESENTATION LEARNING AT SCALE (ICLR-2024)

Reviewer comment (Presentation):
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.

Author response:
**Naming Convention**
It was a hard decision for us. On the one hand, we did want to follow the BERT convention, while we also noticed the recent trend to name the 1B model as the Large model on the [MTEB Leaderboard](https://huggingface.co/spaces/mteb/leaderboard).  However, we are seriously considering your question and trying to find a sweet spot between our current names and CodeSage-Base, CodeSage-Large, and CodeSage-XLarge.

Reviewer comment (Writing):
The paper's title, "Code Representation Learning at Scale," is rather broad and lacks specificity.

Author response:
**Latex Typesetting and Paper Title**
We will improve the formatting as you suggested. We also want to talk about our motivations for the current paper title.

1. Our work is among the first (in addition to StarCoder, to the best of our knowledge) to train the encoder-based embedding model at a large scale of pretraining data and increased model size. We limit our model size to 1B to balance performance and inference cost in many practical settings, e.g., building a retrieval system over billions of examples.
2. More can be different. As reported in Figure 3b of our paper, the dropout-based contrastive learning does not work well with large-scale pretraining data, though it works for UnixCoder and SimCSE, where less than 3M examples are used for training.
3. We indeed hope our work can motivate more researchers to keep pushing the frontiers of code embedding models by advancing the pretraining strategy on a large scale of data, as what has been doing for text embeddings ([MTEB](https://huggingface.co/spaces/mteb/leaderboard)).

We are looking forward to hearing your thoughts and would appreciate it if you have better suggestions.

Reviewer comment (Experiments):
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.

Author response:
**Evaluation on Clone Detection**
We did evaluate BigCloneBenchmark at the beginning of this project. However, we found the task very easy but inefficient to evaluate due to the large data size. Hence, we excluded it and focused on more challenging ones. Presumably, given that CodeSage shows strong performances over previous SOTA models on our chosen tasks across many languages, it will perform on par if not able to largely outperform the baselines. Instead, we dived deep and performed analysis to show the reasons behind the effectiveness of CodeSage (e.g., the effectiveness of MLM/DOBF and CL, performance with scaling), which can drive future research works.

But we are grateful for the suggestion, which motivates us to improve our paper's completeness and overall quality. We summarized our evaluation results on both POJ-104 and BigCloneBenchmark in Table 3 in our general response (2/2) above, where CodeSage-small performs slightly better than UniXCoder. 

We sincerely thank you for the constructive feedback. We have taken your comments into serious consideration. We provided the evaluation results in our general response (2/2) accordingly, which we will also use for the next update of our paper to address your questions comprehensively. Let us know if you have any other comments.

Reviewer comment (Experiments):
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.

Author response:
**The Baseline Models Are Smaller**
One of our goals is to study if larger models would be effective for code representation learning since such models are missing in the literature. Therefore, we urge not to consider that as a weakness of this work.  On the other hand, larger decoders are often not primarily trained to serve retrieval or code understanding works. Those models often yield poor performance in our preliminary evaluations, and we exclude them from comparison as we think it’s unfair to compare with them. However, we appreciate your question, which makes us realize that including these baselines would have made our paper more complete. We summarized the evaluation of CodeGen2.5 (7B), Starcoder (15.5B), and CodeT5+ (16B) in Tables (1a)&(1b) in our general response (2/2) above.

Reviewer comment (Novelty):
The paper is not very clear about the differences between the proposed model and the prior models.

Author response:
Thank you for the feedback. To address your main question, we briefly summarize the fundamental contributions of our work below, which differentiate CodeSage from the existing work. 

1. We are the first to train different-sized code embedding models using a much larger amount of publicly available data and study their effectiveness as an off-the-shelf encoder model in downstream tasks.
2. We innovate an effective two-stage pretraining strategy in the presence of large-scale data, which outperforms the previous SOTA model by significant margins on a wide range of tasks. In particular, 
     * we propose an effective MLM strategy for code, which is different from the existing literature (including CodeBERT and many others) that mainly follows 80-10-10 corruption convention proposed for text in BERT. We dive deep into why 80-10-10 is suboptimal for code through quantitive and qualitative experiments in Figure 2. 
    *  we identify the value of the DOBF objective that is not utilized in prior works for training code embedding models and cross the technical hurdles to make it work for training encoders (see Appendix A.5 for details). We further boost the performance by leveraging our proposed MLM with DOBF to complement each other (see Table 3 in our paper). 
    * we proposed an effective contrastive learning strategy with effective hard positive construction tricks that boost the performance over the baseline (see Figure 3a).  
    * we deep dive into the key ingredients of code representation learning through an exhaustive ablation study, which we believe will help future works advance the performance further.

Reviewer comment (Evaluation):
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?

Author response:
**Why Smaller Models Sometimes Outperform A Larger Model?**

Large models should generally outperform smaller ones on most tasks; we observe such a trend for CodeSage, especially in the zero-shot settings. However, there are no guarantees that larger models are universally better than smaller ones. We can also clearly see this on the [MTEB leaderboard](https://huggingface.co/spaces/mteb/leaderboard) for text; for example, bge-large-en-v1.5 outperforms bge-base-en-v1.5 (https://huggingface.co/BAAI/bge-large-en-v1.5) on average, but the opposite is also reported on some classification tasks and retrieval tasks, with similar findings for e5-large vs. e5-base (https://huggingface.co/intfloat/e5-large) and many other models.

Reviewer comment (Experiments):
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.

Author response:
**Some Classification Benchmarks Ignored**
Instead of going through all benchmarks in CodeXGLUE, we preferred to dive deep and perform analysis to show why CodeSage learns effective code representations. We initially evaluated clone detection, which we found is easy for different models but has high evaluation cost due to the large data size. We excluded it and focused on other challenging classification tasks.  However, we have seriously considered your comments and summarized evaluations on POJ-104 and BigCloneBenchmark (BCB) in Table 2 in our general response (2/2). We can see that CodeSage maintains its strong performance.

Reviewer comment (Experiments):
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. For 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.

Author response:
**Missing Fine-tuned Results on CSN, AdvTest, and CoSQA**
Our primary goal was to build an off-the-shelf embedding model that does not require finetuning in many practical use cases; semantic search is one such case. Strong zero-shot performance would allow wider usage of the embedding models, as annotations for code are rare and very costly to collect. Our evaluation strategy also aligns with the recent trend in the text domain, where strong models have been actively developed to advance the [MTEB Benchmarks](https://huggingface.co/spaces/mteb/leaderboard). We hope our work can motivate a similar trend for code. 

However, we are grateful for your question, which motivates us to better connect with the previous work. We summarized the finetuning results in Table 2 in our general response (2/2), where CodeSage-base (356M) consistently outperforms CodeT5+(770M).  We want to point out that both CodeT5+ 220M and 770M optimize an additional “matching loss” during finetuning (cross-encoder type finetuning of the decoder with pairs selected by the encoder). We hypothesize that if we finetune a copy of UnixCoder or CodeSage on such matching tasks and use them with the original encoders, we can also get a further performance boost for both models.

Reviewer comment (Experiments):
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.

Author response:
**Why Not Code-CPT-001**
We chose Text-Ada002 models per OpenAI’s suggestion on their website and claim that Ada002 is better or at least comparable to CPT-001 on both text and code tasks but at a much cheaper cost. We have validated their claim in tables (1a) and (1b) in our general response (2/2) above.

Reviewer comment (Theory):
Introduction para 1: Don't agree that only encoder models are embedding models. Even Llama, CodeLlama, CodeGen can serve as embedding models.

Author response:
**Missing Comparison with Decoder-only Models** 
As mentioned in our general response (1/2) #1, we tried to avoid unfair comparisons against the larger decoder-only models as they are not primarily trained for discriminative tasks and often yield suboptimal performance. As we can see in Tables (1a)&(1b) in our general response (2/2), CodeGen2.5 (7B), StarCoder (15.5B), and CodeT5+ (16B)-encoder all perform very suboptimal in the zero-shot setting, especially on NL-to-Code search tasks which often require properly designed training objective (e.g., contrastive learning) on the bimodal-data (text-code) to boost the performance.  However, we appreciate your question, which makes us realize we should include these results in our paper for completeness.

Reviewer comment (Experiments):
Question - Why are decoder only models not studied as baselines for the code understanding tasks?

Author response:
**Missing Comparison with Decoder-only Models** 
As mentioned in our general response (1/2) #1, we tried to avoid unfair comparisons against the larger decoder-only models as they are not primarily trained for discriminative tasks and often yield suboptimal performance. As we can see in Tables (1a)&(1b) in our general response (2/2), CodeGen2.5 (7B), StarCoder (15.5B), and CodeT5+ (16B)-encoder all perform very suboptimal in the zero-shot setting, especially on NL-to-Code search tasks which often require properly designed training objective (e.g., contrastive learning) on the bimodal-data (text-code) to boost the performance.  However, we appreciate your question, which makes us realize we should include these results in our paper for completeness.

Reviewer comment (Experiments):
One 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.

Author response:
**Removing One Objective At A Time** 
We partially cover it through different experiments.

1. Figure 4a shows that contrastive learning (stage-II pretraining) helps boost performance on top of stage-I pretraining with DOBF & MLM only. 
2. Figure 5 shows contrastive learning from scratch yields suboptimal performance and cannot benefit from large model sizes. Stage-I pretraining is essential to provide a good starting point for contrastive learning to be effective.
3. Table 3 compares DOBF vs. MLM and effective ways to combine them for the stage-I pretraining.
4. We initially also tried optimizing all three objectives from scratch. However, we observed suboptimal performance that can be explained by #2 above (which indicates that including contrastive learning from scratch is not effective). 

We would like to explore further if we can get computation support for the camera-ready version.

Reviewer comment (Reproducibility):
Do you plan to release the model and/or the code?

Author response:
**Q1:** We plan to release the model checkpoints and the evaluation code.

Reviewer comment (Experiments):
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?

Author response:
**Q2:** Your insights align with our hypothesis that random token replacement could hurt MLM training for multilingual text. However, this requires empirical validation; to the best of our knowledge, the existing multilingual MLM encoders are trained with the 80-10-10 practice. Another hypothesis is that the 80-10-10 corruption convention is proposed by the BERT paper, which pre-trained on a comparatively smaller dataset (BooksCorpus + Wikipedia En) over 40 epochs. In such a setting, corruption can help avoid overfitting and cause no significant performance drop. However, nowadays, scaling up the pretraining data allows us to reduce the training epochs, and the effectiveness of the 80-10-10 convention can change consequently.

Reviewer comment (Reproducibility):
What is the pretraining hardware and compute?

Author response:
**Q3:** We used 128 A100 GPUs for pretraining CodeSage-Small&Base, which takes 3-5 days to finish the two-stage pretraining. We used 256 A100 GPUs for pretraining CodeSage-Large, which takes roughly six days. We will add those details in Appendix A.2.

Reviewer comment (Experiments):
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.

Author response:
**Q4:** We can train even larger models. In this work, we limit the model size to 1B by considering the cost and latency for many practical settings, e.g., providing embeddings for building a retrieval system over billions of examples. However, if we can get enough support on resources, we are also interested in seeing how the performance can be further improved by scaling up both the data and model size, as we can first optimize for performance and then inference cost.

Reviewer comment (Experiments):
For Table 1, are the baseline models also trained on the 9 programming languages being evaluated here?

Author response:
**Q5:** CodeBERT and GraphCodeBERT are trained on 6 languages, UniXCoder (we used UnixCoder-Nine in this paper) and CodeT5+ embedding models are trained on 9 languages, Starencoder is trained on 86 languages. It is unknown how many programming languages are used to train text-embedding-ada-002.