Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MITS: Enhanced Tree Search Reasoning for LLMs via Pointwise Mutual Information

Jiaxi Li, Yucheng Shi, Jin Lu, Ninghao Liu

TL;DR

This work addresses the challenge of efficiently exploring the space of reasoning paths in large language models. It introduces Mutual Information Tree Search (MITS), which uses pointwise mutual information (PMI) to score intermediate reasoning steps, enabling stepwise evaluation without costly look-ahead rollouts. MITS couples PMI-based beam search with entropy-driven dynamic sampling and a weighted voting scheme that fuses PMI with prediction consensus, achieving superior accuracy while maintaining computational efficiency across diverse reasoning benchmarks. The results demonstrate that information-theoretic guidance can substantially improve LLM reasoning performance and robustness, offering a principled framework for test-time reasoning.

Abstract

Tree search has become as a representative framework for test-time reasoning with large language models (LLMs), exemplified by methods such as Tree-of-Thought and Monte Carlo Tree Search that explore multiple reasoning paths. However, it remains difficult to provide instant and reliable quantitative assessments of intermediate reasoning step quality, and extensive path exploration is computationally costly. To address this, we propose Mutual Information Tree Search (MITS), a novel framework that guides reasoning with information-theoretic principles. MITS introduces an effective scoring function based on pointwise mutual information (PMI), which enables step-wise evaluation of reasoning paths and search tree expansion via beam search without expensive look-ahead simulations, achieving superior reasoning performances while maintaining computational efficiency. The framework is complemented by an entropy-based dynamic sampling strategy that adaptively allocates computational resources to uncertain reasoning steps where exploration is most beneficial. For final prediction, MITS employs a weighted voting scheme that combines PMI scores with prediction consensus. Through comprehensive experiments on diverse reasoning benchmarks, MITS consistently surpasses baseline methods, establishing a principled and efficient framework for LLM reasoning.

MITS: Enhanced Tree Search Reasoning for LLMs via Pointwise Mutual Information

TL;DR

This work addresses the challenge of efficiently exploring the space of reasoning paths in large language models. It introduces Mutual Information Tree Search (MITS), which uses pointwise mutual information (PMI) to score intermediate reasoning steps, enabling stepwise evaluation without costly look-ahead rollouts. MITS couples PMI-based beam search with entropy-driven dynamic sampling and a weighted voting scheme that fuses PMI with prediction consensus, achieving superior accuracy while maintaining computational efficiency across diverse reasoning benchmarks. The results demonstrate that information-theoretic guidance can substantially improve LLM reasoning performance and robustness, offering a principled framework for test-time reasoning.

Abstract

Tree search has become as a representative framework for test-time reasoning with large language models (LLMs), exemplified by methods such as Tree-of-Thought and Monte Carlo Tree Search that explore multiple reasoning paths. However, it remains difficult to provide instant and reliable quantitative assessments of intermediate reasoning step quality, and extensive path exploration is computationally costly. To address this, we propose Mutual Information Tree Search (MITS), a novel framework that guides reasoning with information-theoretic principles. MITS introduces an effective scoring function based on pointwise mutual information (PMI), which enables step-wise evaluation of reasoning paths and search tree expansion via beam search without expensive look-ahead simulations, achieving superior reasoning performances while maintaining computational efficiency. The framework is complemented by an entropy-based dynamic sampling strategy that adaptively allocates computational resources to uncertain reasoning steps where exploration is most beneficial. For final prediction, MITS employs a weighted voting scheme that combines PMI scores with prediction consensus. Through comprehensive experiments on diverse reasoning benchmarks, MITS consistently surpasses baseline methods, establishing a principled and efficient framework for LLM reasoning.

Paper Structure

This paper contains 44 sections, 8 equations, 3 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Tree-Search Based Reasoning.
  4. Challenges in Tree Search.
  5. Methodology
  6. Pointwise Mutual Information (PMI) Scoring
  7. Probability Estimation.
  8. PMI Scoring with Incremental Update.
  9. Search Tree Construction
  10. Reasoning Step Generation.
  11. Dynamic Sampling via Entropy Quantification.
  12. Weighted Average Voting for Output Selection
  13. Experiments
  14. Setup
  15. Backbone models and datasets.
  16. ...and 29 more sections

Figures (3)

  • Figure 1: Overview for the MITS approach. The search tree is incrementally constructed by generating reasoning steps conditioned on prior context. Dynamic sampling controls the number of candidates (i.e., nodes) generated at each step. Beam search is applied to prune less promising paths. Weighted average voting is used to aggregate candidates and select the final output.
  • Figure 2: Prompt examples for step-by-step reasoning generation, showing the iterative generation from initial step to subsequent steps, similar to lai2024step.
  • Figure 3: Performance comparison across different rollouts on the ARC-Challenge dataset. MITS not only demonstrates superior performance in early stages with fewer rollouts, but also continues to improve with increased computational budget.