Unlocking Prompt Infilling Capability for Diffusion Language Models

Yoshinari Fujinuma; Keisuke Sakaguchi

cs.CLApr 4, 2026

Unlocking Prompt Infilling Capability for Diffusion Language Models

Authors: Yoshinari Fujinuma, Keisuke Sakaguchi

Organizations: Patronus AI · Tohoku University, RIKEN

Abstract

Masked diffusion language models (dLMs) generate text through bidirectional denoising, yet this capability remains locked for infilling prompts. This limitation is an artifact of the current supervised finetuning (SFT) convention of applying response-only masking. To unlock this capability, we extend full-sequence masking during SFT, where both prompts and responses are masked jointly. Once unlocked, the model infills masked portions of a prompt template conditioned on few-shot examples. We show that such model-infilled prompts match or surpass manually designed templates, transfer effectively across models, and are complementary to existing prompt optimization methods. Our results suggest that training practices, not architectural limitations, are the primary bottleneck preventing masked diffusion language models from infilling effective prompts

Figures & tables

Figure 1: Overview of the prompt infilling procedure and the change in training. (1) Gather few-shot examples with prompt templates and reference responses. (2) A diffusion LM (dLM) infills masked tokens in the prompt template, conditioned on the reference responses. (3) Validate infilled prompts by generating responses across all few-shot examples using either a dLM or an LLM. (4) The best infilled prompt is used for final inference on all inputs.

Model	Prompt	#Shots	EM
LLaDA	“Q: … A: …”	0	70.8
	+ICL	16	73.1
	+Prepend Infilled Tokens	16	69.8
Dream	“Q: … A: …”	0	81.0
	+ICL	16	79.3
	+Prepend Infilled Tokens	16	54.3

Table 1: GSM8K results on public checkpoints (no fine-tuning). In both models, infilled prompts with the publicly available checkpoints decreases the exact match (EM) score (%).

Panel of Figure 2 — Figure 2: Example prompt with infilled tokens using the public LLaDA checkpoint ( LLaDA-8B-Instruct ) on GSM8K. All masked prefix tokens are filled with EOS tokens, confirming the training-inference gap for prompt infilling.

LLaDA
Train	Prompt	Pear.	Spear.	Kend.
None	G-eval	$.397_{.000}$	$.289_{.000}$	$.217_{.000}$
	C-gen	$\underline{.465}_{.000}$	$\underline{.439}_{.000}$	$\underline{.339}_{.000}$
	C-gen Infill	- $.005_{.043}$	- $.013_{.065}$	- $.009_{.044}$
	Judge	$.308_{.000}$	$.336_{.000}$	$.261_{.000}$
	Judge Infill	$.050_{.063}$	$.087_{.069}$	$.064_{.053}$

Table 2: Model performance comparison with different prompting strategies on SummEval dataset. Bold indicates best performance within each model; underline indicates best prompt per (Model, Train) setup.

Panel of Figure 3 — Figure 3: Example of a score rubric with infilled tokens by the model (right) and original score rubric (left). The model replaces descriptions with non-uniform score values (e.g., 1.2, 1.8, 4.8) inferred from few-shot examples, encouraging float score outputs which results in better human correlation.

Train	Pear. $\uparrow$	Spear. $\uparrow$	Kend. $\uparrow$	PPL $\downarrow$
Public	- $.007_{.044}$	$.078_{.058}$	$.068_{.050}$	$7.40$
RO	$.205_{.072}$	$.406_{.091}$	$.168_{.079}$	$7.14$
FS	$\mathbf{.312}_{.013}$	$\mathbf{.548}_{.026}$	$\mathbf{.263}_{.014}$	$\mathbf{3.73}$
FS+RO	$.259_{.039}$	$.525_{.049}$	$.210_{.043}$	$3.78$

Table 3: Diffusion perplexity analysis with LLaDA. (a) BigGen-Bench results using Judge Infill prompt across finetuning stages. PPL: diffusion perplexity with gold responses. Full-sequence masking (FS) achieves best performance with lowest variance and lowest perplexity. (b) Diffusion perplexity (

\downarrow

) on SummEval with gold responses ( [RESULT] {score} ) using finetuned and non-finetuned variants of LLaDA-8B-Instruct.

Panel of Figure 4 — Figure 4: Prompt transfer experiments. (a) GSM8K: infilled prompts achieve higher accuracy with fewer tokens than ICL. (b) SummEval: infilled prompts transfer across LLaDA variants with different training configurations.

Eval Model	Method	Proposal	Pear.	Spear.	Kend.
Prometheus 8x7B	Baseline	–	$.428_{.000}$	$.410_{.000}$	$.372_{.000}$
	COPRO	8x7B	$.465_{.098}$	$.447_{.073}$	$.355_{.061}$
	+ Infill	LLaDA 8B	$\mathbf{.539}_{.020}$	$\mathbf{.490}_{.010}$	$\mathbf{.394}_{.011}$
	GEPA	8x7B	$.377_{.169}$	$.356_{.147}$	$.330_{.136}$
	+ Infill	LLaDA 8B	$.401_{.148}$	$.391_{.133}$	$.315_{.108}$
Prometheus 13B	Baseline	–	$.271_{.000}$	$.248_{.000}$	$.216_{.000}$

