Diffusion Language Models Are Natively Length-Aware

TL;DR

Diffusion Language Models can be made length-aware by estimating output length from latent prompts, enabling dynamic context cropping that reduces computation without harming performance.

Contribution

The paper introduces a zero-shot method to dynamically crop context in diffusion language models based on latent prompt information, improving efficiency.

Findings

Massive FLOPs reduction across tasks

No significant performance degradation

Performance improvement in some tasks

Abstract

Unlike autoregressive language models, which terminate variable-length generation upon predicting an End-of-Sequence (EoS) token, Diffusion Language Models (DLMs) operate over a fixed maximum-length context window for a predetermined number of denoising steps. However, this process is independent of the required response length, resulting in computational waste for the majority of short responses common in reasoning and chat tasks. To address this problem, we conjecture that the latent prompt representation contains sufficient information to estimate the required output length. We provide empirical evidence for this phenomenon and propose a zero-shot mechanism to dynamically crop the context window before generation begins, leading to fewer diffusion steps and substantial computational savings. We evaluate our approach on four benchmarks with diverse tasks -- GSM8K (reasoning), HumanEval (code generation), IfEval (instruction following), and LongFormQA (question answering) -- revealing massive efficiency gains at minimal performance impact. We report significant reductions in FLOPs across all tasks, with no statistically significant performance degradation, and significant performance improvements in 2 out of 4 tasks.