用图神经网络增强代码语义检索,提升复杂代码的匹配准确率。
GNN-Coder: Boosting Semantic Code Retrieval with Combined GNNs and Transformer
- 结合GNN与Transformer,从代码抽象语法树中捕捉结构与语义特征。
- 在CSN数据集上MRR提升1%-10%,零样本场景下CosQA提升20%。
- 提出新型图池化方法和角度均匀性度量,优化嵌入表示区分度。
代码检索是现代软件开发中的关键环节,尤其在大型项目中。现有基于序列的模型常无法充分挖掘代码内在的结构依赖关系,导致在结构复杂的代码片段上表现不佳。本文提出GNN-Coder,一种基于图神经网络(GNN)的新型框架,利用抽象语法树(AST)建模代码结构。首次探索将GNN与Transformer结合以促进语义检索任务,通过捕获代码的结构与语义特征实现性能提升。提出一种专为AST设计的图池化方法,以子节点数量为核心特征,强化树形拓扑关系表达。该设计有效融合序列与层级表示,增强模型对代码结构与语义的建模能力。同时引入均值角间距(Mean Angular Margin, MAM)作为代码嵌入分布均匀性的量化指标,衡量特征可分性。实验表明,所提方法显著提升检索性能:在CSN数据集上MRR提升1%-10%,在CosQA数据集上零样本性能提升20%。
原文摘要 · Abstract (English)
Code retrieval is a crucial component in modern software development, particularly in large-scale projects. However, existing approaches relying on sequence-based models often fail to fully exploit the structural dependencies inherent in code, leading to suboptimal retrieval performance, particularly with structurally complex code fragments. In this paper, we introduce GNN-Coder, a novel framework based on Graph Neural Network (GNN) to utilize Abstract Syntax Tree (AST). We make the first attempt to study how GNN-integrated Transformer can promote the development of semantic retrieval tasks by capturing the structural and semantic features of code. We further propose an innovative graph pooling method tailored for AST, utilizing the number of child nodes as a key feature to highlight the intrinsic topological relationships within the AST. This design effectively integrates both sequential and hierarchical representations, enhancing the model's ability to capture code structure and semantics. Additionally, we introduce the Mean Angular Margin (MAM), a novel metric for quantifying the uniformity of code embedding distributions, providing a standardized measure of feature separability. The proposed method achieves a lower MAM, indicating a more discriminative feature representation. This underscores GNN-Coder's superior ability to distinguish between code snippets, thereby enhancing retrieval accuracy. Experimental results show that GNN-Coder significantly boosts retrieval performance, with a 1\%-10\% improvement in MRR on the CSN dataset, and a notable 20\% gain in zero-shot performance on the CosQA dataset.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。