Tiny Reward Models
作者: Sarah Pan
分类: cs.CL, cs.AI
发布日期: 2025-07-14
备注: 2025 ICML Efficient Systems for Foundation Models Workshop
💡 一句话要点
提出TinyRM,一种小型双向语言模型,用于高效奖励建模。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: 奖励模型 强化学习 人类反馈 小型语言模型 双向模型 低秩适应 模型压缩
📋 核心要点
- 现有奖励模型依赖大型解码器语言模型,推理成本高昂,限制了其在实际场景中的应用。
- TinyRM提出一种小型双向掩码语言模型,结合FLAN提示、DoRA和层冻结等技术,降低模型规模。
- 实验表明,TinyRM在推理和安全偏好建模任务上,性能可与大型模型媲美,同时显著降低计算成本。
📝 摘要(中文)
大型基于解码器的语言模型已成为人类反馈强化学习(RLHF)中奖励建模的主流架构。然而,随着奖励模型越来越多地部署在测试时策略中,它们的推理成本日益受到关注。我们提出了TinyRM,一个小型双向掩码语言模型(MLM)家族,参数最少仅为4亿,但在推理和安全偏好建模任务上,其能力可与大175倍以上的模型相媲美。TinyRM结合了FLAN风格的提示、定向低秩适应(DoRA)和层冻结,尽管使用了明显更少的资源,但在RewardBench上实现了强大的性能。我们的实验表明,小型模型受益于特定领域的调整策略,尤其是在推理方面,轻量级微调方法特别有效。虽然在构建通用模型和对话偏好建模方面仍然存在挑战,但我们的初步结果突出了轻量级双向架构作为高效、可扩展的偏好建模替代方案的潜力。
🔬 方法详解
问题定义:现有基于大型解码器语言模型的奖励模型,在部署时面临高昂的推理成本,限制了其在资源受限环境中的应用。因此,需要开发一种更高效、更轻量级的奖励模型,能够在保证性能的同时,显著降低计算开销。
核心思路:TinyRM的核心思路是利用小型双向掩码语言模型(MLM)作为奖励模型的基础架构。通过结合FLAN风格的提示、定向低秩适应(DoRA)和层冻结等技术,在减少模型参数的同时,保持甚至提升模型在特定任务上的性能。这种设计旨在实现性能与效率之间的平衡。
技术框架:TinyRM的整体框架包括以下几个主要组成部分:1) 小型双向掩码语言模型(MLM)作为基础模型;2) FLAN风格的提示,用于引导模型学习偏好;3) 定向低秩适应(DoRA),用于高效地微调模型;4) 层冻结,用于进一步减少训练参数,防止过拟合。训练流程通常包括预训练和微调两个阶段。
关键创新:TinyRM的关键创新在于将小型双向MLM应用于奖励建模,并结合特定的训练策略,使其在资源受限的情况下也能达到与大型模型相当的性能。与传统的基于解码器的大型语言模型相比,TinyRM在模型规模和推理速度上具有显著优势。DoRA和层冻结的结合使用,使得模型能够以更少的参数进行高效微调。
关键设计:TinyRM的关键设计包括:1) 模型大小:参数量控制在4亿左右;2) 提示工程:采用FLAN风格的提示,提高模型对偏好的理解能力;3) DoRA:使用较小的秩进行低秩分解,减少微调参数;4) 层冻结:冻结部分层,防止过拟合,并减少训练时间;5) 损失函数:使用标准的偏好学习损失函数,例如pairwise ranking loss。
🖼️ 关键图片
📊 实验亮点
TinyRM在RewardBench基准测试中表现出色,在推理和安全偏好建模任务上,性能可与参数量大175倍的模型相媲美。实验结果表明,通过结合FLAN风格的提示、DoRA和层冻结等技术,小型模型能够有效地学习偏好,并在特定领域取得显著的性能提升。尤其是在推理任务上,轻量级微调方法表现出极高的效率。
🎯 应用场景
TinyRM具有广泛的应用前景,包括:1) 资源受限环境下的奖励建模;2) 移动设备或边缘计算平台的AI应用;3) 需要快速推理的在线强化学习系统;4) 作为大型奖励模型的蒸馏目标,进一步提升效率。该研究有助于推动AI技术在更多场景下的应用,并降低AI部署的成本。
📄 摘要(原文)
Large decoder-based language models have become the dominant architecture for reward modeling in reinforcement learning from human feedback (RLHF). However, as reward models are increasingly deployed in test-time strategies, their inference costs become a growing concern. We present TinyRM, a family of small, bidirectional masked language models (MLMs) with as few as 400 million parameters, that rival the capabilities of models over 175 times larger on reasoning and safety preference modeling tasks. TinyRM combines FLAN-style prompting, Directional Low-Rank Adaptation (DoRA), and layer freezing to achieve strong performance on RewardBench, despite using significantly fewer resources. Our experiments suggest that small models benefit from domain-specific tuning strategies, particularly in reasoning, where lightweight finetuning methods are especially effective. While challenges remain in building generalist models and conversational preference modeling, our preliminary results highlight the promise of lightweight bidirectional architectures as efficient, scalable alternatives for preference modeling.