ThoughtProbe: Classifier-Guided LLM Thought Space Exploration via Probing Representations

TL;DR

ThoughtProbe is a framework that uses classifier-guided exploration of LLM reasoning paths at inference time, significantly improving arithmetic reasoning accuracy by efficiently identifying valid reasoning chains.

Contribution

It introduces a novel inference-time exploration method leveraging hidden representations as discriminative signals to guide reasoning path selection in LLMs.

Findings

Achieves significant improvements on arithmetic reasoning benchmarks.

Effectively identifies valid reasoning chains through comprehensive exploration.

Utilizes classifier-guided tree expansion for efficient resource allocation.

Abstract

This paper introduces ThoughtProbe, a novel inference time framework that leverages the hidden reasoning features of Large Language Models (LLMs) to improve their reasoning performance. Unlike previous works that manipulate the hidden representations to steer LLM generation, we harness them as discriminative signals to guide the tree structured response space exploration. In each node expansion, a classifier serves as a scoring and ranking mechanism that efficiently allocates computational resources by prioritizing higher score candidates for continuation. After completing the tree expansion, we collect answers from all branches to form a candidate answer pool. We then propose a branch aggregation method that marginalizes over all supporting branches by aggregating their CoT scores, thereby identifying the optimal answer from the pool. Experimental results show that our framework's comprehensive exploration not only covers valid reasoning chains but also effectively identifies them, achieving significant improvements across multiple arithmetic reasoning benchmarks.