AST-T5: Strukturorientiertes Pretraining für Code-Generierung und -Verständnis

AST-T5: Structure-Aware Pretraining for Code Generation and Understanding

January 5, 2024

Autoren: Linyuan Gong, Mostafa Elhoushi, Alvin Cheung

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) haben bedeutende Fortschritte bei codebezogenen Aufgaben erzielt, doch viele LLMs behandeln Code als einfache Sequenzen und vernachlässigen dessen strukturierte Natur. Wir stellen AST-T5 vor, ein neuartiges Vortrainierungsparadigma, das den Abstract Syntax Tree (AST) für verbesserte Code-Generierung, -Transpilation und -Verständnis nutzt. Mithilfe von dynamischer Programmierung bewahrt unsere AST-Aware Segmentation die Codestruktur, während unser AST-Aware Span Corruption Objective das Modell dazu befähigt, verschiedene Codestrukturen zu rekonstruieren. Im Gegensatz zu anderen Modellen vermeidet AST-T5 komplexe Programmanalysen oder architektonische Änderungen und lässt sich daher nahtlos in jeden Encoder-Decoder-Transformer integrieren. Evaluierungen zeigen, dass AST-T5 durchweg ähnlich große Sprachmodelle bei verschiedenen codebezogenen Aufgaben übertrifft. Die Strukturwahrnehmung macht AST-T5 besonders leistungsstark bei Code-zu-Code-Aufgaben, wobei es CodeT5 im Bugs2Fix-Task um 2 Punkte im Exact-Match-Score und bei der Java-C#-Transpilation in CodeXGLUE um 3 Punkte im Exact-Match-Score übertrifft. Unser Code und Modell sind öffentlich unter https://github.com/gonglinyuan/ast_t5 verfügbar.

English

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.