[论文解读] HELoC: Hierarchical Contrastive Learning of Source Code Representation

Xiao Wang, Qiong Wu+6 more|arXiv (Cornell University)|Mar 27, 2022

Software Engineering Research被引用 4

一句话总结

HELoC 提出了一种用于源代码表示的自监督分层对比学习框架，通过一种新颖的残差自注意力图神经网络（RSGNN）来捕捉局部和全局的抽象语法树（AST）结构。通过预测AST节点层级并利用对比学习优化分层关系，HELoC在多个数据集上的代码分类、克隆检测和聚类任务中均取得了最先进性能。

ABSTRACT

Abstract syntax trees (ASTs) play a crucial role in source code representation. However, due to the large number of nodes in an AST and the typically deep AST hierarchy, it is challenging to learn the hierarchical structure of an AST effectively. In this paper, we propose HELoC, a hierarchical contrastive learning model for source code representation. To effectively learn the AST hierarchy, we use contrastive learning to allow the network to predict the AST node level and learn the hierarchical relationships between nodes in a self-supervised manner, which makes the representation vectors of nodes with greater differences in AST levels farther apart in the embedding space. By using such vectors, the structural similarities between code snippets can be measured more precisely. In the learning process, a novel GNN (called Residual Self-attention Graph Neural Network, RSGNN) is designed, which enables HELoC to focus on embedding the local structure of an AST while capturing its overall structure. HELoC is self-supervised and can be applied to many source code related downstream tasks such as code classification, code clone detection, and code clustering after pre-training. Our extensive experiments demonstrate that HELoC outperforms the state-of-the-art source code representation models.

研究动机与目标

• 为解决在源代码表示中有效学习深层复杂 AST 层次结构的挑战。
• 通过显式建模 AST 节点之间的相邻与非相邻层次关系来改进代码表示。
• 设计一种自监督框架，以在不依赖标注数据的情况下捕捉结构语义。
• 开发一种图神经网络架构，能够在 AST 中平衡局部结构编码与长距离全局依赖建模。
• 在代码分类、克隆检测和聚类等下游软件工程任务中展示卓越性能。

提出的方法

• HELoC 采用一种新颖的分层对比学习目标，将 AST 节点层级预测作为代理任务，以在嵌入空间中保留层次结构。
• 它引入了一种残差自注意力图神经网络（RSGNN），结合图卷积网络用于局部结构建模和自注意力机制用于全局结构建模。
• RSGNN 在内部和外部均引入残差连接，以稳定训练并增强深层 AST 中的特征传播。
• 通过 AST 中的消息传递学习节点表示，父-子关系的路径表示作为 RSGNN 的输入。
• 模型通过对比损失进行自监督预训练，以最大化不同 AST 层级节点之间的距离，同时最小化相似结构上下文节点之间的距离。
• 预训练完成后，RSGNN 编码器在下游任务（如代码分类、克隆检测和聚类）上进行微调。

实验结果

研究问题

• RQ1 自监督对比学习方法能否有效保留代码表示中的 AST 层次结构？
• RQ2 同时建模相邻与非相邻层次关系在多大程度上改善了代码表示？
• RQ3 RSGNN 架构在捕捉局部和全局 AST 结构方面相比标准 GNN 的表现如何？
• RQ4 HELoC 是否在多样化的下游代码理解任务中均达到最先进性能？
• RQ5 HELoC 的各个组件（如节点层级预测和自注意力机制）对整体性能的贡献如何？

主要发现

• 在 GCJ 数据集上，HELoC 在代码分类任务中达到 97.2% 的 F1 值，比之前最先进方法高出 9.0 个百分点。
• 在 OJClone 数据集的代码克隆检测任务中，HELoC 达到 99.6% 的 F1 值，较之前最先进方法提升 10.9 个百分点。
• 在 BCB-ALL 数据集的代码聚类任务中，HELoC 达到 98.0% 的 F1 值，比最强基线高出 11.3 个百分点。
• 消融实验表明，移除节点层级预测、关系优化或自注意力机制均会显著降低性能。
• RSGNN 组件至关重要：若移除它，代码分类任务的 F1 值降至 78.2%，表明其在表示学习中的核心作用。
• HELoC 在所有三个任务和数据集上均表现出一致的优越性，验证了其在捕捉分层 AST 结构方面的有效性。