Enabling Approximate Joint Sampling in Diffusion LMs

Parikshit Bansal; Sujay Sanghavi

cs.CLSep 25, 2025

Enabling Approximate Joint Sampling in Diffusion LMs

Authors: Parikshit Bansal, Sujay Sanghavi

Organizations: UT Austin

Abstract

In autoregressive language models, each token is sampled by conditioning on all the past tokens; the overall string has thus been sampled from the correct underlying joint distribution represented by the model. In contrast, masked diffusion language models generate text by unmasking tokens out of order and potentially in parallel. Generating an overall string sampled from the correct underlying joint distribution would (again) require exactly one token unmasking in every full-model forward pass. The more tokens unmasked in parallel, the further away the string is from the true joint; this can be seen in the resulting drop in accuracy (but, increase in speed). In this paper we devise a way to {\em approximately} sample multiple tokens from the joint distribution in a single full-model forward pass; we do so by developing a new lightweight single-layer ``sampler" on top of an existing large diffusion LM. One forward pass of the full model can now be followed by multiple forward passes of only this sampler layer, to yield multiple unmasked tokens. Our sampler is trained to mimic exact joint sampling from the (frozen) full model. We show the effectiveness of our approximate joint sampling for both pretrained-only (Dream-7B-Base, Llada-7B-Base) and instruction-tuned (Dream-7B-Instruct, Dream-7B-Coder) models on language modeling and math & coding tasks. When four tokens are unmasked for each full-model denoising step, our sampling algorithm achieves a MAUVE score of 0.87 (vs marginal baseline of 0.31) with respect to the true joint distribution.

Figures & tables

Panel of Figure 1 — Figure 1 : In this figure, we use Dream-7B-Base to generate a token string of length 128 tokens starting from an all masked string (kindly refer to Sec 4.2 for details). We report the generative negative log-likelihood (lower is better) and the MAUVE score (higher is better) for the generated strings. We vary the number of tokens sampled per diffusion denoising step from one to four. We observe an increase in gen. NLL and a decrease in the MAUVE score as more tokens are generated in parallel. Our method ADJUST achieves the best NLL and MAUVE (as compared to baseline methods like naive parallel sampling, and energy-based models).

Figure 2 : ADJUST. The figure shows one forward pass of the base diffusion models followed by two passes of the ADJUST sampling module. The tick in the figure denotes the position selection logic (TOP in Algo 2 ). ADJUST module $g$ is a single layer bi-direction attention based diffusion LM. After unmasking a single token ( sat in the figure), the module $g$ conditions the prior output embeddings ${\bm{h}}_{0}$ on the sampled token before proceeding to sample the next token on . Notation in figure follows Algo 2 .

Figure 3

		K=1		K=2		K=4		K=8
		NLL	Mauve	NLL	Mauve	NLL	Mauve	NLL	Mauve
T=0.6	Parallel	1.11	1.00	1.61	0.89	2.21	0.29	2.21	0.02
	EDLM	0.83	0.98	1.33	0.94	1.89	0.31	2.05	0.03
	ADJUST (ours)	1.11	1.00	1.23	0.97	1.64	0.79	1.66	0.67
T=1.0	Parallel	4.25	1.00	5.68	0.76	6.52	0.31	7.10	0.19
	EDLM	3.75	0.96	5.26	0.83	6.18	0.39	6.56	0.16

Table 1 : Unconditional Generation. In this table we compare the NLL and the MAUVE scores for three different methods : Marginal, Joint and Energy-based (EDLM). The base diffusion model is Dream-Base-7B. Across different number of tokens

K

sampled per step, and different sampling temperature, joint sampling consistently outperforms both of the baselines, parallel sampling and EDLM.

Model	Task	K=1	K=2		K=4		K=8
Model	Task	K=1	Parallel	ADJ.	Parallel	ADJ.	Parallel	ADJ.
Llada-Base	GSM	68.94	60.61	66.67	54.55	57.58	19.70	28.79
	MBPP	40.40	28.40	37.40	19.40	29.40	6.60	20.00
	HEval	33.54	25.00	29.88	12.80	23.78	5.49	10.37
Dream-Base	GSM	78.40	71.20	75.60	43.20	64.00	11.60	36.80
Dream-Base	MBPP	48.00	40.80	44.80	33.60	33.60	18.40	19.60

Table 2 : Downstream Evaluation. In this table we report numbers on downstream tasks, GSM8k & MBPP for the Dream-Base (denoted as Base). Since ADJUST for Dream-Instruct (denoted as Inst.) in trained on MetamathQA prompts, we report GSM8k & MATH numbers. Sampling a single token for each denoising step gives the best performance. ADJUST achieves better downstream accuracy tasks across values of tokens sampled in each denoising step

K

