PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiaxi Li, Yucheng Shi, Xiao Huang, Jin Lu, Ninghao Liu

arXiv:2510.03632v2cs.AI

TL;DR

Finding correct reasoning paths is difficult when single-path generation can accumulate errors and exhaustive multi-path search is infeasible. MITS guides tree search with PMI scoring, adaptive sampling, and weighted voting; the paper reports improved benchmark performance and an accuracy/time advantage over MCTS baselines.

  • Problem

    Single-path reasoning can accumulate errors, while generating multiple paths improves solution coverage but exhaustive search is infeasible.

    Source: §1 · p. 1

  • Method

    MITS uses question-specific PMI scoring to guide beam-pruned tree search, entropy-based sampling to allocate compute to uncertain steps, and weighted voting that combines PMI scores with prediction consensus.

    Source: §1 · p. 2, §3.2 · p. 5, §3.3 · p. 5

  • Results

    MITS reports substantial improvements over strong baselines across diverse reasoning benchmarks; on StrategyQA with QWEN2.5-3B, it achieves 68.45% accuracy versus CoT’s 47.34%.

    Source: §1 · p. 2, §4.2 · p. 7

  • Takeaways & Limitations

    On StrategyQA, MITS reports the highest accuracy among the compared methods with moderate computation time, offering a favorable accuracy/time trade-off against MCTS baselines.

    Source: §4.2 · p. 7

Abstract

from arXiv · show

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. 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. The code is available at https://github.com/plusnli/MITS.

1 INTRODUCTION

Multi-path search can improve the chance of finding correct reasoning, but exhaustive exploration is infeasible and existing evaluations can be costly or insufficiently discriminating. MITS uses question-specific PMI scoring and adaptive search to guide reasoning, with experiments reporting gains over strong baselines.

1 INTRODUCTION

  • Generating multiple reasoning paths can improve solution coverage, but the exponentially growing search space makes exhaustive enumeration infeasible.The central challenge is efficiently identifying paths likely to yield correct solutions.Source: §1 · p. 1
  • Existing approaches incur costly forward simulations, provide limited quantitative assessments, or favor generic paths over reasoning tailored to the question.These limitations motivate efficient evaluation that distinguishes question-specific reasoning.Source: §1 · pp. 1–2
  • MITS uses pointwise mutual information (PMI) to score how much a reasoning path’s plausibility increases because of the specific question.Scores are computed as reasoning unfolds, avoiding expensive look-ahead simulations.Source: §1 · p. 2
  • MITS combines PMI scoring, entropy-based dynamic sampling, and weighted voting, and the authors report substantial improvements over strong baselines across diverse benchmarks.Sampling allocates resources to uncertain steps, while voting combines PMI scores with prediction consensus.Source: §1 · p. 2

2 PRELIMINARY

Tree-search reasoning explores multiple intermediate-step paths and selects a plausible solution path. MCTS commonly evaluates intermediate steps through costly future rollouts, motivating a scoring function that assesses reasoning traces without simulated rollouts.

2 PRELIMINARY

  • A search tree represents a question as its root and intermediate reasoning steps as child nodes; a candidate solution path concatenates the question with successive steps.The framework extracts a set of candidate paths and aims to choose the most plausible one as the final answer.Source: §2 · p. 2
  • MCTS typically evaluates intermediate steps with repeated look-ahead rollouts that simulate future trajectories and back-propagate value estimates, incurring substantial computational cost.Prior work also identifies reliable trajectory selection as difficult.Source: §2 · p. 2
  • The section motivates scoring reasoning traces without simulated rollouts to address both computational efficiency and trajectory selection.Source: §2 · p. 2

3 METHODOLOGY

MITS builds a reasoning tree by scoring paths with question-specific information, adapting candidate generation to uncertainty, and using beam search to retain promising paths. It then combines path scores with prediction frequency to select an output.

3.1 POINTWISE MUTUAL INFORMATION (PMI) SCORING

  • PMI scores reasoning paths by how much their probability increases given the question, rewarding question-specific reasoning and penalizing paths likely without the question.The score is estimated from question-conditioned and context-only probabilities; incremental updates measure each new step's information gain beyond preceding steps.Source: §3.1 · pp. 3–4

