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

AI summary

Overview

Research area: Natural Language Processing — discrete diffusion language models (DLMs), with a focus on code infilling (filling in a missing span given surrounding prefix and suffix context).

Technical level: Advanced. The paper assumes familiarity with masked diffusion training objectives, denoising trajectories, and autoregressive fill-in-the-middle (FIM) baselines.

Scope: The paper introduces DreamOn, a framework that adds two special states ([expand] and [delete]) to masked diffusion language models so they can grow or shrink their output during generation instead of requiring a pre-specified fixed-length mask.

What This Paper Is About

Diffusion language models generate by denoising a masked sequence, but they are locked to a fixed-length mask that the user must specify in advance. Because code completions vary widely in length, a mismatched mask size badly hurts accuracy — the authors report an average performance drop of 38% on HumanEval-Infilling when the predefined mask length does not align with the ground truth length. DreamOn's goal is to let the model decide its own output length during denoising, using its own predictions rather than external guidance, and to do so without changing the model architecture.

Key Contributions

1. Native length control for masked diffusion models. DreamOn augments the diffusion process with two length-control states, [expand] and [delete], so the model can autonomously expand or contract the sequence during generation with no architectural changes.
2. A training recipe with minimal deviation from standard diffusion objectives. The authors train on an augmented sequence built via span merging and [delete] insertion, and add a per-token loss weighting scheme to correct the imbalance caused by [delete] tokens contributing disproportionately to the loss.
3. A variable-length inference procedure plus a training-free acceleration trick. Denoising is driven by a per-step unmasking budget n rather than fixed masking schedulers, capped at a maximum length L_max; "deletion broadcasting" removes all trailing [mask] tokens to the right of a predicted [delete], speeding convergence 2.1 times.
4. Empirical validation across models and benchmarks. Built on Dream-7B, DiffuCoder-7B, and DreamCoder-7B, DreamOn delivers an average absolute performance boost of 26.4% over diffusion baselines, approaches oracle-length performance, and matches or exceeds state-of-the-art autoregressive models on HumanEval-Infilling and SantaCoder-FIM.

Main Findings

• Fixed-length masking is a severe bottleneck. DreamCoder-7B's single-line pass@1 on HumanEval-Infilling drops from 93.3 under an oracle target length to an average of 55.3 across fixed mask lengths of 4, 8, 16, 32, and 64; the multi-line average is 26.0 with an oracle of 69.0.
• DreamOn is stable across initial mask lengths. On single-line infilling it scores 88.7, 90.6, 91.0, 91.6, and 92.1 at mask lengths 4, 8, 16, 32, and 64 (average 90.8), close to its own oracle of 91.6. Multi-line rises from 50.2 to 63.8 as mask length grows (average 57.1, oracle 66.6), and SantaCoder-FIM exact match ranges from 75.0 to 79.0 (average 77.4, oracle 82.0).
• Both mechanisms are individually necessary. Removing deletion (w/o Delete) collapses longer mask lengths — single-line drops to 37.8 at mask length 64 (average 67.4) and multi-line averages 45.5. Removing expansion (w/o Expand) collapses short mask lengths — single-line averages 73.4 and multi-line 32.4, driven by weak scores of 25.1 and 5.5 at mask length 4.
• DreamOn beats diffusion baselines by a large margin. Averaged over the nine reported comparisons, gains over the base diffusion models are 26.4% absolute. Dream-7B improves from 48.2 to 88.6 on single-line, 21.9 to 53.3 on multi-line, and 60.3 to 73.8 on SantaCoder-FIM.
• Diffusion models become competitive with leading autoregressive models. DiffuCoder-7B + DreamOn reaches 92.2 single-line and 63.1 multi-line; DreamCoder-7B + DreamOn reaches 92.1 and 63.8. Both exceed Qwen2.5-Coder-7B's 58.7 on multi-line and approach its 92.6 single-line, and exceed Seed-Coder-8B's 59.3 multi-line.
• Loss balancing matters for deletion. Training without down-weighting [delete] loss falls to an 84.6 average pass@1 on single-line infilling (Table 3), confirming the model otherwise overfits to deletion signals. Keeping [delete] tokens in the sequence instead of removing them performs poorly as well (Table 3 average 84.6; the text describes this variant as averaging 85.3%).
• Deletion broadcasting improves efficiency. Disabling it costs 0.6% average performance and makes deletion steps scale roughly linearly with mask length; with broadcasting, deletion takes roughly 1 step on average and total inference steps at mask length 64 fall from 122.8 to 52.4.
• Scheduler mixture matters for expansion. A static merge scheduler alone reaches an 88.9% pass rate for length-4 masks but is limited on longer masks; a dynamic inverse scheduler alone reaches higher performance on longer masks but drops to 82.5% on length-4 masks. A 1:1 mixture works best, peaking at roughly 90.9% average pass@1, and a merge probability of 0.5 gives roughly 90.5%.
• The method generalizes beyond code. On ROCStories, Dream-7B + DreamOn scores Rouge-L of 31.6, 31.4, 31.3, and 30.6 at mask lengths 4, 8, 16, and 32 for narrative infilling, versus 19.2, 29.8, 26.5, and 18.9 for a supervised fine-tuning baseline. Story ending generation follows the same pattern (24.5, 24.6, 24.4, 24.1 versus 16.3, 25.1, 22.4, 16.7).

