[논문 리뷰] Learning Program Representations with a Tree-Structured Transformer

Wenhan Wang, Kechi Zhang외 5명|arXiv (Cornell University)|2022. 08. 18.

Software Engineering Research인용 수 6

한 줄 요약

이 논문은 다중 헤드 어텐션을 활용해 부모-자식 및 형제 노드 간의 의존 관계를 포착하고, 양방향(하향 및 상향) 정보 전파를 통해 프로그램 표현을 학습하는 새로운 트리 구조의 트랜스포머 모델인 Tree-Transformer를 제안한다. 이 방법은 장거리 의존성과 풍부한 노드 수준 표현을 효과적으로 모델링함으로써 프로그램 분류, 버그 로컬라이제이션, 타입 추론 작업에서 기존의 트리 기반 및 그래프 기반 접근 방식을 뛰어넘는 성능을 보인다.

ABSTRACT

Learning vector representations for programs is a critical step in applying deep learning techniques for program understanding tasks. Various neural network models are proposed to learn from tree-structured program representations, e.g., abstract syntax tree (AST) and concrete syntax tree (CST). However, most neural architectures either fail to capture long-range dependencies which are ubiquitous in programs, or cannot learn effective representations for syntax tree nodes, making them incapable of performing the node-level prediction tasks, e.g., bug localization. In this paper, we propose Tree-Transformer, a novel recursive tree-structured neural network to learn the vector representations for source codes. We propose a multi-head attention mechanism to model the dependency between siblings and parent-children node pairs. Moreover, we propose a bi-directional propagation strategy to allow node information passing in two directions, bottom-up and top-down along trees. In this way, Tree-Transformer can learn the information of the node features as well as the global contextual information. The extensive experimental results show that our Tree-Transformer significantly outperforms the existing tree-based and graph-based program representation learning approaches in both the tree-level and node-level prediction tasks.

연구 동기 및 목표

• 기존의 트리 기반 및 그래프 기반 모델이 장거리 의존성을 포착하고 코드에 대한 효과적인 노드 수준 표현을 학습하는 데에 한계가 있음을 해결하기 위해.
• 기존의 순환 신경망에서의 단방향 정보 흐름이 리프 노드 및 저수준 노드의 맥락 학습을 방해함을 극복하기 위해.
• 특히 무작위 순서 불변 모델(예: GNN)이 무시하는 형제 순서를 포함한 프로그램의 순서 구조를 유지하고 활용하기 위해.
• 복잡한 전처리나 그래프 증강이 필요 없이 원본 구문 트리 구조를 유지함으로써 오버헤드를 줄이고 구조적 정확성을 보존하기 위해.
• 향상된 프로그램 표현 학습을 위해 계층적이고 순서가 있는 프로그램 구조를 효과적으로 모델링할 수 있는 트랜스포머 기반 아키텍처를 개발하기 위해.

제안 방법

• Tree-Transformer는 트리 구조에 맞게 조정된 다중 헤드 자기 어텐션 메커니즘을 사용하여, 부모-자식 및 형제 노드 쌍 간의 의존성을 명시적으로 모델링한다.
• 하향식으로는 루트에서 자식 노드로의 맥락을 분배하고, 상향식으로는 리프에서 루트로의 맥락을 집계하는 두 개의 병렬 트랜스포머 유닛을 사용하는 양방향 정보 전파 전략을 도입한다.
• 노드의 형제 순서를 포착하기 위해 위치 인코딩을 통합함으로써, 트리 내에서 서로 다른 구조적 위치를 구분할 수 있도록 한다.
• 어텐션 기반 융합을 통해 자식 표현의 재귀적 집계를 통해 노드 표현을 학습함으로써, 트리 전반에 걸쳐 맥락 정보가 흐르도록 한다.
• 모델은 재귀적이고 트리 구조를 가진 아키텍처를 갖추고 있어, 추상 구문 트리(abstract syntax trees, ASTs)에서 직접 벡터 표현을 끝에서 끝까지 학습할 수 있다.
• 프로그램 분류, 버그 로컬라이제이션, 타입 추론과 같은 최종 작업에서 표준 지도 학습 목표를 통해 모델을 훈련한다.

Figure 1: The overall pipeline of Tree-Transformer.

실험 결과

연구 질문

• RQ1 트랜스포머 기반 아키텍처가 구문 트리 내에서 먼 노드 간의 장거리 의존성을 효과적으로 모델링할 수 있는가, 특히 먼 노드 간의 의존성에 대해?
• RQ2 하향 및 상향의 이중 방향 정보 전파 전략이 단방향 방법보다 노드 수준 표현 학습에 더 나은 성능을 내는가?
• RQ3 위치 인코딩을 통한 형제 순서 통합이 코드의 구문적 구조를 구분하는 데에 모델 능력을 향상시키는가?
• RQ4 다음으로, Tree-Transformer는 최신의 트리 기반 및 그래프 기반 모델과 비교해 프로그램 이해 작업에서 어떻게 성능을 내는가?
• RQ5 복잡한 그래프 구축이나 구조 전처리 없이도 성능을 유지할 수 있는가?

주요 결과

• Tree-Transformer는 다양한 벤치마크에서 기존의 트리 기반 및 그래프 기반 모델을 능가하는 최고 성능을 기록하며 프로그램 분류 작업에서 최고 수준의 성능을 달성했다.
• 모델은 잘못된 연산자 로컬라이제이션 및 복구와 같은 노드 수준 예측 작업에서 뚜렷한 향상을 보이며, 노드 수준에서 뛰어난 표현 학습 능력을 입증했다.
• 양방향 전파 메커니즘은 장거리 의존성 모델링에 효과적이며, 먼 코드 요소 간의 맥락 추론이 필요한 작업에서 성능 향상을 입증했다.
• 형제 순서를 위치 인코딩을 통해 통합함으로써, 문법적 순서에 민감한 작업에서 측정 가능한 성능 향상이 발생했다.
• 추가적인 그래프 구축이나 전처리 없이도 원본 AST 구조를 유지하면서도 뛰어난 성능을 유지함으로써 계산 오버헤드를 줄였다.
• 제거 실험(ablation studies) 결과, 양방향 전파와 다중 헤드 어텐션 구성 요소가 모델 성능 향상에 필수적임을 확인했다.

Figure 2: An example of existing tree-structured and graph neural networks for learning on abstract syntax trees.