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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Guo, Jizhou, Wu, Zhaomin, Yang, Hanchen, Yu, Philip S.
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866909983956795392
author Guo, Jizhou
Wu, Zhaomin
Yang, Hanchen
Yu, Philip S.
author_facet Guo, Jizhou
Wu, Zhaomin
Yang, Hanchen
Yu, Philip S.
contents 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 .
format Preprint
id arxiv_https___arxiv_org_abs_2505_12225
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Mining Intrinsic Rewards from LLM Hidden States for Efficient Best-of-N Sampling
Guo, Jizhou
Wu, Zhaomin
Yang, Hanchen
Yu, Philip S.
Machine Learning
Artificial Intelligence
Computation and Language
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 .
title Mining Intrinsic Rewards from LLM Hidden States for Efficient Best-of-N Sampling
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2505.12225