[Paper Review] HELoC: Hierarchical Contrastive Learning of Source Code Representation

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

Software Engineering Research4 citations

TL;DR

HELoC proposes a self-supervised hierarchical contrastive learning framework for source code representation that leverages a novel Residual Self-attention Graph Neural Network (RSGNN) to capture both local and global AST structures. By predicting AST node levels and optimizing hierarchical relationships through contrastive learning, HELoC achieves state-of-the-art performance in code classification, clone detection, and clustering across multiple datasets.

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.

Motivation & Objective

• To address the challenge of effectively learning deep, complex AST hierarchies in source code representation.
• To improve code representation by explicitly modeling both adjacent and non-adjacent hierarchical relationships between AST nodes.
• To design a self-supervised framework that captures structural semantics without relying on labeled data.
• To develop a GNN architecture capable of balancing local structural encoding and long-range global dependency modeling in ASTs.
• To demonstrate superior performance on downstream software engineering tasks such as code classification, clone detection, and clustering.

Proposed method

• HELoC employs a novel hierarchical contrastive learning objective that treats AST node level prediction as a pretext task to preserve hierarchical structure in the embedding space.
• It introduces a Residual Self-attention Graph Neural Network (RSGNN) that combines graph convolutional networks for local structure and self-attention mechanisms for global structure modeling.
• The RSGNN integrates residual connections both internally and externally to stabilize training and enhance feature propagation across deep ASTs.
• Node representations are learned via message passing in the AST, with path representations from parent-child relationships used as input to the RSGNN.
• The model is pre-trained in a self-supervised manner using contrastive loss to maximize distance between nodes of differing AST levels while minimizing distance for similar structural contexts.
• After pre-training, the RSGNN encoder is fine-tuned on downstream tasks such as code classification, clone detection, and clustering.

Experimental results

Research questions

• RQ1 Can a self-supervised contrastive learning approach effectively preserve AST hierarchy in code representation?
• RQ2 How does modeling both adjacent and non-adjacent hierarchical relationships improve code representation?
• RQ3 To what extent does the RSGNN architecture outperform standard GNNs in capturing local and global AST structures?
• RQ4 Does HELoC achieve state-of-the-art performance across diverse downstream code understanding tasks?
• RQ5 How do individual components of HELoC, such as node level prediction and self-attention, contribute to overall performance?

Key findings

• HELoC achieves 97.2% F1 on code classification (GCJ dataset), surpassing the previous SOTA by 9.0 percentage points.
• On code clone detection (OJClone), HELoC reaches 99.6% F1, a 10.9-point improvement over the prior state-of-the-art.
• In code clustering (BCB-ALL), HELoC attains 98.0% F1, outperforming the strongest baseline by 11.3 points.
• Ablation studies confirm that removing node level prediction, relationship optimization, or the self-attention mechanism significantly degrades performance.
• The RSGNN component is critical: removing it reduces performance to 78.2% F1 on code classification, indicating its central role in representation learning.
• HELoC demonstrates consistent superiority across all three tasks and datasets, validating its effectiveness in capturing hierarchical AST structure.