3.2 SEARCH TREE CONSTRUCTION

  • Entropy-based sampling adjusts how many next-step candidates are generated, while PMI-guided beam search retains only the top-B paths by cumulative score.Entropy thresholds use empirical 33rd and 67th percentiles of generated-step entropies to adapt to task and question difficulty.Source: §3.2 · pp. 4–5

3.3 WEIGHTED AVERAGE VOTING FOR OUTPUT SELECTION

  • Weighted Average Voting reranks the top-K reasoning chains using each chain's PMI multiplied by the frequency of its final prediction.The method combines PMI confidence with prediction consensus to reduce brittleness from spurious query correlations or linguistic-pattern overfitting.Source: §3.3 · p. 5

4 EXPERIMENTS

The experiments compare MITS with prompting, sampling, and tree-search baselines across three reasoning datasets, then examine its computational efficiency and design choices. MITS outperforms baselines across diverse tasks, with experiments also assessing how performance changes with rollout count and other settings.

4.1 SETUP

  • Experiments use four instruction-tuned models across StrategyQA, ARC-Challenge, and CommonsenseQA, comparing training-free CoT, CoT-SC, ToT, RAP, and rStar baselines.The baselines cover single-chain prompting, multiple-chain sampling, and tree-search methods.Source: §4.1 · p. 6, §4.2 · p. 7
  • MITS uses maximum tree depth 10, base sampling Nbase = 3, beam width 32, and defaults evaluator E to generator G.The experiments also compare MITS with MITS-F, which fully expands the search tree without beam search.Source: §4.1 · p. 6

4.2 MAIN RESULTS

  • 68.45% vs 47.34% accuracy on StrategyQA with QWEN2.5-3B, and +5.31% over strongest baseline rStar on ARC-Challenge with QWEN2.5-7B.MITS also maintains strong results across logical, scientific, and commonsense reasoning tasks.Source: §4.2 · p. 7
  • MITS reaches 68.45% accuracy in 64.41s on StrategyQA, versus CoT-SC’s fastest 2.75s at 56.12% accuracy and MCTS runtimes of 203.42s (RAP) and 815.67s (rStar).The experiments measure average wall-clock time from inference start to reasoning completion.Source: Table 2 · p. 7, §4.2 · p. 7

4.3 ABLATION STUDY

  • MITS shows superior accuracy early with fewer rollouts and continues improving as the computational budget increases.The study compares accuracies of CoT-SC, RAP, rStar, and MITS across rollout counts on ARC-Challenge; for CoT-SC and MITS, rollout count equals collected reasoning chains.Source: §4.3 · pp. 7–8, Figure 3 · p. 8
  • With QWEN2.5-3B as generator, accuracy rises from 86.78% with QWEN2.5-1.5B evaluator to 90.61% with QWEN2.5-7B evaluator.More powerful evaluators usually yield better results, which the authors attribute to more accurate probability estimates and more reliable PMI values.Source: Table 3 · p. 8, §4.3 · p. 8
  • ARC-Challenge rises from 84.31% to 89.12% as K increases from 1 to 8; optimal K is 16 on StrategyQA and 32 on ARC-Challenge and CommonsenseQA.Performance stabilizes at larger K values, indicating diminishing returns from including additional reasoning paths.Source: §4.3 · p. 8, Table 5 · p. 9
  • Length-normalized average PMI beats sum aggregation by +2.56% on QWEN2.5-3B and +3.42% on PHI-3.5-MINI, preventing unfair penalties across path lengths.The comparison is on StrategyQA; average aggregation divides cumulative PMI by the number of reasoning steps.Source: §4.3 · pp. 8–9, Table 4 · p. 8

5 RELATED WORK

Related work covers prompting methods for single-round inference and tree-based methods for exploring multiple reasoning paths. MITS is presented as orthogonal to prompting approaches.

5 RELATED WORK

  • Prompting approaches include sub-question decomposition, exemplar selection, higher-order reasoning, and code-integrated prompting for single-round inference.The section describes these methods as orthogonal to MITS.Source: §5 · p. 9
  • Tree-search approaches explore multiple reasoning paths through tree structures, beam search, Monte Carlo Tree Search, and planning applications.Tree-of-Thought is identified as a pioneering tree-search approach.Source: §5 · p. 9
  • Sampling diverse reasoning paths can outperform one-time greedy decoding, with Chain-of-Thought Self-Consistency among the cited approaches.The passage attributes this direction theoretical and empirical support.Source: §5 · p. 9

6 CONCLUSION

