Partial Reasoning in Language Models: Search and Refinement Guided by Uncertainty

TL;DR

This paper presents PREGU, a method that improves reasoning in large language models by monitoring output entropy to trigger localized search and refinement, enhancing multi-step inference accuracy.

Contribution

Introducing PREGU, a novel approach that uses entropy-based uncertainty monitoring to guide partial reasoning and refinement in language models.

Findings

PREGU outperforms baseline methods on multiple reasoning benchmarks.

Entropy effectively signals when to refine reasoning steps.

Localized search improves multi-step inference accuracy.

Abstract

The use of Large Language Models (LLMs) for reasoning and planning tasks has drawn increasing attention in Artificial Intelligence research. Despite their remarkable progress, these models still exhibit limitations in multi-step inference scenarios, particularly in mathematical and logical reasoning. We introduce PREGU (Partial Reasoning Guided by Uncertainty). PREGU monitors the entropy of the output distribution during autoregressive generation and halts the process whenever entropy exceeds a defined threshold, signaling uncertainty. From that point, a localized search is performed in the latent space to refine the partial reasoning and select the most coherent answer, using the Soft Reasoning method. Experiments conducted with LLaMA-3-8B, Mistral-7B, and Qwen2-7B across four reasoning benchmarks (GSM8K, GSM-Hard, SVAMP, and StrategyQA) showed performance greater than or similar to Soft Reasoning, indicating that entropy can serve as an effective signal to trigger selective refinement during reasoning.