Methodology in Plain English

The researchers start from a standard masked diffusion setup: a sequence of tokens is progressively corrupted into [mask] tokens and the model learns to recover them. Their key move is to build an augmented training sequence from each clean example. Some spans of the data are replaced with a special [expand] token, which during generation deterministically turns into two [mask] tokens (growing the sequence). Some positions are marked with a [delete] token, which removes that position from the sequence (shrinking it). Because both special states are always mapped to [mask] in the forward noising process, the ordinary diffusion loss automatically trains the model to predict them too — no new architecture, just two new vocabulary entries.

Which spans become [expand] is controlled by a merging probability. The authors test a static scheduler with a fixed merge probability and a dynamic inverse scheduler whose merge probability falls when there are many masks, then use a 1:1 mix of both. For contraction, they simply append a random number of [delete] tokens (0 to 64) to the middle segment during training. Because each [delete] maps to a single [mask] whereas many [mask] tokens merge into one [expand], they reweight the loss so [delete] predictions collectively count the same as one [mask] prediction.

At inference, the model starts from a masked span and repeatedly predicts tokens at the most confident masked positions, controlled by a per-step unmasking budget rather than a fixed scheduler. Predicted [expand] tokens become two masks; predicted [delete] tokens are removed, and any trailing masks to their right are removed as well. Growth is capped at L_max = 128.

Models were fine-tuned on the education-instruction subset of OpenCoder SFT data — about 110K Python instruction-solution pairs — for 10 epochs with batch size 128, maximum context length 1024, learning rate 1e-5, cosine decay with 10% warmup, roughly 5 hours on 8 H800 GPUs. That compute is 0.15% of pretraining a base model. Inference uses temperature 0.2 and top_p 0.9, with the mask length set to 64 by default for all diffusion models compared.

Why This Matters

Impact on research. The paper removes a structural barrier that has kept diffusion language models out of practical infilling work. It shows that variable-length generation can be obtained through minimal changes to the training objective and no architectural changes, and it provides a controlled comparison against oracle-length performance, which is a useful reference point for the field. It also reinforces the argument that DLMs do not need autoregressive-style FIM workarounds (permuting the target span to the end of the sequence) to compete on infilling.

Real-world applications:

• Code completion in IDEs, where the length of the needed insertion is unknown and varies from a single token to a full function body.
• Fill-in-the-middle refactoring and repair tools, where a developer selects surrounding context and asks a model to supply a missing block.
• Docstring or test generation inserted between existing code sections whose size is not known in advance.
• Narrative or document completion, as shown by the ROCStories experiment, where the natural length of missing content varies by example.

Industry relevance. The reported training cost — roughly 5 hours on 8 H800 GPUs, and 0.15% of base-model pretraining compute — means existing open-weight diffusion models can be adapted rather than retrained. The method is described as model-agnostic and was validated on Dream-7B, DiffuCoder-7B, and DreamCoder-7B at the 7B scale. The deletion broadcasting mechanism also directly targets serving cost by cutting inference steps roughly in half at mask length 64.

Future Directions

• Broaden evaluation beyond code. The authors state their evaluation is currently limited to code infilling tasks requiring variable-length generation, and that future work should assess generalizability to broader applications.
• Develop a more principled inference formulation. The current training and inference procedures rely on heuristics for variable-length generation; the authors call for a more principled formulation for flexible inference in masked diffusion models.
• Richer expansion schemes. The current design uses a single [expand] state that expands into exactly two [mask] tokens, which requires multiple expansion steps when the target is much longer than the initial mask span. The authors suggest multiple expansion factors or an explicit length-prediction head, while noting this enlarges the decision space and requires careful rebalancing of the training objective.
• Investigate the performance gap to oracle. DreamOn approaches but does not reach oracle-length performance on every benchmark — for example, 63.8 versus an oracle of 66.6 on multi-line infilling and 79.0 versus 82.0 on SantaCoder-FIM — leaving open how much of that gap can be closed.

Target Audience

Researchers and practitioners working on diffusion language models, code generation, and infilling; engineers who want to adapt open-weight diffusion models for variable-length completion without architecture changes; and readers interested in how non-autoregressive models can be made competitive with autoregressive models on bidirectional- context tasks. Some background in masked diffusion training objectives and FIM evaluation will make the paper substantially easier to follow.

Authors’ 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.