MaskSearch: A Universal Pre-Training Framework to Enhance Agentic Search Capability

TL;DR

MaskSearch introduces a universal pre-training framework for enhancing agentic search in LLMs by leveraging Retrieval-Augmented Mask Prediction (RAMP). It supports both Supervised Fine-tuning and Reinforcement Learning (via DAPO) and employs curriculum learning to scale difficulty through the number of masked spans. The two-stage approach (RAMP pre-training followed by downstream task fine-tuning) yields substantial gains on open-domain multi-hop QA, with RL achieving larger improvements than SFT and small models benefiting markedly from pre-training. The framework demonstrates strong transferability across model sizes and architectures, and the data-efficient self-evolved distillation strategy further scales training while maintaining data quality.

Abstract

Retrieval-Augmented Language Models (RALMs) represent a classic paradigm where models enhance generative capabilities using external knowledge retrieved via a specialized module. Recent advancements in Agent techniques enable Large Language Models (LLMs) to autonomously utilize tools for retrieval, planning, and reasoning. While existing training-based methods show promise, their agentic abilities are limited by inherent characteristics of the task-specific data used during training. To further enhance the universal search capability of agents, we propose a novel pre-training framework, MaskSearch. In the pre-training stage, we introduce the Retrieval Augmented Mask Prediction (RAMP) task, where the model learns to leverage search tools to fill masked spans on a large number of pre-training data, thus acquiring universal retrieval and reasoning capabilities for LLMs. After that, the model is trained on downstream tasks to achieve further improvement. We apply both Supervised Fine-tuning (SFT) and Reinforcement Learning (RL) for training. For SFT, we combine agent-based and distillation-based methods to generate training data, starting with a multi-agent system consisting of a planner, rewriter, observer, and followed by a self-evolving teacher model. While for RL, we employ DAPO as the training framework and adopt a hybrid reward system consisting of answer rewards and format rewards. Additionally, we introduce a curriculum learning approach that allows the model to learn progressively from easier to more challenging instances based on the number of masked spans. We evaluate the effectiveness of our framework in the scenario of open-domain multi-hop question answering. Through extensive experiments, we demonstrate that MaskSearch significantly enhances the performance of LLM-based search agents on both in-domain and out-of-domain downstream tasks.

Figures (5)

• Figure 1: Overview of MaskSearch, a pre-training framework to incentivize the agentic RAG capabilities of LLMs. Based on the Retrieval-Augmented Mask Prediction (RAMP) task, models can be trained via SFT or RL to acquire generalizable abilities before downstream task training.
• Figure 2: Scaling Performance of SFT with respect to training steps on RAMP. The results obtained at training step 0 align precisely with those of the Direct SFT baseline, which is directly fine-tuned on the 58K CoT trajectories derived from HotpotQA.
• Figure 3: Performance on the dev set while finetuning with varying numbers of masks.
• Figure 4: Exploring RAMP with PPL-based masking strategy. The subfigure (a) illustrates the computation of perplexity for the unmasked spans in each round. The subfigure (b) indicates the effect of the PPL-based masking strategy on downstream test sets along with CL.
• Figure 5: Exploring Different Answer Reward Designs. These figures present the performance of the Qwen-7b model on the development set during RL training. The subfigure (a) illustrates the length of responses generated by the model. The subfigure (b) displays the metric based on token-level recall. The subfigure (c) details the metric assessed using the Qwen-72b model as a judge.