SLIM-RL: 基于风险预算的随机掩码强化学习用于扩散语言模型，无需轨迹切分

# 风险预算约束的随机掩码强化学习：无需轨迹切分的扩散语言模型  
来源：https://arxiv.org/html/2607.00208  

赵瑞康 丹麦技术大学  
王振挺 MBZUAI 基础模型研究所  
韩高 爱荷华州立大学  
韩利光* Red Hat AI 创新实验室 & MIT–IBM Watson AI 实验室  

###### 摘要  

扩散大语言模型 (dLLMs) 的强化学习已主要转向轨迹感知方法。当前最先进的方法 TraceRL 认为随机掩码与模型的推理轨迹不匹配，因此在训练过程中通过将每次 rollout 切片成最多 ⌈K/s⌉ 个轨迹对齐的训练样本来重建轨迹，其成本随块大小 K 增长。我们证明这种不匹配可以在不重建轨迹的情况下得到缓解。我们的方法 SLIM-RL 使用 τ-预算解码器约束每个 rollout 步骤的提交风险，从而降低训练数据中的累积提交风险。在优化过程中，SLIM-RL 在这些风险控制的 rollouts 上使用无迹随机掩码目标函数，该目标函数适配方差缩减工具，结合序列级重要性采样、掩码水平上的确定性正交规则，以及我们引入的一种均值保持、单调递减的逐块掩码调度。在 SDAR-4B 上，SLIM-RL 在块大小 16 时仅使用 TraceRL 训练样本的 0.46× 就达到了其最佳的 MATH500 准确率，在匹配动态采样下，在 MATH500 上比 TraceRL 提升 6.32%，在 GSM8K 上提升 11.05%。在块大小 4 时，4B 的 SLIM-RL 在数学任务上超过了更大的 LLaDA-8B 和 Dream-7B dLLMs，在 MATH500 上超越 LLaDA-8B 达 10.76%，但仍低于自回归模型 Qwen2.5-7B。在代码任务上，它在 MBPP 上比 TraceRL 提升 4.20%，在 HumanEval 上提升 3.65%。τ-预算解码器可无训练地迁移到 LLaDA、Dream 和 SDAR。源代码可在 https://github.com/laolaorkkkkk/SLIM-RL 获取。  

## 1 引言  

扩散大语言模型 (dLLMs) (Nie et al., 2026; Sahoo et al., 2024; Ye et al., 2025) 通过迭代去噪掩码序列来生成文本，同时精炼多个位置，而非严格从左到右。块变体 (Arriola et al., 2025; Cheng et al., 2026) 在自回归解码连续块的同时，在每个块内进行去噪，从而恢复 KV 缓存复用。强化学习已成为提升其推理能力的标准方法 (Zhao et al., 2026; Wang et al., 2025b; Hu et al., 2026; He et al., 2025; Zhu et al., 2025b; Tang et al., 2026; Liu et al., 2026; Zhu et al., 2025a)。早期工作将 RL 应用于随机掩码目标 (Zhao et al., 2026; Yang et al., 2026)，但随机掩码会产生高方差的梯度估计 (Zhu et al., 2025a)，因此近期工作已转向轨迹感知方法 (Huang et al., 2026; Wang et al., 2026)。当前最先进的方法 TraceRL (Wang et al., 2025b) 不在随机掩码上训练，而是在模型的确切解码轨迹上训练，其观点是：在随机掩码下，训练后的目标函数与模型在推理时所遵循的轨迹不匹配。保持轨迹并非没有代价。为了按顺序在每一步解码上进行训练，TraceRL 将一次 rollout 切片成最多 ⌈K/s⌉ 个轨迹对齐的训练样本，每个样本一次前向传播，其中 K 是块大小，s ≥ 1 是收缩因子 (Wang et al., 2025b)。在全保真度 (s=1) 下，一次 rollout 变成最多 K 个样本，因此数据量随 K 增长。提高收缩 s 会限制成本，但会将 s 个连续解码步骤聚合到一个切片中，丢弃其内部顺序，从而用轨迹保真度换取成本。因此，更大的块需要在切片成本和轨迹保真度之间进行权衡，我们将在第 4.5 节中量化这一权衡。问题是，重建确切轨迹是否必要。  

