[논문 리뷰] TreeBERT: A Tree-Based Pre-Trained Model for Programming Language

Xue Jiang, Zhuoran Zheng외 3명|arXiv (Cornell University)|2021. 07. 27.

Software Engineering Research인용 수 1

한 줄 요약

TreeBERT는 코드 내 추상 구문 트리(ASP)를 구성 경로로 표현하고 노드 위치 임베딩을 사용하여 트리 기반 사전 훈련 모델이다. 트리 마스킹 언어 모델링(TMLM)과 노드 순서 예측(NOP)을 통해 프로그래밍 언어 생성을 향상시키며, 코드 요약 및 문서화 작업에서 기존 모델을 능가한다. 또한 미사용된 프로그래밍 언어로의 제로샷 전이도 성공적으로 수행한다.

ABSTRACT

Source code can be parsed into the abstract syntax tree (AST) based on defined syntax rules. However, in pre-training, little work has considered the incorporation of tree structure into the learning process. In this paper, we present TreeBERT, a tree-based pre-trained model for improving programming language-oriented generation tasks. To utilize tree structure, TreeBERT represents the AST corresponding to the code as a set of composition paths and introduces node position embedding. The model is trained by tree masked language modeling (TMLM) and node order prediction (NOP) with a hybrid objective. TMLM uses a novel masking strategy designed according to the tree's characteristics to help the model understand the AST and infer the missing semantics of the AST. With NOP, TreeBERT further extracts the syntactical structure by learning the order constraints of nodes in AST. We pre-trained TreeBERT on datasets covering multiple programming languages. On code summarization and code documentation tasks, TreeBERT outperforms other pre-trained models and state-of-the-art models designed for these tasks. Furthermore, TreeBERT performs well when transferred to the pre-trained unseen programming language.

연구 동기 및 목표

• 사전 훈련 과정에 트리 구조를 가진 AST를 통합하여 프로그래밍 언어 생성을 향상시키기.
• 기존 코드를 위한 사전 훈련 모델에서 트리 구조 활용의 부족을 해결하기.
• 코드 AST의 의미적 및 구문적 구조를 동시에 포착하는 하이브리드 사전 훈련 목표를 개발하기.
• 강력한 구조적 표현을 통해 미사용된 프로그래밍 언어로의 효과적인 제로샷 전이를 가능하게 하기.
• 트리 인식 표현 학습을 통해 코드 요약 및 문서화 생성을 향상시키기.

제안 방법

• 계층적 구조를 유지하고 트리 인식 모델링을 가능하게 하기 위해 코드 AST를 구성 경로로 표현하기.
• AST 내에서 노드의 상대적 위치를 인코딩하기 위해 노드 위치 임베딩 도입하기.
• 결측한 AST 노드를 예측하기 위해 트리 전용 마스킹 전략을 사용하는 트리 마스킹 언어 모델링(TMLM) 목표 설계하기.
• AST 내 노드의 정확한 순서를 예측하여 구문적 제약을 학습하기 위해 노드 순서 예측(NOP) 제안하기.
• TMLM과 NOP를 결합하여 의미적 이해와 구조적 이해를 동시에 최적화하는 하이브리드 사전 훈련 목표 설계하기.
• 다양한 프로그래밍 언어로 구성된 코드 데이터셋을 사전 훈련에 활용하여 언어 간 일반화 능력 향상시키기.

실험 결과

연구 질문

• RQ1 명시적으로 AST 구조를 모델링하는 사전 훈련 모델이 요약 및 문서화와 같은 코드 생성 작업에서 성능 향상을 이룰 수 있는가?
• RQ2 트리 전용 마스킹 전략을 사용한 트리 마스킹 언어 모델링(TMLM)이 AST 의미를 얼마나 효과적으로 포착하는가?
• RQ3 노드 순서 예측(NOP)이 코드 내 구문적 구조 이해에 얼마나 기여하는가?
• RQ4 TreeBERT는 제로샷 전이를 통해 미사용된 프로그래밍 언어로 일반화할 수 있는가?
• RQ5 TreeBERT는 코드 생성 벤치마크에서 기존 사전 훈련 모델 및 최신 기술(SOTA) 방법과 비교해 어떻게 성능을 내는가?

주요 결과

• TreeBERT는 다양한 프로그래밍 언어에서 코드 요약 및 문서화 생성 작업에서 기존 사전 훈련 모델을 능가한다.
• TMLM와 NOP의 하이브리드 사전 훈련 목표는 개별적으로 사용할 경우보다 더 높은 성능을 달성한다.
• TreeBERT는 미사용된 프로그래밍 언어로의 제로샷 전이 성능이 뛰어나 강력한 구조적 일반화 능력을 보여준다.
• TMLM에서 사용된 트리 전용 마스킹 전략은 AST 내 결측한 의미적 내용을 추론하는 데 모델의 능력을 크게 향상시킨다.
• 노드 순서 예측(NOP)은 구문적 제약을 효과적으로 포착하여 모델의 코드에 대한 구조적 이해 능력을 향상시킨다.
• TreeBERT는 코드 생성을 위한 벤치마크 데이터셋에서 최신 기술(SOTA) 성능을 달성하여 트리 기반 사전 훈련의 효과성을 입증한다.