[논문 리뷰] HELoC: Hierarchical Contrastive Learning of Source Code Representation

Xiao Wang, Qiong Wu외 6명|arXiv (Cornell University)|2022. 03. 27.

Software Engineering Research인용 수 4

한 줄 요약

HELoC는 소스 코드 표현을 위한 자기지도형 계층적 대비 학습 프레임워크를 제안하며, 지역적이고 전역적인 AST 구조를 포착하기 위해 새로운 잔차 자기주의 그래프 신경망(RSGNN)을 활용한다. AST 노드의 레벨을 예측하고 대비 학습을 통해 계층적 관계를 최적화함으로써, HELoC는 여러 데이터셋에서 코드 분류, 클론 탐지, 군집화 작업에서 최신 기술 수준(SOTA) 성능을 달성한다.

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에서 국소적 구조 인코딩과 장거리 전역 의존성 모델링을 균형 있게 처리할 수 있는 GNN 아키텍처를 개발하는 것.
• 코드 분류, 클론 탐지, 군집화와 같은 후행 소프트웨어 공학 작업에서 뛰어난 성능을 보여주는 것.

제안 방법

• HELoC는 AST 노드의 레벨 예측을 선행 과제로 삼아 임베딩 공간에서 계층적 구조를 유지하는 데 목적이 있는 새로운 계층적 대비 학습 목적함수를 사용한다.
• 지역적 구조를 위한 그래프 컨볼루션 네트워크와 전역적 구조 모델링을 위한 자기주의 메커니즘을 결합한 잔차 자기주의 그래프 신경망(RSGNN)을 도입한다.
• 내부 및 외부 모두에 잔차 연결을 통합하여 깊은 AST에서 학습 안정성과 특징 전파를 향상시킨다.
• 노드 표현은 AST 내 메시지 전파를 통해 학습되며, 부모-자식 관계에서 유도된 경로 표현이 RSGNN의 입력으로 사용된다.
• 자기지도 학습 방식으로 사전 훈련되며, 대비 손실을 사용하여 서로 다른 AST 레벨의 노드 간의 거리를 최대화하고 유사한 구조적 맥락의 노드 간 거리를 최소화한다.
• 사전 훈련 후, RSGNN 인코더는 코드 분류, 클론 탐지, 군집화와 같은 후행 작업에서 미세조정된다.

실험 결과

연구 질문

• RQ1 자기지도형 대비 학습 접근법이 코드 표현에서 AST 계층을 효과적으로 유지할 수 있는가?
• RQ2 인접 및 비인접한 계층적 관계를 모두 모델링할 경우 코드 표현이 어떻게 향상되는가?
• RQ3 RSGNN 아키텍처가 표준 GNN보다 지역적 및 전역적 AST 구조를 포착하는 데 얼마나 뛰어난가?
• RQ4 HELoC는 다양한 후행 코드 이해 작업에서 최신 기술 수준(SOTA) 성능을 달성하는가?
• RQ5 노드 레벨 예측 및 자기주의 메커니즘과 같은 HELoC의 개별 구성 요소가 전체 성능에 기여하는 정도는 어느 정도인가?

주요 결과

• 코드 분류 작업에서 GCJ 데이터셋 기준으로 97.2%의 F1 스코어를 기록하며, 이는 이전 최신 기술 수준(SOTA)보다 9.0 포인트 높은 성능이다.
• 코드 클론 탐지 작업에서 OJClone 데이터셋 기준으로 99.6%의 F1 스코어를 달성하여, 이전 최신 기술 수준(SOTA)보다 10.9 포인트 향상되었다.
• 코드 군집화 작업에서 BCB-ALL 데이터셋 기준으로 98.0%의 F1 스코어를 기록하며, 가장 강력한 베이스라인보다 11.3 포인트 높은 성능을 보였다.
• 절단 실험 결과, 노드 레벨 예측, 관계 최적화, 또는 자기주의 메커니즘을 제거할 경우 성능이 크게 떨어지는 것으로 확인되었다.
• RSGNN 구성 요소는 핵심적이다: 이를 제거하면 코드 분류 작업에서 F1 스코어가 78.2%로 떨어지며, 표현 학습에서의 중심적 역할을 확인할 수 있다.
• HELoC는 모든 세 가지 작업과 데이터셋에서 일관된 슈퍼리어리티를 보이며, 계층적 AST 구조를 효과적으로 포착할 수 있음을 검증한다.