没有切片的情况下，dLLM-RL 运行仍会在两个独立阶段面临弱点：rollout 生成和优化。Rollouts 使用动态采样 (Wu et al., 2025; Yu et al., 2025) 生成，这是包括 TraceRL 在内的所有 dLLM-RL 方法通用的 rollout 解码器，它会提交每个置信度超过固定阈值 τ 的 token。由于该规则是逐点的，一步解码中如果有很多位置仅略高于 τ，就会过度提交，从而向 rollout 中注入多个错误。在优化阶段，随机掩码导致策略梯度估计高方差 (Zhu et al., 2025a)。  

在这项工作中，我们解决了这两个阶段，并将结果方法称为 SLIM-RL，它结合了第 3 节中发展的三个组件。一个 τ-预算动态解掩码解码器，一种单次前向传播、无需训练的规则，同时用于 rollout 和推理，它限制每一步的累积置信度不确定性。一个方差缩减框架更新策略，它由序列级长度归一化比率 (Zheng et al., 2025)、掩码水平上的确定性正交规则 (Rojas et al., 2025) 和未归一化优势 (Liu et al., 2025) 构成。一种均值保持、单调递减的逐块掩码调度将策略梯度信号集中到最早、条件最丰富的块上。在 SDAR-4B 上，SLIM-RL 在块大小 16 时在数学任务上优于轨迹感知的 TraceRL，在块大小 4 时在数学和代码任务上均优于 TraceRL，并且 4B 模型还超越了更大的 LLaDA-8B 和 Dream-7B 扩散模型，在不到一半的训练数据上达到了 TraceRL 的最佳准确率（表 6）。优势随块大小扩大而增加，第 4 节报告了每个基准的结果。我们的贡献总结如下：  

- • 我们表明，在相等训练成本下，重建确切解码轨迹并非匹配轨迹感知 RL 所必需的。通过限制 rollout 提交风险并降低掩码目标方差，无迹随机掩码在块大小 4 时与轨迹对齐切片相当，在块大小 16 时则更优。  
- • 我们引入了一种 τ-预算动态解掩码解码器，它只提交累积不确定性 ∑ᵢ(1-pᵢ) 保持在校准预算 m(1-τ) 以内的、不确定性最低的较大位置子集，这与动态采样的逐位置阈值形成对比。它无需训练，可迁移到 LLaDA、Dream 和 SDAR。  
- • 我们引入了一种均值保持、单调递减的逐块掩码调度，将策略梯度信号集中到最早、条件最丰富的块上，在 MATH500 上比随机掩码提升 5.91%。  
- • 在 SDAR-4B 上，块大小 16 时，SLIM-RL 在匹配解码下，在 MATH500 和 GSM8K 上分别比 TraceRL 准确率提升 6.32% 和 11.05%，并且仅用 0.46× 的训练数据就达到了 TraceRL 的最佳准确率。块大小 4 时，它在数学和代码任务上均达到比更大的 LLaDA-8B 和 Dream-7B 扩散模型更高的准确率，在 MBPP 上比 TraceRL 提升 4.20%，在 HumanEval 上提升 3.65%。在 SDAR-1.7B 上，SLIM-RL 用 0.76× 的训练数据就达到了 TraceRL 的最佳结果。  

## 2 预备知识  

### 2.1 扩散大语言模型  

#### 前向和反向过程。  
给定一个提示 x 和一个干净响应 y = (y¹, ..., yᴸ)，扩散大语言模型 (dLLM) (Nie et al., 2026; Ye et al., 2025) 定义了一个前向破坏过程和一个学习的反向生成过程。前向过程在掩码水平 t ∈ [0,1] 上独立地用掩码符号 [MASK] 替换每个 token：  
qₜ(yₜ ∣ y) = ∏ᵢ₌₁ᴸ qₜ(yₜⁱ ∣ yⁱ), (1)  
其中每个 token 根据 qₜ(yₜⁱ = [MASK] ∣ yⁱ) = t 和 qₜ(yₜⁱ = yⁱ ∣ yⁱ) = 1-t 独立地掩码，因此 y₀ 是干净响应，y₁ 是完全掩码的。对于反向过程，从 y₁ 开始，模型迭代地预测掩码位置并从 yₜ 重新采样一个噪声较小的 yᵣ (r < t)，最终生成干净的响应（第 3.1 节中的解码器描述了精确的采样和提交规则）。  

