AST-T5: Structure-Aware Pretraining for Code Generation and
  Understanding
Authors: name: Linyuan Gong; name: Mostafa Elhoushi; name: Alvin Cheung
Published: 2024-01-05 06:51:08

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.