Enabling Approximate Joint Sampling in Diffusion LMs

• URL: http://arxiv.org/abs/2509.22738v1
• Date: Thu, 25 Sep 2025 21:23:26 GMT
• Title: Enabling Approximate Joint Sampling in Diffusion LMs
• Abstract summary: Masked diffusion language models generate text by unmasking tokens out of order and potentially in parallel.<n>In this paper we devise a way to em approximately sample multiple tokens from the joint distribution in a single full-model forward pass.
• Score: 16.729589221035074
• License: http://creativecommons.org/licenses/by/4.0/
• 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) and instruction-tuned (Dream-7B-Instruct) 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.

Related papers

• Representation-based Masked Diffusion Model [3.3928176963185237]
  Masked Diffusion Models (MDMs) have emerged as a compelling paradigm for language modeling, offering the capability for efficient parallel text generation.<n>We propose Representation-based Masked Diffusion Model (RMDM), a framework that leverages the text representation to explicitly encode global semantics.
  arXiv  Detail & Related papers  (2026-09-11T03:00:52Z)
• A Ticket from Marginals to Joints: Coupled-Noise Distillation for One-Step Block Generation in Diffusion Language Models [0.6916773850242582]
  We ask whether a discrete masked model can commit an entire block in one pass when its mask embeddings are perturbed by a sampled Gaussian noise field.<n>We propose CONDOR, which trains such a model from scratch without a target-side encoder or an autoregressive teacher.<n>On a controlled TinyStories setting, this recipe yields coherent one-pass continuations that vary with the noise.
  arXiv  Detail & Related papers  (2026-09-06T01:09:14Z)
• Dependency-Guided Parallel Decoding in Discrete Diffusion Language Models [21.699371484195865]
  We propose DEMASK, a lightweight dependency predictor that attaches to the final hidden states of a dLLM.<n> Empirically, DEMASK achieves 1.7-2.2$times$ speedup on Dream-7B while matching or improving accuracy compared to confidence-based and KL-based baselines.
  arXiv  Detail & Related papers  (2026-04-02T22:21:24Z)
• Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models [58.946955321428845]
  This work presents self-rewarding sequential Monte Carlo (SMC)<n>Our algorithm stems from the observation that most existing MDLMs rely on a confidence-based sampling strategy.<n>We introduce the trajectory-level confidence as a self-rewarding signal for assigning particle importance weights.
  arXiv  Detail & Related papers  (2026-02-02T09:21:45Z)
• Diffusion Language Models are Provably Optimal Parallel Samplers [15.981424915336001]
  Diffusion language models (DLMs) have emerged as a promising alternative to autoregressive models.<n>We show that DLMs augmented with a chain-of-thought can simulate any parallel sampling algorithm using an optimal number of sequential steps.
  arXiv  Detail & Related papers  (2025-12-31T18:03:05Z)
• Don't Settle Too Early: Self-Reflective Remasking for Diffusion Language Models [40.902681492117786]
  RemeDi is a mask-based DLM that predicts token distributions and per-token confidence scores at each step.<n>We train a remask-aware pipeline to train this ability, including supervised fine-tuning which teaches the model to detect and remask incorrect tokens.<n>Experiments show that RemeDi achieves the state-of-the-art results among open-source DLMs on multiple datasets.
  arXiv  Detail & Related papers  (2025-09-28T05:39:49Z)
• Sampling from Your Language Model One Byte at a Time [82.71473348639489]
  Tokenization can introduce distortion into the model's generations, known as the Prompt Boundary Problem (PBP)<n>We present an inference-time method to convert any autore LM with a BPE tokenizer into a character-level or byte-level LM.<n>Our method efficiently solves the PBP and is also able to unify the vocabularies of language models with different tokenizers.
  arXiv  Detail & Related papers  (2025-06-17T02:37:04Z)
• Accelerating Diffusion LLMs via Adaptive Parallel Decoding [50.9948753314669]
  We introduce adaptive parallel decoding (APD), a novel method that dynamically adjusts the number of tokens sampled in parallel.<n>APD provides markedly higher throughput with minimal quality degradations on downstream benchmarks.
  arXiv  Detail & Related papers  (2025-05-31T06:10:10Z)
• Text Generation Beyond Discrete Token Sampling [75.96920867382859]
  Mixture of Inputs (MoI) is a training-free method for autoregressive generation.<n>MoI consistently improves performance across multiple models including QwQ-32B, Nemotron-Super-49B, Gemma-3-27B, and DAPO-Qwen-32B.
  arXiv  Detail & Related papers  (2025-05-20T18:41:46Z)
• Reviving Any-Subset Autoregressive Models with Principled Parallel Sampling and Speculative Decoding [55.2480439325792]
  In arbitrary-order language models, it is an open question how to sample tokens in parallel from the correct joint distribution.<n>We find that a different class of models, any-subset autoregressive models (AS-ARMs), holds the solution.<n>We show that AS-ARMs achieve state-of-the-art performance among sub-200M parameter models on infilling benchmark tasks, and nearly match the performance of models 50X larger on code generation.
  arXiv  Detail & Related papers  (2025-04-29T06:33:13Z)
• Exact Byte-Level Probabilities from Tokenized Language Models for FIM-Tasks and Model Ensembles [23.134664392314264]
  Tokenization is associated with many poorly understood shortcomings in language models (LMs)<n>This work studies how tokenization impacts model performance by analyzing and comparing models with their byte-level counterparts.<n>We introduce the Byte-Token Representation Lemma, a framework that establishes a mapping between the learned token distribution and its equivalent byte-level distribution.
  arXiv  Detail & Related papers  (2024-10-11T23:30:42Z)
• Score Mismatching for Generative Modeling [4.413162309652114]
  We propose a new score-based model with one-step sampling. We train a standalone generator to compress all the time steps with the gradient backpropagated from the score network. In order to produce meaningful gradients for the generator, the score network is trained to simultaneously match the real data distribution and mismatch the fake data distribution.
  arXiv  Detail & Related papers  (2023-09-20T03:47:12Z)
• Masked Generative Modeling with Enhanced Sampling Scheme [1.3927943269211591]
  Enhanced Sampling Scheme (ESS) ensures both sample diversity and fidelity. ESS consists of three stages: Naive Iterative Decoding, Critical Reverse Sampling, and Critical Resampling. We demonstrate significant performance gains of ESS in both unconditional sampling and class-conditional sampling.
  arXiv  Detail & Related papers  (2023-09-14T09:42:13Z)
• Extreme Masking for Learning Instance and Distributed Visual Representations [50.152264456036114]
  The paper presents a scalable approach for learning distributed representations over individual tokens and a holistic instance representation simultaneously. We use self-attention blocks to represent distributed tokens, followed by cross-attention blocks to aggregate the holistic instance. Our model, named ExtreMA, follows the plain BYOL approach where the instance representation from the unmasked subset is trained to predict that from the intact input.
  arXiv  Detail & Related papers  (2022-06-09T17:59:43Z)
• ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators [108.3381301768299]
  Masked language modeling (MLM) pre-training methods such as BERT corrupt the input by replacing some tokens with [MASK] and then train a model to reconstruct the original tokens. We propose a more sample-efficient pre-training task called replaced token detection.
  arXiv  Detail & Related papers  (2020-03-23T21:17:42Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.