Left, GSM8K exact match against prompt length for in-context learning and for the prompt infilled by the full-sequence-masking model. Right, SummEval Spearman correlation when each diffusion LM variant uses its own infilled prompt or the transferred FS+RO prompt. Infilled prompts beat 16-shot prompting with a tenth of the tokens, and transfer to other checkpoints.
03
Runs
03
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
Yoshinari Fujinuma, Keisuke Sakaguchi. Unlocking Prompt Infilling Capability for Diffusion Language Models. COLM 2026 Workshop (Workshop on Non-Autoregressive Language Models for Fast and Flexible Text Generation (NonAR-LM), non-archival). arXiv:2604.03677.
@inproceedings{fujinuma2026promptinfill,
title={Unlocking Prompt Infilling Capability for Diffusion Language Models},
author={Yoshinari Fujinuma and Keisuke Sakaguchi},
booktitle={COLM 2026 Workshop, Workshop on Non-Autoregressive Language Models for Fast and Flexible Text Generation (NonAR-LM), non-archival},
year={2026},
eprint={2604.03677},
archivePrefix={arXiv}
}
arXiv 2604.03677
April 4, 2026
COLM 2026 Workshop
Workshop
Workshop on Non-Autoregressive Language Models for Fast and Flexible Text Generation (NonAR-LM), non-archival
