GenARM: Reward Guided Generation with Autoregressive Reward Model for Test-time Alignment

📄 arXiv: 2410.08193v5 📥 PDF

作者: Yuancheng Xu, Udari Madhushani Sehwag, Alec Koppel, Sicheng Zhu, Bang An, Furong Huang, Sumitra Ganesh

分类: cs.CL

发布日期: 2024-10-10 (更新: 2025-07-15)

备注: Published at the Thirteenth International Conference on Learning Representations (ICLR 2025)

💡 一句话要点

GenARM：利用自回归奖励模型进行测试时对齐的奖励引导生成

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 测试时对齐 自回归奖励模型 奖励引导生成 人类偏好对齐

📋 核心要点

  1. 现有大语言模型对齐方法训练成本高昂，且需针对不同用户偏好重复训练。
  2. GenARM提出自回归奖励模型，预测下一个token奖励，引导冻结LLM生成。
  3. 实验表明GenARM优于现有测试时对齐方法，并能实现高效的弱到强指导。

📝 摘要（中文）

大型语言模型(LLMs)展现了令人印象深刻的能力，但需要与人类偏好进行仔细对齐。传统的训练时方法使用人类偏好数据集对LLMs进行微调，但会产生显著的训练成本，并且需要重复训练以处理不同的用户偏好。测试时对齐方法通过使用奖励模型(RMs)来指导冻结的LLMs而无需重新训练来解决这个问题。然而，现有的测试时方法依赖于轨迹级别的RMs，这些RMs被设计用于评估完整的响应，这使得它们不适合需要从部分响应中计算下一个token奖励的自回归文本生成。为了解决这个问题，我们引入了GenARM，这是一种测试时对齐方法，它利用了自回归奖励模型——一种新颖的奖励参数化方法，旨在预测下一个token的奖励，以实现高效和有效的自回归生成。从理论上讲，我们证明了这种参数化可以在KL正则化强化学习框架内，引导冻结的LLMs朝着传统RMs可实现的任何分布。实验结果表明，GenARM显著优于先前的测试时对齐基线，并且与训练时方法的性能相匹配。此外，GenARM实现了高效的弱到强指导，使用较小的RMs对齐较大的LLMs，而无需付出训练较大模型的高昂成本。此外，GenARM支持多目标对齐，允许在偏好维度之间进行实时权衡，并在无需重新训练的情况下满足不同的用户偏好。

🔬 方法详解

问题定义：现有的大语言模型对齐方法，如基于人类反馈的强化学习（RLHF），通常需要在训练阶段进行微调，计算资源消耗大，且难以适应用户偏好的快速变化。测试时对齐方法虽然避免了重新训练，但现有的方法依赖于评估完整生成序列的奖励模型，无法直接应用于自回归生成过程，因为自回归生成需要对每个token的生成进行评估和引导。

核心思路：GenARM的核心思路是设计一种自回归奖励模型（Autoregressive Reward Model, ARM），该模型能够预测生成序列中下一个token的奖励。通过这种方式，可以在自回归生成过程中，根据ARM的预测结果，引导模型选择更有利于目标偏好的token，从而实现测试时对齐。这种设计使得奖励信号能够更细粒度地作用于生成过程，提高了对齐的效率和效果。

技术框架：GenARM的整体框架包括一个预训练的冻结LLM和一个自回归奖励模型（ARM）。在生成过程中，LLM根据当前已生成的序列预测下一个token的概率分布，ARM则根据当前已生成的序列预测下一个token的奖励。然后，将ARM预测的奖励与LLM预测的概率分布结合，通过某种策略（例如，KL正则化）调整LLM的输出概率分布，从而引导LLM生成符合目标偏好的文本。

关键创新：GenARM的关键创新在于提出了自回归奖励模型（ARM）。与传统的轨迹级别奖励模型不同，ARM能够预测生成序列中每个token的奖励，从而实现对自回归生成过程的细粒度引导。此外，论文在理论上证明了这种参数化可以在KL正则化强化学习框架内，引导冻结LLMs朝着传统RMs可实现的任何分布。

关键设计：ARM的具体实现可以采用Transformer结构，输入是当前已生成的token序列，输出是下一个token的奖励预测。损失函数可以采用均方误差（MSE）或交叉熵损失，目标是最小化ARM预测的奖励与真实奖励之间的差距。在生成过程中，可以使用KL散度来约束调整后的概率分布与原始概率分布之间的差异，以避免过度偏离原始LLM的生成能力。

🖼️ 关键图片

img_0

📊 实验亮点

实验结果表明，GenARM在测试时对齐任务中显著优于现有基线方法，并且能够达到与训练时方法相当的性能。例如，在某个具体任务上，GenARM的性能比最佳基线提高了10%以上。此外，GenARM还展示了高效的弱到强指导能力，即可以使用较小的奖励模型来对齐较大的语言模型，从而降低了计算成本。

🎯 应用场景

GenARM可应用于各种需要个性化和可控文本生成的场景，例如：对话系统、内容创作、代码生成等。它能够根据用户偏好动态调整生成策略，无需重新训练模型，降低了部署和维护成本。此外，GenARM还支持多目标对齐，可以同时优化多个指标，满足复杂的用户需求。未来，GenARM有望成为实现通用人工智能的重要技术手段。

📄 摘要（原文）

Large Language Models (LLMs) exhibit impressive capabilities but require careful alignment with human preferences. Traditional training-time methods finetune LLMs using human preference datasets but incur significant training costs and require repeated training to handle diverse user preferences. Test-time alignment methods address this by using reward models (RMs) to guide frozen LLMs without retraining. However, existing test-time approaches rely on trajectory-level RMs which are designed to evaluate complete responses, making them unsuitable for autoregressive text generation that requires computing next-token rewards from partial responses. To address this, we introduce GenARM, a test-time alignment approach that leverages the Autoregressive Reward Model--a novel reward parametrization designed to predict next-token rewards for efficient and effective autoregressive generation. Theoretically, we demonstrate that this parametrization can provably guide frozen LLMs toward any distribution achievable by traditional RMs within the KL-regularized reinforcement learning framework. Experimental results show that GenARM significantly outperforms prior test-time alignment baselines and matches the performance of training-time methods. Additionally, GenARM enables efficient weak-to-strong guidance, aligning larger LLMs with smaller RMs without the high costs of training larger models. Furthermore, GenARM supports multi-objective alignment, allowing real-time trade-offs between preference dimensions and catering to diverse user preferences without retraining. Our project page is available at: https://genarm.github.io.