Free Lunch for Pass@$k$? Low Cost Diverse Sampling for Diffusion Language Models

Sean Lamont; Christian Walder; Paul Montague; Amir Dezfouli; Michael Norrish

cs.CLMar 5, 2026

Free Lunch for Pass@ $k$ ? Low Cost Diverse Sampling for Diffusion Language Models

Authors: Sean Lamont, Christian Walder, Paul Montague, Amir Dezfouli, Michael Norrish

Organizations: Australian National University · Defence Science and Technology Group · Google DeepMind · BIMLOGIQ

Abstract

Diverse outputs in text generation are necessary for effective exploration in complex reasoning tasks, such as code generation and mathematical problem solving. Such Pass@ $k$ problems benefit from distinct candidates covering the solution space. However, traditional sampling approaches often waste computational resources on repetitive failure modes. While Diffusion Language Models have emerged as a competitive alternative to the prevailing Autoregressive paradigm, they remain susceptible to this redundancy, with independent samples frequently collapsing into similar modes. To address this, we propose a training free, low cost intervention to enhance generative diversity in Diffusion Language Models. Our approach modifies intermediate samples in a batch sequentially, where each sample is repelled from the feature space of previous samples, actively penalising redundancy. Unlike prior methods that require retraining or beam search, our strategy incurs negligible computational overhead, while ensuring that each sample contributes a unique perspective to the batch. We evaluate our method on the HumanEval and GSM8K benchmarks using the LLaDA-8B-Instruct model. Our results demonstrate significantly improved diversity and Pass@ $k$ performance across various temperature settings. As a simple modification to the sampling process, our method offers an immediate, low-cost improvement for current and future Diffusion Language Models in tasks that benefit from diverse solution search. We make our code available at https://github.com/sean-lamont/odd.

Figures & tables

Figure 1: Example outputs from LLaDA with standard sampling (left) and our approach (right), both with a temperature of 1. For this example, the base sampler achieved a Pass@16 of 0, while our model found 3 valid solutions. Despite having the same temperature, standard sampling leads to mode collapse over an incorrect approach, whereas our method optimises to select samples distinct from those previously explored.

Figure 2: An overview of our proposed diverse sampling framework, using our orthogonal diversity loss. Given a full batch of $n$ samples, feature vectors $F(\mathbf{x_{i}})$ are extracted for the logits $x_{i}$ for each sample $i\in\{1,\dots,n\}$ . We then loop through each sample, compute an orthogonal basis for all previous sample features $B_{<i}$ , and calculate the projection of the current sample features onto this basis. We then take the norm of the residual with respect to this projection as the diversity objective. Optimising this then moves the sample in a direction diverse compared to previous samples.

GSM8K (200 Problems)
Step Size ( $\alpha$ )	Temperature ( $\theta$ )
Step Size ( $\alpha$ )	0.0	0.5	1.0	1.5	2.0
Baseline (LLaDA)
0	$61.0\pm 0.0$	$74.8\pm 0.2$	$81.3\pm 0.7$	$83.4\pm 0.6$	$76.5\pm 0.4$
LLaDA + DPP
2	$63.5\pm 0.0$	$75.9\pm 0.7$	$82.3\pm 0.7$	$83.9\pm 0.8$	$77.1\pm 0.6$

Table 1: Pass@16 results (mean

\pm

SE over 8 runs) for GSM8K and HumanEval over various temperatures (

\theta

) and step sizes (

\alpha

). Bold values indicate the best result for each

\theta

. LLaDA + DPP (Algorithm 2 ) uses a global DPP objective ( Morshed and Boddeti, 2025 ) . LLaDA + ODD uses our approach in Algorithm 1 .

Approach	GSM8K	HumanEval
Baseline (s)	$22.5\pm 0.3$	$30.8\pm 0.4$
ODD (s)	$23.8\pm 0.5$	$32.0\pm 0.2$
Overhead (%)	+5.8%	+3.9%

Table 2: Wall time comparison for generation between standard LLaDA (Baseline) and ODD (Batch Size 16). We report the execution time per batch (in seconds) averaged over all

\alpha

and temperature settings for each approach, and the relative overhead percentage. Our method introduces minimal latency across both benchmarks, and scales independently of model size.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	$\alpha$ (Repulsion)	Solved (Union)	Coverage
GSM8K	0.0 (Baseline)	198	99.00%
	2.0	193	96.50%
	8.0	192	96.00%
	16.0	195	97.50%
	32.0	195	97.50%
	64.0	195	97.50%

Appendix

Table 3: Cumulative problem coverage across all temperature settings. ODD (

\alpha>0

) significantly improves coverage on HumanEval, finding many problems missed by the baseline.

Figure 3: Average batch diversity (1 - cosine similarity of sentence embeddings) for GSM8K (top) and HumanEval (bottom), comparing ODD to the baseline (

\alpha=0

, as indicated by dashed lines). ODD increases diversity at low temperatures and acts as a coherence filter at high temperatures (evidenced by the improved Pass@16 in Table 1 ), effectively balancing exploration and quality.

Figure 4: Pass@1 vs Pass@16 Pareto frontiers.

\alpha=0

represents standard LLaDA sampling (baseline). Data points represent decreasing temperatures from right (

\theta=0

) to left (

\theta=2

). Top (GSM8K): ODD trades off individual accuracy (shifting left) for superior batch coverage (shifting up). Bottom (HumanEval): For

\alpha\leq 16

, ODD achieves a Pareto improvement, boosting coverage without loss of quality.

Figure 5: Empirical Pass@

k

for ODD over GSM8K (top) and HumanEval (bottom). We observe a consistent improvement relative to baseline for ODD. Given the batch size invariance of ODD, we expect this trend would continue for higher batch sizes.

Figure 6: Computational overhead of ODD compared to standard sampling. Top: Global averages across all runs. Bottom: Scaling behavior demonstrating time and memory overheads across batch size (

B

), sequence length (

S

), and diffusion steps (

T

). ODD incurs less than a

10\%

relative time penalty, with memory scaling independently of the base model size.

Diversity-Aware Sampling

See all 20 papers

A Tale of Two Temperatures: Simple, Efficient, and Diverse Sampling from Diffusion Language Models
Apr 10, 2026Theo X. Olausson, Metod Jazbec, Xi Wang +4Diffusion Language Models Diversity-Aware Sampling
Massachusetts Institute of Technology · UvA-Bosch Delta Lab, University of Amsterdam · Johns Hopkins University +1
Forward-Free Diffusion Language Models
Jun 6, 2026Haotian Sun, Rushi Qiang, Yuqian Zheng +1Diffusion Language Models
Georgia Institute of Technology
LESS Is More: Mutual-Stability Sampling for Diffusion Language Models
Jun 15, 2026Amr Mohamed, Guokan Shang, Michalis VazirgiannisMasked Diffusion Models Diffusion Language Model Decoding
1MBZUAI · 2Ecole Polytechnique

Apr 10, 2026J/K move · Enter open · S save

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

Theo X. Olausson, Metod Jazbec, Xi Wang, Armando Solar-Lezama+3

Massachusetts Institute of Technology · UvA-Bosch Delta Lab, University of Amsterdam · Johns Hopkins University+1

Diffusion Language Models Diversity-Aware Sampling

Open paper AskPDF arXiv

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.

Abstract

Diversity-Aware Sampling

Explore similar work

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

Forward-Free Diffusion Language Models

LESS Is More: Mutual-Stability Sampling for Diffusion Language Models