SIRS：单次查询实现掩码扩散模型高效拒绝采样

掩码扩散模型（Masked Diffusion Models, MDMs）通过并行解掩多个 token 来加速生成，但这种独立采样的方式忽略了 token 之间的相关性，导致输出分布出现偏差。Anari 等人（2026）利用拒绝采样修正了这一偏差，将对数接受概率分解为 token 级项之和，却带来一个新的代价：每解掩一个 token 就需要一次 score 查询。在并行度较高的场景下，这一开销会迅速累积。

针对这一矛盾，最新研究提出了 Single-Increment Rejection Sampler（SIRS），试图用单次查询完成对数接受概率的估计，从而在保持无偏性的同时大幅降低查询成本。

并行掩码扩散的偏差从何而来

掩码扩散模型的生成过程可以理解为：从全掩码序列出发，每一轮并行地解掩一批 token，直到所有位置都被填充。这种“块内并行”策略显著减少了前向次数，但代价是采样时假设块内各 token 相互独立。

真实的数据分布中，同一块内的 token 往往存在强依赖关系——比如自然语言里相邻词之间的语法和语义约束。忽略这些依赖，等价于用边缘分布的乘积去近似联合分布，输出自然会产生偏差。Anari 等人的工作从理论上刻画了这一偏差，并给出基于拒绝采样的修正方案：把联合接受概率拆解成每个 token 的贡献之和，再逐一评估。

掩码扩散模型（MDM） 是扩散模型在离散序列上的一种变体。与连续空间扩散通过逐步加噪、去噪不同，MDM 的“加噪”过程是把序列中的 token 逐步替换为特殊的 [MASK] 标记，训练目标是让模型学会根据当前部分可见的序列，预测被掩位置上的原始 token。生成时从全掩码序列出发，反复调用模型预测各位置的条件分布并采样填充，这一过程称为“解掩（unmask）”。由于每个位置的条件分布可以同时计算，MDM 天然支持并行解码，这一点与自回归模型逐个生成 token 的方式形成鲜明对比。

Score 查询 在 MDM 语境中指调用模型前向计算，得到给定当前已解掩上下文时各掩码位置的条件分布或对数概率。拒绝采样需要知道提议分布与目标分布的比值，因此每评估一次接受概率就要查询一次模型。Anari 等人的方案把整块接受概率拆成 token 级项之和，每项都依赖模型在该位置的输出，于是块内有多少个待解掩 token，就需要多少次查询。在块很大时，查询次数会线性增长，构成主要瓶颈。

SIRS 的核心思路：一次查询估计整个增量

SIRS 的关键洞察在于：与其为块内每个 token 分别查询 score，不如随机采样一个增量（increment），用这一次查询的结果去估计整个块的对数接受概率。具体来说，模型每次前向会返回块内所有剩余 token 的条件分布，SIRS 利用这一信息对增量做 Rao-Blackwellization——在不需要额外查询的前提下，对块内 token 求平均，从而降低估计方差。

方差降低直接带来一个实际好处：裁剪（clipping）减少。拒绝采样中，接受概率的估计若波动过大，就需要更激进的裁剪来保证数值稳定，而裁剪会引入额外的分布偏差。SIRS 通过更稳的估计，把裁剪控制在更小的范围内。

随后，SIRS 借助 Chen 等人（2026）提出的 Bernoulli factory，将一个服从 Poisson 分布的裁剪估计值转换为逐 token 的接受/拒绝决策。这一转换保证了采样过程在统计上的正确性，而不是简单的启发式近似。

Rao-Blackwellization 是一种经典的方差缩减技术，得名于统计学家 C. R. Rao 与 David Blackwell。其基本思想是：如果某个估计量依赖于一个可以条件求期望的辅助随机变量，那么用该辅助变量关于充分统计量的条件期望代替原估计量，可以在不引入偏差的前提下降低方差。用公式表达，对任意随机变量 X 和 Y，有 Var(E[X|Y]) ≤ Var(X)，即对 Y 取条件期望只会让方差变小或不变。