#### 块扩散。  
块扩散 (Arriola et al., 2025; Cheng et al., 2026) 将响应分割成大小 K 的 B 个块（最后一个块可能更小）。解码按块顺序进行：对于每个块 b ∈ {1,...,B}，模型在潜在变量 zb∼Unif[0,1] 的引导下，将块 b 的初始完全掩码序列去噪到其最终形式。下一个块在解码完当前块后开始，因此因果依赖关系得以保持。为了生成块 b，模型使用条件上下文 x ∪ {y₁,...,y_{b-1}}，通过多个去噪步骤，将块 b 内的 Token 从 [MASK] 转换为干净 Token。在推理时，标准方法是在每个去噪步骤中提交所有 pᵢ > τ 的位置（动态采样 (Wu et al., 2025; Yu et al., 2025)）。第 3.2 节描述了 τ-预算解码器，它通过硬预算来约束这种提交。  

### 2.2 用于 dLLMs 的 GRPO  

对于每个提示 x ∼ D，旧策略采样 G 个响应 {yⱼ}ⱼ₌₁ᴳ ∼ π_θ_old(· ∣ x)，并获得标量奖励 {rⱼ}ⱼ₌₁ᴳ，标准化优势 Âⱼ = (rⱼ - mean{rᵢ}) / std{rᵢ}。标准 GRPO (Shao et al., 2024) 定义了 token 级重要性比率 ρⱼᵏ(θ) = π_θ(yⱼᵏ ∣ x, y_{<k}) / π_θ_old(yⱼᵏ ∣ x, y_{<k})，并优化以下目标：  

L_GRPO(θ) = ¹/ᴳ Σⱼ₌₁ᴳ ¹/|yⱼ| Σₖ₌₁^{|yⱼ|} [min(ρⱼᵏ(θ) Âⱼ, clip(ρⱼᵏ(θ), 1-ε, 1+ε) Âⱼ) - β D_KL(π_θ || π_ref)], (2)  

其中 clip 将比率限制在 [1-ε, 1+ε] 内，β 项惩罚偏离参考策略。后续工作 (Zhao et al., 2026; Yang et al., 2026; Wang et al., 2025b) 将 GRPO 适应于 dLLMs，但采用随机的 Token 级掩码，这使得 (2) 中的梯度估计具有高方差 (Zhu et al., 2025a)。  

## 3 方法  

SLIM-RL 由三个组件组成：一个 τ-预算解码器（第 3.1 节），它通过限制每一步的不确定性来提供更安全的 rollouts；一个方差缩减的随机掩码目标函数（第 3.2 节），它联合使用序列级比率、确定性正交规则和未归一化优势；以及一个均值保持的掩码调度（第 3.3 节），它优先考虑早期块的梯度信号。  

### 3.1 τ-预算解码器：受控提交  

在块的每个去噪步骤 s，模型为每个掩码位置 i 生成置信度 pᵢ（例如，logits 上的 softmax 概率）。标准的动态采样提交所有位置，在这些位置 pᵢ > τ，其中 τ 是固定阈值。激进的阈值（小 τ）会每步提交更多 token，但会注入错误，这些错误会通过后续去噪传播。保守阈值保留了质量，但需要更多精炼步骤，从而削弱了 dLLM 的并行性优势。  

