Source-linked AI summary
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
TL;DR
DLMs are hindered by fixed-length masks that degrade code infilling when completion lengths differ. DREAMON adds prediction-driven expansion and deletion states without architectural changes, achieving competitive benchmark performance and near-oracle results while remaining limited to code infilling and heuristic length control.
Problem
Fixed-length masks prevent DLMs from dynamically determining output length and sharply degrade code infilling when mask and target lengths mismatch.
Method
DREAMON augments diffusion with [expand] and [delete] states that adapt sequence length from model predictions without architectural changes.
Results
DREAMON achieves a 26.4% average absolute improvement over diffusion baselines, matches oracle-length performance, and reaches competitive or superior performance to autoregressive models across infilling benchmarks.
Takeaways & Limitations
DREAMON removes the fixed-length bottleneck for practical variable-length code infilling while preserving the diffusion framework.
Takeaways & Limitations
Evaluation is limited to code infilling, while the method relies on heuristics and a binary expansion state requiring multiple steps for much longer completions.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Diffusion language models support flexible, any-order infilling, but fixed-length masks make code completion highly sensitive to mismatched target lengths. DREAMON addresses this bottleneck with prediction-driven length adaptation and reports strong infilling performance across benchmarks.
- DLMs use multi-step denoising for flexible, any-order generation and are inherently suited to infilling.
- A fixed-length mask prevents DLMs from dynamically determining output length, which is especially damaging when code completion lengths vary.
- 38% average performance drop occurs on HumanEval-Infilling when the predefined mask length mismatches the ground-truth length.
- DREAMON dynamically expands or contracts masked sequences using [expand] and [delete] states without architectural modifications.
- 26.4% average absolute improvement over diffusion baselines is reported across multiple infilling benchmarks, with performance matching oracle length and approaching or reaching autoregressive results.
2 PRELIMINARY
Masked diffusion progressively corrupts token sequences and learns reverse transitions to reconstruct them. The framework can use continuous-time parameterization and trains with a weighted loss evaluated on masked positions.
- Masked diffusion replaces tokens with an absorbing [mask] state according to a noise schedule that approaches a fully masked sequence.
- The generative model starts from x_T and applies parameterized reverse transitions that approximate the posterior toward x_0.
- Continuous-time parameterization t ∈ [0, 1] avoids bias from predefined discretization over diffusion time steps.
- The training objective evaluates loss only on masked positions and weights it using a time-dependent term derived from the noise schedule.
3 METHOD
DREAMON extends masked diffusion with adaptive length control, allowing sequences to expand or contract during generation without architectural changes. Its training augments standard diffusion with special states, while inference uses model predictions to adjust length dynamically.
- Framework: DREAMON introduces [expand] and [delete] states so the model can increase or decrease sequence length during generation.[expand] replaces one token with two [mask] tokens, while [delete] removes a token.
- Training: Training constructs augmented sequences by merging masked spans into [expand] and inserting [delete] tokens.The augmented sequence is then processed through masked diffusion so both special states become prediction targets.
- Training: DREAMON downweights [delete] predictions to compensate for their greater frequency relative to [expand] targets.The normalization preserves consistent loss magnitude across sequences with different numbers of deletions.
- Inference: During inference, the model predicts masked positions, selectively remasks tokens by entropy, and applies [expand] or [delete] predictions to change sequence length.The procedure uses an initial length, maximum length, unmasking budget, and sampling temperature.
- Implementation: DREAMON combines static and dynamic inverse span-merging schedulers during training to control how masked spans are merged.The dynamic inverse scheduler reduces merging when more [mask] tokens are present.
- Implementation: Deletion broadcasting accelerates inference by removing subsequent masked tokens when the model predicts [delete].This training-free adaptation reduces unnecessary computation without sacrificing generation quality.
4 EXPERIMENTS
The experiments evaluate DREAMON on code-infilling benchmarks against autoregressive and diffusion baselines. DREAMON improves over diffusion baselines and reaches competitive performance with autoregressive models, including on multi-line infilling.
- Setup: The study fine-tunes Dream-7B, DiffuCoder-7B, and DreamCoder-7B on about 110K Python instruction-solution pairs.The data come from the education-instruction subset of OpenCoder SFT data.
- Setup: Evaluation covers HumanEval-Infilling single-line and multi-line subsets and the Python subset of Santacoder-FIM.The reported metrics are pass@1 for HumanEval-Infilling and exact match for Santacoder-FIM.
- Results: 26.4% average absolute improvement over diffusion baselines is achieved by DREAMON across the reported infilling benchmarks.The result is presented as evidence of a model-agnostic enhancement over fixed-length diffusion models.
- Results: DREAMON-equipped DiffuCoder-7B and DreamCoder-7B match Qwen2.5-Coder-7B and surpass it on the more challenging multi-line infilling benchmark.The comparison is reported against leading autoregressive models.
- Results: Table 1 compares pass@1 and exact match for open-source autoregressive and diffusion model baselines.Best and second-best diffusion-model results are distinguished in the table.
- Results: Table 2 compares infilling performance across diffusion-model designs using oracle target length as a reference.An AST parser is used for the dagger-marked exact-match results to normalize large syntactic differences.
5 ANALYSIS
The analysis shows that DREAMON’s performance depends on combining expansion and deletion, while scheduler and broadcasting choices affect robustness and inference efficiency across mask lengths.
- Performance Breakdown: DREAMON achieves near-oracle performance across a wide range of initial mask lengths, with gains from combining expansion and contraction.It remains stable on single-line infilling and SantaCoder-FIM regardless of the initial mask length.
- Ablation on Length Control: Removing deletion sharply harms longer masks through over-generation, whereas removing expansion harms short masks by preventing longer completions.Expansion also provides smaller benefits on long masks through finer-grained length adjustment.
- Expansion Mechanism Design: 88.9% pass rate is achieved by the static scheduler on length-4 masks, but its performance is limited on longer masks.The static scheduler uses a fixed merge probability.
- Expansion Mechanism Design: The dynamic inverse scheduler performs better on longer masks but drops to 82.5% on length-4 masks, while a 1:1 mixture balances performance across mask lengths.It merges less when more mask tokens are present.
- Deletion Mechanism Design: 84.6% average pass@1 results when loss balancing is removed, while persistent deletion tokens yield 85.3% average performance.Disabling deletion broadcasting reduces performance by 0.6% on average, especially for overly long masks.
- Efficiency Analysis: Deletion broadcasting reduces inference steps from 122.8 to 52.4 at mask length 64 without affecting output quality.It converts token-by-token deletion into a length-prediction action and accelerates generation by 2.1×.
6 RELATED WORK
Related work positions DREAMON among diffusion, non-autoregressive, and variable-length generation methods. Its distinction is native length control for masked diffusion with minimal training changes and no architectural modification.
- Code Infilling with Autoregressive Models: Autoregressive code infilling commonly relocates the target span to the sequence end so decoder-only models can generate it left to right.This workaround preserves standard autoregressive architectures while accommodating bidirectional infilling context.
- Discrete Diffusion Language Models: Discrete diffusion models generate token sequences through iterative bidirectional denoising rather than strictly left-to-right decoding.This line of work forms the broader modeling basis for DREAMON.
- Non-autoregressive Models with Length Control: Prior non-autoregressive methods address variable length with separate length predictors, latent alignment marginalization, or edit operations.These strategies target sequence expansion or contraction through mechanisms distinct from DREAMON’s masked-diffusion integration.
- Non-autoregressive Models with Length Control: Recent methods jointly model edits, token states, positions, or insertion and unmasking rates to support variable-length generation.The cited approaches include Edit flows, DDOT, and FlexMDM.
- DREAMON: DREAMON adds native length control to masked diffusion with minimal additional training and no architectural modifications.This design directly adapts pretrained diffusion language models and avoids multistage inference pipelines.
7 CONCLUSIONS
The conclusion presents DREAMON as a no-architecture-change framework for dynamic length control and reports near-oracle, competitive infilling performance. It also identifies limited evaluation scope, heuristic procedures, and expansion-step inefficiency as open boundaries.
- 7 CONCLUSIONS: DREAMON uses [delete] and [expand] states to dynamically control sequence length without architectural changes.The model learns expansion and contraction from auxiliary length-control states based on its own predictions.
- 7 CONCLUSIONS: DREAMON approaches oracle-length performance and achieves competitive results with state-of-the-art autoregressive models.The conclusion frames these results as evidence for more practical and flexible diffusion language models.
- Limitations: Evaluation is limited to code infilling, leaving DREAMON’s generalizability to broader applications for future work.The authors also identify heuristic training and inference procedures as a limitation.
- Limitations: A single [expand] state deterministically creates two [mask] tokens, requiring multiple expansion steps when completions greatly exceed the initial mask span.Richer expansion factors or an explicit length-prediction head could reduce denoising iterations but enlarge the decision space.
ETHICS STATEMENT
The ethics statement reports no human-subject, private-information, or human-evaluation components and relies on publicly available benchmark datasets.
- ETHICS STATEMENT: The study uses publicly available benchmarks, reports no human subjects or private information, and includes no human evaluation.The authors state that dataset usage licenses were followed.
- ETHICS STATEMENT: The authors report no foreseeable risks of misuse or societal harm from the method.
REPRODUCIBILITY STATEMENT
The authors provide detailed experimental setups and hyperparameter configurations to support reproducibility. LLMs were used mainly for polishing text, not for core research activities.
- Detailed experimental setups and hyperparameter configurations are provided in §4.
- LLMs were used primarily to correct grammar and improve clarity.
- LLMs did not significantly contribute to idea generation, experimental design, or substantive manuscript writing.
B GENERALIZABILITY BEYOND CODE INFILLING
DREAMON is evaluated beyond code infilling on ROCStories using narrative infilling and prefix-guided settings. Across varying initial mask lengths, it outperforms the baseline and maintains stable generation quality.
- DREAMON and an SFT baseline are evaluated on ROCStories for narrative infilling and prefix-guided generation.
- The baseline degrades when the initialized mask length mismatches the natural length of missing content, especially at lengths 4 and 32.
- DREAMON achieves superior Rouge-L performance across varying initial mask lengths in both evaluation settings.
- Generation quality remains consistent regardless of initial mask length, supporting generalizability to variable-length natural-language tasks.
C ABLATION FOR HYPERPARAMETERS
The ablations examine scheduler mixtures and DREAMON’s two complementary length-control states. Performance is evaluated across multiple initial mask lengths using average pass@1, while expansion and deletion respectively address under- and over-length generations.
- Scheduler hyperparameters: 90.9% Pass@1 is reached with a balanced 1:1 mixture of static and dynamic inverse schedulers for single-line infilling.
- Evaluation setup: Performance is computed as average pass@1 across mask lengths 4, 8, 16, 32, and 64.
- Length adaptation: The [expand] state enables generation longer than the initial mask, preventing truncation of comprehensive sequences.
- Length adaptation: The [delete] mechanism removes redundant mask tokens to approach the predicted sequence length and streamline generation.