Learning Explainable Dense Reward Shapes via Bayesian Optimization

本文提出一种通过Bayesian Optimization学习解释性密集奖励形状的方法，以解决RLHF中奖励稀疏问题，实现token级信用分配优化，提升训练效率和性能，同时保持最优政策不变。

Reinforcement Learning, Reward Shaping, Token-Level Credit Assignment, Explainability Methods, Bayesian Optimization, Policy Training

Ryan Koo, Ian Yang, Vipul Raheja, Mingyi Hong, Kwang-Sung Jun, Dongyeop Kang

University of Minnesota, Georgia Institute of Technology, Grammarly, University of Arizona

Generated by grok-3-mini-latest

Background Problem

当前强化学习从人类反馈（RLHF）管道中，大型语言模型（LLM）的对齐通常通过对序列分配标量奖励，并使用最终token作为整个序列质量的代理指标，这导致反馈信号稀疏且token级信用分配不佳。研究的起点是解决RL中奖励稀疏性问题，特别是RLHF应用中常见的挑战，例如中间状态缺乏细粒度反馈，可能会导致代理学习效率低下和不稳定。具体问题包括：奖励信号带宽低，无法有效评估中间token的质量；收集token级人类标注成本高且主观性强；现有奖励整形方法（如基于注意力机制）可能与token贡献不直接相关。因此，本文通过将奖励整形框架化为优化问题，旨在利用解释性方法生成密集奖励，以提供更细粒度的反馈。

Method

核心思想： 本文将奖励整形视为优化问题，旨在通过解释性方法（如SHAP和LIME）估计token级贡献，并使用双层优化框架学习奖励整形函数的参数，以生成密集奖励而不改变最优政策。具体地，解释性方法属于可加特征归因函数家族，确保政策不变性。
如何实现： 首先，使用解释性方法（如SHAP值或LIME）从奖励模型中估计每个token的贡献，形成一个token级分数矩阵。然后，通过Bayesian Optimization（BO）在外部优化层中采样权重向量 $w$ ，以构建密集奖励函数。具体公式为： $r'(s,a) = w^{\top} \mathcal{E} \cdot r(s,a)$ 其中 $\mathcal{E}$ 是解释性分数矩阵， $r(s,a)$ 是原始稀疏奖励。BO使用高斯过程作为代理模型，通过获取函数（如log Noisy Expected Improvement）平衡探索和利用，采样权重。内部优化层使用PPO算法训练策略 $\pi_{\theta}$ ，优化加权奖励。整个过程确保了密集奖励的计算高效，且在不修改原始奖励模型的情况下，提供细粒度反馈。
主要步骤： (1) 初始化BO模型；(2) 采样权重并计算解释性分数；(3) 使用加权组合生成密集奖励；(4) 通过PPO迭代更新策略；(5) 使用验证集性能更新BO模型。

Experiment

实验设置： 本文在HH-RLHF（关注帮助性）和Ultrafeedback数据集上进行实验，使用PPO作为策略梯度算法。基线包括：监督微调（SFT）模型、稀疏奖励RLHF、注意力-based信用分配等。实验设计考虑了BO维度的变化（例如d=2、3、4），并限制BO迭代次数（m=25）以控制计算开销。数据集被分割为训练集和验证集（90%/10%），每个BO迭代中随机采样子集进行训练和验证。目的是评估密集奖励是否加速学习、提高稳定性，并避免奖励过拟合。
为什么这样设计： 实验设置合理地平衡了计算复杂性和泛化能力，通过比较不同解释性方法（如SHAP、LIME）和BO优化的组合，验证了方法的鲁棒性。使用开源基准如AlpacaEval-2和MTBench评估泛化性能，确保结果不仅在测试集上表现良好，还能在分布外数据上保持优势。
结果： 结果显示，解释性密集奖励方法显著改善了性能：例如，在HH-RLHF上，SHAP-based方法在保持平均奖励的同时，降低了价值函数损失，提高了AlpacaEval-2的胜率（例如从48.7%到57.62%）。BO优化进一步提升了奖励分配平衡，加速了训练（如图4所示，密集奖励早早达到高奖励）。与稀疏奖励相比，避免了奖励过拟合问题（胜率在分布外基准上更高）。然而，高维度BO（如d=4）可能由于迭代限制而未充分收敛，表明需要更多样本以探索复杂空间。整体而言，实验结果符合预期，证明了方法的有效性和优越性。

Further Thoughts

这个方法巧妙地将解释性技术与Bayesian Optimization结合，展示了在RLHF中利用不确定性估计来优化奖励形状的潜力；未来可以探索将其扩展到多模态任务中，例如结合视觉或语音数据，以处理更复杂的序列决策问题；此外，考虑到解释性方法（如SHAP）的噪声敏感性，或许可以整合元学习框架来适应不同领域的数据分布，类似于一些工作（如AlphaPO）中对奖励函数形状的动态调整，从而进一步提升泛化能力和鲁棒性。