MaskSearch: Training LLMs for Expert-Level Search Capabilities

MaskSearch is a pretraining framework designed to sharpen a model's fundamental search capabilities. Rather than relying on fine-tuning for narrow, specific tasks, it introduces a Retrieval Augmented Mask Prediction (RAMP) objective. Under this framework, the model learns to reconstruct masked segments of pretraining data by actively utilizing search tools. This process establishes a robust foundation for retrieval and reasoning, which translates into superior performance for large language model (LLM) agents.

The generation of training data combines agentic workflows with knowledge distillation. The process begins with a multi-agent system—consisting of a Planner, a Rewriter, and an Observer—to coordinate complex search paths. Subsequently, a self-evolving teacher model refines these outputs to create high-quality reasoning trajectories for the student model.

Empirical results highlight the effectiveness of this approach. MaskSearch consistently outperforms existing methods for LLM-based search agents across both in-domain and out-of-domain tasks.

The following table details the evaluation across several open-domain QA benchmarks:

*✗ indicates no specific pre-training or post-training in that category. SFT = Supervised Fine-Tuning. RL = Reinforcement Learning. RAG-PE and Agent-PE are baseline prompting methods.

Installation and Setup

Before running the framework, replace the placeholder keys in the following files with your Qwen (DashScope) and Google Search credentials:

• src/RAMP/model.py
• src/multi_agent/model.py
• src/multi_agent/web_news_get.py

Set your API keys as follows:

DASHSCOPE_API_KEY = "YOUR_API_KEY"
GOOGLE_API_KEY = "YOUR_API_KEY"

Install the necessary dependencies:

pip install -r requirements.txt

Generate RAMP QA Data

Generate the QA dataset using a Wikipedia dump. Point the --corpus argument to your local Wikipedia directory.

python gen_qa.py \
    --model "$model" \
    --corpus "Wikipedia Directory" \
    --output_path "output_path"

Build CoT Trajectories

Run the multi-agent pipeline to construct Chain-of-Thought (CoT) reasoning steps. This generates the SFT dataset. If necessary, configure your data paths within src/multi_agent/dataset.py.

python cot_construct.py \
    --model "$model" \
    --dataset "dataset" \
    --output_path "output_path"

Training

Once the data is prepared, begin the training process. For SFT, follow the standard LLaMA-Factory workflow. For RL training, refer to the configurations used in Search-R1 or ZeroSearch.