💬 NLP

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

DreamOn is a novel diffusion framework that overcomes the fixed-length generation limitation of existing Diffusion Language Models by introducing dynamic length control states, enabling flexible, variable-length code infilling that achieves performance comparable to state-of-the-art autoregressive models without requiring architectural changes.

Original authors: Zirui Wu, Lin Zheng, Zhihui Xie, Jiacheng Ye, Jiahui Gao, Shansan Gong, Yansong Feng, Zhenguo Li, Wei Bi, Guorui Zhou, Lingpeng Kong

Published 2026-02-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zirui Wu, Lin Zheng, Zhihui Xie, Jiacheng Ye, Jiahui Gao, Shansan Gong, Yansong Feng, Zhenguo Li, Wei Bi, Guorui Zhou, Lingpeng Kong

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to fill in a blank space in a story or a piece of code. You have the beginning (the prefix) and the end (the suffix), and you need to write what goes in the middle.

For a long time, the smartest computer programs (called Autoregressive Models) have done this by writing one word at a time, from left to right, like a person typing a sentence. They are very good at it.

Then, a new type of program called a Diffusion Language Model (DLM) arrived. Think of these like a sculptor. Instead of writing word-by-word, they start with a block of clay where the middle part is completely hidden (masked). They chip away at the hidden parts, refining the shape until the perfect words appear. This is great because they can fix mistakes in any order, not just left-to-right.

The Problem: The "Fixed-Size Canvas" Trap
However, these sculptors had a major flaw. They could only work on a block of clay that was exactly the size you told them to use.

If you gave them a small block (say, 4 hidden words) but the answer needed 20 words, they would run out of space and give you a broken, incomplete answer.
If you gave them a huge block (say, 64 hidden words) but the answer only needed 4 words, they would get confused and fill the extra space with nonsense or repetitive garbage.

The paper calls this the "Fixed-Size Canvas" problem. The model is stuck waiting for you to guess the exact right size of the hole before it starts working. If you guess wrong, the result is bad.

The Solution: DREAMON
The authors introduce DREAMON, a new way to teach these sculptors how to change the size of their clay block while they are working.

They added two special "magic commands" to the model's vocabulary:

[expand]: "I need more space! Turn this one hidden spot into two hidden spots."
[delete]: "I have too much space! Remove this hidden spot entirely."

How It Works (The Analogy)
Imagine the model is drawing a picture inside a frame.

Old Way: You tell the model, "Draw inside this 10x10 square." If the picture needs to be 20x20, the model is squished. If it needs to be 2x2, the model is wasted space.
DREAMON Way: You give the model a frame, but you let it hold a remote control. As it draws, if it realizes the picture is getting too big for the frame, it presses [expand] and the frame magically grows. If it realizes the picture is tiny and the frame is huge, it presses [delete] and the frame shrinks.

The model does this entirely on its own, based on what it thinks the answer should look like. It doesn't need you to tell it the length beforehand.

The Results
The researchers tested this on coding tasks (filling in missing parts of computer programs).

Before DREAMON: The diffusion models were terrible at this if the "hole" size didn't match the answer perfectly. Their performance dropped by about 38% just because the size was slightly off.
After DREAMON: The models became as good as the best "word-by-word" writers (the autoregressive models). They could handle holes of any size, growing or shrinking the answer until it was just right.
Efficiency: They also added a trick called "Deletion Broadcasting." If the model decides to delete a huge chunk of the frame, it doesn't do it one tiny piece at a time; it deletes the whole chunk instantly, making the process much faster.

In Summary
DREAMON takes a powerful but rigid tool (Diffusion Models) and gives it the flexibility to stretch and shrink its output. It removes the need for humans to guess the perfect answer length, allowing these models to finally compete with the best code-writing AI we have today, but with the added benefit of being able to fix things in any order.

Technical Summary: DREAMON – Diffusion Language Models for Code Infilling Beyond Fixed-Size Canvas

1. Problem Statement

Diffusion Language Models (DLMs) offer a promising alternative to autoregressive models by enabling flexible, any-order generation, which is theoretically ideal for code infilling tasks (generating content between a prefix and suffix). However, current DLMs suffer from a critical limitation: they require a pre-specified, fixed-length masked sequence for generation.

This constraint creates a "fixed-size canvas" bottleneck. In code infilling, the ideal completion length varies significantly across examples. When the predefined mask length does not match the ground-truth solution length, performance degrades drastically. The paper reports that Dream-Coder-7B suffers an average performance drop of 38% on the HumanEval-Infilling benchmark when the mask length mismatches the target. Specifically:

Too few masks: The model lacks sufficient space to generate meaningful code, leading to incomplete solutions.
Too many masks: The model tends to over-generate unnecessary or incorrect code snippets to fill the excess space.

