Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

StepSearch: Igniting LLMs Search Ability via Step-Wise Proximal Policy Optimization

Ziliang Wang, Xuhui Zheng, Kang An, Cijun Ouyang, Jialu Cai, Yuhang Wang, Yichao Wu

TL;DR

StepSearch tackles the challenge of efficient multi-hop QA by introducing step-wise proximal policy optimization (StePPO) with token-level supervision and two reward channels to guide iterative retrieval. A MuSiQue-based data augmentation pipeline generates sub-question trajectories, enabling rich, stepwise learning signals. Across four multi-hop QA benchmarks and with 3B–7B Qwen models, StepSearch achieves state-of-the-art improvements using only 19k training examples, demonstrating faster convergence and better retrieval fidelity than prior RL-based search methods. The approach offers a robust, data-efficient path to enhancing retrieval-augmented LLMs, with potential to extend to larger models and multimodal tasks.

Abstract

Efficient multi-hop reasoning requires Large Language Models (LLMs) based agents to acquire high-value external knowledge iteratively. Previous work has explored reinforcement learning (RL) to train LLMs to perform search-based document retrieval, achieving notable improvements in QA performance, but underperform on complex, multi-hop QA resulting from the sparse rewards from global signal only. To address this gap in existing research, we introduce StepSearch, a framework for search LLMs that trained with step-wise proximal policy optimization method. It consists of richer and more detailed intermediate search rewards and token-level process supervision based on information gain and redundancy penalties to better guide each search step. We constructed a fine-grained question-answering dataset containing sub-question-level search trajectories based on open source datasets through a set of data pipeline method. On standard multi-hop QA benchmarks, it significantly outperforms global-reward baselines, achieving 11.2% and 4.2% absolute improvements for 3B and 7B models over various search with RL baselines using only 19k training data, demonstrating the effectiveness of fine-grained, stepwise supervision in optimizing deep search LLMs. Our code will be released on https://github.com/Zillwang/StepSearch.

StepSearch: Igniting LLMs Search Ability via Step-Wise Proximal Policy Optimization

TL;DR

StepSearch tackles the challenge of efficient multi-hop QA by introducing step-wise proximal policy optimization (StePPO) with token-level supervision and two reward channels to guide iterative retrieval. A MuSiQue-based data augmentation pipeline generates sub-question trajectories, enabling rich, stepwise learning signals. Across four multi-hop QA benchmarks and with 3B–7B Qwen models, StepSearch achieves state-of-the-art improvements using only 19k training examples, demonstrating faster convergence and better retrieval fidelity than prior RL-based search methods. The approach offers a robust, data-efficient path to enhancing retrieval-augmented LLMs, with potential to extend to larger models and multimodal tasks.

Abstract

Efficient multi-hop reasoning requires Large Language Models (LLMs) based agents to acquire high-value external knowledge iteratively. Previous work has explored reinforcement learning (RL) to train LLMs to perform search-based document retrieval, achieving notable improvements in QA performance, but underperform on complex, multi-hop QA resulting from the sparse rewards from global signal only. To address this gap in existing research, we introduce StepSearch, a framework for search LLMs that trained with step-wise proximal policy optimization method. It consists of richer and more detailed intermediate search rewards and token-level process supervision based on information gain and redundancy penalties to better guide each search step. We constructed a fine-grained question-answering dataset containing sub-question-level search trajectories based on open source datasets through a set of data pipeline method. On standard multi-hop QA benchmarks, it significantly outperforms global-reward baselines, achieving 11.2% and 4.2% absolute improvements for 3B and 7B models over various search with RL baselines using only 19k training data, demonstrating the effectiveness of fine-grained, stepwise supervision in optimizing deep search LLMs. Our code will be released on https://github.com/Zillwang/StepSearch.

Paper Structure

This paper contains 21 sections, 13 equations, 5 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Data Augmentation Pipeline
  4. Train LLM with Search Actions
  5. StepSearch
  6. Type 1 Reward: Global Signal
  7. Type 2 Reward: Search Step
  8. Experiment
  9. Dataset and Evaluation Metrics
  10. Baselines
  11. Training Details
  12. Main Results
  13. Further Analysis
  14. Different RL Comparison
  15. Ablation Study
  16. ...and 6 more sections

Figures (5)

  • Figure 2: Step-wise search involves interactive rounds, with information gain being rewarded and redundancy penalised. Each interaction evaluates thinking and searching behaviour based on the retrieved results, with the final answer being used as the basis for global rewards.
  • Figure 3: Data pipeline for generating the corresponding search query for the Q&A intermediate process.
  • Figure 4: Overview of StepSearch. At each step, the model issues queries to an external engine and receives snippets. Search Step Reward score, combining information gain and redundancy penalty, are applied to tokens within each round, while the global reward, based on final answer accuracy and keyword hit rate, is applied at the last token. Retrieved content is masked during training to isolate the model’s generative parameters.
  • Figure 5: Comparison of StePPO vs. PPO.
  • Figure 6: Training dynamics of correctness, response length, and search-key scores in ablation experiments. 'w/o' represent ‘with out’ while ‘ow’ for 'only with', 'rp' stands for 'redundancy penalty','ig' for 'information gain', 'ss' is 'step score' ('ig' + 'rp') and 'skr' means the global reward 'search key reward'