mits-pointwise-mutual-information-tree-search

name	mits-pointwise-mutual-information-tree-search
title	MITS: Enhanced Tree Search Reasoning for LLMs via Pointwise Mutual Information
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2510.03632
keywords	["Tree Search","Reasoning","Pointwise Mutual Information","LLM Decoding","Test-Time Scaling"]
description	Score reasoning paths using PMI to identify question-specific relevant steps without rollout simulations, enabling efficient tree search 12× faster than MCTS while improving reasoning accuracy.

Technique: PMI-Based Efficient Tree Search for LLM Reasoning

Large language models generate multiple reasoning trajectories to solve complex problems, but evaluating which paths are most promising typically requires expensive rollout simulations or external verifiers. MITS uses Pointwise Mutual Information (PMI) to score reasoning steps directly: how much does a particular step increase the likelihood of being correct specifically for this question?

Unlike generic step quality scores, PMI captures question-specific relevance. A mathematical notation definition might be high-quality universally but irrelevant to an algebra problem about finances. By computing PMI(question; step), MITS filters out generic patterns and focuses computation on steps that genuinely inform the specific problem.

Core Concept

MITS operates on three key innovations:

1. PMI Scoring: Each reasoning step S receives a score measuring "how much a reasoning path's plausibility increases because of the specific question," computed as log p(S|q) / p(S).

2. Incremental Computation: PMI is calculated as reasoning unfolds, not recomputed from scratch. Each new step contributes its conditional probability given prior steps and the question, enabling efficient computation without look-ahead simulations.

3. Entropy-Based Adaptive Allocation: High-uncertainty steps receive more candidate exploration; certain steps require fewer attempts. This optimizes overall search budget automatically.

Architecture Overview

• Question Input: Problem statement for which reasoning is needed
• LLM Generation: Generate candidate reasoning steps (chains or trees)
• PMI Computation: Score each step using conditional probability estimates
• Adaptive Sampling: Concentrate search effort on high-entropy decision points
• Voting: Combine top-scored paths using weighted frequency and PMI consensus
• Final Output: Aggregated answer from best-performing reasoning paths

Implementation Steps

Compute the base and conditional probabilities needed for PMI calculation. These can be extracted from LLM logits during generation.

def compute_pmi_score(step_text, question, prior_steps, lm):
    """
    Compute PMI(question; step) measuring question-specific relevance.

    Args:
        step_text: The reasoning step to score
        question: The problem statement
        prior_steps: List of previous steps in the chain
        lm: Language model with log-probability access

    Returns:
        pmi_score: log p(step|question, prior_steps) - log p(step|prior_steps)
    """
    
    context_with_q =  + \
                     .join(prior_steps) + 

    log_prob_given_question = lm.get_log_prob(
        context_with_q, step_text
    )

    
    context_without_q = .join(prior_steps) + 
    log_prob_given_context = lm.get_log_prob(
        context_without_q, step_text
    )

    
    pmi_score = log_prob_given_question - log_prob_given_context

     pmi_score