Improving Code Search with Multi-Modal Momentum Contrastive Learning

Contrastive learning has recently been applied to enhancing the BERT-based pre-trained models for code search. However, the existing end-to-end training mechanism cannot sufficiently utilize the pre-trained models due to the limitations on the number and variety of negative samples. In this paper, we propose MoCoCS, a multi-modal momentum contrastive learning method for code search, to improve the representations of query and code by constructing large-scale multi-modal negative samples. MoCoCS increases the number and the variety of negative samples through two optimizations: integrating multi-batch negative samples and constructing multi-modal negative samples. We first build momentum contrasts for query and code, which enables the construction of large-scale negative samples out of a mini-batch. Then, to incorporate multi-modal code information, we build multi-modal momentum contrasts by encoding the abstract syntax tree and the data flow graph with a momentum encoder. Experiments on CodeSearchNet with six programming languages demonstrate that our method can further improve the effectiveness of pre-trained models for code search.