CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking

We thank the reviewer for their valuable feedback, and we address below their concerns and comments:

Weakness 1: This paper's contribution is mainly incremental. It improves data quality through filtering and negative sample mining methods on the Stack V2 dataset, which has been extensively explored in previous research work [1,2]. This paper only makes minor improvements based on these papers.

Response:
We would like to clarify that our contributions extend significantly beyond incremental improvements:

1. Release of largest public high-quality contrastive code data: We introduce CoRNStack, which is ten times larger than CodeSearchNet—the largest existing public contrastive code training dataset. Also, we are the first to apply consistency filtering to enhance data quality for contrastive fine-tuning of code embedding models. This results in CoRNStack having significantly higher query-positive correctness compared to existing contrastive code datasets like CodeSearchNet and CosQA (see Table 2 in the paper). To further highlight the impact of CoRNStack’s scale and quality, we finetuned Arctic-Embed, a text embedding model, for one epoch separately on CoRNStack and CodeSearchNet. To specifically show the benefit of CoRNStack’s quality, we also report results for fine-tuning on 2 million randomly sampled datapoints from CoRNStack, the same amount of data as CodeSearchNet. The results, shown in the table below, clearly illustrate the improvement in code retrieval performance from finetuning on CoRNStack with both a 2M subset and the full 21M examples. We also added this comparison to Section A.5 of the Appendix.

Moreover, unlike state-of-the-art open-source code embedding models, such as CodeSage and Jina-Code-v2 , which only provide model weights without releasing their training data, CoRNStack is publicly available at the following anonymized link. This establishes CoRNStack as a standardized high-quality dataset for enabling future research in code retrieval and reranking.

2. First to finetune LLMs as listwise code rerankers: While listwise reranking has gained popularity in text retrieval [1], its application in code ranking has been largely unexplored, due to the lack of high-quality contrastive bimodal (text, code) data for fine-tuning. In section 3.2 of our paper, we have demonstrated how CoRNStack, with the curated set of (text, code) pairs that involve positives and negatives, can be leveraged to finetune instruction-tuned code generation models as effective listwise code rerankers. Tables 6 and 8 in the paper show significant performance improvement from using our code reranker on top of results from our code retriever.

3. Demonstrate the benefit of improved code retrieval & reranking for real-world software development: Our work specifically addresses the gap between existing code ranking models and their practical application in complex software development tasks like repository-level function localization. Prior work, such as CodeRAG-Bench [2], hasn’t seen much success from leveraging code embedding models on complex tasks such as bug localization on SWE-Bench-Lite. In this paper, we show that this gap is largely due to a lack of high-quality contrastive data. We demonstrate that combining our code retriever & reranker setup, which were finetuned on CoRNStack, shows significant improvements in repository-level bug localization over Agentless [3], a top-performing automated software development framework.

[1] RankZephyr: Effective and Robust Zero-Shot Listwise Reranking is a Breeze!; Pradeep et al.

[2] “Coderag-bench: Can retrieval augment code generation?”; Wang et al., arXiv 2024

[3] “Agentless: Demystifying llm-based software engineering agents”; Xia et al., arXiv 2024

Weakness 2: The paper lacks a more detailed comparison and analysis between the final obtained and original datasets beyond just quantity (such as case studies and other statistical conclusions).

Response:

To assess the accuracy of the text-code pairings in CoRNStack, we conducted an automatic evaluation as described in lines 174–180 and shown in Table 2 of the original paper. In short, Qwen2.5-Coder-7B-Instruct, a code generation model, was prompted to judge whether the positive code snippet fully answers the corresponding query. This was performed on 10,000 randomly sampled pairs from CoRNStack, the original Stack v2, and other contrastive code training datasets like CosQA and CodeSearchNet (CSN), across three random seeds. The results indicate that CoRNStack has a significantly higher mean percentage of correct pairings. We further provide below the language-wise correctness for the mean numbers reported in the paper and add this table to section A.3 of the revised version.

Weakness 3: In the writing of the paper, the construction of the dataset and the harmful sample mining methods (especially the curriculum learning technique, which is not model-independent) are coupled. In my opinion, these are two parts: one part is the construction of the dataset, and the other part is the harmful sample mining method designed by the author. I believe these two parts should be more clearly distinguished in the writing.

Response:

We sincerely thank the reviewer for this valuable suggestion. We agree that the curriculum learning component, which is related to the model training process, can be separated from the dataset construction description. We have updated the paper accordingly in the revised version. Specifically, we have moved the curriculum learning details to Section 3.1 (changes highlighted in blue). We hope that this revision adequately addresses the reviewer’s concern.

We appreciate the reviewer's positive feedback on our paper's clarity, motivation, and significant results across various scenarios, and we address here the concerns in the experimental section that they have suggested to further strengthen our work:

