Elamur Research⌕
☰
home / papers / 2026 / ReProbe: Efficient Test-Time Scaling of Multi-Step Reasoning by Probing Internal States of Large Language Models

ReProbe: Efficient Test-Time Scaling of Multi-Step Reasoning by Probing Internal States of Large Language Models

Paper · 2026 · Jingwei Ni, Ekaterina Fadeeva, Tianyi Wu et al.

Summary

Test-time scaling improves LLM reasoning by sampling several candidate reasoning steps or full chains and keeping the most promising ones. It needs a step-level scorer, usually a Process Reward Model (PRM). PRMs are separate 1.5–8B models that are costly to train and run, and they often fail to transfer beyond math. The authors propose ReProbe, a small transformer-based probe of under 10M parameters. It sits on a frozen target LLM and estimates whether each reasoning step is correct from the model's own internal signals: attention weights, top-K logits, or hidden states from all layers. Step-level training labels come either from DeepSeek-R1 or from the target model itself, so no human labels or Monte-Carlo rollouts are needed. The probes were tested with Best-of-N and beam search on math (in-domain) and on planning and QA tasks (out-of-domain). They matched or beat PRMs up to 810× their size and generalized especially well out of domain. This suggests that LLM internal states carry usable confidence signals for cheap, general step verification.

Key points

  • Two feature sets were compared:
    • Attn+Logit: attention over the 5 preceding tokens across all layers, plus top-K candidate logits.
    • Hidden states taken from every layer.
  • Benchmarks:
    • In-domain math: MATH, GSM8k, ProofNet.
    • Out-of-domain planning: Trips, Meetings, Calendar.
    • Out-of-domain QA: StrategyQA, ScienceQA.
    • Training questions come from the PRM800K training set, which allows a fair comparison with PRMs such as Qwen2.5-Math-7B-PRM800k.
  • Size versus performance: ReProbes beat PRMs up to 150× larger and stay competitive with PRMs up to 810× larger.
  • Complementarity with PRMs: combining PRM and ReProbe scores improves step-level PR-AUC on Qwen3-8B. For example, on MATH the score rises from .586 (PRM1 alone) to .613.
  • Scaling and limitations:
    • More training data and more diverse data both help.
    • Annotating 32K trajectories with DeepSeek-R1 cost over 2000 USD.
    • Each probe works only for the specific LLM it was trained on.
    • Code, probes and annotations are released at reprobe.github.io.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓National University of SingaporeSingapore · Singapore180 works · 34,304 citations🎓The University of MelbourneMelbourne · Australia100 works · 22,618 citations🎓University of ZurichZurich · Switzerland78 works · 12,260 citations🎓ETH ZurichZurich · Switzerland72 works · 5,176 citations🎓Mohamed bin Zayed University of Artificial IntelligenceAbu Dhabi · United Arab Emirates124 works · 5,676 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works

🧭 Topics 3Research subjects and languages the work is about.Language models & pretraining · 3,320 worksReasoning · 1,426 worksInterpretability · 130 works

Cite

Cite as Jingwei Ni, Ekaterina Fadeeva, Tianyi Wu, Mubashara Akhtar, Jiaheng Zhang, Elliott Ash, Markus Leippold, Timothy Baldwin, See-Kiong Ng, Artem Shelmanov and Mrinmaya Sachan. 2026. ReProbe: Efficient Test-Time Scaling of Multi-Step Reasoning by Probing Internal States of Large Language Models. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.acl-long.536

