AST-Monte Carlo Diffusion: Structure-Aware Super Data Learning for Code

TL;DR: Instead of masking random tokens, mask meaningful code chunks from the AST so the model learns program structure—even when it rereads the same files. First experiment: inject AST-guided corruption plus frequency-aware weighting into a code DLM and compare to AR coders under repeated Python data; hypothesis is improved syntactic correctness and pass@k at fixed unique tokens.

Research Question: Can syntax-aware, structure-preserving corruption amplify the DLM advantage in code modeling under scarce unique code, leading to fewer syntax errors and higher task accuracy than AR models?

Hypothesis: AST-guided, span-level denoising creates stronger implicit augmentation than token-level masking, improving long-range consistency and compositional generalization. Coupled with block diffusion (for variable-length code) and any-order modeling, this will surpass matched AR coders with the same unique data and compute.

Experiment Plan: - Method: Use TreeDiff-style AST spans for corruption targets; bias sampling toward rare constructs (frequency-informed). Train block diffusion variants to allow flexible sequence lengths and KV caching.

• Data: 5–10B unique Python tokens (as in Ni et al.), repeated to reach ~1T–1.5T token compute; evaluate on HumanEval/MBPP, syntax error rates, and fill-in-the-middle coding.
• Baselines: AR coder trained with matched compute/data; standard MDLM without AST guidance; DiffuLLaMA-style adaptation from an AR coder as an additional baseline.
• Measures: pass@1/pass@k, compilation rate, functional correctness, and generation latency.
• Expected: Significant gains in compilation rate and pass@k vs. AR and token-level MDLM; better fill-in-the-middle due to bidirectional denoising; competitive inference efficiency via block-wise sampling.

References: 1. Ni, J., Liu, Q., Dou, L., Du, C., Wang, Z., Yan, H., Pang, T., & Shieh, M. (2025). Diffusion Language Models are Super Data Learners.
2. Zeng, Y., Cao, J., Li, Z., Chen, Y., Ren, T., Xiang, D., Wu, X., Gao, S., & Yu, T. (2025). TreeDiff: AST-Guided Code Generation with Diffusion LLMs. arXiv.org.
3. Arriola, M., Gokaslan, A., Chiu, J. T., Yang, Z., Qi, Z., Han, J., Sahoo, S., & Kuleshov, V. (2025). Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models. International Conference on Learning Representations.
4. Gong, S., Agarwal, S., Zhang, Y., Ye, J., Zheng, L., Li, M., An, C., Zhao, P., Bi, W., Han, J., Peng, H., & Kong, L. (2024). Scaling Diffusion Language Models via Adaptation from Autoregressive Models. International Conference on Learning Representations.
5. Ye, J., Zheng, Z., Bao, Y., Qian, L., & Gu, Q. (2023). Diffusion Language Models Can Perform Many Tasks with Scaling and Instruction-Finetuning. arXiv.org.