用抽象语法树指导扩散模型生成代码,提升语法正确性。
TreeDiff: AST-Guided Code Generation with Diffusion LLMs
- 基于抽象语法树选择性掩码关键节点,而非随机掩码
- 相比随机掩码方法,代码生成准确率提升13.3%
- 适合需要高语法正确性的代码生成任务
代码生成在实际应用中日益重要,但基于扩散的大型语言模型仍面临挑战。与自由文本不同,代码要求严格的语法结构,微小的结构错误即可能导致程序无法执行。现有扩散模型依赖随机标记进行数据污染,导致两个关键缺陷:在迭代去噪过程中缺乏对语法边界的感知,且无法捕捉程序正确性所需的长程分层依赖关系。为此,我们提出TreeDiff,一种结合抽象语法树(AST)结构先验的语法感知扩散框架。不采用随机掩码,而是有选择地掩码属于关键AST节点的标记。通过将污染过程与代码底层结构对齐,该方法促使模型内化编程语言的组合特性,从而生成符合语法规则、保留长程依赖的程序。实验表明,该方法相较随机掩码训练方式实现13.3%的相对性能提升,验证了利用代码结构先验在代码生成任务中的有效性。
原文摘要 · Abstract (English)
Code generation is increasingly critical for real-world applications. Still, diffusion-based large language models continue to struggle with this demand. Unlike free-form text, code requires syntactic precision; even minor structural inconsistencies can render a program non-executable. Existing diffusion-based large language models rely on random token masking for corruption, leading to two key failures: they lack awareness of syntactic boundaries during the iterative denoising process, and they fail to capture the long-range hierarchical dependencies essential for program correctness. We propose TreeDiff to address both issues. Specifically, we propose a syntax-aware diffusion framework that incorporates structural priors from Abstract Syntax Tree (AST) into the corruption process. Instead of masking individual tokens at random, we selectively mask tokens belonging to key AST nodes. By aligning the corruption process with the underlying structure of code, our method encourages the model to internalize the compositional nature of programming languages, enabling it to reconstruct programs that respect grammatical boundaries and capture long-range dependencies. Our method achieves a 13.3% relative improvement over the random masking training method, demonstrating its effectiveness in code generation task by leveraging underlying structures.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。