EntRGi：基于熵感知的离散扩散语言模型奖励引导方法 EntRGi: Entropy Aware Reward Guidance for Diffusion Language Models
用模型预测熵自适应混合软/硬token嵌入，解决离散扩散LLM奖励引导的梯度-可靠性矛盾
前置知识
离散扩散语言模型（dLLM）
离散扩散语言模型是一类非自回归的文本生成模型，与传统GPT式的从左到右逐token生成不同，dLLM从一个完全被mask的序列出发，在多个去噪步骤中并行地逐步恢复真实token。典型的mask diffusion设定中，词汇表包含K个"真实"token和一个特殊的mask token $m$。在每个去噪步骤 $t$，模型对当前所有被mask的位置输出一个概率分布 $p_\theta^l(z_t)$，然后按照预设的选择逻辑（如选择熵最低的若干位置）进行unmask。这个"unmask and commit"模式意味着一旦某个token被确定，它在后续步骤中保持不变。
本文的方法EntRGi是专门为dLLM设计的，理解dLLM的生成过程（mask→去噪→commit）是理解整个方法的前提，尤其是如何在每个去噪步骤中对当前被mask位置的logits进行奖励引导更新。
奖励引导（Reward Guidance）
奖励引导是一种在推理时或训练时利用下游奖励模型的反馈来引导生成过程的技术。在连续扩散模型中（如图像生成），奖励模型的输出对模型输出是可微的，可以直接通过梯度反传来调整生成方向，使得生成结果获得更高奖励。这种方法不需要重新训练基础模型，只需要一个预训练好的奖励模型即可实现推理时的可控生成。在语言模型领域，奖励模型通常是经过微调的语言模型，输出标量分数来衡量生成质量。
本文的核心贡献就是将奖励引导从连续扩散模型迁移到离散扩散语言模型，但面临离散token无法直接求梯度的根本挑战，理解奖励引导的范式是理解本文动机的关键。
直通估计器（Straight-Through Estimator, STE）
STE是一种在离散变量上进行梯度估计的技术，最早由Bengio等人（2013）和Jang等人（2017）提出。其核心思想是：在前向传播中使用离散的采样值（如从分布中采样的hard token），但在反向传播时将梯度"直通"地传回采样前的连续分布。用数学表达就是 $e_l^{\text{APS}} = e_l^{\text{soft}} + \text{sg}(e_l^{\text{hard}} - e_l^{\text{soft}})$，其中 $\text{sg}$ 是stop-gradient算子。前向计算得到 $e_l^{\text{hard}}$（离散token的嵌入），但梯度通过 $e_l^{\text{soft}}$（连续期望嵌入）传播。
STE是本文对比基线APS方法的核心机制，也是EntRGi方法中在高熵位置使用的技术。理解STE的工作原理及其引入的近似误差 $E_l = \|e_l^{\text{hard}} - e_l^{\text{soft}}\|$，是理解EntRGi如何改进APS的关键。
词汇误差与近似误差
论文定义了两个关键误差量。词汇误差 $D_l = \min_k \|e_l - E_R^k\|$ 衡量的是送给奖励模型的输入嵌入 $e_l$ 与奖励模型词汇表中最接近的真实token嵌入 $E_R^k$ 之间的距离，反映了奖励模型看到的输入是否"合理"。近似误差 $E_l = \|e_l^{\text{hard}} - e_l^{\text{soft}}\|$ 衡量的是STE中离散采样token与连续期望之间的差距，反映了梯度传播的准确性。这两个误差存在此消彼长的关系：期望方法有零词汇误差但高近似误差（高熵时），STE方法有零词汇误差但高近似误差。
这两个误差的矛盾是本文要解决的核心技术难题，EntRGi的核心创新就是通过熵感知的自适应插值来同时控制这两个误差。
预测熵（Predictive Entropy）
在dLLM中，每个被mask的位置 $l$ 在去噪步骤 $t$ 会输出一个概率分布 $q_l$，该分布的Shannon熵 $H(q_l) = -\sum_k q_l^k \log q_l^k$ 反映了模型对该位置应该填什么token的不确定性。熵越高表示模型越不确定（分布越均匀），熵越低表示模型越确信（分布越集中于少数token）。在去噪早期（$t$ 接近 $T$），大多数位置的熵较高；随着去噪进行，越来越多位置的熵降低。最大可能熵为 $\log K$，其中 $K$ 是词汇表大小。
EntRGi的核心创新就是利用这个预测熵作为插值权重 $w_l = H(q_l) / \log K$，在低熵（模型确信）时偏向连续软嵌入，在高熵（模型不确定）时偏向离散硬嵌入，从而自适应地平衡梯度准确性和奖励可靠性。
研究动机
离散扩散语言模型（dLLM）在奖励引导方面面临一个根本性的困境。在连续扩散模型中（如图像生成），可以直接通过奖励模型反传梯度来引导生成，因为连续输出天然可微。但dLLM的输出是离散token，无法直接求梯度。现有两种解决方案各有严重缺陷：第一种是"期望方法"（Murata et al., 2024; Tae et al., 2025），将离散分布的期望软嵌入 $e_l^{\text{soft}} = \sum_{k=1}^K q_l^k E_R^k$ 送给奖励模型。但奖励模型只在真实token嵌入上训练过，从未见过这种"软"输入，尤其在高熵时，期望嵌入会远离任何真实token嵌入，导致奖励估计不可靠，词汇误差 $D_l$ 很大。第二种是APS方法（Rout et al., 2025c），使用STE在前向传播中采样hard token给奖励模型（保证 $D_l = 0$），但梯度通过soft embedding传播。这引入了近似误差 $E_l = \|e_l^{\text{hard}} - e_l^{\text{soft}}\|$，在高熵位置尤为严重——因为采样的hard token与期望soft embedding差距很大，梯度方向不准确。简而言之，现有方法要么牺牲奖励模型的可靠性，要么牺牲梯度的准确性，无法兼得。
本文的目标是本文的目标是设计一种新的奖励引导机制，能够同时保证奖励模型的可靠性（输入是奖励模型能理解的真实token嵌入）和梯度传播的准确性（梯度方向能正确反映奖励函数的变化方向）。具体来说，作者希望在dLLM的每个去噪步骤中，对当前被mask的位置计算出一组能最大化下游奖励的logits更新，使得最终生成的文本能获得更高奖励。此外，作者还希望将这种推理时的奖励引导能力扩展到后训练（post-training）阶段，通过自蒸馏的方式生成高质量的奖励引导样本，进而微调模型本身。
与已有工作不同的是，本文的独特切入角度在于认识到：在dLLM的去噪过程中，不同token位置的不确定性是不同的，而且同一个位置在不同去噪步骤中的不确定性也在变化。这是一个被所有现有方法忽视的关键信号。期望方法和APS方法都对所有位置采用相同的处理策略——前者全部用软嵌入，后者全部用硬嵌入——没有考虑到模型自身对每个位置的"信心"程度。本文提出利用模型自身的预测熵 $H(q_l)$ 作为自适应信号：当模型对某个位置很确信时（低熵），用软嵌入更安全，因为软嵌入已经接近某个真实token；当模型对某个位置不确定时（高熵），必须用硬嵌入，否则奖励模型看到的输入会严重偏离训练分布。这个看似简单的直觉——"让模型的不确定性指导如何向奖励模型展示自己的输出"——是本文最核心的创新点。
核心方法
EntRGi的核心直觉可以用一个日常类比来理解：想象你在向一位严格的产品经理（奖励模型）展示你的设计方案（生成的token序列）。当你对自己的方案很确信时，可以直接展示详细的设计草图（软嵌入），因为PM能从中看到你想表达的东西；但当你还在探索多种可能性时（高熵），最好先确定一个具体的方案再展示（硬嵌入），否则PM会因为看到模糊的、不伦不类的设计而给出不准确的反馈。从技术路线来说，EntRGi在每个去噪步骤中对每个被mask的位置 $l$ 计算一个插值权重 $w_l = H(q_l) / \log K$，然后构造奖励模型的输入为 $$e_l^{\text{EntRGi}} = e_l^{\text{soft}} + \text{sg}[w_l(e_l^{\text{hard}} - e_l^{\text{soft}})]$$ 这是一个关于 $e_l^{\text{soft}}$ 的凸组合，其中 $w_l = 0$ 退化为期望方法，$w_l = 1$ 退化为APS。通过让权重随熵自适应变化，EntRGi在每个位置自动选择最合适的策略，从而在保持梯度流的同时确保奖励模型看到合理的输入。
EntRGi与已有方法的本质区别在于引入了"位置级"和"时间步级"的自适应机制。期望方法和APS都是对所有位置、所有时间步采用统一策略的"一刀切"方法。EntRGi的核心创新是认识到：在去噪早期（$t$ 接近 $T$），大多数位置的熵很高，此时应该用STE（高 $w_l$）确保奖励模型看到合理输入；在去噪后期，大部分token已确定，剩余位置的熵较低，此时应该用软嵌入（低 $w_l$）获得更准确的梯度。即使在同一个时间步内，已获得大量上下文的位置（如句末的标点）熵可能很低，而语义上开放的位置（如形容词选择）熵可能较高，EntRGi能对它们分别处理。从数学上可以严格证明，EntRGi的近似误差 $E_l^{\text{EntRGi}} = w_l \cdot E_l^{\text{APS}} < E_l^{\text{APS}}$（只要 $w_l < 1$），而词汇误差 $D_l^{\text{EntRGi}} \leq (1 - w_l) \|e_l^{\text{hard}} - e_l^{\text{soft}}\|$，在低熵时很小，高熵时 $w_l \to 1$ 保证 $D_l \to 0$。因此EntRGi严格优于APS。
方法步骤详情
EntRGi的完整推理时算法（Algorithm 1）如下：首先初始化一个全mask序列 $z_T = m^L$，mask位置集合 $M_T = [L]$。然后对每个去噪步骤 $t = T, T-1, \ldots, 1$：(1) 将当前序列 $z_t$ 输入dLLM，获得所有被mask位置的概率分布 $q^{M_t} = \{q_l, l \in M_t\}$；(2) 对每个被mask位置，计算期望软嵌入 $e_l^{\text{soft}} = \sum_{k=1}^K q_l^k E_R^k$；(3) 从 $q_l$ 中采样得到 $x_l \sim q_l$，得到hard token嵌入 $e_l^{\text{hard}} = E_R[x_l]$；(4) 计算每个位置的熵权重 $w_l = H(q_l) / \log K$；(5) 构造奖励模型输入：对被mask位置 $l \in M_t$，$e_l^{\text{EntRGi}} = e_l^{\text{soft}} + \text{sg}[w_l(e_l^{\text{hard}} - e_l^{\text{soft}})]$；对已unmask位置，直接用token嵌入；(6) 通过奖励模型 $R(e^{\text{EntRGi}})$ 的梯度对logits $q_l$（$l \in M_t$）进行梯度上升更新，重复 $M$ 步内循环；(7) 根据更新后的logits，按照预设选择逻辑 $U(q^{M_t})$ unmask若干token并commit，其余位置保持mask。最终返回奖励引导后的完整序列 $z_0$。后训练算法RGRL（Algorithm 2）则是在每个训练步中采样prompt，用Algorithm 1生成 $N$ 个奖励引导的补全，然后更新模型参数 $\theta$ 以增加这些高质量补全的似然。
技术新颖性
EntRGi的技术新颖性体现在多个层面。首先，在离散扩散语言模型的奖励引导中引入自适应插值机制是全新的——之前的方法要么全部用软嵌入（期望方法），要么全部用硬嵌入（APS），从未考虑过根据模型的不确定性进行动态选择。其次，将dLLM自身的预测熵作为自适应信号是非常自然但之前无人提出的想法——这个信号完全来自模型本身，不需要额外的计算或启发式。第三，通过严格的数学分析证明了EntRGi的近似误差在每个熵水平上都严格低于APS，给出了理论保证。第四，RGRL是第一个利用奖励梯度反馈（而非纯标量奖励）进行dLLM后训练的工作，通过自蒸馏奖励引导的补全来微调模型，相比标准的diffu-GRPO等RL方法能提供更密集的监督信号。最后，论文还处理了dLLM与奖励模型tokenizer不匹配的实际问题（如LLaDA与Llama/Qwen的重叠率只有45%-55%），对不重叠的token设零嵌入，这在实际应用中非常重要。
实验结果
论文在两个7B+参数的dLLM（Dream-v0-Instruct-7B和LLaDA-8B-Instruct）、三个评估基准（Reward-Bench-2、JudgeBench、RM-Bench）、四个Skywork奖励模型（0.6B到4B参数）上进行了全面实验。在测试时自适应设置中，EntRGi在Reward-Bench-2上Top@1得分3.07（τ=0.1）和3.91（τ=0.7），分别超过APS的2.95和3.62；LMUnit外部评判分数从4.19提升到4.22。特别是在高采样温度（τ=0.7）下，EntRGi的优势更明显，而APS明显退化，这验证了高熵场景下STE的局限性。在RM-Bench上，EntRGi的Top@1从APS的4.72提升到4.90（τ=0.1），5.11提升到5.70（τ=0.7），提升幅度达11.5%。在tokenizer不匹配的LLaDA-8B上，EntRGi同样保持优势，如在Reward-Bench-2上Top@1达6.40（vs APS的6.32）和2.80（vs 2.35）。在后训练实验中，RGRL在WildChat-IF数据集上用Dream模型实现了高达70%的相对提升（绝对+0.90），显著优于diffu-GRPO。RGRL-EntRGi在所有设置中均获得最高持续增益。在不同奖励模型规模（0.6B→1.7B→4B）的实验中，EntRGi在每个规模上都优于APS，且随奖励模型增大，性能差距保持稳定。内循环优化步数 $M$ 的分析表明，$M = 3$-$4$ 是最佳范围，超过此范围可能出现reward hacking。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Reward-Bench-2（测试时自适应，Dream-v0-7B，τ=0.1） | Top@1 | EntRGi: 3.07±0.22 | APS: 2.95±0.21 | +4.1% |
| Reward-Bench-2（测试时自适应，Dream-v0-7B，τ=0.7） | Top@1 | EntRGi: 3.91±0.30 | APS: 3.62±0.27 | +8.0% |
| RM-Bench（测试时自适应，Dream-v0-7B，τ=0.7） | Top@1 | EntRGi: 5.70±0.12 | APS: 5.11±0.14 | +11.5% |
| JudgeBench（测试时自适应，Dream-v0-7B，τ=0.7） | Top@1 | EntRGi: 2.44±0.06 | APS: 1.87±0.14 | +30.5% |
| Reward-Bench-2（后训练，Dream-v0-7B，WildChat-IF） | 平均奖励 | RGRL-EntRGi: +0.90 | diffu-GRPO: 基线 | 70%相对提升 |
| Reward-Bench-2（测试时自适应，LLaDA-8B，tokenizer不匹配，τ=0.7） | Top@1 | EntRGi: 6.40±0.33 | APS: 6.32±0.48 | +1.3%，方差更小 |
| LMUnit评分（Dream-v0-7B，三个基准平均，τ=0.7） | LMUnit | EntRGi: 4.06 | APS: 4.05 | 在所有奖励模型规模上均一致领先 |
局限与改进
论文承认的局限性包括：(1) 所有基于梯度的奖励引导方法（包括EntRGi和RGRL）都需要可微的奖励模型，这限制了方法的适用范围——许多实际部署的奖励模型可能是黑盒API或不可微的模型；(2) 反传梯度通过奖励模型引入了额外的计算开销，在某些设置下可能抵消样本效率的提升——论文在4个后训练设置中只有1个实现了wall-clock时间的加速；(3) 最优内循环步数 $M$ 是数据集依赖的（Figure 5显示不同数据集在 $M$ = 3-4时最优），需要针对具体任务调参；(4) 当 $M$ 过大时会出现reward hacking，LMUnit分数在 $M > 4$ 后崩溃。此外，我观察到：(5) EntRGi的提升幅度在低熵位置（去噪后期）与APS趋同（Figure 3），其主要优势体现在早期高熵阶段，这意味着如果dLLM本身有很好的初始去噪质量，EntRGi的边际收益可能有限；(6) 在tokenizer严重不匹配的情况下（如LLaDA仅45%-55%重叠），不重叠token的梯度完全丢失，这可能导致优化信号不完整。
独立分析的弱点
尽管EntRGi取得了显著改进，仍存在几个值得深入分析的弱点。首先，熵权重 $w_l = H(q_l) / \log K$ 的设计虽然直觉合理，但缺乏对该特定函数形式的理论最优性证明——论文在附录C.3中消融了替代加权机制，但没有给出选择 $H(q_l) / \log K$ 而非其他映射（如sigmoid变换）的充分理由。改进方向：可以探索基于奖励模型对输入分布敏感度的自适应权重学习。其次，内循环的 $M$ 步梯度更新在每步去噪中都是固定的，但实际上不同时间步的最优 $M$ 可能不同——早期可能需要更多更新（因为不确定性高），后期可能一步就够。改进方向：可以根据熵的变化动态调整 $M$。第三，方法目前只在Skywork系列奖励模型上验证，这些模型本身是基于Qwen/Llama架构微调的，与其他类型奖励模型（如基于对比学习的、基于Bradley-Terry模型的）的兼容性未被探索。改进方向：在更多样化的奖励模型架构上测试。第四，论文的后训练实验（RGRL）在lmsys-chat-1m数据集上的增益相对较小（如Dream上仅从-0.28提升到+0.33），暗示在已有较好初始质量的数据上，密集奖励梯度信号的优势减弱。改进方向：探索在更有挑战性的数据分布上应用RGRL。
未来方向
论文作者提出了四个未来方向：(1) 选择性地仅在最具信息量的去噪步骤中应用奖励梯度反馈，以更好地权衡样本效率和计算开销——这可以借鉴自适应计算的思想，在熵变化最剧烈的步骤给予更多关注；(2) 扩展到多目标奖励组合，即同时优化安全性、有帮助性、事实性等多个维度，这在实际对齐场景中非常重要；(3) 开发更原则性的tokenizer不匹配处理方案，而不是简单地将不匹配token设为零嵌入；(4) 将可微奖励模型的密集反馈与RLVR（Reinforcement Learning with Verifiable Rewards）风格的方法结合。基于已有成果，还可以延伸以下方向：(5) 将EntRGi的自适应插值思想应用到连续扩散模型的离散化变体中；(6) 研究在few-shot或in-context learning场景下，无需额外奖励模型，直接利用任务指标的梯度进行推理时引导；(7) 探索EntRGi与Best-of-N采样的组合策略，在计算预算充裕时获得更好的trade-off。
复现评估
论文的可复现性较好。代码已在GitHub开源（https://atutej.github.io/entrgi-rgrl），实验中使用的基础模型（Dream-v0-Instruct-7B、LLaDA-8B-Instruct）和奖励模型（Skywork系列，0.6B/1.7B/4B）均为公开可下载的模型。数据集方面，Reward-Bench-2、JudgeBench、RM-Bench均为公开基准，后训练使用的WildChat-IF和lmsys-chat-1m也是公开数据集。论文提供了详细的附录（Appendix B包含实验设置，Appendix C包含消融实验），包括超参数设置（引导尺度 $\eta$、内循环步数 $M$、采样温度 $\tau$）。在算力需求方面，实验在两个7B+参数模型上进行，内循环需要 $M$ 步奖励模型前向+反向传播，因此复现需要多GPU资源。论文对所有实验报告了标准差（如Table 1中的 ± 值），便于统计显著性评估。总体而言，给定足够的GPU资源和论文提供的细节，复现应该是可行的。
论文图表