Weakness 1: The motivation for selecting Arctic-Embed-M as the initial encoder for training, instead of other models, is not clearly explained.

Response:

Our motivation for using Arctic-Embed-M as the initial encoder is based on the hypothesis that code retrievers can benefit, especially for text-code retrieval tasks, from pretraining on supervised text ranking data, which is typically abundant. To validate this hypothesis, we performed the following experiment:

We selected three ~130M parameter text and code encoders, specifically Arctic-Embed-M and Nomic-Embed as the text encoders due to their strong performance on text retrieval benchmarks like MTEB, and CodeSage-Small as the code encoder due to its overall performance on code retrieval benchmarks. We evaluated these models on code retrieval tasks before and after finetuning on CoRNStack for one epoch. The results are presented in the table below and also in Section A.4 of the revised paper.

Firstly, we see that finetuning on CoRNStack significantly improves code retrieval performance for all models. Moreover, we observe that a stronger text embedding model (Artic-Embed-M in this case) leads to better code retrieval performance after finetuning on CoRNStack, even outperforming the code-pretrained CodeSage-Small in most cases. These results highlight that CoRNStack, being large-scale and high-quality, can be leveraged to finetune text encoders into performant code retrievers. Therefore, we selected Arctic-Embed-M as it provides a strong foundation from supervised text ranking pretraining, which, when combined with fine-tuning on CoRNStack, leads to superior code retrieval performance.

Weakness 2: The setup in section 4.2.1 lacks clarity. For example, the statement “a sampling strategy with varying window sizes (between 3 to 10) and random shuffling leads to 250k training instances” could be elaborated more for better clarity.

Response:

To clarify, we follow the data augmentation methodology established in prior work [1], incorporating techniques to create a more diverse and challenging training setup in order to obtain a more robust trained reranker. Specifically, the augmentation process consists of two key components:

1. Variable Window Sizes: For each training instance, a random subset of candidate code snippets (ranging from 3 to 10 candidates) is sampled from the original ranked list to pass as input to the teacher model. This introduces variability in the input size and diversifies complexity of the reranking task, ensuring that the reranker model encounters a broader range of scenarios during training.
2. Random Shuffling: To enhance the model's generalization ability across different document orders—beyond the default order provided by the retriever— random permutations are applied to the sampled windows. This technique has demonstrated effectiveness in traditional text reranking tasks [1], and we extend it to code reranking.

We add these details to Section A.8 in the Appendix of the revised paper.

[1] RankZephyr: Effective and Robust Zero-Shot Listwise Reranking is a Breeze!; Pradeep et al.

Hello Reviewer bBWX,

We would like to thank you for the thoughtful and constructive feedback on our submission.

As the rebuttal period comes to a close, we are following up to inquire if our responses have addressed the concerns and would appreciate you considering raising your final rating or providing further feedback/concerns. We are looking forward to your response!

Thanks, ICLR 2025 Conference Paper 12599 Authors

We thank the reviewer for the valuable feedback and comments. We address the reviewer’s concerns and questions below:

Weakness 1: The process of handling negative samples in Figure 1 needs to be explained in more detail and more directly.

Response: Our hard negative mining strategy is divided into two stages:

1. An offline stage that is coupled with consistency filtering and computes a corpus-level similarity score matrix S between all text and code instances, which is used to filter false negatives.
2. An online stage that uses a softmax-based sampling strategy with a curriculum based on the pre-computed matrix S to progressively select diverse, challenging negatives during the contrastive finetuning.

The retriever uses the negatives from the online stage to incorporate a curriculum learning strategy while training. The reranker, on the other hand, is trained using top-M negatives based on S from the offline stage, since it is relies on the ranking ordering supervision for these negatives provided by the teacher model, which is not feasible to obtain in an online fashion. We have made the distinction between the online and offline stages more clear in the revised version of the paper. The offline and online stages are described in more detail in lines 203-207 of Section 2.3, 224-235 of Section 3.1, and 267-269 of Section 3.2.

Question 1: The quality of the new dataset needs to be explained, especially how accurate the pairings are between text queries and code.

Response:

To assess the accuracy of the text-code pairings in CoRNStack, we conducted an automatic evaluation as described in lines 174–180 and shown in Table 2 of the original paper. In short, Qwen2.5-Coder-7B-Instruct, a code generation model, was prompted to judge whether the positive code snippet fully answers the corresponding query. This was performed on 10,000 randomly sampled pairs across three random seeds from CoRNStack, the original Stack v2, and other contrastive code training datasets like CosQAand CodeSearchNet (CSN). The results indicate that CoRNStack has a significantly higher mean percentage of correct pairings. We further provide below the language-wise correctness for the mean numbers reported in the paper and add this table to section A.3 of the revised version.

Question 2: It should also be clarified whether the code data in the dataset has been checked manually or tested with examples, as this is important for the quality of the code retrieval tasks.

Response:

