Learning Explainable Dense Reward Shapes via Bayesian Optimization 通过贝叶斯优化学习可解释的密集奖励形状
https://arxiv.org/pdf/2504.16272
摘要
当前用于大语言模型(LLM)对齐的基于人类反馈的强化学习(RLHF)流程通常为整个序列分配一个标量奖励,并以最后一个token作为整个序列质量的代理指标。然而,这种方法导致反馈稀疏以及次优的token级别信用分配。在本研究中,我们将奖励塑形(reward shaping)视为一个专注于token级别信用分配的优化问题。我们提出了一种利用可解释性方法(如SHAP和LIME)从奖励模型中估计每个token奖励的奖励塑形函数。为了学习该塑形函数的参数,我们采用了一种双层优化框架,结合贝叶斯优化与策略训练,以应对token级奖励估计中的噪声。实验结果表明,在token级别奖励归因上实现更好的平衡,能够在下游任务中取得优于基线方法的性能,并在训练过程中更快地找到最优策略。此外,我们在理论上证明:对于属于特征加性归因函数(feature additive attribution functions)类别的可解释性方法,其保持了原始奖励下的最优策略不变。代码已公开发布。
1 引言
强化学习(RL)中的一个基本挑战源于奖励信号的稀疏性:反馈信号通常仅在轨迹结束时提供,关于中间状态几乎没有或完全没有评估信息。因此,这种限制导致了对某些中间状态是否有利存在认知空白,因为智能体缺乏必要的细粒度反馈来判断哪些动作对最终结果是有益的。在近年来许多基于人类反馈的强化学习(RLHF)应用中也观察到了类似的挑战,其中稀疏奖励十分普遍(Zheng et al., 2023b;Chaudhari et al., 2024)。
考虑到语言模型的序列特性以及RLHF中通常优化的token级别的价值函数(Zhong et al., 2025;Rafailov et al., 2024),确定各个token对整体奖励的贡献至关重要。问题在于,奖励通常是在序列生成结束后才给予的,表现为一个代表整个序列质量的标量值,已知这种做法不稳定(Razin et al., 2024;Engstrom et al., 2020),并且编码的是低带宽信号,无法帮助判断中间token的相对质量。因此,为解决这一局限,若能为策略提供token级别的奖励以实现细粒度反馈,将是大有裨益的(Wu et al., 2023;Xie et al., 2024)。
然而,在自回归生成过程中为每个token打分计算成本高昂,而在监督学习设置下收集大量精细的人类标注既昂贵又容易产生高度分歧。因此,一些研究探索了“奖励塑形”技术,将稀疏奖励转化为密集奖励(Sutton & Barto, 2018;Ng et al., 1999),从而促进更高效、更可解释的优化过程,并对中间决策提供更精细的控制。例如,最近的一些尝试考察了“过程奖励”(process rewards)(Lightman et al., 2023;Uesato et al., 2022),即为思维链(chain-of-thought)生成提供中间反馈,或直接利用奖励模型的注意力图(attention map)(Chan et al., 2024)来重新分配奖励。然而,这些方法较为复杂,往往需要高质量的人类反馈,且每个token上的注意力可能并不直接对应其输出的解释性(Jain & Wallace, 2019)。
在本研究中,我们提出一种两阶段方法来使稀疏奖励变得密集:(1)使用可解释性方法构建密集奖励信号;(2)使用贝叶斯优化学习由可解释性得分构成的新奖励塑形函数的权重。我们将奖励塑形建模为一个双层优化问题:在上层,我们优化塑形奖励函数的系数;在下层,我们学习相应的最优策略,如图1所示。首先,我们使用LIME(Ribeiro et al., 2016)或SHAPley值(Lundberg & Lee, 2017)等可解释性技术来估计token级别的贡献。其次,由于这些可解释性方法已知对噪声敏感(Li et al., 2020),我们将其视为具有不确定性的估计,并试图学习其上的最优加权方案。然而,由于奖励景观的复杂性,穷举所有可能的权重配置在计算上是不可行的。为此,我们将该问题视为黑箱优化问题,并采用贝叶斯优化来学习最佳奖励系数,这是一种对噪声目标函数具有天然鲁棒性的方法(Fröhlich et al., 2020;Daulton et al., 2022)。
我们证明,可解释性方法提供了一种自然的方式,可以从奖励模型中提取更多信息以实现奖励密集化,并满足基于势能的奖励塑形(potential-based reward shaping)(Ng et al., 1999)的条件,从而保证策略不变性——即使用原始奖励函数时的最优策略在塑形后保持不变。此外,我们引入了一种新的优化框架,整合了多种来源的token级信息,并表明贝叶斯优化有助于学习最佳的奖励塑形函数。实证结果显示,与稀疏奖励相比,使用可解释性方法能积极影响RL训练过程,加快学习速度并使价值函数更新更加稳定。我们还表明,加入贝叶斯优化以合理地塑造奖励,相较于其简单版本设置,在下游任务上的生成质量有所提升。
2 预备知识
在本节中,我们首先介绍经典的基于人类反馈的强化学习(RLHF)方法,并为本文所提出的方法建立相关的理论基础,涵盖可解释性技术与贝叶斯优化的基本概念。
因此,接收到的信号非常稀疏(Zhong 等人,2025;Chan 等人,2024),模型除了 KL 正则化项之外,没有关于中间令牌生成的相关反馈。因此,我们感兴趣的是如何设计令牌级奖励以提供高质量、密集的信号,使学习更容易,并可能提高性能。
2.2 作为令牌级奖励的可解释性
我们首先引入一个通用定义,用于估计文本环境中令牌级信息的可解释性。为了解释复杂黑盒函数 f(即我们的奖励模型)的预测,通常使用一个小的线性模型 g 来局部近似它,涉及提示和完成。这里,我们只考虑加性特征归因函数(Ribeiro 等人,2016;Lundberg & Lee,2017)的家族,使得原始奖励的最优策略保持不变(Ng 等人,1999):
2.3 奖励塑形的贝叶斯优化
3 将可解释性奖励塑形建模为双层优化问题
在本节中,我们讨论优化设置以及用于从每个token级别的得分中寻找最佳奖励形状的初步算法。首先,我们讨论奖励形状的公式化方法,即通过可解释性方法在局部近似token级别的奖励,从而计算出密集奖励,并获得更细粒度的反馈。接着,我们讨论双层优化问题,在该框架中,我们将贝叶斯优化(BO)与经典的RLHF问题建模为一个嵌套目标,以同时优化奖励形状和策略。
3.1 计算奖励形状
在优化过程的每一步中,从上层步骤采样得到权重 w 后,我们通过计算所有 token 级别得分的加权线性组合来塑造奖励,并将该奖励在整个序列上进行广播:
3.2 贝叶斯优化与RLHF作为嵌套问题
传统上,下层问题是一个辅助性问题,旨在帮助求解上层问题(Zhang et al., 2023)。相比之下,本文中我们不仅关注上层问题的解,也同样关注由上层解所决定的下层问题的解。换句话说,我们采用如下的双层优化设置:
4 实验
任务。在本节中,我们通过下游任务实证验证了所提出方法的有效性。具体而言,我们在两个不同的单轮对话数据集上训练各个模型:1)HH-RLHF(有帮助性)(Bai et al., 2022)和 2)Ultrafeedback(Cui et al., 2024)数据集。对于 HH-RLHF 实验,我们采用 OpenLLaMA 系列模型,使用经过指令微调的 7B 模型作为监督微调(SFT)模型,并采用 Dong 等人(2023)提出的 3B 参数奖励模型,设置方式遵循 Chan 等人(2024)的方案。对于 Ultrafeedback 实验,我们使用 LLAMA-3.2-Instruct 3B(Team, 2024)作为 SFT 模型,并在 Ultrafeedback 的偏好数据上对 LLAMA-3.2-Instruct 1B 进行微调,作为奖励模型。
为了验证每种方法的有效性,我们首先考虑奖励模型在每个数据集测试集分割上的平均整体奖励(average holistic reward)。此外,我们还在开放基准如 AlpacaEval-2(Dubois et al., 2025)和 MT-Bench(Zheng et al., 2023a)上进行评估,以检验我们的方法是否相比仅使用稀疏奖励优化能够达到更好的局部最优解,同时避免奖励过拟合。
基线方法。为了衡量我们方法的改进效果,我们考虑了若干基线:1)SFT 模型,用于校准改进程度;2)在稀疏奖励上进行 RLHF 训练;3)基于注意力的信用分配方法(Chan et al., 2024);4)在使用贝叶斯优化进行奖励形状优化之前,单独使用每种可解释性方法的结果。
训练设置。在所有实验中,我们采用 PPO(近端策略优化)作为主要的策略梯度算法。由于 PPO 的计算复杂度较高,同时为了运行足够数量的贝叶斯试验,我们必须减少每次试验的训练样本数量,以避免过高的计算开销。因此,我们并未在每次贝叶斯优化(BO)试验中遍历整个数据集,而是进行 m = 25 次贝叶斯优化试验,每次试验包含少量训练轮次(即 10 轮),每轮批大小为 8(共 80 个样本),总计使用 2000 个样本。
为了初始采样奖励塑形的权重,我们在前五次试验中采用 Sobol 采样,以构建高斯过程(GP)模型的先验,之后的试验则从该 GP 模型中进行采样。
在每次 BO 迭代后,我们在验证集(validation split 4)上评估训练好的模型,并将平均验证奖励作为效用值,用于更新 BO 的代理函数(surrogate function)。为了继续训练,我们采用模型检查点机制,从表现最好的检查点恢复训练,而不是从头开始完整的训练循环。因此,我们最多仅对数据集 D₁:ₙ 遍历两次(一次用于 BO,一次用于完整的 PPO 训练)。为保证随机性,我们在每次贝叶斯试验中从训练数据集中随机采样一个子集用于 PPO 训练,同时在每次验证步骤中也从验证集中随机采样一个子集。
我们在附录 C 和 D 中详细说明了模型设置与超参数,并讨论了计算复杂度。
4.1基于解释的奖励在多大程度上优化了RLHF目标?
我们首先分析基于可解释性方法的奖励在优化标准RLHF目标中的影响。
首先,我们并不期望仅通过可解释性方法就能带来显著的改进。由于我们将可解释性得分构建为一种基于势能的奖励塑形(potential-based reward shaping)变换,因此并不能保证一定能找到更优的局部最优解(Ng et al., 1999),但精心选择的奖励形状有可能改善强化学习的学习复杂度(Fu et al., 2025;Gupta et al., 2022)。最重要的是,我们关注的是:相比于仅仅使用稀疏的标量奖励,可解释性方法是否能够带来任何有意义的改进或优化。
利用可解释性进行信用分配有助于探索 在应用任何贝叶斯优化之前,我们手动选择一组权重组合,参照 Chan 等人(2024)的做法,对token级别的得分赋予更高权重,设 w=0.8。总体而言,基于可解释性的方法在训练早期就达到了较高的平均奖励,其表现与图4(a)中观察到的基于注意力的基线相当。此外,从图4(b)可以看出,基于可解释性的方法相比稀疏奖励实现了更稳定的更新。特别是,可解释性方法生成的密集奖励显著降低了PPO的价值头(value head)损失,这表明我们对状态到长期回报的估计更为准确,从而在训练过程中实现了更稳定的策略更新。
基于可解释性的奖励有助于避免奖励过拟合 在基于人类反馈的强化学习中,一个主要风险是奖励过拟合。这种情况尤其令人担忧,因为奖励模型仅是人类偏好的代理,过度优化其得分反而可能损害在真实任务上的表现(Gao et al., 2022)。为了研究这一问题,我们在开放基准测试上评估了各基线方法及我们提出的方法在指令遵循和多轮对话能力方面的表现,这些测试数据并不完全属于训练数据的分布范围。从表1可以看出,基于可解释性的方法在HH-RLHF测试集上的表现并未显著超越基线;然而,它们在开放基准测试中取得了明显更高的胜率,表明我们并未严重过拟合奖励函数,同时保持了良好的泛化性能。值得注意的是,对于HH-RLHF模型,稀疏奖励方法在AlpacaEval2上的表现甚至不如SFT模型。我们推测这是由于在奖励函数上发生了过拟合,导致在与“有帮助性”无直接关联的数据分割上性能下降。为验证这一点,我们仅在AlpacaEval2的“有帮助性”子集上将该方法与SFT模型进行对比,发现其胜率有所提升,被偏好选择的比例达到 56.34% ± 4.03。
4.2 贝叶斯优化是否有助于更好地平衡token级别的奖励?
接下来,我们评估贝叶斯优化(BO)在采样最优权重以构建基于token级别可解释性得分的新型密集奖励函数方面的效果。我们考虑了由SHAP、LIME以及奖励模型注意力图所得到的多种token级别得分的组合。具体而言,我们的目标是评估BO是否能够有效权衡这些不同来源的token级别重要性,从而在策略优化过程中生成更具信息量的奖励信号。
然而,由于可用试验次数有限,我们承认BO的收敛性无法得到保证,尤其是在奖励权重组合的搜索空间具有高维性且可能非凸的情况下(Loeppky et al., 2012;Snoek et al., 2012)。在这种场景下,探索与利用之间的平衡问题进一步加剧了优化过程的复杂性,特别是当面对噪声大且非平稳的奖励曲面时。
尽管存在这些局限性,表1和图3(a)表明,由贝叶斯优化(BO)采样的权重相较于基线方法,普遍带来了下游任务性能的提升,并且在BO优化的密集奖励引导下的模型往往能在训练中获得更高的奖励。
探索-利用权衡 为了更好地理解BO的行为,我们在图5和图3(c)中分析了权重分布以及密集奖励形状在各次试验中的演变过程。特别地,我们注意到:最终将所有得分组合(输入维度为 d = 4)的方案并未优于更简单的组合。这在预期之中,因为随着搜索空间复杂性的增加,我们也应相应增加试验次数;然而由于计算资源的限制,我们保持试验次数不变。理想情况下,如果BO成功识别出最优权重,那么任何得分子集的组合都应具有不低于完整组合的性能下界;也就是说,只要权重得当,引入更多得分不应损害性能。然而,d = 4 时性能的下降表明,BO可能未能充分探索或利用奖励空间,可能过早收敛到了次优区域。从图5的所有试验中可以看出,在前五次使用Sobol采样后,BO倾向于在探索与利用之间持续切换,这表明可能需要更多的探索。
贝叶斯优化与信用分配 此外,我们检查了在Ultrafeedback(Cui et al., 2024)数据集的一个代表性样例上奖励形状在优化前后的变化情况。该任务要求模型从三个句子中选择两个能够与未完成故事自然衔接的句子。在图6中,我们展示了多种奖励塑形方法对SFT模型生成结果的影响,其中每个token根据其在计算塑形权重后获得的奖励值进行着色。特别地,几乎所有塑形方法仍然对终止token赋予较高重要性。我们推测,由于奖励模型在训练时是在终止token上进行反向传播的,因此为了最大化奖励函数,我们无法忽略该位置的作用。具体而言,我们观察到SHAP与LIME的组合能够突出模型选择第一和第二个句子的推理依据。类似的现象也出现在包含所有token级别得分的 d = 4 组合中。
然而,我们发现当标量奖励在最终的塑形得分中被赋予零权重时,分数的分配变得更加均匀。
当前可解释性方法的局限性 尽管贝叶斯优化(BO) 能够利用多种可解释性信号之间的互补效应,但单个的机制性可解释方法(例如模型的内在注意力)从根本上仍然是次优的。如图6所示,这些方法往往无法与人类对相关特征的感知精确对齐(Bereska & Gavves, 2024)。未来研究的一个有前景的方向是引入更细粒度且与人类判断对齐的信号,这些信号可能来源于显式的token级别标注(Hayati et al., 2021),或来自阅读文本时的认知指标,例如眼动轨迹(De Langis & Kang, 2022)。
5 相关工作
奖励塑形(Reward Shaping)是强化学习(RL)中的一种技术,通过引入更具信息量的奖励信号来补充传统上稀疏的奖励(Hu et al., 2020)。早期关于基于势能的奖励塑形(Potential-Based Reward Shaping, PBRS)的研究表明,奖励塑形可以在保留最优策略的同时减少训练时间,或引入领域知识(Hu et al., 2020;Cao et al., 2024a)。然而,挑战仍然存在,因为通常很难保证所添加的奖励在任意给定任务中都是有益的(Hu et al., 2020)。此外,我们并非首个从双层优化视角研究奖励塑形或RLHF的工作。例如,参数化奖励塑形的双层优化方法(BiPaRS)(Hu et al., 2020)自适应地利用给定的奖励函数,同时优化策略本身以及奖励函数的塑形权重。Shen 等人(2024)也从惩罚函数的角度将RLHF建模为一种双层优化形式。其他奖励塑形方法则侧重于可适应性以克服局限性,例如 AlphaPO(Gupta et al., 2025)在简单策略优化(SimPO)基础上引入一个新的标量参数 α,表明奖励函数的形状可以影响策略的似然位移。
Token-Level RLHF 另一些方法旨在通过在token级别分配奖励来改进奖励信号,以捕捉生成文本中特定词语、短语(Yoon et al., 2024)或话语结构(Kim et al., 2025)的细微差异。例如,Token-Level Continuous Reward(TLCR)使用token级别的偏好判别器(Yoon et al., 2024);RLMEC 利用最小编辑约束生成token级别的监督信号(Chen et al., 2024b);DRLC 则利用大语言模型(LLM)在响应中识别密集的正负标签(Cao et al., 2024b)。也有研究探索直接使用大语言模型本身来提供密集奖励信号(Huang et al., 2024)。这些工作的共同特点是采用对比式训练:将初始输出与修改后的输出进行比较,通过分析差异获得细粒度的奖励信号(Guo et al., 2023)。诸如token级PPO(TPPO)(Ouyang et al., 2024)和token级DPO(TDPO)(Zeng et al., 2024)等方法进一步支持了在token级别进行微调的思路。
贝叶斯优化(Bayesian Optimization, BO)近年来也被用于增强大语言模型的训练与微调(Austin et al., 2024;Liu et al., 2024;Yang et al.;Gao et al., 2024;Agarwal et al., 2025;Kristiadi et al., 2024;Opsahl-Ong et al., 2024;Chen et al., 2024a)。BO通过不断加入新信息或候选解,迭代估计最优的后验预测分布(Agarwal et al., 2025)。一些近期工作将BO与LLM结合,用于在离散搜索空间中进行选择(Kristiadi et al., 2024;Opsahl-Ong et al., 2024)。另一些研究探索使用BO来优化LLM的提示(prompt),例如BOPRO(Agarwal et al., 2025)、InstructZero(Chen et al., 2024a),利用贝叶斯优化探索提示的搜索空间。BO也可用于预训练阶段,例如通过寻找检查点融合的最优权重(Liu et al., 2024)。
此外,BO还可用于不确定性估计(Yang et al.)、偏差缓解(通过胜率校准)(Gao et al., 2024)以及引导式查询生成(Austin et al., 2024)。这些任务最终都得益于BO在优化过程中对不确定性的建模能力,从而展现出对噪声的强鲁棒性(Opsahl-Ong et al., 2024)。
6 结论
在本研究中,我们表明可解释性方法为重塑奖励以实现token级别信用分配提供了一种有效途径。此外,我们的研究结果表明,贝叶斯优化(BO)是一种有前景的工具,可用于学习并平衡不同来源的token级别信息,以进行奖励塑形。结合两者,我们证明:在贝叶斯框架下优化得到的、基于可解释性的密集奖励形状,相较于稀疏奖励,能够通过加速学习过程和对价值函数提供更稳定的更新,积极影响强化学习的训练效果。此外,通过可解释方法对token进行合理归因,不仅在下游任务上优于稀疏奖励,而且在理论上保持了策略不变性,从而忠实于原始奖励模型。
未来,我们可以考虑将上下文信息引入BO的权重中,以动态地塑造token级别的奖励。例如,可以学习一组与策略模型生成结果中token嵌入相对应的高维权重,从而动态计算奖励形状,而非当前采用的静态加权奖励形状。
原文链接:https://arxiv.org/pdf/2504.16272
热门跟贴