MITS uses PMI-guided tree search, entropy-based dynamic sampling, and weighted average voting. Experiments report stronger performance than several baselines across diverse reasoning datasets, alongside computational efficiency.

6 CONCLUSION

  • MITS substantially outperforms several strong baseline methods across diverse reasoning datasets while providing a computationally efficient alternative to traditional tree search.The conclusion says these results support information-theoretic guidance across diverse reasoning tasks.Source: §6 · p. 9
  • MITS combines PMI-based reasoning-path evaluation, entropy-based adaptive compute allocation, and weighted average voting for answer selection.The conclusion identifies these as its three key innovations.Source: §6 · p. 9
  • The paper reports extensive ablation studies intended to provide analysis and insights for advanced LLM reasoning techniques.Source: §6 · p. 9

A.1 DETAILED DERIVATION FOR PMI SCORING

The appendix derives a recurrence for PMI scoring and details an entropy-based sampling adjustment that scales sampling intensity according to deviation from a moderate entropy range.

A.1 DETAILED DERIVATION FOR PMI SCORING

  • The appendix derives the PMI recurrence from Equation 2 and presents the resulting recurrence relation in Equation 3.The displayed score term subtracts the unconditional log probability of a step from its question-conditioned log probability.Source: §A.1 · p. 14

A.2 FORMULA FOR DYNAMIC SAMPLING

  • Dynamic sampling uses proportional control to adjust sample counts according to how far current entropy falls outside the moderate range.The base sample count is 3, and dynamic sampling is enabled only when m ≥ 10; the adjustment is bounded.Source: §A.2 · p. 14
  • The entropy adjustment is scaled relative to the range between Hlow and Hhigh, with Vp serving as the proportional gain.The appendix describes this normalization as making the controller robust to different entropy scales.Source: §A.2 · p. 14

B.1 HYPERPARAMETERS

The appendix reports MITS inference and search-tree settings, then provides case studies with example prompts and reasoning paths for StrategyQA and ARC-Challenge.

B.1 HYPERPARAMETERS

  • The search-tree settings use 3 candidate next steps, maximum depth 10, a 512-token limit per candidate, and beam width 32.The hyperparameter search also covers decoding temperature, top-k, and top-p; Table 6 is titled “Hyperparameter Configuration for MITS.”Source: §B.1 · p. 14, Table 6 · p. 15

C CASE STUDY

  • The case studies use QWEN2.5-3B as both generator and evaluator on StrategyQA, ARC-Challenge, and CommonsenseQA.The appendix records each node’s index, depth, and reasoning-step content, and computes PMI from the difference between two log probabilities.Source: C CASE STUDY · p. 15

C.1 STRATEGYQA

  • For CAS number 8009-03-8, identified as TDI, the example concludes it is not specifically harmful to rashes for the majority of people, so the rash-harm statement is false.The cited reasoning steps are at depths 1, 3, and 4; the answer is [B. False].Source: §C.1 · pp. 15–16

C.2 ARC-CHALLENGE

  • In the ARC-Challenge example, the selected safety measure is labeling the new ammonia container clearly.The reasoning path contrasts this choice with recycling the old bottle, heating the container, or refrigerating it.Source: §C.2 · p. 17

C.3 COMMONSENSEQA

The CommonsenseQA example asks where a weasel could access chicken eggs and shows a reasoning path that selects the barn. The path assigns PMI scores to successive reasoning steps, from contextualizing the question to stating the answer.

C.3 COMMONSENSEQA

  • The example question asks where a weasel could get into chicken eggs, with barn among the listed options.The prompt requests step-by-step reasoning and a final answer in square brackets.Source: §C.3 · p. 17
  • The path first frames the question as asking where eggs are stored and scores that context at PMI 50.72.The passage reports log p(S|q) = -86.52 and log p(S) = -137.24.Source: §C.3 · p. 18
  • The option analysis favors the barn as a place where agricultural products, including eggs, are commonly stored, scoring PMI 83.32.The passage reports log p(S|q) = -127.25 and log p(S) = -210.57.Source: §C.3 · p. 18
  • The path identifies the barn as the most likely answer, then concludes with [B. barn], with PMI scores of 19.61 and 22.88, respectively.The most-likely-answer step has PMI 19.61; the final-answer step has PMI 22.88.Source: §C.3 · p. 18
Loading 2510.03632v2…