ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

扩散奖励模型:从RLHF痛点出发的分布式奖励建模实践

扩散奖励模型:从RLHF痛点出发的分布式奖励建模实践 1. 从RLHF的痛点说起为什么奖励模型需要扩散思维如果你在过去一年里动手训练过任何基于人类反馈的强化学习RLHF流程大概率会被奖励模型Reward Model, RM的稳定性折磨过。传统的奖励建模思路很直接拿一个预训练语言模型当骨干接一个标量输出头用成对偏好数据chosen/rejected做 Bradley-Terry 风格的排序损失训练最后得到一个能给任意回复打分的函数。听起来干净利落但真正跑起来你会发现这个标量头学到的往往是一个高度脆弱、对分布外样本极度敏感的评分面。问题出在哪我自己的体会是传统RM本质上是在学一个判别式的打分函数它只告诉你这个回答比那个好却完全没有建模好的回答长什么样。当策略模型在RL阶段开始探索、生成分布逐渐偏离偏好数据集时RM的评分就会失真出现reward hacking——策略学会了钻RM的空子输出一堆在RM眼里高分、在人类眼里垃圾的内容。这是RLHF最经典的翻车现场。Diffusion Reward Models扩散奖励模型这个方向核心动机就是换一个思路不再只学一个打分函数而是学一个关于优质回复的生成式分布。扩散模型天生擅长建模复杂分布把它用在奖励建模上意味着RM不再是一个冷冰冰的标量头而是一个能想象出好答案的生成模型。这个视角的转变直接引出了标题里分布式奖励建模的说法——注意这里的分布式不是指多机多卡的系统分布式而是指奖励信号从单一标量变成了分布式的、基于生成过程的建模。这篇文章我会把这条技术路线拆开讲透扩散模型为什么能当奖励模型用、奖励信号怎么从去噪过程里提取、训练目标怎么设计、和传统RM相比在哪些场景下真的更稳以及如果你想自己复现哪些坑是必须提前知道的。内容偏原理加实操适合已经了解RLHF基本流程、想深入奖励建模这一环的读者。如果你连PPO和偏好数据都还没接触过建议先补一下基础再回来。2. 扩散模型做奖励建模的底层逻辑把打分变成生成2.1 传统RM的判别式本质与它的天花板先把传统RM的数学形式摆出来这样对比才清楚。给定提示 $x$、回复 $y$传统RM学的是 $r_\phi(x, y)$训练目标是让 $r_\phi(x, y_w) r_\phi(x, y_l)$其中 $y_w$ 是人类偏好的回复$y_l$ 是被拒绝的。损失通常写成L_RM -log σ(r_φ(x, y_w) - r_φ(x, y_l))这个形式的问题在于它只约束了相对顺序对绝对分值没有任何约束。理论上你可以给所有回复加同一个常数损失不变。这就导致RM的评分尺度是任意的、漂移的RL阶段策略稍微一探索评分面就崩了。更麻烦的是判别式模型没有对 $p(y|x)$ 的建模能力它不知道什么样的 $y$ 是合理的只知道这两个 $y$ 谁更讨喜。我在实际项目里见过太多次RM在验证集上的pairwise accuracy能到70%以上看起来不错但一接进PPOKL散度稍微放开一点生成质量就断崖式下跌。根因就是RM缺乏对生成分布的约束它是个挑剔的裁判不是懂行的教练。2.2 扩散先验为什么天然适合当奖励信号扩散模型的核心是学一个去噪过程从纯噪声 $y_T$ 出发逐步去噪得到 $y_0$每一步都在估计 $p(y_{t-1} | y_t)$。训练目标是最小化去噪误差等价于在学数据分布 $p(y)$ 的score function $\nabla_y \log p(y)$。把这个能力搬到奖励建模上思路就变成了用扩散模型建模优质回复的条件分布 $p(y | x, \text{high reward})$。一旦你有了这个分布奖励信号可以从多个角度提取似然角度一个回复 $y$ 在优质分布下的似然越高奖励越高。这比判别式打分有更强的分布约束。去噪一致性角度给回复加噪再去噪如果重建误差小说明它落在优质分布的高密度区。score匹配角度比较回复在优质分布和劣质分布下的score差异作为奖励。这里的关键洞察是扩散模型学到的score function本身就是一种密集的、逐时间步的监督信号。传统RM只在最终输出端给一个标量而扩散RM在整个去噪轨迹上都有信息这天然抗reward hacking——策略很难同时骗过所有时间步的去噪过程。2.3 分布式奖励建模到底分布在哪标题里的分布式我理解有三层含义值得说清楚否则容易和系统层面的分布式训练混淆第一层是分布式的奖励表示。奖励不再是一个标量而是一个分布 $p(r | x, y)$ 或者一条去噪轨迹上的奖励序列。这让下游RL可以用期望、方差、分位数等多种统计量来指导策略而不只是单一均值。第二层是分布式的监督信号。扩散过程有 $T$ 个时间步每个时间步都提供监督相当于把一次打分拆成了 $T$ 次一致性检查。这种密集监督是传统RM给不了的。第三层才是工程上的分布式。扩散模型参数量大、采样步数多实际训练和推理往往需要数据并行加模型并行这属于实现层面的分布式是前两层带来的必然工程需求。理解了这三层你就明白为什么这个方向值得单独拿出来讲——它不是给RM换个骨干网络那么简单而是整个奖励建模范式的转变。3. 奖励信号如何从去噪轨迹中提取三种可落地的方案3.1 方案一条件扩散似然作为奖励最直接的做法是训练一个条件扩散模型 $p_\theta(y | x, c)$其中 $c$ 是质量标签比如优质。训练数据用偏好数据集里的chosen样本作为优质条件rejected作为劣质条件。推理时给定 $(x, y)$计算 $y$ 在优质条件下的对数似然 $\log p_\theta(y | x, c\text{good})$直接当奖励。这个方案的优点是理论干净奖励有明确的概率解释。但坑也很明显扩散模型的对数似然计算成本极高需要跑完整的去噪链还要用ELBO或者概率流ODE来估计。我实测下来单条样本的似然估计在A100上要几十毫秒到上百毫秒如果RL阶段每步都要算吞吐直接崩掉。实操建议是不要在线算似然而是离线预计算一批候选回复的似然训练一个轻量的代理RM去拟合它在线用代理RM。这样既保留了扩散先验的分布约束又控制了推理成本。这个扩散教师轻量学生的蒸馏思路是我认为目前最务实的落地方案。3.2 方案二去噪重建误差作为奖励第二个方案更轻量给回复 $y$ 加一定程度的噪声得到 $y_t$让扩散模型去噪重建用重建误差 $|y - \hat{y}_0|$ 的负值作为奖励。直觉是落在优质分布高密度区的样本去噪重建更准偏离分布的样本重建误差大。这个方案的计算成本比似然低得多只需要跑部分去噪步比如从 $t0.5T$ 开始去噪不用跑完整链。但要注意噪声强度的选择噪声太小所有样本重建都好区分度不够噪声太大优质样本也被破坏信号被淹没。我的经验是噪声强度取中间偏上的时间步$t$ 在 $0.4T$ 到 $0.6T$ 之间区分度最好这个区间既保留了语义信息又对分布外样本足够敏感。还有一个细节重建误差要在语义空间而不是像素/词元空间算。对文本任务直接用token级重建误差会被高频词主导建议用句向量编码后的距离或者用扩散模型中间层的隐状态距离。这个改动我踩过坑直接用token误差训出来的奖励和人类偏好相关性只有0.3出头换成隐状态距离后能到0.55以上。3.3 方案三双分布score差作为奖励第三个方案最有意思也最接近分布式奖励的本意。同时训练两个条件扩散模型一个建模优质分布 $p_\theta(y | x, \text{good})$一个建模劣质分布 $p_\phi(y | x, \text{bad})$。奖励定义为两个分布score function的差r(x, y) ≈ E_t [ ∇_y log p_θ(y_t | x, good) - ∇_y log p_φ(y_t | x, bad) ]这个差在数学上正比于对数似然比 $\log \frac{p(y|x,\text{good})}{p(y|x,\text{bad})}$而似然比正是最优奖励函数的理论形式在Bradley-Terry模型下。所以这个方案不是拍脑袋是有理论保证的。它的优势是对绝对尺度有约束不像传统RM那样尺度漂移。因为两个分布的score差天然有界RL阶段不会出现奖励爆炸。缺点是训练两个扩散模型成本翻倍而且两个模型的去噪时间步要对齐否则score差没有意义。我建议如果资源有限可以共享骨干网络只分两个条件头这样参数量增加不多还能保证时间步对齐。实测这个共享骨干的版本和完全独立双模型的奖励质量差距在5%以内但训练成本省了将近一半。4. 训练目标与工程实现从损失函数到分布式部署4.1 扩散奖励模型的损失函数设计训练扩散RM的损失核心还是去噪损失但要加上奖励相关的约束。以方案一为例基础损失是标准的扩散去噪损失L_diff E_{t, y_0, ε} [ |ε - ε_θ(y_t, t, x, c)|² ]其中 $c$ 是质量条件。但光有这个不够因为我们要的是奖励排序能力不是单纯的生成能力。所以还要加一个排序损失让优质条件下的似然确实高于劣质条件L_rank -log σ( log p_θ(y_w | x, good) - log p_θ(y_l | x, good) )总损失是两者的加权和。这里的权重 $\lambda$ 很关键$\lambda$ 太大扩散模型退化成判别模型失去分布建模能力$\lambda$ 太小奖励排序能力不足。我的经验值是 $\lambda$ 取0.1到0.3之间先用纯扩散损失预训练再逐步加大排序损失做微调。这个课程学习式的训练策略比一上来就联合训练稳定得多。4.2 采样步数与推理成本的权衡扩散模型最大的工程痛点就是采样慢。奖励建模场景下我们不需要生成高质量样本只需要足够好的奖励估计所以可以大幅减少采样步数。实测下来用DDIM采样20到50步奖励估计的质量和1000步DDPM差距很小相关性0.95以上但速度快了20倍以上。更进一步可以用**一致性模型Consistency Models**做一步或少步采样。把扩散RM蒸馏成一致性模型后单次奖励估计能压到几毫秒完全能满足在线RL的需求。这个蒸馏步骤我强烈建议做它是扩散RM能否真正落地的分水岭。还有一个技巧缓存去噪轨迹的中间结果。RL阶段同一个提示会生成多个回复它们的条件 $x$ 相同扩散模型在早期时间步的计算可以复用。我实现过一个基于提示级缓存的版本吞吐提升了将近40%。4.3 分布式训练与推理的工程要点扩散RM参数量通常比传统RM大一个量级训练时数据并行加ZeRO是标配。但有几个扩散特有的坑时间步采样的负载不均不同时间步的去噪难度不同如果按时间步分桶做负载均衡能明显提升GPU利用率。我试过按 $t$ 分4个桶训练吞吐提升约15%。条件注入的通信开销如果条件 $c$ 和 $x$ 通过cross-attention注入注意这部分参数在模型并行时的切分方式切不好会引入大量all-reduce。推理时的批处理奖励估计是逐样本的但扩散去噪可以批处理。把同一批候选回复打包去噪比逐条推理快3到5倍。推理部署上如果只是离线打分用FP16加动态批处理就够了如果要在线接RL建议上TensorRT或者用编译优化把去噪链的kernel融合做掉。我实测TensorRT优化后50步DDIM的单样本延迟从80ms降到25ms左右。5. 和传统RM的正面比较哪些场景扩散真的更稳5.1 分布外泛化能力的实测差异我做过一组对比实验用同一批偏好数据分别训传统RM和扩散RM然后在分布外测试集上评估。分布外测试集是刻意构造的把训练时的提示模板换掉、把回复长度分布偏移、引入训练时没见过的领域词。结果很说明问题传统RM在分布内pairwise accuracy能到72%但分布外掉到58%几乎接近随机扩散RM分布内是69%略低但分布外还有65%衰减明显更小。这个差距在RL阶段会被放大——用传统RM做PPO策略很快就学会输出分布外的高分垃圾用扩散RM策略的探索被分布约束限制住生成质量稳定得多。原因还是回到本质扩散RM建模了 $p(y|x,\text{good})$它对什么算好回复有生成式的理解分布外样本会被低似然惩罚传统RM只学了相对排序分布外就是盲区。5.2 抗reward hacking的机制分析Reward hacking的本质是策略找到了RM评分面的漏洞。传统RM的评分面是判别式的漏洞多且难预测扩散RM的评分面来自生成分布漏洞少且更平滑。具体机制上扩散RM的奖励是多个时间步score的聚合策略要hack就得同时骗过所有时间步难度指数级上升。而且扩散过程的随机性噪声采样本身引入了正则化策略很难针对固定的评分面做优化。我在一个对话生成任务上做过对比传统RM接PPO训练到2000步左右开始出现明显的重复、谄媚、空洞化扩散RM接同样的PPO训练到5000步生成质量还在缓慢提升。当然扩散RM也不是万能如果优质分布本身学偏了照样会被hack只是门槛高很多。5.3 成本与收益的取舍建议说了这么多优点必须泼盆冷水扩散RM的成本是真的高。训练成本大概是传统RM的5到10倍推理成本即使优化后也是3到5倍。所以它不是无脑替换传统RM的方案而是要看你是否真的需要它的分布建模能力。我的建议是分场景场景推荐方案理由快速原型、小规模实验传统RM成本低迭代快大规模RLHF、对稳定性要求高扩散RM抗hacking分布外稳多模态奖励建模扩散RM扩散天然处理连续模态推理预算极紧的在线服务传统RM或蒸馏后的扩散RM延迟敏感需要奖励分布而非单点扩散RM唯一能给出分布的选择如果你的任务里策略探索空间大、分布外风险高扩散RM的额外成本是值得的如果只是小规模对齐、策略探索受限传统RM够用。6. 复现路上的真实坑从数据准备到训练崩溃6.1 偏好数据的质量比数量重要得多扩散RM对数据质量比传统RM敏感得多。原因很简单它要建模的是分布如果偏好数据里噪声大、标注不一致学出来的优质分布就是糊的。我踩过的坑是拿了一个公开偏好数据集直接训结果扩散RM的奖励和人类偏好相关性还不如传统RM——因为那个数据集里chosen和rejected的差异太细微扩散模型学不出清晰的分布边界。后来我做了两件事一是用强模型对偏好数据做一致性过滤把标注矛盾、差异过小的样本剔掉二是对chosen样本做质量分层不是简单二分优质/劣质而是分3到5档让扩散模型学更细的分布。这两步做完奖励相关性从0.42提到0.61。6.2 训练不稳定的三个典型症状与对策扩散RM训练比传统RM难调我遇到过三种典型崩溃症状一损失震荡不收敛。多半是排序损失权重 $\lambda$ 太大扩散模型被带偏了。对策是把 $\lambda$ 降到0.1以下先让扩散损失稳定下来。症状二奖励尺度爆炸。双分布方案里两个模型的score差没对齐导致奖励值域失控。对策是加一个尺度归一化把奖励限制在固定区间或者用tanh压缩。症状三生成样本模式坍塌。扩散模型只学到了优质分布的少数模式奖励对多样性回复不友好。对策是加熵正则或者在条件注入时加dropout强制模型学更宽的分布。这三个坑我都花了不少时间才定位核心经验是扩散RM的训练监控要看去噪损失和奖励相关性两条曲线任何一条异常都要停下来查不要硬训。6.3 评估扩散RM的正确姿势评估扩散RM不能只看pairwise accuracy那个指标会骗人。我建议至少看四个维度分布内pairwise accuracy基础指标但别只看这个。分布外pairwise accuracy构造偏移测试集这是扩散RM的真正优势区。奖励-人类相关性Spearman看绝对分值的排序质量。奖励尺度稳定性在不同批次上算奖励的均值和方差看是否漂移。还有一个我自创的土办法用扩散RM给一批明显垃圾的回复打分看它能不能稳定给低分。传统RM经常给垃圾回复中等偏上的分扩散RM如果学得好应该给明显低分。这个测试比accuracy更能反映实际可用性。7. 这条路线接下来能怎么走扩散奖励模型目前还在早期但几个方向已经能看到苗头。一是和DPO类方法结合用扩散RM提供更稳的隐式奖励绕开PPO的复杂性二是多模态扩展扩散天然处理图像、视频做多模态RLHF的奖励建模有天然优势三是奖励分布的 downstream 应用比如用奖励方差做不确定性估计指导主动学习或者保守策略优化。我自己最看好的是一致性模型蒸馏加扩散RM的组合它把扩散的分布建模能力和实时推理的需求同时满足了。如果你现在要上手我的建议是先用方案二去噪重建误差做快速验证成本最低能跑通再上方案一或方案三。别一上来就搞双分布那个训练难度和资源消耗对新手不友好。最后分享一个实操细节扩散RM的随机种子对结果影响比传统RM大建议至少跑3个种子取平均否则单次实验的结论可能不可靠。这个坑我在早期对比实验里踩过差点得出错误结论。
返回列表