Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models¶
会议: NeurIPS 2026
arXiv: 2602.01849
OpenReview: https://openreview.net/forum?id=gFBldFk335
代码: https://github.com/Algolzw/self-rewarding-smc
领域: LLM 其他
关键词: 掩码扩散语言模型, 序贯蒙特卡洛, 自奖励, 推理阶段扩展, 粒子重采样
一句话总结¶
针对掩码扩散语言模型依赖贪婪置信度重掩码导致的短视局部最优与路径坍缩问题,本文提出自奖励序贯蒙特卡洛(SR-SMC)算法,将轨迹累积置信度严格推导为粒子重要性权重,无需额外训练或外部判别器即可在推理阶段引导扩散过程向全局高置信度解收敛。
研究背景与动机¶
离散掩码扩散语言模型(MDLMs)作为自回归生成的重要替代范式,通过向前向加掩码过程注入离散噪声并在反向过程中逐步反掩码恢复文本,展现出高并行度和双向上下文建模的独特优势。然而,当前主流 MDLMs(如 MDLM、BD3-LMs、LLaDA 与 Dream 等)普遍依赖基于步级置信度的贪婪反掩码策略,即在每个去噪时间步仅保留预测概率最高的 token,而将其余低置信度位置重新置为掩码。这种局部贪婪机制极其脆弱:一旦早期解码步生成了局部的幻觉或次优 token,该错误会在后续反向步中持续级联传播,导致生成路径陷入单调坍缩,丧失了在更高质量文本空间全局探索的可能性。
现有的推理阶段扩展(Inference-time scaling)尝试多借助外部奖励模型(如毒性抑制器、格式约束器或偏好模型)来倾斜扩散轨迹。但这类方法高度依赖针对特定下游任务人工调优的外部判别器,无法泛化到通用离散文本生成场景;同时,单纯增加单轨迹去噪步数也会受制于自回归或贪婪重掩码的近视陷阱。扩散模型在去噪过程中固有的局部置信度分布与全局轨迹质量之间存在显著割裂,如何从离散扩散的反向转移机制本身出发,构建一种通用的、无外部奖励依赖的推理性轨迹搜索框架,成为释放离散扩散大模型潜能的核心难题。
本文的切入视角是回归扩散反向过程的概率测度本质:既然离散扩散模型本身就能在每个时间步给出其预测分布的精确后验似然,那么未掩码序列的累积似然即可直接作为衡量生成轨迹优劣的内在指标。核心 idea:将掩码扩散语言模型的反向解码重构为费曼-卡茨(Feynman-Kac)路径测度模型,利用粒子扩散动力学在离散状态空间并行探索多条反掩码轨迹,并以轨迹级累积置信度自发作为粒子重要性权重进行序贯蒙特卡洛(SMC)重采样,实现完全自奖励的推理期搜索扩展。
方法详解¶
整体框架¶
自奖励序贯蒙特卡洛(SR-SMC)的整体输入为初始的全掩码离散序列与预训练离散扩散语言模型,输出为最终完成全部反掩码的确定性高质量文本。该算法维护 \(N\) 个并行演化的扩散采样过程(称为粒子)。在反向去噪过程由 \(t=T\) 倒退至 \(t=0\) 的每个时间步中,粒子系统依次经历自适应重采样(Resample)、局部转移传播(Propagate)与重要性权重更新(Re-weight)三阶段:系统监测粒子权重退化程度并自适应复制高潜力轨迹,在传播步依照扩散模型的预测分布与重掩码策略更新未掩码子集,最终利用新接受 token 的联合概率在轨迹级别累积自奖励得分。整个生成流程图如下所示:
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["全掩码初始化序列<br/>N 个全 [MASK] 粒子"] --> B["自适应重采样与退化检测<br/>基于有效样本量 ESS 动态触发"]
B --> C["扩散转移核局部传播<br/>Gumbel-Max 采样与重掩码策略"]
C --> D["轨迹级自奖励权重更新<br/>累积接受 token 联合置信度"]
D -->|未到达终点 t > 1| B
D -->|到达终点 t = 0| E["全局最优粒子选择<br/>输出累积权重最大的文本序列"]
关键设计¶
1. 费曼-卡茨路径测度与扩散转移核提议:为无奖励搜索提供理论立足点 针对传统采样中缺乏全局轨迹评价测度的缺陷,本文在离散扩散路径空间引入费曼-卡茨测度系统。设目标未归一化路径测度满足递归关系 \(\tilde{\pi}_{t-1}(\mathbf{x}_{t-1:T}) = \tilde{\pi}_{t}(\mathbf{x}_{t:T}) K_t(\mathbf{x}_t, \mathbf{x}_{t-1}) G_{t-1}(\mathbf{x}_t, \mathbf{x}_{t-1})\),其中转移核 \(K_t\) 严格服从掩码扩散模型包含低置信度重掩码的离散逆向转移: $\(K_t(\mathbf{x}_t, \mathbf{x}_{t-1}) = \prod_{j=1}^L p_\theta(\mathbf{x}_{t-1}(j) \mid \mathbf{x}_t)\)$ 当将 SMC 的逆向马尔可夫提议分布 \(q_{t-1}(\mathbf{x}_{t-1} \mid \mathbf{x}_t)\) 直接选定为扩散转移核 \(K_t(\mathbf{x}_t, \mathbf{x}_{t-1})\)(即 Bootstrap SMC 架构)时,权重公式分子与分母中的高维序列转移核完全相互对消。这种选择在数学上保证了粒子演化与预训练扩散模型的逆动力学自然对齐,不仅避免了设计复杂提议核的繁琐开销,且直接证明了以模型预测概率作为序贯自奖励是概率测度推演的必然结果而非经验凑合。
2. 轨迹级自奖励权重更新:消除步级贪婪带来的短视偏差 为了度量粒子当前轨迹的可信度,本文将势能函数 \(G_{t-1}(\mathbf{x}_t, \mathbf{x}_{t-1})\) 定义为当前步骤中所有被接受(未被重新置为掩码)的新 token 的模型预测联合概率。从而推导得出粒子在时间步 \(t-1\) 的增量重要性权重仅取决于本次反掩码集合 \(\mathcal{S}_t\) 中各 token 的预测置信度乘积: $\(\tilde{w}_{t-1}(\mathbf{x}_{t-1:T}) = \prod_{j \in \mathcal{S}_t} \mathbf{c}_t(j) = \prod_{j \in \mathcal{S}_t} p_\theta(\hat{\mathbf{x}}_0(j) \mid \mathbf{x}_t)\)$ 由于重要性权重在反向时间步上递推相乘,粒子累积的总权重等价于完整生成文本中所有生成 token 在各自被确定时刻的置信度连乘积。相比传统采样只在单步依靠 \(\mathbf{c}_t(j)\) 做局部 hard 截断,SR-SMC 允许粒子在早期即使持有部分不确定的局部 token 也能存活,并依赖全局轨迹累积置信度在后续步骤中修正局部偏见,通过全局概率乘积挑选出整体最具可信度的输出序列。
3. 自适应重采样与随机性控制:平衡粒子多样性与计算开销 频繁进行多粒子全量重采样会导致粒子多样性过早崩溃(Path Degeneracy),并引入不必要的排序与复制通信开销。为此,算法引入有效样本量(Effective Sample Size, ESS)监控机制,在每一步计算: $\(\mathrm{ESS} = \frac{1}{\sum_{i=1}^N (w_t^i)^2}\)$ 仅当粒子系统权重的方差急剧增大、\(\mathrm{ESS} < N/2\) 时才显式触发多项式重采样,淘汰落后轨迹并克隆优质粒子;在分块自回归扩散大模型(dLLMs)中则自然对齐至块边界触发。此外,在离散 token 传播采样时引入受温度参数 \(\tau\) 控制的 Gumbel-Max 采样,既保证了粒子探索不同语义分支所需的离散随机扰动,又通过重采样消除了低确定性搜索路径带来的离散退化。
损失函数 / 训练策略¶
SR-SMC 是纯粹的推理阶段扩展(Inference-time scaling)算法,完全不需要任何模型微调或参数更新,直接无缝插拔于任意预训练好的离散掩码扩散架构之上。在测试时,扩散步数 \(T\) 与粒子数 \(N\) 构成可灵活调配的推理计算预算;对于大语言扩散模型(dLLMs,如 LLaDA-1.5 与 Dream-7B),SR-SMC 原生兼容 KV Cache 缓存加速与块级并行解码策略(Block-wise decoding),在几乎不增加端到端延时量级的前提下将并行显卡算力转化为显著的生成质量提升。
实验关键数据¶
主实验¶
论文首先在通用语言建模任务(OpenWebText 预训练,评估 GPT2-Large 生成困惑度 Gen. PPL 及函数评估次数 NFEs)以及扩散大语言模型代码/数学推理任务(GSM8K、MATH、HumanEval、MBPP)上全面验证 SR-SMC。主实验对比展现了显著的质量飞跃:
| 模型 / 架构 | 评测基准 / 序列长度 | 本文方法 (SR-SMC) | 原始基线 | 提升幅度 |
|---|---|---|---|---|
| MDLM (Diffusion) | OpenWebText (L=1024 Gen. PPL ↓) | 25.8 (4K NFEs) | 46.8 (1K NFEs) | -21.0 PPL (相对改善 44.9%) |
| BD3-LMs (\(L'=4\)) | OpenWebText (L=1024 Gen. PPL ↓) | 16.1 (4K NFEs) | 25.7 (1K NFEs) | -9.6 PPL (逼近 AR 模型的 14.1) |
| BD3-LMs (\(L'=4\)) | OpenWebText (L=2048 Gen. PPL ↓) | 15.1 (8K NFEs) | 23.6 (2K NFEs) | -8.5 PPL |
| LLaDA-1.5 (dLLM) | GSM8K 5-shot (\(L=512\) Acc ↑) | 82.0% | 80.4% | +1.6% |
| LLaDA-1.5 (dLLM) | MATH 4-shot (\(L=512\) Acc ↑) | 45.4% | 41.4% | +4.0% |
| LLaDA-1.5 (dLLM) | HumanEval 0-shot (\(L=512\) Acc ↑) | 41.5% | 35.4% | +6.1% |
| LLaDA-1.5 (dLLM) | 4 项基准平均准确率 (\(L=512\)) | 53.0% | 49.4% | +3.6% |
| Dream-7B (dLLM) | HumanEval 0-shot (\(L=512\) Acc ↑) | 53.7% | 45.1% | +8.6% |
| Dream-7B (dLLM) | MBPP 3-shot (\(L=256\) Acc ↑) | 48.6% | 42.0% | +6.6% |
| Dream-7B (dLLM) | 4 项基准平均准确率 (\(L=512\)) | 56.4% | 51.8% | +4.6% |
消融实验¶
粒子数目 \(N\) 是控制推理阶段计算量与探索广度的核心超参数。下表总结了在生成长度 \(L=256\) 条件下,扩展粒子规模对 LLaDA-1.5 与 Dream-7B 推理表现的影响(其中 \(N=1\) 等价于标准贪婪并行解码下界):
| 模型 | 粒子数配置 (\(N\)) | GSM8K (5-shot) | MATH (4-shot) | HumanEval (0-shot) | MBPP (3-shot) | 综合平均分 (Acc %) |
|---|---|---|---|---|---|---|
| LLaDA-1.5 | \(N=1\) (贪婪基线) | 79.8 | 38.2 | 38.7 | 40.4 | 49.3 |
| LLaDA-1.5 | \(N=2\) | 81.7 | 41.8 | 39.0 | 44.2 | 51.7 |
| LLaDA-1.5 | \(N=3\) | 80.7 | 41.8 | 39.0 | 44.2 | 51.4 |
| LLaDA-1.5 | \(N=4\) (默认配置) | 80.7 | 41.8 | 41.5 | 44.2 | 52.1 (+2.8) |
| Dream-7B | \(N=1\) (贪婪基线) | 76.2 | 41.6 | 47.9 | 42.0 | 51.9 |
| Dream-7B | \(N=2\) | 76.3 | 43.2 | 52.4 | 43.6 | 53.9 |
| Dream-7B | \(N=3\) | 77.9 | 45.6 | 54.6 | 44.2 | 55.6 |
| Dream-7B | \(N=4\) (默认配置) | 78.0 | 45.2 | 53.7 | 48.6 | 56.4 (+4.5) |
关键发现¶
- 逆袭事件(Particle Overtake)证实了全局探索的有效机制:作者对分块生成中 640 个 block 进行统计,发现在各任务中平均有高达 27.6%(MATH 达 30.8%,HumanEval 达 30.6%)的 block 在结束时得分最高的胜出粒子并不是该 block 开始时概率最大的主导粒子。这证明粒子系统并非在贪婪轨迹周围做无意义的微小扰动,而是成功保护了初始不显山露水但具备更高全局演化潜力的生成分支,避免了单轨迹贪婪解码不可逆转的错误锁定。
- 极端随机温度下的卓越鲁棒性:在低采样温度下,标准并行解码基线(如 Dream-7B)在代码和数学任务上因确定性重复陷入极端退化(\(\tau=0.1\) 时 MBPP 准确率暴跌至 1.80%);而 SR-SMC 依靠粒子重加权与重采样,能够在 \(\tau \in [0.0, 1.0]\) 全区间维持高准确率(\(\tau=0.1\) 时稳定在 43.4%),极大拓宽了离散扩散生成的超参安全裕度。
- 质量提升不以牺牲多样性为代价:生成文本的词元熵分析显示,SR-SMC 在困惑度大幅下降的同时,生成熵(\(L=2048\) 时达 5.41~5.49)几乎完全保持在原始人类文本的熵水平(5.44),表明该算法并未坍缩至低熵的退化模式。
亮点与洞察¶
- 将扩散反向核直接视作提议分布使得增量权重完美化简:数学推导极其优雅,利用 Bootstrap SMC 结构令复杂的反向转移核直接相消,证明了“当前步接受 token 预测概率积”即为严格的无偏增量重要性比率,赋予了启发式“置信度加权”坚实的统计物理学理论基础。
- 无需训练和外置奖励即可实现推理算力转化为性能收益:打破了过去离散扩散模型只能靠拉长去噪步数或调用笨重外部判别器做引导的局限,提供了一种通用、即插即用且能充分释放现代 GPU 并行批处理特性的推理性扩展(test-time scaling)新维度。
- 粒子竞争与逆袭机制具有广泛迁移价值:该方法可直接迁移至各类离散非自回归生成任务(如分子序列离散生成、离散代码修复、机器翻译等),为解决掩码自编码器迭代去噪中的级联误差累积提供了通用的即插即用范式。
局限与展望¶
- 推理 FLOPs 与显存随粒子数线性增加:运行 \(N\) 个并行粒子意味着模型正向传播的函数评估次数(NFEs)扩大至 \(N\) 倍,在大参数量扩散模型上长文本生成时的显存峰值开销相对显著。
- 置信度自奖励对模型内生幻觉缺乏外部纠偏能力:当扩散模型自身对某个错误事实具有高度一致的“自信”时,轨迹累积置信度依然可能偏好该高置信度的幻觉轨迹,模型似然与真实世界事实性之间仍存在固有鸿沟。
- 未来方向可探索前瞻性(Look-ahead)与扭曲(Twisted)提议分布:当前提议核直接采用先验反向扩散核,未来若能引入轻量前瞻机制预测未来去噪风险,可进一步减少低效粒子的分支探索,提升计算效率。
相关工作与启发¶
- vs Particle Gibbs for DLMs (Dang et al., 2025): 依赖外部特定奖励模型提供引导并在轨迹间做复杂的 Gibbs 刷新;本文 SR-SMC 完全摆脱外部 reward,直接从扩散内生似然抽取信号,适用于无奖励定义的任意通用文本生成。
- vs Standard Greedy Remasking (Sahoo et al., 2024 / MDLM): 标准方法仅在每一步进行单轨迹 hard 截断,无法挽救早期决策偏差;本文通过多粒子并行与重加权机制,保留了多条潜在分支,有效打破了离散扩散的近视解码陷阱。
- vs Fast-dLLM (Wu et al., 2025b): Fast-dLLM 侧重于利用并行解码与 KV Cache 在单轨迹下压缩推理时延;本文 SR-SMC 则在此高效底层基础之上构建上层粒子探索机制,将并发解码吞吐量直接转化为采样精度的飞跃。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次从严格的 Feynman-Kac 路径测度出发将序贯蒙特卡洛与离散掩码扩散结合,推导简洁优美且免外置奖励。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 种代表性扩散架构(MDLM/BD3/LLaDA/Dream)及困惑度、长文本生成、代码生成与数学推理等全维度验证,粒子消融与逆袭分析极为扎实。
- 写作质量: ⭐⭐⭐⭐⭐ 动机清晰直击痛点,数学推导详尽自洽,实验图表信息量极其丰富。
- 价值: ⭐⭐⭐⭐⭐ 为掩码扩散语言模型的推理阶段扩展(Test-time scaling)奠定了重要的基石性范式。