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

TL;DR

This work analyzes RLVR’s tendency to extinguish rare yet valid reasoning paths due to a high-pass selection effect, termed the Normalization Squeeze. It introduces Amortized Reasoning Tree Search (ARTS), a decoupled Proposer-Verifier framework that trains a flow-based verifier via Flow Matching to guide inference-time search, preserving long-tail reasoning without altering the base model. Theoretical and empirical results on MATH-500 show ARTS matching the performance of fully finetuned baselines (BoN@16 of 74.6–74.7%) and uniquely recovering hard, sparse traces where RLVR collapses, demonstrating robustness in high-entropy, combinatorial reasoning tasks. The findings highlight a practical trade-off between internalizing reasoning into weights and leveraging robust inference-time search, suggesting a scalable path for solving complex, long-tail reasoning tasks while maintaining model diversity.

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.

Figures (3)

• Figure 1: (a) Visualizing the "Normalization Squeeze". Base Model (Gray): Covers both common and rare valid reasoning paths. RL Policy (Red): Acts as a High-Pass Filter, amplifying the dominant mode while suppressing rare traces to near zero (Extinction Zone). (b) Visualization of Capability Extinction. We track the relative log-likelihood evolution ($\Delta \log_{10} P(\tau)$) of valid reasoning traces over 300 GRPO steps. While dominant Mode Traces (Blue) are consistently reinforced, Rare (Red) and Endangered (Yellow) traces suffer systematic decay despite being fully correct. This confirms that standard RLVR acts as a High-Pass Likelihood Filter, preserving only the high-redundancy patterns while purging the model's capability to solve hard, long-tail problems. ARTS (Green): Decouples generation from verification, using inference-time search to recover these extinguished paths and unlock latent long-tail capacity.
• Figure 2: Visualizing Sparse Deep Fork (SDF). Instead of exhaustive tree search, SDF constructs a cost-effective DAG by stochastically branching (Orange) from the base model's main trajectory (Blue). This structure efficiently exposes the verifier to local decision boundaries and alternative outcomes—such as discovering a correct path ($R=10$) diverging from an incorrect one—without the exponential cost of full expansion.
• Figure 3: Performance on the "Extinction Set". We compare recovery rates on the hardest subset where the RL policy (GRPO) fails completely (0%). While standard verifiers (PRM/RM) suppress rare traces below the random sampling baseline (Dashed Line), ARTS uniquely achieves positive recovery. This confirms that Flow Matching preserves valid reasoning signals in the long tail where standard discriminative objectives collapse.

Theorems & Definitions (2)

• Proposition 1: Probability Decay under Sparse Updates
• proof