Unlocking the Potential of Diffusion Language Models through Template Infilling
Junhoo Lee, Seungyeon Kim, Nojun Kwak
TL;DR
This work addresses the limitation of diffusion language models relying on prefix-based conditioning by introducing Template Infilling (TI), which defines an explicit structural template with masked segments to guide generation. To adapt template granularity to content needs, Dynamic Segment Allocation (DSA) expands mask lengths based on generation confidence, enabling problem-adaptive inference. Across mathematical reasoning and code generation benchmarks, TI achieves an average improvement of $17.01\%$ points over baselines and remains effective in multi-token generation, suggesting substantial gains from exploiting DLM bidirectional conditioning. The results indicate TI as a new direction for structured, template-based conditioning in diffusion generation, with potential impacts on efficiency and the broader design of conditioning strategies beyond autoregressive prompts.
Abstract
Diffusion Language Models (DLMs) have emerged as a promising alternative to Autoregressive Language Models, yet their inference strategies remain limited to prefix-based prompting inherited from the autoregressive paradigm. In this paper, we propose Template Infilling (TI), a tailored conditioning methodology for DLMs' generation process. Unlike conventional prefix prompting, TI first generates a structural template for the target response, then fills in the masked segments. To enhance the flexibility of this structural control, we introduce Dynamic Segment Allocation (DSA), which adaptively adjusts segment lengths based on generation confidence. We demonstrate the effectiveness of our approach on mathematical reasoning and code generation benchmarks, achieving consistent improvements of 17.01$\%$p over baseline. Furthermore, we show that TI provides additional advantages in multi-token generation settings, enabling effective speedup while maintaining generation quality.