[논문 리뷰] AST-T5: Structure-Aware Pretraining for Code Generation and Understanding

Linyuan Gong, Mostafa Elhoushi외 1명|arXiv (Cornell University)|2024. 01. 05.

Software Engineering Research인용 수 5

한 줄 요약

AST-T5는 AST 인식 사전학습을 통해 인코더-디코더 LMs를 개선하며, AST 인식 세분화 및 AST 인식 subtree 마스킹을 사용해 아키텍처 변경 없이 코드 생성, 트랜스파일링 및 이해를 향상시킨다.

ABSTRACT

Large language models (LLMs) have made significant advancements in code-related tasks, yet many LLMs treat code as simple sequences, neglecting its structured nature. We introduce AST-T5, a novel pretraining paradigm that leverages the Abstract Syntax Tree (AST) for enhanced code generation, transpilation, and understanding. Using dynamic programming, our AST-Aware Segmentation retains code structure, while our AST-Aware Span Corruption objective equips the model to reconstruct various code structures. Unlike other models, AST-T5 avoids intricate program analyses or architectural changes, so it integrates seamlessly with any encoder-decoder Transformer. Evaluations show that AST-T5 consistently outperforms similar-sized LMs across various code-related tasks. Structure-awareness makes AST-T5 particularly powerful in code-to-code tasks, surpassing CodeT5 by 2 points in exact match score for the Bugs2Fix task and by 3 points in exact match score for Java-C# Transpilation in CodeXGLUE. Our code and model are publicly available at https://github.com/gonglinyuan/ast_t5.

연구 동기 및 목표

• 코드를 시퀀스가 아닌 구조화된 것으로 간주하려는 동기 부여, 코드 관련 작업을 향상시키기 위함.
• 복잡한 프로그램 분석 없이도 작동하는 구조 인식 사전학습 패러다임 개발.
• 기존의 인코더-디코더 트랜스포머와 드롭인 대체로 매끄럽게 통합 가능하도록 지원.
• 생성, 트랜스파일링, 이해 벤치마크 전반에서 개선을 시연.

제안 방법

• Tree-sitter를 사용하여 구조를 포착하기 위해 코드를 AST로 파싱.
• AST-인식 세분화를 적용하여 최소한의 AST 구조 파손으로 코드를 분할.
• span 손상 중 AST 하위트리를 마스킹하기 위해 AST-인식 하위트리 손상 적용.
• 구조 인식 마스킹으로 증강된 단일 사전학습 목표(span corruption)로 학습.
• 추가 헤드나 아키텍처 없이 Vanilla T5와의 호환성 유지.

실험 결과

연구 질문

• RQ1 사전학습 중 AST-인식 세분화가 그리디 세분화보다 코드 구조를 더 잘 보존하는가?
• RQ2 표준 span corruption과 비교했을 때 AST-인식 하위트리 마스킹이 생성 및 트랜스파일 품질을 향상시키는가?
• RQ3 구조 인식 수정이 적용된 단일 사전학습 목표로 코드 이해 및 생성 작업에서 경쟁력 있는 성능을 달성할 수 있는가?
• RQ4 HumanEval, MBPP, Bugs2Fix 및 Java-C# 트랜스파일에서 AST-T5의 성능은 CodeT5 및 CodeT5+에 비해 어떤가?

주요 결과

모델	HumanEval	Concode	Bugs2Fix	Java-C#	Clone	Defect	Avg
T5	5.2	18.3	21.2/13.8	65.5/68.4	96.9	64.1	44.2
+ AST. Segmentation	7.2	20.2	22.5/15.1	66.3/69.3	98.3	65.9	45.7
+ AST-Aware Subtree Corrupt	9.6	22.1	23.3/16.5	67.3/72.2	98.6	66.0	47.0
+ Mask 25% (AST-T5)	14.0	22.9	23.8/16.1	68.9/72.3	98.6	65.8	47.9
+ Mask 50%	14.3	22.0	21.9/15.0	66.5/70.1	97.1	64.2	46.4

• AST-인식 세분화는 그리디 세분화보다 더 일관된 학습 청크를 생성하여 AST 구조 손상을 줄인다.
• AST-인식 하위트리 손상은 일반 span corruption에 비해 생성 및 트랜스파일 지표를 향상시킨다.
• 마스크 비율을 25%로 올리면 이해에 큰 영향 없이 생성이 향상되며; 50%는 수익이 감소한다.
• AST-T5는 크기에 비해 여러 작업에서 CodeT5 및 CodeT5+와 같은 규모의 모델을 능가하는 강력한 성과를 달성한다.
• 트랜스파일링 및 복제 탐지에서 AST-인식은 기준선에 비해 현저한 우위를 제공한다.
• HumanEval 및 MBPP에서 더 큰 모델과 경쟁하며 매개변수 효율성을 강조한다.