ThoughtProbe: Classifier-Guided LLM Thought Space Exploration via Probing Representations
This work addresses the challenge of efficient and effective reasoning in LLMs, representing an incremental advancement over existing methods for steering generation.
The paper tackles the problem of improving reasoning performance in Large Language Models by introducing ThoughtProbe, a framework that uses hidden reasoning features to guide tree-structured exploration and branch aggregation, achieving significant improvements on multiple arithmetic reasoning benchmarks.
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.