# TreeDiff: AST-Guided Code Generation with Diffusion LLMs

arXiv: https://arxiv.org/abs/2508.01473 · HTML: https://mlsanity.com/paper/2508.01473/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 49 references that are in the map.
- First posted on arXiv 2 Aug 2025.

## Builds on (most cited first, 30 of 31)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Code Llama: Open Foundation Models for Code](https://mlsanity.com/paper/2308.12950/) (2023)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [A Survey on Large Language Models for Code Generation](https://mlsanity.com/paper/2406.00515/) (2024)
- [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://mlsanity.com/paper/2503.09573/) (2025)
- [Demystifying Long Chain-of-Thought Reasoning in LLMs](https://mlsanity.com/paper/2502.03373/) (2025)
- [DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generation](https://mlsanity.com/paper/2506.20639/) (2025)
- [code2seq: Generating Sequences from Structured Representations of Code](https://mlsanity.com/paper/1808.01400/) (2018)
- [OpenCodeReasoning: Advancing Data Distillation for Competitive Coding](https://mlsanity.com/paper/2504.01943/) (2025)
- [Abstract Syntax Networks for Code Generation and Semantic Parsing](https://mlsanity.com/paper/1704.07535/) (2017)
- [RewriteLM: An Instruction-Tuned Large Language Model for Text Rewriting](https://mlsanity.com/paper/2305.15685/) (2023)
- [Recent Advances in Large Langauge Model Benchmarks against Data Contamination: From Static to Dynamic Evaluation](https://mlsanity.com/paper/2502.17521/) (2025)
- [AST-T5: Structure-Aware Pretraining for Code Generation and Understanding](https://mlsanity.com/paper/2401.03003/) (2024)
- [StudentEval: A Benchmark of Student-Written Prompts for Large Language Models of Code](https://mlsanity.com/paper/2306.04556/) (2023)
- [Uncertainty-Guided Chain-of-Thought for Code Generation with LLMs](https://mlsanity.com/paper/2503.15341/) (2025)
- [CodeFusion: A Pre-trained Diffusion Model for Code Generation](https://mlsanity.com/paper/2310.17680/) (2023)
- [CodeJudge-Eval: Can Large Language Models be Good Judges in Code Understanding?](https://mlsanity.com/paper/2408.10718/) (2024)
- [Reassessing Layer Pruning in LLMs: New Insights and Methods](https://mlsanity.com/paper/2411.15558/) (2024)
- [Dynamic Benchmarking of Reasoning Capabilities in Code Large Language Models Under Data Contamination](https://mlsanity.com/paper/2503.04149/) (2025)
- [Bridging the Editing Gap in LLMs: FineEdit for Precise and Targeted Text Modifications](https://mlsanity.com/paper/2502.13358/) (2025)
- [Chain-of-Thought Matters: Improving Long-Context Language Models with Reasoning Path Supervision](https://mlsanity.com/paper/2502.20790/) (2025)
- [PPM: Automated Generation of Diverse Programming Problems for Benchmarking Code Generation Models](https://mlsanity.com/paper/2401.15545/) (2024)
- [From LLM-anation to LLM-orchestrator: Coordinating Small Models for Data Labeling](https://mlsanity.com/paper/2506.16393/) (2025)

## Built on by (most cited first, 11 of 11)

- [LaDiR: Latent Diffusion Enhances LLMs for Text Reasoning](https://mlsanity.com/paper/2510.04573/) (2025)
- [Exploring the Potential of Diffusion Large Language Models in Code Generation](https://mlsanity.com/paper/2509.11252/) (2025)
- [Lookahead-then-Verify: Reliable Constrained Decoding for Diffusion LLMs under Context-Free Grammars](https://mlsanity.com/paper/2602.00612/) (2026)
- [LaDi-RL: Latent Diffusion Reasoning Prevents Entropy Collapse in Reinforcement Learning](https://mlsanity.com/paper/2602.01705/) (2026)
- [WavefrontDiffusion: Dynamic Decoding Schedule for Improved Reasoning](https://mlsanity.com/paper/2511.19473/) (2025)
- [From AR to Diffusion: Efficiently Adapting Large Language Models with Strictly Causal and Elastic Horizons](https://mlsanity.com/paper/2605.27387/) (2026)
- [DiffuTester: Accelerating Unit Test Generation for Diffusion LLMs via Mining Structural Pattern](https://mlsanity.com/paper/2509.24975/) (2025)
- [AnCoder: Anchored Code Generation via Discrete Diffusion Models](https://mlsanity.com/paper/2602.17688/) (2026)
- [LLMs + Security = Trouble](https://mlsanity.com/paper/2602.08422/) (2026)
- [Beyond Execution: Static-Analysis Rewards and Hint-Conditioned Diffusion RL for Code Generation](https://mlsanity.com/paper/2605.17174/) (2026)
- [Beyond the Autoregressive Horizon: A Comprehensive Survey of Diffusion Models, World Modelling, and State Space Models for Code](https://mlsanity.com/paper/2606.23690/) (2026)

## Abstract (arXiv)

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.