[论文解读] AST-T5: Structure-Aware Pretraining for Code Generation and Understanding

Linyuan Gong, Mostafa Elhoushi+1 more|arXiv (Cornell University)|Jan 5, 2024

Software Engineering Research被引用 5

一句话总结

AST-T5 引入针对 AST 的感知预训练，用于编码器-解码器语言模型，结合 AST 感知分割和 AST 感知子树遮蔽以在不改变架构的情况下提升代码生成、转译与理解。

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.

研究动机与目标

• 将代码视为结构化对象而非仅仅是一个序列的动机，以改进与代码相关的任务。
• 开发一个不需要复杂程序分析的结构感知预训练范式。
• 实现与现有编码器-解码器 Transformer 的无缝集成，作为简单替换。
• 在生成、转译和理解基准上展示改进。

提出的方法

• 使用 Tree-sitter 将代码解析为 AST 以捕捉结构。
• 应用 AST-Aware Segmentation 将代码分割，同时尽量减少 AST 结构中断。
• 应用 AST-Aware Subtree Corruption 在跨度损坏时遮蔽 AST 子树。
• 使用单一预训练目标（跨度损坏）并结合结构感知遮蔽进行训练。
• 保持与 Vanilla T5 的兼容性，无需额外的头部或架构。

实验结果

研究问题

• RQ1 AST 感知分割在预训练阶段是否比贪婪分割更好地保留代码结构？
• RQ2 与标准跨度损坏相比，AST 感知子树遮蔽是否改善生成和转译质量？
• RQ3 单一的预训练目标结合结构感知修改是否能够在代码理解和生成任务上实现具有竞争力的性能？
• RQ4 相较于 CodeT5 和 CodeT5+，AST-T5 在 HumanEval、MBPP、Bugs2Fix 以及 Java-C# 转译任务上的表现如何？

主要发现

Model	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-Aware Segmentation 比 Greedy Segmentation 产生更连贯的训练块，减少 AST 结构中断。
• AST-Aware Subtree Corruption 相较于原生跨度损坏在生成和转译指标上有提升。
• 将遮罩比例提高到 25% 在对理解影响不大的前提下提升生成；50% 时收益递减。
• AST-T5 在同等规模下取得了强劲结果，优于 CodeT5 和 CodeT5+ 等同规模模型在多任务上的表现。
• 在转译和克隆检测方面，AST 感知相对于基线具有显著优势。
• AST-T5 在 HumanEval 和 MBPP 上与更大模型相比具备参数效率的竞争力。