Adaptive Uncertainty-Aware Tree Search for Robust Reasoning

TL;DR

This paper introduces Uncertainty-Aware Tree Search (UATS), a novel method that improves reasoning in large language models by estimating uncertainty and dynamically allocating computational resources, especially for out-of-distribution samples.

Contribution

The paper presents a new uncertainty-aware search strategy that combines Monte Carlo Dropout and reinforcement learning to enhance reasoning robustness in LLMs.

Findings

UATS reduces errors on out-of-distribution reasoning tasks.

Theoretical analysis shows sublinear regret with UATS.

Empirical results demonstrate improved performance over baseline methods.

Abstract

Inference-time reasoning scaling has significantly advanced the capabilities of Large Language Models (LLMs) in complex problem-solving. A prevalent approach involves external search guided by Process Reward Models (PRMs). However, a fundamental limitation of this framework is the epistemic uncertainty of PRMs when evaluating reasoning paths that deviate from their training distribution. In this work, we conduct a systematic analysis of this challenge. We first provide empirical evidence that PRMs exhibit high uncertainty and unreliable scoring on out-of-distribution (OOD) samples. We then establish a theoretical framework proving that while standard search incurs linear regret accumulation, an uncertainty-aware strategy can achieve sublinear regret. Motivated by these findings, we propose Uncertainty-Aware Tree Search (UATS), a unified method that estimates uncertainty via Monte Carlo Dropout and dynamically allocates compute budget using a reinforcement learning-based controller. Extensive experiments demonstrate that our approach effectively mitigates the impact of OOD errors.