T: Progressive Block Scaling for Masked Diffusion Language Models Through Trajectory Aware Reinforcement Learning
Paper · 2026 · Hanchen Xia, Baoyou Chen, Yutang Ge et al.
Summary
Masked diffusion language models (MDMs) usually decode in blocks. Larger blocks allow more tokens to be finalized in parallel, but math reasoning accuracy tends to drop as block size grows, especially for smaller models. The authors propose T⋆, a curriculum built on the trajectory-aware RL method TraceRL. It starts from an AR-initialized SDAR checkpoint with a small block size. The block size is then doubled stage by stage (4 → 8 → 16 → 32), and the denoising policy is re-optimized at each stage. Every update step applies TraceRL twice: once with the standard block partition and once with a partition shifted by half a block, which reduces dependence on fixed boundaries. Experiments on SDAR-1.7B-Chat and SDAR-4B-Chat cover MATH500, GSM8K and AIME24. T⋆ generally matches or beats both the base checkpoints and direct large-block TraceRL. Direct TraceRL, by contrast, suffers sudden training collapses. The learned unmasking order stays mostly monotone but does not revert to strict left-to-right decoding. This makes T⋆ a practical way to gain decoding parallelism without giving up much reasoning ability.
Key points
- Training setup: 8K math problems (difficulty levels 3–5) from OpenR1-Math. Each step samples 128 problems with 16 rollouts each. Training ran on 8×H200 with AdamW (lr 1e-6) and a KL penalty of β = 0.01. Results are reported as Pass@3, with checkpoints selected on MATH500.
- MATH500 results (SDAR-1.7B, base / TraceRL / T⋆):
- B=8: 55.9 / 60.2 / 63.4
- B=16: 52.4 / 54.4 / 59.8
- B=32: 54.2 / 54.1 / 59.0
- Baseline instability: direct TraceRL fell from about 56% to the low 40s during the B=8 stage, and to about 30% late in the B=16 stage. The authors link this to likelihood drift under noisy, low-confidence rollouts.
- Parallelism (tokens per forward, TPF): TPF rises from 2.91 to 3.38 to 3.74 as B goes from 8 to 16 to 32. This means roughly 14–22% fewer forward passes, and the gain is preserved after RL training.
- Decoding order: the LocalStrict metric for T⋆ stays between 0.73 and 0.85, compared with 1.0 for AR models. The limitations section notes that blocks of 64 or more were not tested because the inference engine becomes unstable at those sizes.
Institutions
Fields, topics and concepts
📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksComputer Vision and Pattern Recognition · 3,486 works
🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,894 works
Cite
Cite as Hanchen Xia, Baoyou Chen, Yutang Ge, Guojiang Zhao and Siyu Zhu. 2026. T: Progressive Block Scaling for Masked Diffusion Language Models Through Trajectory Aware Reinforcement Learning. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers).
@inproceedings{xia2026t,
title = {{T: Progressive Block Scaling for Masked Diffusion Language Models Through Trajectory Aware Reinforcement Learning}},
author = {Hanchen Xia, Baoyou Chen, Yutang Ge, Guojiang Zhao and Siyu Zhu},
booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)},
year = {2026},
doi = {10.18653/v1/2026.acl-short.67},
url = {https://doi.org/10.18653/v1/2026.acl-short.67}
}References
- 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. In International Conference on Learning Representations. find
- Art of Problem Solving. 2024a. 2024 AIME i. AoPS Wiki. Accessed: 2026-01-05. find
- Art of Problem Solving. 2024b. 2024 AIME ii. AoPS Wiki. Accessed: 2026-01-05. find
- Maciej Besta, Nils Blach, Ales Kubicek, Robert Gerstenberger, Michal Podstawski, Lukas Gianinazzi, Joanna Gajda, Tomasz Lehmann, Hubert Niewiadomski, Piotr Nyczyk, and Torsten Hoefler. 2024. Graph of thoughts: Solving elaborate problems with large language models. In Thirty-Eighth AAAI Conference on… find
- Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter,… — Language Models are Few-Shot Learners
- Shuang Cheng, Yihan Bian, Dawei Liu, Yuhua Jiang, Yihao Liu, Linfeng Zhang, Wenhai Wang, Qipeng Guo, Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the…
- Chang Gao, Chujie Zheng, Xiong-Hui Chen, Kai Dang, Shixuan Liu, Bowen Yu, An Yang, Shuai Bai, Jingren Zhou, and Junyang Lin. 2025. Soft adaptive policy optimization. arXiv preprint arXiv:2511.20347. — [arXiv:2511.20347](https://arxiv.org/abs/2511.20347) find
- 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. find
- Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021. Measuring mathematical problem solving with the MATH dataset. In Advances in Neural Information Processing Systems.
- Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2023a. Let’s verify step by step. arXiv preprint arXiv:2305.20050. Defines the nonstandard MATH-500 evaluation split released in the PRM800K repository.
- Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2023b. LetâĂŹs verify step by step. find
- Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. 2019. Roberta: A robustly optimized BERT pretraining approach.
- Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T. Chiu, Alexander Rush, and Volodymyr Kuleshov. 2024. find
- Simple and effective masked diffusion language models. In Advances in Neural Information Processing Systems.
- John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and Pieter Abbeel. 2015. High-dimensional continuous control using generalized advantage estimation. find
- Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. LLaMA: Open and efficient foundation language models.
- Yinjie Wang, Ling Yang, Bowen Li, Ye Tian, Ke Shen, and Mengdi Wang. 2025. Revolutionizing reinforcement learning framework for diffusion large language models. find
- Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao, and Karthik Narasimhan. 2023. Tree of thoughts: Deliberate problem solving with large language models. In Advances in Neural Information Processing Systems. find
- Yao Yao, Zuchao Li, and Hai Zhao. 2024. GoT: Effective graph-of-thought reasoning in language models. In Findings of the Association for Computational Linguistics: NAACL 2024, pages 2901–2921, Mexico City, Mexico. Association for Computational Linguistics. find
- Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. 2025. Dream 7b: Diffusion large language models.
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/xia-2026-t-star-progressive-block-scaling/