A Tale of Two Temperatures: Simple, Efficient, and Diverse Sampling from Diffusion Language Models

TL;DR

This paper introduces a simple and efficient method to increase diversity in sampling from diffusion language models by using tempered confidence heuristics, improving exploration and downstream performance.

Contribution

The work proposes tempered heuristics for sampling that enhance diversity without sacrificing computational efficiency, bridging the gap between existing methods.

Findings

Tempered heuristics close the exploration gap (pass@k) between confidence-based and autoregressive sampling.

The approach outperforms existing methods when controlling for computational cost.

Increased diversity benefits downstream post-training and test-time scaling.

Abstract

Much work has been done on designing fast and accurate sampling for diffusion language models (dLLMs). However, these efforts have largely focused on the tradeoff between speed and quality of individual samples; how to additionally ensure diversity across samples remains less well understood. In this work, we show that diversity can be increased by using softened, tempered versions of familiar confidence-based remasking heuristics, retaining their computational benefits and offering simple implementations. We motivate this approach by introducing an idealized formal model of fork tokens and studying the impact of remasking on the expected entropy at the forks. Empirically, the proposed tempered heuristics close the exploration gap (pass@k) between existing confidence-based and autoregressive sampling, hence outperforming both when controlling for cost (pass@NFE). We further study how the increase in diversity translates to downstream post-training and test-time compute scaling. Overall, our findings demonstrate that simple, efficient, and diverse sampling from dLLMs is possible.