Table 4: Effect of LLaDA infilling on COPRO and GEPA prompts across Prometheus models (8x7B v2.0 and 13B v1.0) on SummEval (16-shot). Infilling improves both methods, but is far more effective on COPRO: variance largely drops (

.073\rightarrow.010

) while GEPA barely changes (

.147\rightarrow.133

). COPRO + Infill on 8x7B achieves the best result overall.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Full prompt transfer evaluation on SummEval using LLaDA models trained with different configurations (mean

\pm

std), showing Pearson, Spearman, and Kendall correlations. Baseline results from Table 2 showing Judge Infilled prompts results across different training stages. The transferred prompt improves on Public and RO models showing that the infilled prompts transfer across different LLaDA variants.

Hyperparameter	Value
Learning rate (LLaDA)	$1\times 10^{-5}$
Learning rate (Dream)	$2\times 10^{-6}$
LR schedule	Cosine decay
Batch size (LLaDA)	1 (32 grad. accum. steps)
Batch size (Dream)	2 (64 grad. accum. steps)
Training epochs (FS)	8

Appendix

Table 5: Training hyperparameters for full-sequence masking (FS) and response-only masking (RO).

Parameter	Value
Diffusion steps	$128$
Response length	$128$
Temperature (LLaDA)	$0.8$
Temperature (Dream)	$0.1$

Appendix

Table 6: Inference configuration.

Figure 6: GSM8K exact match rate vs. prompt length for Dream public checkpoint. The infilled prompt (red star) achieves the best exact match accuracy (81.3%) with only 463 tokens on average, outperforming all in-context learning approaches while using significantly fewer tokens than 8-shot (2117 tokens) and 16-shot (4169 tokens) configurations.

Model	Recall	F1	Label Acc
LLaDA (Public)	$.560_{.002}$	$.268_{.001}$	$.435_{.010}$
LLaDA (FS)	$.556_{.001}$	$.279_{.001}$	$.550_{.003}$
LLaDA (RO)	$.533_{.001}$	$.163_{.001}$	$.503_{.013}$
LLaDA (FS+RO)	$.547_{.001}$	$.211_{.000}$	$.604_{.014}$

Appendix

Table 7: HoVer few-shot evaluation results (mean

\pm

std over 3 runs). Public: no additional training.

Diffusion Language Models

Diffusion language models generate and refine text through iterative denoising, enabling bidirectional context and parallel token updates.

See all 443 papers

LED: Improving Masked Diffusion Language Models by Lookahead Editing
Oct 16, 2026Yuyan Zhou, Zeyuan Yang, Yuxuan Song +5Diffusion Language Model Decoding Masked Diffusion Models
Department of Computer Science and Engineering, The Hong Kong University of Science and Technology · ByteDance Seed · Institute for AI Industry Research (AIR), Tsinghua University +1
DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling
May 22, 2026Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3Diffusion Language Models Diffusion Language Model Distillation
NVIDIA
LogicDiff: Logic-Guided Denoising Improves Zero-Shot Reasoning in Masked Diffusion Language Models
Mar 24, 2026Shaik AmanMasked Diffusion Models Diffusion Language Model Decoding
Independent Researcher, Nellore, Andhra Pradesh, India

Oct 16, 2026J/K move · Enter open · S save

LED: Improving Masked Diffusion Language Models by Lookahead Editing

Yuyan Zhou, Zeyuan Yang, Yuxuan Song, Zhicheng Cai+4

Department of Computer Science and Engineering, The Hong Kong University of Science and Technology · ByteDance Seed · Institute for AI Industry Research (AIR), Tsinghua University+1

Diffusion Language Model Decoding Masked Diffusion Models Parallel Decoding for Diffusion Language Models Diffusion Model Fine-Tuning

Open paper AskPDF arXiv

Diffusion large language models (dLLMs) have emerged as a promising alternative to autoregressive language models, enabling non-autoregressive text generation through parallel decoding. Existing dLLMs can be broadly categorized into two families: masked diffusion models (MDMs) and uniform diffusion models (UDMs) In this work, we theoretically show that MDMs are easier to optimize due to their mask-based transition schedule but have lower computational expressivity, whereas UDMs enable iterative state revision and greater expressivity but are substantially harder to train. This trade-off raises a natural question: can we obtain the expressivity of iterative revision without giving up the simple, easy-to-optimize transition schedule of MDMs? Motivated by this analysis, we propose Lookahead Editing (\ours{}), which augments masked diffusion decoding with a rewritable proposal state. At each decoding step, the model performs several lookahead updates to iteratively refine future predictions in this proposal state before committing tokens to the masked sequence. This design introduces a rewritable proposal that uses additional compute to improve expressivity while keeping the simple mask-based transition schedule of MDMs. We fine-tune Qwen2.5-7B and Qwen3-8B to obtain LED-7B and LED-8B, respectively. Across eight benchmarks, LED SFT yields average gains of

+11.4

and

+7.7

points over standard MDM SFT. Moreover, using SFT alone, LED-8B outperforms LLaDA2.1-16B Q-mode, which is trained with continued pre-training, supervised fine-tuning, and reinforcement learning.

Abstract

Diffusion Language Models

Explore similar work

LED: Improving Masked Diffusion Language Models by Lookahead Editing

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

LogicDiff: Logic-Guided Denoising Improves Zero-Shot Reasoning in Masked Diffusion Language Models