MGS3: A Multi-Granularity Self-Supervised Code Search Framework

2025-05-30

Knowledge Discovery and Data Mining

AI Summary

Existing code search methods primarily focus on function-level alignment, overlooking the reusability of finer-grained code fragments—such as blocks and statements—thereby limiting cross-granularity retrieval performance. To address this, we introduce MGCodeSearchNet, the first multi-granularity code search dataset covering functions, blocks, and statements. We further propose MGS3, a novel framework featuring a Hierarchical Multi-Granularity Representation (HMGR) module. HMGR integrates syntax-aware hierarchical encoding, granularity-consistent positive sample construction, and intra-function negative sampling to enable self-supervised contrastive learning for cross-granularity alignment. MGS3 is compatible with mainstream pre-trained code models and achieves state-of-the-art performance across all granularities on multi-granularity benchmarks—outperforming prior approaches by +12.7% average Mean Reciprocal Rank (MRR) at both block- and statement-level retrieval.

Technology Category

Natural Language Processing: Code Generation / Program Synthesis from Natural LanguageSearch and Optimization: Distributed SearchMachine Learning: Large Multimodal Models (LMMs)

Application Category

Search and Retrieval-Augmented AI: Multilingual and cross-lingual Web searchGraph Algorithms and Modeling for the Web: Querying, indexing, and retrieval in Web-related graphsWeb Mining and Content Analysis: Mining multimedia, multimodal, multilingual, cross-lingual Web data

Abstract

In the pursuit of enhancing software reusability and developer productivity, code search has emerged as a key area, aimed at retrieving code snippets relevant to functionalities based on natural language queries. Despite significant progress in self-supervised code pre-training utilizing the vast amount of code data in repositories, existing methods have primarily focused on leveraging contrastive learning to align natural language with function-level code snippets. These studies have overlooked the abundance of fine-grained (such as block-level and statement-level) code snippets prevalent within the function-level code snippets, which results in suboptimal performance across all levels of granularity. To address this problem, we first construct a multi-granularity code search dataset called MGCodeSearchNet, which contains 536K+ pairs of natural language and code snippets. Subsequently, we introduce a novel Multi-Granularity Self-Supervised contrastive learning code Search framework (MGS3). First, MGS3 features a Hierarchical Multi-Granularity Representation module (HMGR), which leverages syntactic structural relationships for hierarchical representation and aggregates fine-grained information into coarser-grained representations. Then, during the contrastive learning phase, we endeavor to construct positive samples of the same granularity for fine-grained code, and introduce in-function negative samples for fine-grained code. Finally, we conduct extensive experiments on code search benchmarks across various granularities, demonstrating that the framework exhibits outstanding performance in code search tasks of multiple granularities. These experiments also showcase its model-agnostic nature and compatibility with existing pre-trained code representation models.

Problem

• Improving code search across multiple granularity levels
• Addressing neglect of fine-grained code snippets in existing methods
• Enhancing self-supervised contrastive learning for code retrieval

Innovation

• Hierarchical Multi-Granularity Representation module
• Self-Supervised contrastive learning framework
• Multi-Granularity code search dataset