Unlocking the Potential of Diffusion Language Models through Template Infilling
This addresses the problem of inefficient inference for researchers and practitioners using DLMs, offering a novel method that is incremental but provides strong specific gains.
The paper tackles the limited inference strategies of Diffusion Language Models (DLMs) by proposing Template Infilling (TI), a conditioning methodology that generates structural templates before filling masked segments, with Dynamic Segment Allocation (DSA) for adaptive segment lengths. The result is consistent improvements of 17.01 percentage points over baselines on mathematical reasoning and code generation benchmarks, with additional advantages in multi-token generation speedup.
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.