@inproceedings{ni2026reprobe,
  title = {{ReProbe: Efficient Test-Time Scaling of Multi-Step Reasoning by Probing Internal States of Large Language Models}},
  author = {Jingwei Ni, Ekaterina Fadeeva, Tianyi Wu, Mubashara Akhtar, Jiaheng Zhang, Elliott Ash, Markus Leippold, Timothy Baldwin, See-Kiong Ng, Artem Shelmanov and Mrinmaya Sachan},
  booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  year = {2026},
  doi = {10.18653/v1/2026.acl-long.536},
  url = {https://doi.org/10.18653/v1/2026.acl-long.536}
}

References

55 references, 22 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Marah Abdin, Sahaj Agarwal, Ahmed Awadallah, Vidhisha Balachandran, Harkirat Behl, Lingjiao Chen, Gustavo de Rosa, Suriya Gunasekar, Mojan Javaheripi, Neel Joshi, et al. 2025. Phi-4-reasoning technical report. arXiv preprint arXiv:2504.21318. find
  2. AURORA. 2025. Aurora: Automated training framework of universal process reward models via ensemble prompting and reverse verification. find
  3. Amos Azaria and Tom Mitchell. 2023. The internal state of an LLM knows when it’s lying. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 967–976, Singapore. Association for Computational Linguistics.
  4. Zhangir Azerbayev, Bartosz Piotrowski, Hailey Schoelkopf, Edward W. Ayers, Dragomir Radev, and Jeremy Avigad. 2023. Proofnet: Autoformalizing and formally proving undergraduate-level mathematics. Preprint, arXiv:2302.12433.
  5. Yung-Sung Chuang, Linlu Qiu, Cheng-Yu Hsieh, Ranjay Krishna, Yoon Kim, and James R. Glass. 2024. Lookback lens: Detecting and mitigating contextual hallucinations in large language models using only attention maps. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Proces…
  6. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021. Training verifiers to solve math word problems. Preprint, arXiv:2110.14168.
  7. DeepSeek-AI. 2025. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. Preprint, arXiv:2501.12948.
  8. Ekaterina Fadeeva, Aleksandr Rubashevskii, Artem Shelmanov, Sergey Petrakov, Haonan Li, Hamdy Mubarak, Evgenii Tsymbalov, Gleb Kuzmin, Alexander Panchenko, Timothy Baldwin, et al. 2024. Fact-checking the output of large language models via token-level uncertainty quantification. In Findings of the A…
  9. Ekaterina Fadeeva, Roman Vashurin, Akim Tsvigun, Artem Vazhentsev, Sergey Petrakov, Kirill Fedyanin, Daniil Vasilev, Elizaveta Goncharova, Alexander Panchenko, Maxim Panov, Timothy Baldwin, and Artem Shelmanov. 2023. LM-polygraph: Uncertainty estimation for language models. In Proceedings of the 202… find
  10. Yichao Fu, Xuewei Wang, Yuandong Tian, and Jiawei Zhao. 2025. Deep think with confidence. Preprint, arXiv:2508.15260. find
  11. Yarin Gal et al. 2016. Uncertainty in deep learning. find
  12. Mor Geva, Daniel Khashabi, Elad Segal, Tushar Khot, Dan Roth, and Jonathan Berant. 2021. Did Aristotle Use a Laptop? A Question Answering Benchmark with Implicit Reasoning Strategies. Transactions of the Association for Computational Linguistics (TACL). find
  13. Teofilo F. Gonzalez. 1985. Clustering to minimize the maximum intercluster distance. Theoretical Computer Science, 38:293–306. find
  14. Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, et al. 2025. Deepseek-r1 incentivizes reasoning in llms through reinforcement learning. Nature, 645(8081):633–638. find
  15. Jinwen He, Yujia Gong, Zijin Lin, Cheng’an Wei, Yue Zhao, and Kai Chen. 2024a. LLM factoscope: Uncovering LLMs’ factual discernment through measuring inner states. In Findings of the Association for Computational Linguistics: ACL 2024, pages 10218–10230, Bangkok, Thailand. Association for Computatio… find
  16. Jujie He, Tianwen Wei, Rui Yan, Jiacai Liu, Chaojie Wang, Yimeng Gan, Shiwen Tu, Chris Yuhao Liu, Liang Zeng, Xiaokun Wang, Boyang Wang, Yongcong Li, Fuxiang Zhang, Jiacheng Xu, Bo An, Yang Liu, and Yahui Zhou. 2024b. Skywork-o1 open series. https://huggingface.co/Skywork. ↗
  17. Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021. Measuring mathematical problem solving with the math dataset. Preprint, arXiv:2103.03874.
  18. Xiaolin Hu, Qinghua Zhou, Bogdan Grechuk, and Ivan Y. Tyukin. 2025. Stepproof: Step-by-step verification of natural language mathematical proofs. Preprint, arXiv:2506.10558. — [arXiv:2506.10558](https://arxiv.org/abs/2506.10558) find
  19. HuggingFaceH4. 2025. Qwen2.5-math-1.5b-instruct-prm-0.2. Hugging Face Model Hub. Fine-tuned version of Qwen/Qwen2.5-Math-1.5B-Instruct on the prm800k-trl-dedup dataset, trained with PRM and TRL. Safetensors format. arXiv:2211.14275.
  20. Saurav Kadavath, Tom Conerly, Amanda Askell, Tom Henighan, Dawn Drain, Ethan Perez, Nicholas Schiefer, Zac Hatfield-Dodds, Nova DasSarma, Eli Tran-Johnson, Scott Johnston, Sheer El-Showk, Andy Jones, Nelson Elhage, Tristan Hume, Anna Chen, Yuntao Bai, Sam Bowman, Stanislav Fort, Deep Ganguli, Danny… find
  21. Zhewei Kang, Xuandong Zhao, and Dawn Song. 2025. Scalable best-of-n selection for large language models via self-certainty. In Advances in Neural Information Processing Systems. find
  22. Lorenz Kuhn, Yarin Gal, and Sebastian Farquhar. 2023. Semantic uncertainty: Linguistic invariances for uncertainty estimation in natural language generation. In The Eleventh International Conference on Learning Representations. find
  23. Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2024. Let’s verify step by step. In The Twelfth International Conference on Learning Representations.
  24. Zhen Lin, Shubhendu Trivedi, and Jimeng Sun. 2024. Generating with confidence: Uncertainty quantification for black-box large language models. find
  25. Chengwu Liu, Ye Yuan, Yichun Yin, Yan Xu, Xin Xu, Zaoyu Chen, Yasheng Wang, Lifeng Shang, Qun Liu, and Ming Zhang. 2025. Safe: Enhancing mathematical reasoning in large language models via retrospective step-aware formal verification. In Proceedings of the 63rd Annual Meeting of the Association for… find
  26. Pan Lu, Swaroop Mishra, Tony Xia, Liang Qiu, Kai-Wei Chang, Song-Chun Zhu, Oyvind Tafjord, Peter Clark, and Ashwin Kalyan. 2022. Learn to explain: Multimodal reasoning via thought chains for science question answering. In The 36th Conference on Neural Information Processing Systems (NeurIPS). find
  27. LM-polygraph. Transactions of the Association for Computational Linguistics, 13:220–248. find
  28. Liangchen Luo, Yinxiao Liu, Rosanne Liu, Samrat Phatale, Meiqi Guo, Harsh Lara, Yunxuan Li, Lei Shu, Yun Zhu, Lei Meng, Jiao Sun, and Abhinav Rastogi. 2024. Improve mathematical reasoning in language models by automated process supervision. Preprint, arXiv:2406.06592.
  29. Roman Vashurin, Maiya Goloburda, Albina Ilina, Aleksandr Rubashevskii, Preslav Nakov, Artem Shelmanov, and Maxim Panov. 2025b. Cocoa: A minimum bayes risk framework bridging confidence and consistency for uncertainty quantification in llms. In The Thirty-ninth Annual Conference on Neural Information… find
  30. Andrey Malinin and Mark J. F. Gales. 2021. Uncertainty estimation in autoregressive structured prediction. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. find
  31. Shentong Mo and Miao Xin. 2024. Tree of uncertain thoughts reasoning for large language models. In ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 12742–12746. IEEE. find
  32. Artem Vazhentsev, Ekaterina Fadeeva, Rui Xing, Gleb Kuzmin, Ivan Lazichny, Alexander Panchenko, Preslav Nakov, Timothy Baldwin, Maxim Panov, and Artem Shelmanov. 2025a. Unconditional truthfulness: Learning unconditional uncertainty of large language models. In Proceedings of the 2025 Conference on E…
  33. Artem Shelmanov, Ekaterina Fadeeva, Akim Tsvigun, Ivan Tsvigun, Zhuohan Xie, Igor Kiselev, Nico Daheim, Caiqi Zhang, Artem Vazhentsev, Mrinmaya Sachan, Preslav Nakov, and Timothy Baldwin. 2025a. A head to predict and a head to question: Pre-trained uncertainty quantification heads for hallucination… find
  34. Artem Vazhentsev, Lyudmila Rvanova, Ivan Lazichny, Alexander Panchenko, Maxim Panov, Timothy Baldwin, and Artem Shelmanov. 2025b. Token-level density-based uncertainty quantification methods for eliciting truthfulness of large language models. In Proceedings of the 2025 Conference of the Nations of…
  35. Artem Shelmanov, Maxim Panov, Roman Vashurin, Artem Vazhentsev, Ekaterina Fadeeva, and Timothy Baldwin. 2025b. Uncertainty quantification for large language models. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 5: Tutorial Abstracts), pages 3– 4,… find
  36. Peiyi Wang, Lei Li, Zhihong Shao, Runxin Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. 2024. Math-shepherd: Verify and reinforce LLMs step-by-step without human annotations. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Paper…
  37. Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. 2024. Scaling llm test-time compute optimally can be more effective than scaling model parameters. arXiv preprint arXiv:2408.03314.
  38. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. 2022. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837.
  39. Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. 2025. Scaling llm test-time compute optimally can be more effective than scaling parameters for reasoning. In International Conference on Representation Learning, volume 2025, pages 10131–10165. find
  40. Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins. 2022. Solving math word problems with process- and outcome-based feedback. arXiv preprint arXiv:2211.14275.
  41. Roman Vashurin, Ekaterina Fadeeva, Artem Vazhentsev, Lyudmila Rvanova, Daniil Vasilev, Akim Tsvigun, Sergey Petrakov, Rui Xing, Abdelrahman Sadallah, Kirill Grishchenkov, Alexander Panchenko, Timothy Baldwin, Preslav Nakov, Maxim Panov, and Artem Shelmanov. 2025a. Benchmarking uncertainty quantifica…
  42. Wei Xiong, Hanning Zhang, Nan Jiang, and Tong Zhang. 2024. An implementation of generative prm. find
  43. Wei Xiong, Wenting Zhao, Weizhe Yuan, Olga Golovneva, Tong Zhang, Jason Weston, and Sainbayar Sukhbaatar. 2025. Stepwiser: Stepwise generative judges for wiser reasoning. Preprint, arXiv:2508.19229. — [arXiv:2508.19229](https://arxiv.org/abs/2508.19229) find
  44. Hang Yan, Fangzhi Xu, Rongman Xu, Yifei Li, Jian Zhang, Haoran Luo, Xiaobao Wu, Luu Anh Tuan, Haiteng Zhao, Qika Lin, et al. 2025. Mur: Momentum uncertainty guided reasoning for large language models. arXiv preprint arXiv:2507.14958.
  45. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. 2025. Qwen3 technical report. arXiv preprint arXiv:2505.09388.
  46. Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. 2023. Tree of thoughts: Deliberate problem solving with large language models. Advances in neural information processing systems, 36:11809–11822. find
  47. Jin Zhou, Charles Staats, Wenda Li, Christian Szegedy, Kilian Weinberger, and Yuhuai Wu. 2024. Don't trust: Verify – grounding llm quantitative reasoning with autoformalization. In International Conference on Representation Learning, volume 2024, pages 2544– 2563. find
  48. Kuo Zhou and Lu Zhang. 2025. Step-wise formal verification for llm-based mathematical problem solving. Preprint, arXiv:2505.20869. — [arXiv:2505.20869](https://arxiv.org/abs/2505.20869) find
  49. Zhangyue Yin, Qiushi Sun, Qipeng Guo, Zhiyuan Zeng, Xiaonan Li, Junqi Dai, Qinyuan Cheng, Xuanjing Huang, and Xipeng Qiu. 2024. Reasoning in flux: Enhancing large language models reasoning through uncertainty-aware adaptive guidance. In Proceedings of the 62nd Annual Meeting of the Association for C… find
  50. Jinghan Zhang, Xiting Wang, Fengran Mo, Yeyang Zhou, Wanfu Gao, and Kunpeng Liu. 2025a. Entropy-based exploration conduction for multi-step reasoning. In Findings of the Association for Computational Linguistics: ACL 2025, pages 3895–3906, Vienna, Austria. Association for Computational Linguistics. find
  51. Yanzhao Zhang, Mingxin Li, Dingkun Long, Xin Zhang, Huan Lin, Baosong Yang, Pengjun Xie, An Yang, Dayiheng Liu, Junyang Lin, Fei Huang, and Jingren Zhou. 2025b. Qwen3 embedding: Advancing text embedding and reranking through foundation models. arXiv preprint arXiv:2506.05176.
  52. Zhenru Zhang, Chujie Zheng, Yangzhen Wu, Beichen Zhang, Runji Lin, Bowen Yu, Dayiheng Liu, Jingren Zhou, and Junyang Lin. 2025c. The lessons of developing process reward models in mathematical reasoning. Preprint, arXiv:2501.07301.
  53. Jian Zhao, Runze Liu, Kaiyan Zhang, Zhimu Zhou, Junqi Gao, Dong Li, Jiafei Lyu, Zhouyi Qian, Biqing Qi, Xiu Li, and Bowen Zhou. 2025. Genprm: Scaling test-time compute of process reward models via generative reasoning. arXiv preprint arXiv:2504.00891. — [arXiv:2504.00891](https://arxiv.org/abs/2504.00891) find
  54. Chujie Zheng, Zhenru Zhang, Beichen Zhang, Runji Lin, Keming Lu, Bowen Yu, Dayiheng Liu, Jingren Zhou, and Junyang Lin. 2025. ProcessBench: Identifying process errors in mathematical reasoning. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long…
  55. Huaixiu Steven Zheng, Swaroop Mishra, Hugh Zhang, Xinyun Chen, Minmin Chen, Azade Nova, Le Hou, Heng-Tze Cheng, Quoc V. Le, Ed H. Chi, and Denny Zhou. 2024. Natural plan: Benchmarking llms on natural language planning. Preprint, arXiv:2406.04520. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/ni-2026-reprobe-efficient-test-time-scaling/