Enhancing Test-Time Scaling of Large Language Models with Hierarchical Retrieval-Augmented MCTS

2025/07/08 by Dou, Alex ZH, Zhongwei Wan, Dongfei Cui +14 · 1 citation
Computer Science · #Computation and Language (cs.CL) #FOS: Computer and information sciences #Multimodal Machine Learning Applications #Natural Language Processing Techniques #Topic Modeling

paper · pdf · doi:10.48550/arxiv.2507.05557

openalex publication_date 2025/07/08 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Test-time scaling has emerged as a promising paradigm in language modeling, leveraging additional computational resources at inference time to enhance model performance. In this work, we introduce R2-LLMs, a novel and versatile hierarchical retrieval-augmented reasoning framework designed to improve test-time scaling in large language models (LLMs) without requiring distillation from more advanced models to obtain chain-of-thought (CoT) training data. R2-LLMs enhances inference-time generalization by integrating dual-level retrieval-based in-context learning: (1) At the coarse level, our approach extracts abstract templates from complex reasoning problems and retrieves similar problem-answer pairs to facilitate high-level in-context learning; (2) At the fine level, during Monte Carlo Tree Search (MCTS), R2-LLMs efficiently retrieves analogous intermediate solution steps from reference mathematical problem datasets, refining step-wise reasoning with the aid of a process reward model (PRM) for scoring. R2-LLMs is a robust hierarchical reasoning-augmentation method that enhances in-context-level reasoning while seamlessly integrating with step-level tree search methods. Utilizing PRM, it refines both candidate generation and decision-making for improved reasoning accuracy. Empirical evaluations on the MATH500, GSM8K, and OlympiadBench-TO datasets achieve substantial relative improvement with an increase of up to 16% using LLaMA-3.1-8B compared to the baselines, showcasing the effectiveness of our approach in complex reasoning tasks.

Citations

• SRPO: Enhancing Multimodal LLM Reasoning via Reflection-Aware Reinforcement Learning
• SwingArena: Competitive Programming Arena for Long-context GitHub Issue Solving
• Enhancing Text-to-Image Diffusion Transformer via Split-Text Conditioning
• Beyond Distillation: Pushing the Limits of Medical LLM Reasoning with Minimalist Rule-Based RL
• PhyX: Does Your Model Have the "Wits" for Physical Reasoning?
• SVD-LLM V2: Optimizing Singular Value Truncation for Large Language Model Compression
• Knowledge-enhanced Multimodal ECG Representation Learning with Arbitrary-Lead Inputs
• MEDA: Dynamic KV Cache Allocation for Efficient Multimodal Long-Context Inference
• Recent Advances in Large Langauge Model Benchmarks against Data Contamination: From Static to Dynamic Evaluation
• ParallelComp: Parallel Long-Context Compressor for Length Extrapolation
• ReasonFlux: Hierarchical LLM Reasoning via Scaling Thought Templates
• rStar-Math: Small LLMs Can Master Math Reasoning with Self-Evolved Deep Thinking
• BoostStep: Boosting mathematical capability of Large Language Models via improved single-step reasoning
• Beyond Examples: High-level Automated Reasoning Paradigm in In-Context Learning via MCTS
• ClinicalBench: Can LLMs Beat Traditional ML Models in Clinical Prediction?
• Autoregressive Models in Vision: A Survey
• NeuroClips: Towards High-fidelity and Smooth fMRI-to-Video Reconstruction
• Can Medical Vision-Language Pre-training Succeed with Purely Synthetic Data?
• Rewarding Progress: Scaling Automated Process Verifiers for LLM Reasoning
• UNComp: Can Matrix Entropy Uncover Sparsity? -- A Compressor Design from an Uncertainty-Aware Perspective
• LLaMA-Berry: Pairwise Optimization for O1-like Olympiad-Level Mathematical Reasoning
• Famba-V: Fast Vision Mamba with Cross-Layer Token Fusion
• Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters
• The Llama 3 Herd of Models
• Large Language Monkeys: Scaling Inference Compute with Repeated Sampling
• Evolver: Chain-of-Evolution Prompting to Boost Large Multimodal Models for Hateful Meme Detection
• Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies
• LiteSearch: Efficacious Tree Search for LLM
• LOOK-M: Look-Once Optimization in KV Cache for Efficient Multimodal Long-Context Inference
• Step-level Value Preference Optimization for Mathematical Reasoning
• ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search
• Improve Mathematical Reasoning in Language Models by Automated Process Supervision
• Structure-based Drug Design Benchmark: Do 3D Methods Really Dominate?
• SVD-LLM: Truncation-aware Singular Value Decomposition for Large Language Model Compression
• Zero-Shot ECG Classification with Multimodal Learning and Test-time Clinical Knowledge Enhancement
• MathScale: Scaling Instruction Tuning for Mathematical Reasoning
• OlympiadBench: A Challenging Benchmark for Promoting AGI with Olympiad-Level Bilingual Multimodal Scientific Problems
• BioXP-0.5B: Explainable Medical-AI via RL-GRPO
• Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations
• CritiqueLLM: Towards an Informative Critique Generation Model for Evaluation of Large Language Model Generation
• The Mystery of In-Context Learning: A Comprehensive Survey on Interpretation and Analysis
• KCTS: Knowledge-Constrained Tree Search Decoding with Token-Level Hallucination Detection
• Language Agent Tree Search Unifies Reasoning Acting and Planning in Language Models
• Alphazero-like Tree-Search can Guide Large Language Model Decoding and Training
• Don't throw away your value model! Generating more preferable text with Value-Guided Monte-Carlo Tree Search decoding
• MAmmoTH: Building Math Generalist Models through Hybrid Instruction Tuning
• ETP: Learning Transferable ECG Representations via ECG-Text Pre-training
• Let's Verify Step by Step
• Med-UniC: Unifying Cross-Lingual Medical Vision-Language Pre-Training by Diminishing Bias
• Reasoning with Language Model is Planning with World Model
• CRITIC: Large Language Models Can Self-Correct with Tool-Interactive Critiquing
• Self-Refine: Iterative Refinement with Self-Feedback
• Planning with Large Language Models for Code Generation
• A Survey on In-context Learning
• G-MAP: General Memory-Augmented Pre-trained Language Model for Domain Tasks
• Solving Quantitative Reasoning Problems with Language Models
• Making Large Language Models Better Reasoners with Step-Aware Verifier
• Large Language Models are Zero-Shot Reasoners
• Self-Consistency Improves Chain of Thought Reasoning in Language Models
• Training Verifiers to Solve Math Word Problems
• Measuring Mathematical Problem Solving With the MATH Dataset
• Large Scale Many-Objective Optimization Driven by Distributional Adversarial Networks
• Many-Objective Estimation of Distribution Optimization Algorithm Based on WGAN-GP
• Scaling Laws for Neural Language Models
• Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks
• MathQA: Towards Interpretable Math Word Problem Solving with Operation-Based Formalisms
• The Internet of Things in the Era of Generative AI: Vision and Challenges

Cited by

• Less Diverse, Less Safe: The Indirect But Pervasive Risk of Test-Time Scaling in Large Language Models
• SRPO: Enhancing Multimodal LLM Reasoning via Reflection-Aware Reinforcement Learning

Related

• Reasoning Beyond Limits: Advances and Open Problems for LLMs
• Guided by Gut: Efficient Test-Time Scaling with Reinforced Intrinsic Confidence
• Rethinking the Unsolvable: When In-Context Search Meets Test-Time Scaling
• The Art of Scaling Test-Time Compute for Large Language Models
• Can 1B LLM Surpass 405B LLM? Rethinking Compute-Optimal Test-Time Scaling