在 SIRS 中，“增量”是随机采样的辅助变量，而模型前向返回的块内各 token 条件分布恰好提供了对增量取条件期望所需的全部信息。因此 SIRS 可以在不为每个 token 单独查询的前提下，对增量做解析平均，得到方差更低的接受概率估计。这解释了为什么单次查询即可获得足够稳定的估计：方差下降不是靠增加样本量，而是靠解析地消去一部分随机性。

Bernoulli factory 是一个把某种概率分布的随机样本转换为另一种分布样本的算法框架。在拒绝采样中，我们需要根据估计出的接受概率 p 生成一个以概率 p 接受、以概率 1-p 拒绝的伯努利决策，但 p 本身往往是估计值而非精确值，且可能以某种变换形式出现。Bernoulli factory 能仅利用对 p 的“黑箱”查询（如判断 p 是否大于某个阈值），构造出精确服从参数 p 的伯努利分布，不需要知道 p 的解析形式。

SIRS 中，裁剪后的接受概率估计服从 Poisson 分布，无法直接当作接受概率使用。Chen 等人（2026）的 Bernoulli factory 恰好能处理这一类输入，把它精确转换为逐 token 的接受/拒绝决策。这一步骤是采样无偏性的关键：如果只是把估计值直接当作接受概率，得到的分布会有额外偏差；而通过 Bernoulli factory 转换后，只要估计量本身无偏，最终的接受决策在统计上就是正确的。

理论保证：KL 散度有界

SIRS 给出了明确的误差界：数据分布与 SIRS 输出分布之间的 KL 散度，被一个裁剪项加上模型的训练误差所界定。

这个界的意义在于它把采样误差与模型本身的质量解耦开来——采样算法引入的额外偏差只与裁剪范围有关，不会随块大小或并行度失控增长。论文进一步对比了两种情形：

• 精确 score 下，并行解码需要支付块内依赖的全部代价；
• SIRS 只需支付超出裁剪范围的那部分依赖代价。

换句话说，只要裁剪范围设置得当，SIRS 的实际开销可以远低于朴素的并行拒绝采样。

KL 散度（Kullback-Leibler divergence） 是衡量两个概率分布差异的非对称指标，定义为 D_KL(P||Q) = Σ_x P(x) log(P(x)/Q(x))，取值恒非负，当且仅当两分布几乎处处相等时为零。在生成模型的误差分析中，KL 散度常被用来刻画采样算法输出分布与真实数据分布之间的差距：散度越小，采样结果在信息论意义上越接近目标分布。

SIRS 给出的界把 KL 散度分解为两部分：一部分来自裁剪操作引入的偏差，另一部分来自模型 score 估计本身的训练误差。这种分解的价值在于，它表明采样算法引入的额外误差与块大小、并行度等超参数无关，只受裁剪阈值控制。因此可以在不牺牲并行效率的前提下，通过收紧裁剪范围把采样误差压到任意小——代价是更多的拒绝和重采样，形成一种可控的精度与效率权衡。

对掩码扩散加速的意义

掩码扩散模型近年来在文本生成、代码补全等任务上受到关注，其并行解码能力是相对自回归模型的主要效率优势。但并行度越高，块内依赖被忽略得越严重，质量与速度的权衡就越尖锐。

SIRS 的价值在于它没有回避这一权衡，而是用更精细的统计工具把代价压到理论下限附近。单次查询的估计方式让查询成本与块大小解耦，Rao-Blackwellization 则在不增加成本的前提下进一步压缩方差。对于需要在推理阶段大批量采样的应用场景，这种查询效率的提升可能比单纯增加并行度更有实际价值。

当然，SIRS 的效果依赖于 score 估计的质量和裁剪阈值的选取，这两者在实际部署中如何调优，仍是工程落地需要回答的问题。