τ-预算解码器不是提交所有置信度高于 τ 的位置，而是提交累积置信度不确定性保持在已校准预算内的最大位置子集。形式上，设 M_s 为步骤 s 时块内的掩码位置集合，并设  
A_s = {i ∈ M_s : pᵢ > τ} (4)  
为标准动态采样在步骤 s 会提交的位置集合。我们为每个候选分配置信度不确定性 uᵢ = 1 - pᵢ，并定义集合 S ⊆ A_s 的累积步骤级不确定性为 U(S) = ∑_{i∈S} uᵢ = ∑_{i∈S} (1-pᵢ)。给定一个风险预算调度 m ≥ 1，我们设置步骤级预算 B_s = m(1-τ)。我们按不确定性升序排序候选：u_{i₁} ≤ u_{i₂} ≤ ... ≤ u_{i_{|A_s|}}，并提交累积不确定性保持在预算内的最大前缀：  

k_s = max{k: ∑_{r=1}^{k} u_{i_r} ≤ B_s}, (5)  
C_s = {i₁, ..., i_{k_s}}.  

参见图 1 说明：单个去噪步骤上的动态采样与 τ-预算解掩码对比。  
两者从相同的部分去噪块和相同的超越阈值候选集 A_s 开始，即置信度超过 τ=0.9 的三个位置。  
*(A) 动态采样提交每个候选，最终确定所有三个，包括置信度为 0.94 的位置。  
*(B) τ-预算为每个候选分配不确定性 uᵢ=1-pᵢ，按升序排序，并仅提交累积不确定性保持在预算 B_s=m(1-τ)=0.10 内的最大前缀。剩余位置保持掩码，用于后续去噪步骤。  
如果 A_s = ∅，我们提交单个置信度最高的掩码位置。  
图 1 对比了 τ-预算与动态采样在单个去噪步骤上的表现，算法 2 给出了完整的单次前向传播解码器。  

#### 解释：校准的错误提交预算。  
设 WrongCommit_s = ∑_{i∈C_s} 1{ŷᵢ ≠ yᵢ*} 统计与正确 token 不一致的已提交位置。如果最大置信度近似等于 argmax 正确的概率，则 E[WrongCommit_s] ≈ ∑_{i∈C_s} (1-pᵢ) = U(C_s) ≤ B_s，因此预算限制了一个校准的代理量，即每步预期错误提交数。逐点阈值则提交所有 A_s，其累积不确定性随接受 token 数量增长，因此多个略高于阈值的位置可能共同注入较大的步骤级风险。校准假设仅近似成立，因此我们将其视作动机，而 τ-预算 rollouts 每步提交的预期错误 token 少于动态采样（附录图 5）。  

### 3.2 训练目标：方差缩减的随机掩码  

收集 rollouts 后，SLIM-RL 不重建 token 级解码顺序。随机掩码用于 RL 的本质弱点是高方差的策略梯度。方差的第一个来源是重要性比率的粒度。token 级比率将一个响应级优势附加到多个独立裁剪的逐位置比率上。如 GSPO (Zheng et al., 2025) 所指出的，这种 token 级重要性权重会引入高方差训练噪声，并可能随响应长度累积。因此，我们使用序列级比率，它针对给定掩码水平的整个响应计算。  

对于一个掩码水平 t，令 ỹⱼ(t) 为 yⱼ 的随机掩码版本，令 Uⱼ(t) 表示其掩码位置，并令 Mⱼ(t) = |Uⱼ(t)|。归一化的序列去噪对数得分为：  

ℓ_θⱼ(t) = ¹/ᴹⱼ(ᵗ) ∑_{i∈Uⱼ(t)} log p_θ(yⱼⁱ ∣ x, ỹⱼ(t)), (6)  

其中 yⱼⁱ 是掩码位置 i 处的干净 token。对应的序列级去噪比率是掩码位置上逐位置去噪比率的几何平均：  

ρⱼ^(t)(θ) = exp(ℓ_θⱼ(t) - ℓ_θ_oldⱼ(t)) = exp(¹/ᴹⱼ(ᵗ) ∑_{i∈Uⱼ(t)} log[p_θ(yⱼⁱ ∣ x, ỹⱼ(t)) / p_θ_old(yⱼⁱ ∣ x, ỹⱼ(t))]). (7)  

### 3.3 确定性正交规则与未归一化优势  

方差的第二个来源是掩码