DreamOn: Diffusion Language Models For Code Infilling Beyond Fixed-size Canvas
Zirui Wu, Lin Zheng, Zhihui Xie, Jiacheng Ye, Jiahui Gao, Shansan Gong, Yansong Feng, Zhenguo Li, Wei Bi, Guorui Zhou, Lingpeng Kong
Abstract
Diffusion Language Models (DLMs) present a compelling alternative to autoregressive models, offering flexible, any-order infilling without specialized prompting design. However, their practical utility is blocked by a critical limitation: the requirement of a fixed-length masked sequence for generation. This constraint severely degrades code infilling performance when the predefined mask size mismatches the ideal completion length. To address this, we propose DreamOn, a novel diffusion framework that enables dynamic, variable-length generation. DreamOn augments the diffusion process with two length control states, allowing the model to autonomously expand or contract the output length based solely on its own predictions. We integrate this mechanism into existing DLMs with minimal modifications to the training objective and no architectural changes. Built upon Dream-Coder-7B and DiffuCoder-7B, DreamOn achieves infilling performance on par with state-of-the-art autoregressive models on HumanEval-Infilling and SantaCoder-FIM and matches oracle performance achieved with ground-truth length. Our work removes a fundamental barrier to the practical deployment of DLMs, significantly advancing their flexibility and applicability for variable-length generation. Our code is available at https://github.com/DreamLM/DreamOn.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 62a927fa-e2de-46cd-9996-20df5fd3d283Cited by top-tier papers9
- Any-Order Flexible Length Masked DiffusionJaeyeon Kim, Cheuk Lee Kit, Carles Domingo-Enrich, Yilun Du et al.ICLR 2026 · 51 citations
- Scaling Behavior of Discrete Diffusion Language ModelsDimitri von Rütte, Janis Fluri, Omead Pooladzandi, Bernhard Schölkopf et al.ICLR 2026 · 34 citations
- Constrained Decoding of Diffusion LLMs with Context-Free GrammarsNiels Mündler, Jasper Dekoninck, Martin VechevICLR 2026 · 18 citations
- Watermarking Diffusion Language ModelsThibaud Gloaguen, Robin Staab, Nikola Jovanović, Martin VechevICLR 2026 · 13 citations
- Rainbow Padding: Mitigating Early Termination in Instruction-Tuned Diffusion LLMsBumjun Kim, Dongjae Jeon, Dueun Kim, Wonje Jeung et al.ICLR 2026 · 11 citations
Builds on20
- Structured Denoising Diffusion Models in Discrete State-SpacesJacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow et al.NeurIPS 2021 · 2,256 citations
- Large Language Diffusion ModelsShen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang et al.NeurIPS 2025 · 949 citations
- Simple and Effective Masked Diffusion Language ModelsSubham S. Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan et al.NeurIPS 2024 · 929 citations
- Argmax Flows and Multinomial Diffusion: Learning Categorical DistributionsEmiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré et al.NeurIPS 2021 · 782 citations
- Simplified and Generalized Masked Diffusion for Discrete DataJiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet et al.NeurIPS 2024 · 693 citations
Related papers
- Beyond Fixed: Training-Free Variable-Length Denoising for Diffusion Large Language ModelsJinsong Li, Xiaoyi Dong, Yuhang Zang, Yuhang Cao et al.ICLR 2026 · 34 citations
- WavefrontDiffusion: Dynamic Decoding Schedule for Improved ReasoningHaojin Yang, Rui Hu, Zequn Sun, Rui Zhou et al.ICLR 2026 · 5 citations
- Block Diffusion: Interpolating Between Autoregressive and Diffusion Language ModelsMarianne Arriola, Aaron Gokaslan, Justin T. Chiu, Zhihan Yang et al.ICLR 2025
- Soft-Masked Diffusion Language ModelsMichael Hersche, Samuel Moor-Smith, Thomas Hofmann, Abbas RahimiICLR 2026 · 15 citations
- Scaling Diffusion Language Models via Adaptation from Autoregressive ModelsShansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye et al.ICLR 2025