We acknowledge that the reviewer raises a very interesting point. With CoRNStack being a large-scale training dataset containing millions of data points, any validation must be automated, since it is not feasible to manually check the quality of each code snippet.

To ensure the syntactic correctness of the code data in CoRNStack, we followed the approach in [1] by using the Tree-sitter parsing toolkit to filter out any functions that cannot be parsed into a syntax tree during data extraction from The Stack v2.

For evaluating the functional correctness, we employed a heuristic validation method described in our response to Question 1. Specifically, we utilized a code language model to judge whether a given positive code snippet fully answers its corresponding query. Our results, shown in Table 2 of the paper, demonstrate that CoRNStack has a significantly higher mean percentage correctness than existing code datasets.

Another way of validating the functional correctness of the code data would be to generate test cases using a combination of LLMs and/or human annotators. However, considering the scale of the dataset, this approach would be computationally intensive to generate diverse test cases that cover all edge cases and would require substantial resources to set up the necessary environments for execution-based validation. Nevertheless, we agree that this is a valuable area for future exploration and include a discussion of this in lines 500-505 of the revised paper.

[1] “GraphCodeBert: Pretraining Code Representations with Data Flow”; Guo et al, ICLR 2021

Hello Reviewer MeoD,

We would like to thank you for the thoughtful and constructive feedback on our submission.

As the rebuttal period comes to a close, we are following up to inquire if our responses have addressed the concerns and would appreciate you considering raising your final rating or providing further feedback/concerns. We are looking forward to your response!

Thanks, ICLR 2025 Conference Paper 12599 Authors

We would like to thank the reviewer for appreciating the paper with a positive score. Here, we address the reviewer’s concerns and questions:

Weakness 1: There is no analysis of the reason for the improved performance of code rerankers. I guess this is also due to the hard negative samples?

Response: We attribute the improvement in performance to the finetuning of LLMs for the listwise code reranking task. As Table 6 in the paper shows, the instruction-tuned code generation model, when directly used for listwise code reranking, has considerably lower performance. However, it improves after being finetuned for this task, using the data made possible by getting a relevance ordering on the negatives present within CoRNStack.

Weakness 2: As stated by the authors, one of the motivations for constructing this dataset is that existing datasets contain many false negative samples. In [1], the authors also discussed this problem and proposed methods to address the problem during model fine-tuning, which I think should also be discussed in the paper.

Response: We sincerely thank the reviewer for highlighting this relevant work. The paper mentioned introduces the Soft-InfoNCE loss to reduce the impact of false negatives with contrastive training when using InfoNCE loss, by weighting negative samples based on their similarity to the query. We view this work as complementary to ours, as the mined hard negatives from CoRNStack could be seamlessly integrated into the Soft-InfoNCE loss, in a similar fashion to that described in lines 236–248 of our paper, to further improve effectiveness. We appreciate the reviewer bringing up this work and will incorporate a discussion into the final version.

Weakness 3: Although the authors discussed the problems of existing datasets in the introduction and some preliminaries of code retrieval in the paper, I still think it is better to write a related work section to make them clearer. Since code retrieval is not an extremely hot topic, many readers may not know many details of previous works in this area.

Response: We thank the reviewer for this suggestion. We are actively working on drafting a related work section and will add it into the revised paper by the camera-ready deadline at the latest.

Question 1: Why do you use pre-trained text embedding models for dual consistency filtering? What about using pre-trained code retrieval models pre-trained on previous code retrieval datasets?

Response: We use the pre-trained code embedding model Jina-Code-v2 as the proxy embedding model N (line 166 in the paper) for dual consistency filtering. We selected this embedding model for filtering because of its performance on code retrieval benchmarks like CodeSearchNet, AdvTest, and CoIR. We have clarified this detail in footnote 4 of the revised paper.

Question 2: Fine-tuning models with contrastive learning is a common and traditional method for code retrieval tasks. Recently, there has been a new paradigm named Generation-Augmented Retrieval (GAR) framework proposed for code retrieval tasks as well [1, 2]. I think this framework makes the fine-tuning of code retrieval models somewhat unnecessary. What's your opinion on the necessity of fine-tuning code retrieval models?

Response: We thank the reviewer for highlighting the Generation-Augmented Retrieval (GAR) framework. We acknowledge that GAR introduces an innovative approach by using a generation model to expand queries with exemplar code snippets, addressing issues of short or ambiguous queries. However, the retrieval component in GAR still depends on a model trained with contrastive learning to match the expanded queries with relevant code snippets effectively. Therefore, we believe that fine-tuning code retrieval models remains essential to enhance their ability to understand code semantics and improve retrieval accuracy. Additionally, GAR may introduce additional latency due to the generation step, whereas using finetuned retrieval models directly provides faster results. In summary, finetuning code retrieval models is still necessary for optimal performance, even when employing GAR techniques.