[论文解读] CoCoSoDa: Effective Contrastive Learning for Code Search

Ensheng Shi, Wenchao Gub+6 more|arXiv (Cornell University)|Apr 7, 2022

Topic Modeling被引用 9

一句话总结

CoCoSoDa 提出了一种用于代码搜索的对比学习框架，结合了软数据增强和基于动量的负样本队列，以提升代码与查询的表征学习效果。该方法在六种编程语言上实现了最先进性能，平均 MRR 相较于 CodeBERT、GraphCodeBERT 和 UniXcoder 分别提升了 13.3%、10.5% 和 5.9%。

ABSTRACT

Code search aims to retrieve semantically relevant code snippets for a given natural language query. Recently, many approaches employing contrastive learning have shown promising results on code representation learning and greatly improved the performance of code search. However, there is still a lot of room for improvement in using contrastive learning for code search. In this paper, we propose CoCoSoDa to effectively utilize contrastive learning for code search via two key factors in contrastive learning: data augmentation and negative samples. Specifically, soft data augmentation is to dynamically masking or replacing some tokens with their types for input sequences to generate positive samples. Momentum mechanism is used to generate large and consistent representations of negative samples in a mini-batch through maintaining a queue and a momentum encoder. In addition, multimodal contrastive learning is used to pull together representations of code-query pairs and push apart the unpaired code snippets and queries. We conduct extensive experiments to evaluate the effectiveness of our approach on a large-scale dataset with six programming languages. Experimental results show that: (1) CoCoSoDa outperforms 14 baselines and especially exceeds CodeBERT, GraphCodeBERT, and UniXcoder by 13.3%, 10.5%, and 5.9% on average MRR scores, respectively. (2) The ablation studies show the effectiveness of each component of our approach. (3) We adapt our techniques to several different pre-trained models such as RoBERTa, CodeBERT, and GraphCodeBERT and observe a significant boost in their performance in code search. (4) Our model performs robustly under different hyper-parameters. Furthermore, we perform qualitative and quantitative analyses to explore reasons behind the good performance of our model.

研究动机与目标

• 通过增强的数据增强和负样本挖掘，有效利用对比学习以提升代码搜索性能。
• 解决现有对比学习方法在代码搜索中的局限性，如静态增强和负样本多样性不足。
• 使预训练模型能够在无需大量微调的情况下学习到更鲁棒且对齐的代码-查询表征。
• 验证所提方法在多种预训练模型（包括 RoBERTa、CodeBERT、GraphCodeBERT 和 UniXcoder）上的泛化能力。
• 在不同超参数设置和零样本设置下，验证框架的鲁棒性与有效性。

提出的方法

• 提出四种软数据增强（SoDa）技术，通过动态掩码或替换代码标记为其类型，生成语义相似的正样本。
• 采用动量编码器和动量更新的队列，为每个小批量维持一个大规模且一致的负样本集合，提升对比学习的稳定性。
• 应用多模态对比学习，对齐成对代码-查询实例的表征，同时拉远未配对的代码片段与查询之间的表征。
• 将对比学习目标集成到多种预训练模型（包括 RoBERTa、CodeBERT、GraphCodeBERT 和 UniXcoder）的训练中，以提升其代码搜索性能。
• 采用双对比学习策略，结合模态内（代码-代码、查询-查询）和模态间（代码-查询）的对比目标，增强表征对齐。
• 通过在不进行微调的情况下评估预训练模型，将框架适配至零样本评估，证明其仅通过预训练即可学习到鲁棒表征的能力。

实验结果

研究问题

• RQ1 通过类型替换或掩码代码标记的软数据增强是否能提升代码搜索中的表征学习？
• RQ2 使用基于动量的负样本队列是否能增强对比学习在代码搜索中的判别能力？
• RQ3 相较于单模态对比学习，多模态对比学习（代码与查询之间）在表征对齐与性能提升方面有多显著？
• RQ4 CoCoSoDa 在不同预训练模型（包括非代码专用模型如 RoBERTa）上的有效性如何？
• RQ5 CoCoSoDa 是否能在无需微调的情况下实现优异性能，表明其学习表征的鲁棒性与泛化能力？

主要发现

• CoCoSoDa 在 CodeSearchNet 基准上实现最先进性能，超越 18 个 SOTA 基线模型，六种编程语言上的平均 MRR 达到 0.738。
• CoCoSoDa 相较于 CodeBERT、GraphCodeBERT 和 UniXcoder，平均 MRR 分别提升 13.3%、10.5% 和 5.9%，展现出显著的性能增益。
• 在零样本评估中，CoCoSoDa 的性能超越微调后的 CodeBERT、CodeT5 和 GraphCodeBERT，平均 MRR 提升超过 50%，在 PHP 上最高提升达 70%。
• 消融实验证实，软数据增强与基于动量的负样本队列均为关键组件，各自对性能提升有显著贡献。
• 该方法在不同预训练模型（RoBERTa、CodeBERT、GraphCodeBERT、UniXcoder）上均持续提升性能，表明其与多种架构的正交性与泛化能力。
• 该模型在不同超参数设置下表现出鲁棒性，在各种配置中均保持强性能。