AST-T5: Structuurbewuste voorafgaande training voor codegeneratie en -begrip

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

January 5, 2024

Auteurs: Linyuan Gong, Mostafa Elhoushi, Alvin Cheung

cs.AI

Samenvatting

Grote taalmodellen (LLMs) hebben aanzienlijke vooruitgang geboekt in code-gerelateerde taken, maar veel LLMs behandelen code als eenvoudige sequenties en negeren de gestructureerde aard ervan. Wij introduceren AST-T5, een nieuw pretrainingsparadigma dat gebruikmaakt van de Abstract Syntax Tree (AST) voor verbeterde codegeneratie, transpilatie en begrip. Door middel van dynamisch programmeren behoudt onze AST-Aware Segmentation de codestructuur, terwijl ons AST-Aware Span Corruption-objectief het model in staat stelt om verschillende codestructuren te reconstrueren. In tegenstelling tot andere modellen vermijdt AST-T5 ingewikkelde programma-analyses of architectuurwijzigingen, waardoor het naadloos integreert met elke encoder-decoder Transformer. Evaluaties tonen aan dat AST-T5 consistent beter presteert dan vergelijkbaar grote taalmodellen in diverse code-gerelateerde taken. Structureel bewustzijn maakt AST-T5 bijzonder krachtig in code-naar-code taken, waarbij het CodeT5 overtreft met 2 punten in de exacte matchscore voor de Bugs2Fix-taak en met 3 punten in de exacte matchscore voor Java-C# Transpilatie in CodeXGLUE. Onze code en model zijn publiekelijk beschikbaar op https://github.com/gonglinyuan/ast_t5.

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.