Existing solutions often rely on specialized prompting or workarounds that disrupt natural contextual structures, but DLMs have lacked a native mechanism to dynamically adjust output length.

2. Methodology: DREAMON

The authors propose DREAMON, a framework that enables dynamic, variable-length generation in discrete diffusion models without architectural modifications. The core innovation is the introduction of two length-control states integrated into the diffusion process.

2.1 Core Mechanism: Special States

DREAMON augments the vocabulary with two special sentinel tokens:

[expand]: When predicted, this token is deterministically expanded into two [mask] tokens, effectively increasing the sequence length and providing more room for generation.
[delete]: When predicted, this token removes the corresponding position from the sequence, effectively contracting the length.

2.2 Training Procedure

The training process involves data augmentation and loss re-weighting:

Augmented Sequence Construction ( $z_0$ ): From the original input sequence ( $x_0$ ), the authors construct an augmented sequence $z_0$ . Random token spans are merged into [expand] tokens, and [delete] tokens are inserted. This simulates the need for length adjustment during the forward diffusion process.
Forward Process: Standard masked diffusion is applied to $z_0$ . Crucially, both [expand] and [delete] states are mapped to [mask] during the forward noising process, ensuring they contribute to the learning signal.
Loss Weighting: A significant imbalance exists where [delete] targets are more frequent than [expand] targets (since one [delete] replaces one token, while one [expand] represents a merge of multiple tokens). To prevent the model from overfitting to deletion signals, a loss weighting scheme is introduced. The loss contribution of [delete] predictions is down-scaled so that the total weight of all [delete] predictions in a sequence equals the weight of a single standard [mask] prediction.

2.3 Inference Procedure

During inference, the model performs standard denoising but with dynamic length control:

Adaptive Trajectory: The model predicts all masked positions. If [expand] is predicted, the position is replaced by two [mask] tokens. If [delete] is predicted, the token is removed.
Deletion Broadcasting: To improve efficiency when the initial mask is much larger than the target, the authors introduce a heuristic: if the model predicts a [delete], all subsequent tokens to the right (if they are all [mask]) are immediately eliminated. This transforms a token-by-token deletion process into a length prediction action, significantly reducing inference steps.
Constraints: A maximum output length ( $L_{max}$ ) is enforced to prevent unbounded growth.

3. Key Contributions

Dynamic Length Control: DREAMON is the first framework to enable native, variable-length generation in masked diffusion models using only two special tokens ([expand] and [delete]) without changing the model architecture.
Minimal Training Overhead: The method requires only minimal modifications to the training objective (auxiliary states and loss weighting) and can be applied to existing pre-trained DLMs (e.g., Dream-Coder-7B, DiffuCoder-7B) with negligible additional compute (0.15% of pretraining cost).
Performance Recovery: The approach allows diffusion models to achieve performance on par with state-of-the-art autoregressive models, effectively removing the fixed-length bottleneck.

4. Experimental Results

The authors evaluated DREAMON on HumanEval-Infilling (single-line and multi-line) and SantaCoder-FIM (Python subset).

Baseline Comparison: Standard diffusion models (Dream-7B, DiffuCoder-7B, DreamCoder-7B) lagged significantly behind autoregressive baselines (e.g., Qwen2.5-Coder-7B, Deepseek-Coder-6.7B) due to fixed-length constraints.
DREAMON Performance:
- DREAMON achieved an average absolute performance boost of 26.4% over diffusion baselines.
- On HumanEval-Infilling, DREAMON-enhanced models (e.g., DiffuCoder-7B + DREAMON) reached 92.2 Pass@1 (single-line) and 63.1 Pass@1 (multi-line), surpassing or matching top autoregressive models like Qwen2.5-Coder-7B.
- DREAMON variants achieved performance near-oracle levels (performance when the initial mask length matches the ground truth), demonstrating that the model can successfully adapt to the correct length regardless of the initial mask size.
Ablation Studies:
- Removing either [expand] or [delete] caused significant performance drops, confirming both mechanisms are necessary.
- Loss balancing was critical; without it, performance dropped by ~7% on average due to overfitting to deletion.
- Deletion broadcasting reduced total inference steps by roughly 50% (e.g., from 122.8 to 52.4 steps at mask length 64) without sacrificing quality.

5. Significance and Claims

The paper claims that DREAMON removes a fundamental barrier to the practical deployment of Diffusion Language Models. By enabling dynamic length adjustment, DREAMON significantly advances the flexibility and applicability of DLMs for variable-length generation tasks, particularly code infilling.

The authors position their work as a step toward making DLMs competitive with autoregressive models in real-world scenarios where output length is unknown. They note that while their current design uses a single expansion factor (one-to-two), future work could explore richer expansion vocabularies or explicit length-prediction heads to further optimize efficiency. The work emphasizes that these gains are achieved with no architectural changes to the underlying transformer, preserving the simplicity of the original diffusion framework.