Improving Variable-Length Generation in Diffusion Language Models via Length Regularization

TL;DR

This paper introduces LR-DLLM, a length-regularized inference method for diffusion language models that improves variable-length generation by explicitly modeling and correcting length bias, leading to more reliable and flexible text generation.

Contribution

It proposes a novel length regularization framework that enables diffusion language models to accurately determine and adapt generation length without retraining.

Findings

Achieves 51.3% Pass@1 on HumanEvalInfilling with unknown lengths.

Outperforms DreamOn with a +13.4% improvement.

Demonstrates effectiveness across four languages with a 14.3% gain.

Abstract

Diffusion Large Language Models (DLLMs) are inherently ill-suited for variable-length generation, as their inference is defined on a fixed-length canvas and implicitly assumes a known target length. When the length is unknown, as in realistic completion and infilling, naively comparing confidence across mask lengths becomes systematically biased, leading to under-generation or redundant continuations. In this paper, we show that this failure arises from an intrinsic lengthinduced bias in generation confidence estimates, leaving existing DLLMs without a robust way to determine generation length and making variablelength inference unreliable. To address this issue, we propose LR-DLLM, a length-regularized inference framework for DLLMs that treats generation length as an explicit variable and achieves reliable length determination at inference time. It decouples semantic compatibility from lengthinduced uncertainty through an explicit length regularization that corrects biased confidence estimates. Based on this, LR-DLLM enables dynamic expansion or contraction of the generation span without modifying the underlying DLLM or its training procedure. Experiments show that LRDLLM achieves 51.3% Pass@1 on HumanEvalInfilling under fully unknown lengths (+13.4% vs. DreamOn) and 51.5% average Pass@1 on four-language McEval (+14.3% vs. DreamOn).