Figure 3 : Throughput evaluation. In this figure we report tokens produced per second for unconditional (leftmost) and downstream evaluation (middle) on GSM8k for Dream-Base-7B (Sec 4.4 ). Recall that DPD is a greedy (temperature=0) sampling based method and hence is missing from unconditional generation figure which requires sampling to generate text from base model’s underlying distribution. For both evaluation settings, ADJUST demonstrates gain in performance for a given throughput compared to naive parallel sampling. Additionally, while DPD is limited to a smaller range of tokens produced per second, ADJUST achieves higher throughput (>80 toks/sec) while outperforming both parallel sampling and DPD in that regime. The rightmost figure shows comparison with numbers with tokens per pass on the x-axis

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4 : In this figure we compare adaptive parallel decoding (working left-to-right) to naive parallel decoding and ADJUST . Sampling $K>1$ tokens per diffusion step, for left-to-right decoding shows a steeper increase in NLL (left) than using the base model’s unmasking logic. Hence APD generation is not suited for such open-ended tasks and under-performs. The right figure shows GSM8k evaluation for high temperature sampling (temperature=1.0) for Dream-Instruct model. Left-to-right generation is naturally good for math reasoning tasks, as shown by higher accuracy under APD when compared to naive/joint for $K=1$ tokens per step. APD leads to a quicker decrease in accuracy when compared to baseline and our method.

		K=1		K=2		K=3		K=4
		NLL	Mauve	NLL	Mauve	NLL	Mauve	NLL	Mauve
T=0.6	Parallel	1.11	1.00	1.61	0.89	2.18	0.57	2.21	0.29
	EDLM	0.83	0.98	1.33	0.94	1.85	0.54	1.89	0.31
	ADJUST (ours)	1.11	1.00	1.23	0.97	1.54	0.82	1.64	0.79
T=1.0	Parallel	4.25	1.00	5.68	0.76	6.14	0.48	6.52	0.31
	EDLM	3.75	0.96	5.26	0.83	5.76	0.58	6.18	0.39

Appendix

Table 4 : In this table we compare the NLL and the MAUVE scores for three different methods : Marginal, Joint and Energy-based (EDLM). Across different number of tokens

K

sampled per step, and different sampling temperature, joint sampling consistently outperforms both of the baselines, parallel sampling and EDLM.

Model	W2S(E)	W2S(M)	W2S(H)	Copy	Ins(idx)	Ins(rand)	Rm(idx)
Dream (x1)	82.0	87.0	69.0	73.0	30.0	42.0	72.0
Parallel (x4)	81.0	77.0	63.0	55.0	12.0	13.0	57.0
ADJUST (x4)	84.0	85.0	69.0	63.0	21.0	20.0	62.0
Model	Rm(rand)	Rpl(idx)	Rpl(rand)	Reverse	Shuffle	Sort
Dream (x1)	62.0	51.0	52.0	68.0	62.0	21.0
Parallel (x4)	36.0	35.0	1.0	27.0	24.0	13.0

Appendix

Table 5 : Evaluation on ParallelBench using Dream-Base-7B model

	T=0.6			T=1.0
	Spec=1	Spec=2	Spec=3	Spec=1	Spec=2	Spec=3
Tok/sec	35.57	52.87	71.00	35.56	52.91	70.97

Appendix

Table 6 : We report the runtime of the EDLM algorithm (with Qwen-2.5-0.5B) for experimental setup in Table 1 in Tokens/sec. These numbers are comparable to those in figure 1. In this table we report the negative log-likelihood numbers on decoding with EDLM. We fix the generation length to be 128 tokens and nucleus sampling to be 0.95, and temperature to be 1.0. Here we use Qwen2.5-7B as the verifier for EDLM (Big), i.e., the same model which is used to evaluate responses, while using Qwen2.5-0.5B as the verifier for EDLM (Small). We can see identical results because at each position we sample only two options both of which are identically rated by both the small and the big models

Figure 5 : This figure shows the architecture used for our draft model

Explore similar work

LESS Is More: Mutual-Stability Sampling for Diffusion Language Models
Jun 15, 2026Amr Mohamed, Guokan Shang, Michalis VazirgiannisDiffusion Language Models Diffusion Models
1MBZUAI · 2Ecole Polytechnique
How to Train Your Latent Diffusion Language Model Jointly With the Latent Space
May 8, 2026Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov +4Diffusion Language Models Latent Diffusion Model
Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia +2
DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling
May 22, 2026Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3Diffusion Language Models Dataset Distillation
NVIDIA

Jun 15, 2026, cs.CLJ/K move · Enter open · S save

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

Amr Mohamed, Guokan Shang, Michalis Vazirgiannis

1MBZUAI · 2Ecole Polytechnique

Diffusion Language Models Diffusion Models Accelerated Sampling Adaptive Sampling

Open paper AskPDF arXiv

Diffusion large language models (dLLMs) offer a promising alternative to autoregressive decoding by iteratively refining masked sequences, enabling parallel token updates and bidirectional conditioning. Their practical efficiency, however, is limited by sampling procedures that execute a fixed number of reverse denoising steps selected before decoding, spending computation on already-stable positions and sometimes committing unstable ones too early. We present \textsc{LESS}, a training-free, model-agnostic adaptive sampler that treats token commitment as an online stopping problem. \textsc{LESS} implements mutual-stability sampling through a joint stability rule that makes a masked position eligible for unmasking only when its top-1 prediction has high confidence, its top-1 token persists across recent reverse steps, and its predictive distribution is stable under top-

K

inter-step Jensen--Shannon divergence. We evaluate \textsc{LESS} on Dream-7B, LLaDA-8B, and LLaDA-1.5-8B, covering full-sequence diffusion and semi-autoregressive blockwise sampling regimes, across seven benchmarks spanning general knowledge, math, and code. \textsc{LESS} improves average accuracy over strong training-free adaptive samplers while using

72.1\%

fewer reverse steps than fixed-budget decoding. Since each reverse step requires a Transformer forward pass, these step-count reductions translate into fewer forward evaluations, lower measured wall-clock latency, and lower estimated inference compute.

Abstract

Explore similar work

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

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling