Elamur Research⌕
☰
home / papers / 2026 / TreeDiff: AST-Guided Code Generation with Diffusion LLMs

TreeDiff: AST-Guided Code Generation with Diffusion LLMs

Paper · 2026 · Yiming Zeng, Jinghan Cao, Zexin Li et al.

Summary

Diffusion-based language models (DLLMs) generate text by iteratively un-masking a corrupted sequence. They do this poorly for code, partly because training masks tokens uniformly at random, ignoring program grammar. As a result, models often produce syntactically broken drafts and miss long-range relations such as scope and control flow. TreeDiff replaces uniform masking in the code part of each training example with corruption guided by the Abstract Syntax Tree (AST): tokens belonging to important syntactic nodes are preferentially masked. Each training sequence has three parts, handled differently:

  • the natural-language prompt stays visible;
  • a <think> reasoning trace receives ordinary random masking;
  • the code receives AST-weighted masking.

The authors fine-tune LLaDA 8B with LoRA on roughly 150K long code-reasoning samples. They evaluate on HumanEval, HumanEval+, MBPP and MBPP+, where TreeDiff beats random masking across inference lengths. Its largest relative gain, 13.3%, comes on HumanEval+. Qualitative trajectories suggest the model builds a structural skeleton early in denoising. The work matters because it shows structural priors can help narrow the gap between diffusion and autoregressive coders.

Key points

  • Base model: LLaDA 8B, adapted via LoRA on 8 NVIDIA A100 GPUs, since full fine-tuning on 4096-token traces did not fit in memory. Inference traces were capped at 1,024 tokens because of latency.
  • Training data: about 150K long code-reasoning samples. Evaluation used HumanEval, HumanEval+, MBPP and MBPP+.
  • Main result: up to 13.3% relative improvement over the random-masking baseline, on HumanEval+.
  • Masking granularity matters. Masking whole contiguous AST subtrees (span masking, in the spirit of AST-T5) removes too much context. On MBPP at T = 512 it scores 39.7%, versus 41.5% for random masking, whereas node-level AST masking stays more robust.
  • Case study on HumanEval/6 (nested parenthesis depth): TreeDiff produces stack-based logic such as stack.pop() by step 110 of 256. The baseline instead falls back on a shallow counter heuristic.
  • The authors promise to release code and checkpoints upon acceptance.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓National University of SingaporeSingapore · Singapore180 works · 34,304 citations🎓University of PittsburghPittsburgh · United States59 works · 22,324 citations🎓University of California, RiversideRiverside · United States7 works · 339 citations🎓University of ConnecticutStorrs · United States11 works · 357 citations🎓Florida State UniversityTallahassee · United States8 works · 363 citations

Also on the byline: San Francisco State University.

 Fields, topics and concepts

📚 Shelves 6The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksInformation Systems · 1,437 worksSoftware · 92 works

🧭 Topics 2Research subjects and languages the work is about.Language models & pretraining · 3,320 worksCode & programming languages · 129 works

Cite

Cite as Yiming Zeng, Jinghan Cao, Zexin Li, Yiming Chen, Tao Ren, Zhuochun Li, Dawei Xiang, Xidong Wu, Shangqian Gao and Tingting Yu. 2026. TreeDiff: AST-Guided Code Generation with Diffusion LLMs. Proceedings of the First Workshop on Structured Understanding, Retrieval, and Generation in the LLM Era (SURGeLLM 2026).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.surgellm-1.5

