Learning Goal-Conditioned Representations for Language Reward Models
作者: Vaskar Nath, Dylan Slack, Jeff Da, Yuntao Ma, Hugh Zhang, Spencer Whitehead, Sean Hendryx
分类: cs.CL
发布日期: 2024-07-18 (更新: 2024-10-23)
💡 一句话要点
提出目标条件对比学习方法，提升语言模型奖励模型的性能和可控性。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 奖励模型 对比学习 目标条件化 语言模型 人工反馈强化学习
📋 核心要点
- 现有方法在利用离线数据或自监督目标改进表征学习方面取得了显著成果，但其在语言模型的人工反馈强化学习（RLHF）中的收益尚不明确。
- 本文提出了一种对比式的目标条件奖励模型训练方法，通过优化轨迹上未来状态表征的相似性，提升奖励模型的性能和可控性。
- 实验结果表明，该方法在多个基准测试中显著提高了奖励模型的性能，并实现了对生成内容在有用性和复杂性等方面的细粒度控制。
📝 摘要（中文）
本文提出了一种对比式的、目标条件化的方法来训练奖励模型（RM），通过增加采样到的偏好轨迹上未来状态的表征相似性，并减少随机采样的非偏好轨迹上的相似性。该目标显著提高了RM的性能，在MATH和GSM8k等具有挑战性的基准测试中，AUROC提高了高达0.09。这些发现也适用于通用对齐，在Helpful-Harmless数据集上，准确率提高了2.3%。除了提高奖励模型性能外，这种训练RM表征的方式还提高了可控性，因为它允许我们评估一个动作实现特定目标状态的可能性。利用这一洞察力，我们发现可以通过丢弃可能最终进入“不正确”状态的轨迹来过滤掉高达55%的生成token，从而显著节省成本。此外，我们发现这些表征可以通过调节期望的未来目标状态来执行细粒度控制。例如，我们表明，使用我们的方法引导Llama 3模型生成有用的内容，比经过监督微调的基线提高了9.6%的有用性。同样，引导模型生成复杂的内容，比基线提高了21.6%的复杂性。总的来说，我们发现以这种对比式的、目标条件化的方式训练RM可以显著提高性能并实现模型的可控性。
🔬 方法详解
问题定义：现有语言模型的人工反馈强化学习（RLHF）中，如何有效地利用离线数据或自监督学习来提升奖励模型（RM）的性能和可控性是一个关键问题。传统的表征学习方法在传统强化学习中表现良好，但在RLHF中收益不明确。现有奖励模型在理解和区分不同目标状态的能力上存在不足，限制了模型的可控性。
核心思路：本文的核心思路是通过对比学习的方式，训练一个目标条件化的奖励模型。具体来说，模型学习区分偏好轨迹和非偏好轨迹，并学习将相似的目标状态映射到相似的表征空间。通过这种方式，奖励模型能够更好地理解不同目标状态的语义信息，从而提高奖励预测的准确性和模型的可控性。
技术框架：该方法主要包含以下几个阶段：1）数据收集：收集包含偏好轨迹和非偏好轨迹的数据集。2）表征学习：使用对比学习的目标函数，训练奖励模型学习目标条件化的表征。具体来说，对于偏好轨迹，增加未来状态表征的相似性；对于非偏好轨迹，减少未来状态表征的相似性。3）奖励预测：使用学习到的表征，预测给定状态的奖励值。4）模型控制：利用学习到的表征，评估一个动作实现特定目标状态的可能性，从而实现对生成内容的细粒度控制。
关键创新：该方法最重要的创新点在于将对比学习和目标条件化学习相结合，用于训练奖励模型。与传统的奖励模型训练方法相比，该方法能够更好地学习目标状态的语义信息，从而提高奖励预测的准确性和模型的可控性。此外，该方法还提出了一种新的模型控制方法，通过评估动作实现目标状态的可能性，实现对生成内容的细粒度控制。
关键设计：在对比学习的目标函数中，使用了InfoNCE损失函数来最大化偏好轨迹上未来状态表征的互信息，并最小化非偏好轨迹上未来状态表征的互信息。目标状态的表征通过一个神经网络编码器获得。在模型控制方面，使用学习到的表征来计算一个动作实现特定目标状态的概率，并根据该概率来过滤生成的token。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在MATH和GSM8k等具有挑战性的基准测试中，AUROC提高了高达0.09。在Helpful-Harmless数据集上，准确率提高了2.3%。通过过滤掉可能导致错误结果的token，可以节省高达55%的生成成本。使用该方法引导Llama 3模型生成有用的内容，比经过监督微调的基线提高了9.6%的有用性。引导模型生成复杂的内容，比基线提高了21.6%的复杂性。
🎯 应用场景
该研究成果可应用于各种需要人工反馈的语言模型训练场景，例如对话系统、文本生成、代码生成等。通过提高奖励模型的性能和可控性，可以生成更符合人类偏好和目标的内容。此外，该方法还可以用于降低生成成本，通过过滤掉可能导致错误结果的token，减少不必要的计算资源消耗。未来，该方法可以进一步扩展到多模态场景，例如图像生成和视频生成。
📄 摘要（原文）
Techniques that learn improved representations via offline data or self-supervised objectives have shown impressive results in traditional reinforcement learning (RL). Nevertheless, it is unclear how improved representation learning can benefit reinforcement learning from human feedback (RLHF) on language models (LMs). In this work, we propose training reward models (RMs) in a contrastive, $\textit{goal-conditioned}$ fashion by increasing the representation similarity of future states along sampled preferred trajectories and decreasing the similarity along randomly sampled dispreferred trajectories. This objective significantly improves RM performance by up to 0.09 AUROC across challenging benchmarks, such as MATH and GSM8k. These findings extend to general alignment as well -- on the Helpful-Harmless dataset, we observe $2.3\%$ increase in accuracy. Beyond improving reward model performance, we show this way of training RM representations enables improved $\textit{steerability}$ because it allows us to evaluate the likelihood of an action achieving a particular goal-state (e.g., whether a solution is correct or helpful). Leveraging this insight, we find that we can filter up to $55\%$ of generated tokens during majority voting by discarding trajectories likely to end up in an "incorrect" state, which leads to significant cost savings. We additionally find that these representations can perform fine-grained control by conditioning on desired future goal-states. For example, we show that steering a Llama 3 model towards helpful generations with our approach improves helpfulness by $9.6\%$ over a supervised-fine-tuning trained baseline. Similarly, steering the model towards complex generations improves complexity by $21.6\%$ over the baseline. Overall, we find that training RMs in this contrastive, goal-conditioned fashion significantly improves performance and enables model steerability.