DreamOn: Diffusion Language Models For Code Infilling Beyond Fixed-size Canvas

TL;DR

DreamOn addresses the fixed-length mask bottleneck in diffusion language models for code infilling by introducing two length-control states, [expand] and [delete], that enable dynamic, end-to-end length adaptation without architectural changes. It augments the diffusion process with an auxiliary augmented latent and trains on augmented transitions, using a balanced loss to learn when to expand or delete tokens. During inference, DreamOn denoises a configurable number of masked tokens per step, applying expansions and deletions to adjust sequence length with a maximum cap $L_{max}$, and employs span-merging schedulers and deletion broadcasting to accelerate convergence. Empirically, DreamOn achieves competitive or superior performance to state-of-the-art autoregressive models on HumanEval-Infilling and SantaCoder-FIM, closely matching oracle-length results and delivering substantial improvements over baseline diffusion models. The approach significantly broadens the practical applicability of DLMs for variable-length generation with minimal training or architectural changes.

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.

Figures (7)

• Figure 1: Example of DreamCoder-7B failing at code infilling due to the length mismatch between masked input and ground truth. Incorrect and correct code is marked in red and green. With too few masked tokens, diffusion models lack sufficient room for meaningful code infilling. Too many masks cause over-generation of unnecessary code snippet (e.g., depth $>$ 0 that is incorrect).
• Figure 2: Overview of the augmented diffusion process. Top: the forward augmentation-and-noising procedure maps the input sequence $\boldsymbol{\mathbf{x}}_0$ to an augmented latent $\boldsymbol{\mathbf{z}}_0$ containing [expand] and [delete] states, and then applies a standard masked diffusion process over $\boldsymbol{\mathbf{z}}_0$ to obtain $\boldsymbol{\mathbf{z}}_t$ and eventually $\boldsymbol{\mathbf{z}}_T$. Bottom: a single denoising step where [mask] positions in $\boldsymbol{\mathbf{z}}_t$ can be predicted as either regular tokens or special states; [expand] deterministically expands into two [mask] tokens, while [delete] will remove the corresponding position, yielding a new sequence $\boldsymbol{\mathbf{z}}_{t-1}$ with a different length from $\boldsymbol{\mathbf{z}}_t$.
• Figure 3: Ablation on merging rate scheduler design choices.
• Figure 4: Average generation steps of DreamCoder + DreamOn on multi-line infilling subset.
• Figure 5: Performance on single-line subset of HumanEvalInfilling-FIM with different hyperparameters during training. The performance is computed as the average pass@1 with mask length 4, 8, 16, 32 and 64.