@inproceedings{zeng2026treediff,
  title = {{TreeDiff: AST-Guided Code Generation with Diffusion LLMs}},
  author = {Yiming Zeng, Jinghan Cao, Zexin Li, Yiming Chen, Tao Ren, Zhuochun Li, Dawei Xiang, Xidong Wu, Shangqian Gao and Tingting Yu},
  booktitle = {Proceedings of the First Workshop on Structured Understanding, Retrieval, and Generation in the LLM Era (SURGeLLM 2026)},
  year = {2026},
  doi = {10.18653/v1/2026.surgellm-1.5},
  url = {https://doi.org/10.18653/v1/2026.surgellm-1.5}
}

References

31 references, 14 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Wasi Uddin Ahmad, Sean Narenthiran, Somshubra Majumdar, Aleksander Ficek, Siddhartha Jain, Jocelyn Huang, Vahid Noroozi, and Boris Ginsburg. 2025. Opencodereasoning: Advancing data distillation for competitive coding. Preprint, arXiv:2504.01943.
  2. Alfred V. Aho, Monica S. Lam, Ravi Sethi, and Jeffrey D. Ullman. 2006. Compilers: Principles, Techniques, and Tools (2nd Edition). Addison-Wesley. find
  3. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhai… — Evaluating Large Language Models Trained on Code
  4. Linyuan Gong, Mostafa Elhoushi, and Alvin Cheung. 2024. Ast-t5: structure-aware pretraining for code generation and understanding. In Proceedings of the 41st International Conference on Machine Learning, ICML’24. JMLR.org. find
  5. Shansan Gong, Ruixiang Zhang, Huangjie Zheng, Jiatao Gu, Navdeep Jaitly, Lingpeng Kong, and Yizhe Zhang. 2025. Diffucoder: Understanding and improving masked diffusion models for code generation. Preprint, arXiv:2506.20639. — [arXiv:2506.20639](https://arxiv.org/abs/2506.20639) find
  6. Uri Alon, Shaked Brody, Omer Levy, and Eran Yahav. 2019. code2seq: Generating sequences from structured representations of code. In International Conference on Learning Representations. ArXiv:1808.01400. find
  7. Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, et al. 2025. Deepseek-r1 incentivizes reasoning in llms through reinforcement learning. Nature, 645(8081):633–638. find
  8. Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, and Volodymyr Kuleshov. 2025. Block diffusion: Interpolating between autoregressive and diffusion language models. arXiv preprint arXiv:2503.09573. — [arXiv:2503.09573](https://arxiv.org/abs/2503.09573) find
  9. Vincent J. Hellendoorn, Petros Maniatis, Rishabh Singh, Charles Sutton, and David Bieber. 2020. Global relational models of source code. In International Conference on Learning Representations. find
  10. Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. 2021a. Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems.
  11. Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. 2021b. Program synthesis with large language models. arXiv preprint arXiv:2108.07732.
  12. Haolin Chen, Shiyu Wang, Can Qin, Bo Pang, Zuxin Liu, Jielin Qiu, Jianguo Zhang, Yingbo Zhou, Zeyuan Chen, Ran Xu, Shelby Heinecke, Silvio Savarese, Caiming Xiong, Huan Wang, and Weiran Yao. 2025. Coda: Coding lm via diffusion adaptation. Preprint, arXiv:2510.03270. — [arXiv:2510.03270](https://arxiv.org/abs/2510.03270) find
  13. Jonathan Ho, Ajay Jain, and Pieter Abbeel. 2020. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems. find
  14. Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2022. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations.
  15. Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori B. Hashimoto. 2022. Diffusion-lm improves controllable text generation. In Advances in Neural Information Processing Systems.
  16. Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang, and Lingming Zhang. 2023. Is your code generated by chat-GPT really correct? rigorous evaluation of large language models for code generation. In Thirty-seventh Conference on Neural Information Processing Systems. find
  17. Iulian Neamtiu, Jeffrey S. Foster, and Michael Hicks. 2005. Understanding source code evolution using abstract syntax tree matching. In Proceedings of the 2005 International Workshop on Mining Software Repositories, MSR ’05, page 1–5, New York, NY, USA. Association for Computing Machinery. find
  18. Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. 2025. Large language diffusion models. Preprint, arXiv:2502.09992.
  19. OpenAI. 2022. Introducing chatgpt. https://openai.com/blog/chatgpt. Accessed: 2025-07-24. ↗
  20. Terence Parr. 2013. The definitive antlr 4 reference. find
  21. Maxim Rabinovich, Mitchell Stern, and Dan Klein. 2017. Abstract syntax networks for code generation and semantic parsing. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics. ArXiv:1704.07535. find
  22. Pengcheng Yin and Graham Neubig. 2017. A syntactic neural model for general-purpose code generation. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics. ArXiv:1704.01696.
  23. Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. 2025. d1: Scaling reasoning in diffusion large language models via reinforcement learning. In The Thirty-ninth Annual Conference on Neural Information Processing Systems. find
  24. Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, Yifan Du, Chen Yang, Yushuo Chen, Zhipeng Chen, Jinhao Jiang, Ruiyang Ren, Yifan Li, Xinyu Tang, Zikang Liu, Peiyu Liu, Jian-Yun Nie, and Ji-Rong Wen. 2026. A survey of l…
  25. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21(1).
  26. Subham Sekhar Sahoo, Marianne Arriola, Aaron Gokaslan, Edgar Mariano Marroquin, Alexander M Rush, Yair Schiff, Justin T Chiu, and Volodymyr Kuleshov. 2024. Simple and effective masked diffusion language models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.
  27. Mukul Singh, José Cambronero, Sumit Gulwani, Vu Le, Carina Negreanu, and Gust Verbruggen. 2023. Code-Fusion: A pre-trained diffusion model for code generation. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 11697–11708, Singapore. Association for Com… find
  28. Gemini Team, Rohan Anil, et al. 2023. Gemini: A family of highly capable multimodal models. arXiv preprint arXiv:2312.11805.
  29. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoc V Le, and Denny Zhou. 2022. Chain of thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems.
  30. Zhihui Xie, Jiacheng Ye, Lin Zheng, Jiahui Gao, Jingwei Dong, Zirui Wu, Xueliang Zhao, Shansan Gong, Xin Jiang, Zhenguo Li, and Lingpeng Kong. 2025. Dream-coder 7b: An open diffusion language model for code. Preprint, arXiv:2509.01142. — [arXiv:2509.01142](https://arxiv.org/abs/2509.01142) find
  31. Edward Yeo, Yuxuan Tong, Morry Niu, Graham Neubig, and Xiang Yue. 2025. Demystifying long chain-of-thought reasoning in llms. arXiv preprint arXiv:2502.03373.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zeng-2026-treediff-ast-guided-code-generation/