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

TL;DR

The paper tackles the inefficiency of traditional reward modeling for Best-of-N sampling in LLMs by mining intrinsic signals from the model's own hidden states. It introduces SWIFT, a compact token-level reward model that linearly combines per-token hidden-state information to produce a path-level reward with far fewer parameters than text-based baselines. Empirical results across math, code, and reasoning benchmarks show SWIFT achieves higher accuracy than open-source baselines and delivers orders-of-magnitude gains in time, FLOPs, and memory, while also supporting logit-only training for closed-source models. SWIFT can further augment external reward models, generalizes across tasks, and maintains strong performance when using only late-layer signals, offering a practical, scalable approach for efficient LLM post-training.

Abstract

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 .

Figures (9)

• Figure 1: An illustration comparing traditional reward model and SWIFT.
• Figure 2: Cross-validation accuracy of the PCA+LDA pipeline for predicting correctness at each layer of Llama-3.1-8B. The results demonstrate that hidden states contain information about reasoning correctness with linear representation.
• Figure 3: Comparison of average time and FLOPs per sample for SWIFT and baselines, averaged across different datasets and task-performing models with 1-sigma error bars using standard error of the mean. The y-axis is plotted in log-scale. It shows that SWIFT achieves orders-of-magnitude higher efficiency than baselines. See Appendix Figure \ref{['fig:time_and_flops_full_1']} and Figure \ref{['fig:time_and_flops_full_2']} for full details.
• Figure 4: SWIFT has positive scaling with increased number of training samples and with the number of reasoning paths for inference. The result is averaged across datasets and task-performing models.
• Figure 5: Average performance of combined SWIFT and external reward models across different datasets and task-performing models, for both rank selection and scaled averaging. See Appendix Figure \ref{['fig:reward_combination_full_1']} and Figure \ref{['fig:reward_combination_full_2']} for full details.