mlsanityData to 14 Oct 2026

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

arXiv 2505.12225 ↗First posted 18 May 2025cs.LG cs.AI cs.CLcited by 9

On the map: LLM alignment / RLHF / preference optimisation → Process reward models for reasoning

Abstract (arXiv)

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 .

Key facts. Cited by 9 papers in the map (data through 14 Oct 2026). Placed in Process reward models for reasoning, one of 10 directions in LLM alignment / RLHF / preference optimisation. Builds on 37 of its 76 references that are in the map. First posted on arXiv 18 May 2025.

Builds on

37 of 76 references.

Closest ancestors · 2

Same direction, newest first

  1. 2024Advancing LLM Reasoning Generalists with Preference Trees2 Apr
  2. 2023Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations14 Dec

Other references · 35

  1. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
  2. 2023GPT-4 Technical Reportcited by 18,396
  3. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 13,028
  4. 2021Training Verifiers to Solve Math Word Problemscited by 10,109
  5. 2023Direct Preference Optimization: Your Language Model is Secretly a Reward Modelcited by 8,065
  6. 2020Scaling Laws for Neural Language Modelscited by 6,586
  7. 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,944
  8. 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,371
  9. 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,100
  10. 2022Locating and Editing Factual Associations in GPTcited by 2,691
  11. 2023Representation Engineering: A Top-Down Approach to AI Transparencycited by 1,336
  12. 2022Discovering Latent Knowledge in Language Models Without Supervisioncited by 767
23 more
  1. 2023The Internal State of an LLM Knows When It's Lyingcited by 659
  2. 2017Program Induction by Rationale Generation : Learning to Solve and Explain Algebraic Word Problemscited by 650
  3. 2022DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scalecited by 407
  4. 2023UltraFeedback: Boosting Language Models with Scaled AI Feedbackcited by 406
  5. 2024Skywork-Reward: Bag of Tricks for Reward Modeling in LLMscited by 327
  6. 2024PKU-SafeRLHF: Towards Multi-Level Safety Alignment for LLMs with Human Preferencecited by 261
  7. 2024From r to Q^*: Your Language Model is Secretly a Q-Functioncited by 242
  8. 2024AnyEdit: Mastering Unified High-Quality Image Editing for Any Ideacited by 241
  9. 2025Reasoning Models Know When They're Right: Probing Hidden States for Self-Verificationcited by 159
  10. 2025Can 1B LLM Surpass 405B LLM? Rethinking Compute-Optimal Test-Time Scalingcited by 136
  11. 2023WebGLM: Towards An Efficient Web-Enhanced Question Answering System with Human Preferencescited by 100
  12. 2024Noise Contrastive Alignment of Language Models with Explicit Rewardscited by 85
  13. 2023Search-in-the-Chain: Interactively Enhancing Large Language Models with Search for Knowledge-intensive Taskscited by 73
  14. 2024Robust Preference Optimization through Reward Model Distillationcited by 67
  15. 2025Open Problems in Machine Unlearning for AI Safetycited by 51
  16. 2024Logits of API-Protected LLMs Leak Proprietary Informationcited by 40
  17. 2024Calibrating Reasoning in Language Models with Internal Consistencycited by 28
  18. 2023Learning to Relate to Previous Turns in Conversational Searchcited by 26
  19. 2024TRAD: Enhancing LLM Agents with Step-Wise Thought Retrieval and Aligned Decisioncited by 23
  20. 2022UnifieR: A Unified Retriever for Large-Scale Retrievalcited by 19
  21. 2023GLM-Dialog: Noise-tolerant Pre-training for Knowledge-grounded Dialogue Generationcited by 9
  22. 2024Would I Lie To You? Inference Time Alignment of Language Models using Direct Preference Headscited by 5
  23. 2024Exploring Domain Robust Lightweight Reward Models based on Router Mechanismcited by 1

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 3; 2026: 6.

9 papers.

Similar, not cited

Close in wording, no citation either way

  1. 2024Automated Rewards via LLM-Generated Progress Functions
  2. 2026Scaling Reward Modeling without Human Supervision
  3. 2025Sailing by the Stars: A Survey on Reward Models and Learning Strategies for Learning from Rewards
  4. 2025Confidence as a Reward: Transforming LLMs into Reward Models
  5. 2026Consolidating Rewarded Perturbations for LLM Post-Training
  6. 2026PRISM: A Unified Framework for Post-Training LLMs Without Verifiable Rewards
  7. 2024LASeR: Learning to Adaptively Select Reward Models with Multi-Armed Bandits
  8. 2026WiSparse: Boosting LLM Inference Efficiency with Weight-Aware Mixed Activation Sparsity