arXiv:2505.12225cs.LGcs.AI2025-05KDD被引 11

用大模型内部状态直接生成奖励,让高效采样不再依赖大量标注数据。

Mining Intrinsic Rewards from LLM Hidden States for Efficient Best-of-N Sampling

  • 从大模型隐藏状态中挖掘内在奖励信号,无需外部文本奖励模型。
  • 在MATH数据集上比基线高12.7%准确率,参数量不足其0.005%。
  • 轻量高效,可适配闭源模型,适合追求效率的部署场景。

Best-of-N采样能显著提升大语言模型性能,但通常依赖于大量文本奖励模型,这些模型计算成本高且需海量标注数据训练,造成显著的数据与效率瓶颈。本文提出SWIFT(Simple Weighted Intrinsic Feedback Technique),一种轻量级方法,直接从大模型自身隐藏状态中学习奖励函数。该方法在词元嵌入层面使用简单线性层,有效区分偏好与非偏好生成结果,避免了复杂的文本建模。在标准基准测试中,SWIFT在MATH数据集上比EurusRM-7B高出12.7%准确率,同时仅需不到0.005%的参数量。其强可扩展性、支持通过logit访问适配部分闭源模型,以及可与传统奖励模型结合进一步提效,展现出显著实用价值,推动更高效的数据驱动式大模型后训练。代码已开源:https://github.com/aster2024/SWIFT。

原文摘要 · Abstract (English)

Best-of-N sampling is a powerful method for improving Large Language Model (LLM) performance, but it is often limited by its dependence on massive, text-based reward models. These models are not only computationally expensive but also data-hungry, requiring extensive labeled datasets for training. This creates a significant data challenge, as they overlook a rich, readily available data source: the LLM's own internal hidden states. To address this data and efficiency gap, we introduce SWIFT (Simple Weighted Intrinsic Feedback Technique), a novel and lightweight method that learns a reward function directly from the rich information embedded in LLM hidden states. Operating at the token embedding level, SWIFT employs simple linear layers to effectively distinguish between preferred and dispreferred generations, eliminating the need for computationally intensive text-based modeling. Extensive experiments on standard benchmarks show that SWIFT outperforms existing baselines (12.7% higher accuracy than EurusRM-7B on MATH dataset) while using less than 0.005% of their parameters. Its robust scalability, compatibility with certain closed-source models via logit access, and ability to combine with traditional reward models for additional performance highlight SWIFT's practical value and contribution to more efficient data-driven LLM post-training. Our code is available at https://github.com/aster2024/SWIFT .

奖励模型高效采样隐状态利用

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。