Amortized Reasoning Tree Search: Decoupling Proposal and Decision in Large Language Models

TL;DR

This paper introduces ARTS, a novel approach that decouples reasoning and verification in large language models, effectively addressing the suppression of rare correct reasoning paths and improving performance on complex tasks.

Contribution

The paper proposes ARTS, which separates generation from verification using a Flow Matching objective, enhancing reasoning robustness without altering the base model.

Findings

Achieves 74.6% performance on MATH-500 benchmark.

Recovers performance on long-tail reasoning tasks where RL collapses.

Matches fully fine-tuned policies without modifying the generative backbone.

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) has established itself as the dominant paradigm for instilling rigorous reasoning capabilities in Large Language Models. While effective at amplifying dominant behaviors, we identify a critical pathology in this alignment process: the systematic suppression of valid but rare (low-likelihood under the base model distribution) reasoning paths. We theoretically characterize this phenomenon as a "Normalization Squeeze," where the interplay between mode-seeking policy gradients and finite sampling acts as a high-pass likelihood filter, driving the probability of rare correct traces to statistical extinction. To counteract this collapse without discarding the base model's latent diversity, we propose Amortized Reasoning Tree Search (ARTS). Unlike standard approaches that force internalization via parameter updates, ARTS prioritizes deliberation by decoupling generation from verification. We introduce a Flow Matching objective that repurposes the verifier to estimate the conservation of probability flow, enabling robust navigation through sparse, high-entropy search spaces where traditional discriminative objectives fail. Extensive experiments on the MATH-500 benchmark demonstrate that ARTS achieves a performance of 74.6% (BoN@16), effectively matching fully fine-tuned policies (74.7%) without modifying the generative backbone. Crucially, on the long-tail subset where coupled RL optimization collapses to 0% pass@k, ARTS uniquely recovers significant performance, suggesting that disentangling verification from generation offers a more robust pathway for solving complex reasoning tasks.