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

📄 arXiv: 2401.03003v4 📥 PDF

作者: Linyuan Gong, Mostafa Elhoushi, Alvin Cheung

分类: cs.SE, cs.CL, cs.LG

发布日期: 2024-01-05 (更新: 2024-06-23)

备注: 15 pages; ICML 2024: https://icml.cc/virtual/2024/poster/33601

🔗 代码/项目: GITHUB

💡 一句话要点

提出AST-T5以提升代码生成与理解能力

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 抽象语法树 代码生成 结构感知 深度学习 Transformer 代码理解 机器学习

📋 核心要点

  1. 现有的代码生成模型往往忽视代码的结构特性，导致生成质量不高。
  2. AST-T5通过引入抽象语法树（AST）来增强代码生成和理解能力，保持代码结构信息。
  3. 实验结果显示，AST-T5在多个代码相关任务中表现优异，超越了CodeT5等基线模型。

📝 摘要（中文）

大型语言模型（LLMs）在代码相关任务上取得了显著进展，但许多LLMs将代码视为简单序列，忽视了其结构特性。我们提出了AST-T5，这是一种新颖的预训练范式，利用抽象语法树（AST）来增强代码生成、转译和理解。通过动态编程，我们的AST感知分割保留了代码结构，而AST感知跨度腐蚀目标使模型能够重构各种代码结构。与其他模型不同，AST-T5避免了复杂的程序分析或架构更改，因此可以无缝集成到任何编码-解码Transformer中。评估结果表明，AST-T5在各种代码相关任务中始终优于同等规模的语言模型。

🔬 方法详解

问题定义：现有的代码生成和理解模型通常将代码视为简单的序列，忽略了代码的结构特性，导致生成的代码质量和理解能力不足。

核心思路：AST-T5通过引入抽象语法树（AST）来增强模型的结构感知能力，利用AST进行代码的分割和重构，从而提高代码生成和理解的效果。

技术框架：AST-T5的整体架构基于编码-解码Transformer，主要包括AST感知分割模块和AST感知跨度腐蚀目标模块，前者用于保留代码结构，后者用于重构代码。

关键创新：AST-T5的核心创新在于其结构感知机制，避免了复杂的程序分析和架构更改，使其能够与现有的Transformer架构无缝集成。

关键设计：在模型设计中，采用动态编程技术进行AST感知分割，并设计了特定的损失函数以优化模型在代码重构任务中的表现。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

AST-T5在多个代码相关任务中表现优异，特别是在Bugs2Fix任务中比CodeT5提高了2个点的准确匹配分数，在Java-C#转译任务中提高了3个点，显示出其在代码到代码任务中的强大能力。

🎯 应用场景

AST-T5在代码生成、转译和理解等领域具有广泛的应用潜力，能够帮助开发者更高效地生成高质量代码，提升代码的可维护性和可读性。此外，该模型的结构感知能力也为未来的自动化编程工具和智能编程助手奠定了基础。

📄 摘要（原文）

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.