Enabling Approximate Joint Sampling in Diffusion LMs

2025/09/25 by Parikshit Bansal, Sujay Sanghavi, Bansal, Parikshit +1 · 1 citation
Engineering · #Acoustic Wave Resonator Technologies #Computation and Language (cs.CL) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Ultrasonics and Acoustic Wave Propagation

paper · pdf · doi:10.48550/arxiv.2509.22738

openalex publication_date 2025/09/25 · openalex created_date 2025/10/19 · openalex updated_date 2026/07/28

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.

Citations

• Encoder-Decoder Diffusion Language Models for Efficient Training and Inference
• ParallelBench: Understanding the Trade-offs of Parallel Decoding in Diffusion LLMs
• Self-Speculative Masked Diffusions
• Dream 7B: Diffusion Large Language Models
• Seed Diffusion: A Large-Scale Diffusion Language Model with High-Speed Inference
• Your LLM Knows the Future: Uncovering Its Multi-Token Prediction Potential
• Mercury: Ultra-Fast Language Models Based on Diffusion
• Esoteric Language Models: A Family of Any-Order Diffusion LLMs
• Accelerating Diffusion LLMs via Adaptive Parallel Decoding
• Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding
• Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking
• Anchored Diffusion Language Model
• Reviving Any-Subset Autoregressive Models with Principled Parallel Sampling and Speculative Decoding
• EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test
• Large Language Diffusion Models
• Path Planning for Masked Diffusion Model Sampling
• Think While You Generate: Discrete Diffusion with Planned Denoising
• Discrete Copula Diffusion
• Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling
• Learning Harmonized Representations for Speculative Sampling
• Simple and Effective Masked Diffusion Language Models
• Simplified and Generalized Masked Diffusion for Discrete Data
• EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty
• Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution
• Fast Inference from Transformers via Speculative Decoding
• Training and Inference on Any-Order Autoregressive Models the Right Way
• Structured Denoising Diffusion Models in Discrete State-Spaces
• Divergence Frontiers for Generative Models: Sample Complexity, Quantization Effects, and Frontier Integrals

Cited by

• Learning Unmasking Policies for Diffusion Language Models

Related

• Fast and Fluent Diffusion Language Models via Convolutional Decoding and Rejective Fine-tuning
• Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking
• Simple and Effective Masked Diffusion Language Models
• Beyond Autoregression: Fast LLMs via Self-Distillation Through Time
• Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling