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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Dou, Alex ZH, Wan, Zhongwei, Cui, Dongfei, Wang, Xin, Xiong, Jing, Lin, Haokun, Tao, Chaofan, Yan, Shen, Zhang, Mi
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913932106530816
author Dou, Alex ZH
Wan, Zhongwei
Cui, Dongfei
Wang, Xin
Xiong, Jing
Lin, Haokun
Tao, Chaofan
Yan, Shen
Zhang, Mi
author_facet Dou, Alex ZH
Wan, Zhongwei
Cui, Dongfei
Wang, Xin
Xiong, Jing
Lin, Haokun
Tao, Chaofan
Yan, Shen
Zhang, Mi
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2507_05557
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Enhancing Test-Time Scaling of Large Language Models with Hierarchical Retrieval-Augmented MCTS
Dou, Alex ZH
Wan, Zhongwei
Cui, Dongfei
Wang, Xin
Xiong, Jing
Lin, Haokun
Tao, Chaofan
Yan, Shen
Zhang, Mi
Computation and Language
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.
title Enhancing Test-Time Scaling of Large Language Models with Hierarchical Retrieval-Augmented MCTS
topic Computation and Language
url https://arxiv.org/abs/2507.05557