Saved in:
Bibliographic Details
Main Authors: Zeng, Yiming, Cao, Jinghan, Li, Zexin, Chen, Yiming, Ren, Tao, Li, Zhuochun, Xiang, Dawei, Wu, Xidong, Gao, Shangqian, Yu, Tingting
Format: Preprint
Published: 2025
Subjects:
Online Access:https://arxiv.org/abs/2508.01473
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908749071908864
author Zeng, Yiming
Cao, Jinghan
Li, Zexin
Chen, Yiming
Ren, Tao
Li, Zhuochun
Xiang, Dawei
Wu, Xidong
Gao, Shangqian
Yu, Tingting
author_facet Zeng, Yiming
Cao, Jinghan
Li, Zexin
Chen, Yiming
Ren, Tao
Li, Zhuochun
Xiang, Dawei
Wu, Xidong
Gao, Shangqian
Yu, Tingting
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2508_01473
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle TreeDiff: AST-Guided Code Generation with Diffusion LLMs
Zeng, Yiming
Cao, Jinghan
Li, Zexin
Chen, Yiming
Ren, Tao
Li, Zhuochun
Xiang, Dawei
Wu, Xidong
Gao, Shangqian
Yu, Tingting
Computation and Language
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.
title TreeDiff: AST-Guided Code Generation with Diffusion LLMs
topic Computation and Language
url https://arxiv.org/abs/2508.01473