arXiv cs.CL· Parikshit Bansal, Sujay Sanghavi·· 11 小时前AI 评分34
让扩散语言模型实现近似联合采样
Enabling Approximate Joint Sampling in Diffusion LMs
AI 导读
研究者提出一种在扩散语言模型中近似联合采样的方法：在已有大型扩散 LM 之上训练一个轻量单层“采样器”，使一次全模型前向传播后可接多次采样器前向，从而并行解掩多个 token。
正文 · 原文
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.
| Comments: | Accepted to NeurIPS 2026 |
| Subjects: | Computation and Language (cs.CL); Machine Learning (cs.LG) |
| Cite as: | arXiv:2509.22738 [cs.CL] |
| (or arXiv:2509.22738v4 [cs.CL] for this version) | |
| https://doi.org/10.48550/arXiv.2509.22738 arXiv-issued DOI via DataCite |
Submission history
From: Parikshit Bansal [view email]
[v1]
Thu, 25 Sep 2025 21:23:26 UTC (344 KB)
[v2]
Mon, 2 Feb 2026 17:02:01 UTC (342 KB)
[v3]
Thu, 24 Sep 2026 10:48:07 UTC (1,447 KB)
[v4]
Thu, 5 Nov 2026 23:41:05 UTC (1,631 KB)
来源：arXiv cs.CL · arxiv.org