Source-linked AI summary
Predict, Don't Iterate: Efficient Adaptive-Length Infilling for Diffusion Language Models
Haobo Xu, Sirui Chen, Yuanchen Bei, Lingjie Chen, Yuchen Yan, Dongqi Fu, Jingrui He, Hanghang Tong
TL;DR
DLM infilling requires choosing a fixed span length before decoding, although the correct length depends on the content and prior adaptive methods are sensitive to initialization and inefficient. PILL predicts the length without a preset, decodes nearby candidates in parallel, and selects a coherent result with two extra forward passes. Across five DLMs and eight benchmarks, it improves over the strongest baseline on code and text while running 1.82× faster, though scaling to substantially larger backbones remains unverified.
Problem
DLM infilling needs a fixed length before decoding even though the correct span length is unknown and prior adaptive methods are sensitive to initial length and costly.
Method
PILL predicts target length from a mask token’s bidirectional hidden state, decodes candidate lengths in parallel, and selects one with a post-hoc coherence score.
Results
+4.8 average pass rate on code and +6.0 BLEU-2 on text, while PILL runs 1.82× faster than the strongest baseline across five DLMs and eight benchmarks.
Takeaways & Limitations
PILL provides preset-free variable-length infilling for fixed-length DLMs with substantially lower inference overhead than prior adaptive-length baselines.
Takeaways & Limitations
Whether the length probe and post-hoc selection scale to substantially larger diffusion backbones remains to be verified.
Abstract
from arXiv · showhide
Diffusion language models (DLMs) have emerged as a promising alternative to the auto-regressive paradigm. With bidirectional attention and any-order generation, DLMs naturally fit infilling tasks, which require generating a middle span conditioned on both the prefix and the suffix. However, infilling is sensitive to the length of the span, while DLMs require the length to be fixed before generation. Although prior studies extend DLMs to dynamic lengths, they still suffer from two limitations. (i) Sensitivity to initial length. These methods require a preset length to initialize the search and are highly sensitive to this initial length, often yielding suboptimal results. (ii) Inference inefficiency. They either insert length-changing operations during generation or repeatedly search for an appropriate length using multi-step denoising confidence, both of which introduce substantial extra forward passes and computational cost. Therefore, we propose PILL (Probing-based InfiLling with preset-Length-free decoding), an efficient infilling method for DLMs that requires no preset initial length and adds far fewer extra forward passes than baselines, substantially reducing inference time. Experiments show that, across five DLMs spanning different families, architectures, and training recipes on eight infilling benchmarks, PILL improves over the strongest baseline by +4.8 average pass rate on code and +6.0 BLEU-2 on text, while running 1.82x faster than that baseline. The code is available at https://github.com/Hsu1023/PILL.
1 Introduction
DLMs suit infilling because bidirectional attention conditions generation on both prefix and suffix, but their fixed decoding length is unknown and highly consequential. PILL predicts lengths without a preset, decodes candidates in parallel, and selects among them efficiently, improving benchmark performance.
- Bidirectional attention makes DLMs well suited to infilling because they condition middle-span generation on both prefix and suffix.
- Infilling correctness is highly sensitive to span length, while DLMs require that length before decoding and keep it fixed.
- Existing adaptive-length methods depend on preset initial lengths and add substantial inference cost through length-changing operations or repeated confidence searches.
- PILL predicts the target length from a single mask token’s bidirectional hidden state instead of searching for it.
- PILL expands its estimate into candidate lengths, decodes them in parallel, and selects the final span using a post-hoc coherence score.
- +4.8 pass rate on code and +6.0 BLEU-2 on text, with PILL running 1.82× faster than the strongest baseline across five DLMs and eight benchmarks.
2 Related Work
Related work adapts DLM infilling through canvas expansion, explicit length changes, fine-tuning, or confidence-based length search. These approaches involve additional costs or sensitivity to the preset initial length, motivating PILL’s alternative.
- DAEDAL expands the generation canvas during inference but only appends tokens at the sequence end.
- DreamOn introduces explicit length-changing operations but requires fine-tuning the DLM, potentially compromising general-purpose capabilities.
- FlexMDM and DDOT also introduce additional fine-tuning costs.
- CAL searches for suitable lengths using early-step denoising confidence but incurs extra forward passes and remains sensitive to initial length.
3 Preliminaries
DLMs iteratively denoise masked sequences with bidirectional attention, but infilling requires fixing a masked region length before decoding even though the target length is unknown.
- DLMs corrupt sequences by independently replacing tokens with a mask and train a network to reconstruct masked tokens.
- At inference, a prefix and fixed target length initialize a masked canvas that is iteratively denoised by unmasking high-confidence positions.
- Infilling adds a suffix as a right-side condition, while the gold span length remains unknown at inference time.
4 Method
PILL decouples length prediction from content generation through probing, parallel candidate decoding, and one-pass post-hoc selection. Its design uses shared context, slot-wise attention, and interpolated position ids to preserve candidate alignment and reduce scoring cost.
- 4 Method: PILL treats infilling as coupled length determination and content generation, then decouples them through a three-stage design.Stage I predicts a target length, Stage II generates nearby candidates, and Stage III selects one.
- 4.2 Length Probing: A lightweight probe predicts target length from hidden states produced by inserting one mask token between the prefix and suffix.The representation combines nearby prefix states, the mask state, and nearby suffix states; a three-layer MLP maps it to a scalar estimate.
- 4.3 Multi-slot Parallel Decoding: PILL expands the predicted length ˆL into 2r + 1 nearby candidates and decodes them in parallel using separate slots.Each slot is isolated by a slot-wise attention mask while sharing the prefix, middle mask, and suffix context.
- 4.3 Multi-slot Parallel Decoding: Interpolated position ids place each candidate across a shared prefix-to-suffix interval, keeping the suffix anchored consistently across lengths.This avoids position-id gaps associated with padding-based alignment and addresses under-generation caused by extra or misaligned positions.
- 4.4 Post-hoc Selection: Because DLMs cannot use ordinary next-token likelihood for committed outputs, PILL scores candidate and suffix tokens with a single carefully masked forward pass.The selection combines internal coherence and suffix alignment, while sparse suffix selection reduces unnecessary scoring.
5 Experiments
The experiments apply PILL across five DLMs and eight infilling benchmarks covering code and natural-language tasks. The model set spans families, architectures, and training recipes.
- Models: PILL is evaluated on five diffusion language models spanning the LLaDA and Dream families.The models include dense and MoE architectures and both base and instruction-tuned training recipes.
- Datasets: The evaluation covers code and natural-language infilling benchmarks with disjoint datasets.The listed benchmark groups include Python, other programming languages, and text tasks.
LLaDA-8B-Base
On LLaDA-8B-Base, PILL achieves strong infilling quality and efficiency by predicting candidate lengths, decoding them in parallel, and selecting among them post hoc. Ablations show that length probing, interpolated slot positions, and combined coherence scoring are important design choices.
- Main Results: PILL improves over CAL by +4.8 pass rate on code and +6.0 BLEU-2 on text across eight benchmarks.The comparison spans five DLMs and both code and natural-language infilling.
- Efficiency: PILL adds only 7% wall-clock time over fixed-length decoding, compared with +94% for CAL and +117% for DAEDAL.It is 1.82× faster than CAL while achieving 6.6 points higher accuracy on HumanEval-S.
- Ablations: PILL’s oracle score keeps rising with candidate radius, but its performance saturates beyond r = 2.Post-hoc Pass@1 peaks at m = 4 scored suffix tokens and around α = 0.5, combining inner coherence with suffix alignment.
- Length Probing: PILL avoids preset-length sensitivity, while CAL’s MAE rises from 3.56 at l=4 to 24.13 at l=32 and Acc@1 falls from 0.58 to 0.02.PILL achieves MAE 2.50 without tuning and reaches 86.4% Acc@5.
- Slot Position-ID Design: Interpolated slot positions outperform left- and right-padding, reaching Avg Pass@1 51.6 versus 28.2/31.6 and Oracle Pass@1 81.9 versus 73.2/74.2.The gap persists under oracle selection, indicating structural generation errors from padding.
6 Conclusion
PILL converts fixed-length DLMs into adaptive-length infillers through length probing, parallel candidate decoding, and post-hoc coherence selection. It requires no preset length, adds only two forward passes, and outperforms prior adaptive-length baselines across models and tasks while running faster.
- 6 Conclusion: PILL turns fixed-length DLMs into adaptive-length infillers through three lightweight stages.The stages predict target length, decode multiple candidates in parallel, and select using a post-hoc coherence score.
- 6 Conclusion: PILL requires no preset length and adds only two forward passes.Its backbone remains frozen while the method performs probing and scoring around parallel candidate decoding.
- 6 Conclusion: PILL outperforms prior adaptive-length baselines across five DLMs on code and text infilling while running substantially faster.
Limitations
PILL is evaluated across DLMs and infilling benchmarks, but its applicability to more complex editing settings and substantially larger backbones remains unresolved.
- Nested spans and repository-level code patches remain practical editing scenarios for future work.
- Experiments use DLMs with roughly 7B/8B parameters on a single GPU.
- Whether the length probe and post-hoc selection scale to substantially larger diffusion backbones remains to be verified.
A.1 Details of the Length Probe
The length probe is a lightweight, backbone-specific regressor trained on auxiliary masked-span examples, with evidence that it can generalize beyond its training domain.
- The probe is implemented as a 3-layer MLP with hidden dimensions of 512 and 128, using GeLU activation.
- Each backbone uses a separate probe trained on 196k code and text infilling samples.Samples come from Py150, LeetCode, CodeContests, and C4.
- The probe predicts masked-span length from self-supervised examples containing prefix, masked span, and suffix contexts.The original masked-span token length provides supervision, without using evaluation instances or test cases.
- The training recipe uses AdamW regression with MSE loss, learning rate 1 × 10^-3, weight decay 1 × 10^-4, batch size 16, dropout 0.1, and up to 50 epochs.
- PILL outperforms DreamOn across all tested length configurations despite the text tasks being out-of-distribution for its code-trained probe.The comparison covers both datasets.
A.2 Per-Stage Overhead at a Fixed Length
The fixed-length overhead analysis isolates PILL’s added runtime from decoding-length effects by comparing it with a matched-length fixed-length baseline.
- PILL’s runtime is decomposed against a single fixed-length decoding at the same predicted length ˆL.This controlled baseline is denoted Basic Dec.
- The matched-length comparison isolates overhead from the probe and subsequent stages rather than from differences in decoding length.This differs from the end-to-end comparison, which averages backbone decoding over preset lengths L ∈ {4, 8, 16, 32}.
- Table 6 reports per-stage overhead at a fixed length.
A.3 Full Results across Initial Lengths
Across initial-length, backbone, multi-span, robustness, and efficiency evaluations, PILL avoids preset-length tuning while maintaining strong infilling quality and low overhead. Its adaptive probing remains effective across datasets and structurally complex spans, although candidate coverage and selection remain error sources.
- Initial-length sensitivity: Baseline quality varies non-monotonically with initial length, and the best length differs across Human-S, Java, and Wikitext.CAL falls from 70.86 at L=4 to 50.82 at L=32 on Human-S and from 22.69 to 5.35 on Wikitext.
- Multi-span infilling: PILL jointly handles multiple disjoint gaps, improving pass rate by 11.79 points over fixed-length decoding and 4.86 points over CAL.On multi-span MBPP, runtime is 6.70s versus 6.60s for the backbone and 2.6× faster than CAL.
- Robustness: Length probing remains effective under structural complexity, with 81.4% of high-complexity Java and 70.2% of high-complexity C/C++ examples within two tokens of the gold length.Greater complexity moderately increases MAE and reduces near-exact prediction accuracy.
- Probe-training data: PILL surpasses CAL with 10k probe-training examples, while gains become modest beyond 50k samples.On Human-S, PILL reaches 67.96 versus 64.74 for CAL with 10k examples.
- Comparison with oracle initialization: PILL outperforms oracle best-initialized CAL on four of six benchmarks, despite requiring no preset initial length.CAL’s optimal initialization varies from L=8 to L=32 across datasets, and oracle selection uses unavailable test-set performance.