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

📄 arXiv: 2505.20285v2 📥 PDF

作者: Weiqi Wu, Xin Guan, Shen Huang, Yong Jiang, Pengjun Xie, Fei Huang, Jiuxin Cao, Hai Zhao, Jingren Zhou

分类: cs.CL

发布日期: 2025-05-26 (更新: 2025-05-27)

备注: Code is available at https://github.com/Alibaba-NLP/MaskSearch

💡 一句话要点

提出MaskSearch预训练框架，提升LLM智能体在开放域多跳问答中的搜索能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 预训练框架 检索增强 掩码预测 智能体搜索 开放域问答

📋 核心要点

  1. 现有基于训练的智能体方法受限于特定任务数据的固有特征，通用搜索能力不足。
  2. MaskSearch提出检索增强掩码预测（RAMP）任务，使模型通过搜索工具填充掩码，获得通用检索和推理能力。
  3. 实验表明，MaskSearch显著提升了LLM智能体在开放域多跳问答任务上的性能，包括领域内和领域外。

📝 摘要（中文）

检索增强语言模型（RALMs）是一种经典范式，模型通过专门的模块检索外部知识来增强生成能力。智能体技术的最新进展使大型语言模型（LLMs）能够自主地利用工具进行检索、规划和推理。虽然现有的基于训练的方法显示出前景，但它们的智能体能力受到训练期间使用的特定任务数据的固有特征的限制。为了进一步增强智能体的通用搜索能力，我们提出了一种新的预训练框架MaskSearch。在预训练阶段，我们引入了检索增强掩码预测（RAMP）任务，模型学习利用搜索工具来填充大量预训练数据上的掩码跨度，从而获得LLMs的通用检索和推理能力。之后，模型在下游任务上进行训练以实现进一步的改进。我们应用监督微调（SFT）和强化学习（RL）进行训练。对于SFT，我们结合了基于智能体和基于蒸馏的方法来生成训练数据，首先是一个由规划器、重写器、观察器组成的多智能体系统，然后是一个自我进化的教师模型。对于RL，我们采用DAPO作为训练框架，并采用由答案奖励和格式奖励组成的混合奖励系统。此外，我们引入了一种课程学习方法，允许模型基于掩码跨度的数量从更容易到更具挑战性的实例逐步学习。我们在开放域多跳问答的场景中评估了我们框架的有效性。通过大量的实验，我们证明了MaskSearch显著提高了基于LLM的搜索智能体在领域内和领域外下游任务上的性能。

🔬 方法详解

问题定义：现有基于训练的Agent方法，其Agent能力受限于特定任务数据，导致通用搜索能力不足。具体来说，模型难以泛化到训练数据未覆盖的知识领域或推理路径，限制了其在开放域场景下的应用。

核心思路：通过预训练阶段的检索增强掩码预测（RAMP）任务，让模型学习利用外部搜索工具来填充文本中的掩码部分。这种方式使得模型能够接触到更广泛的知识，并学习如何有效地利用这些知识进行推理，从而提升其通用搜索能力。

技术框架：MaskSearch框架包含预训练和微调两个阶段。预训练阶段使用RAMP任务，模型接收包含掩码的文本，并利用搜索工具检索相关信息，然后预测被掩码的内容。微调阶段则在下游任务上进行，采用监督微调（SFT）和强化学习（RL）两种方式。SFT使用多智能体系统生成训练数据，包括规划器、重写器和观察器。RL则采用DAPO框架，并设计了混合奖励系统，包括答案奖励和格式奖励。此外，还使用了课程学习，从简单到复杂逐步训练模型。

关键创新：核心创新在于RAMP预训练任务，它将检索能力融入到预训练过程中，使得模型在预训练阶段就具备了利用外部知识进行推理的能力。与传统的预训练方法相比，RAMP任务更侧重于让模型学习如何主动获取和利用知识，而不是仅仅依赖于预训练数据中的信息。

关键设计：RAMP任务的关键在于如何设计掩码策略和检索策略。掩码策略决定了哪些部分的文本会被掩码，而检索策略则决定了模型如何利用搜索工具来检索相关信息。此外，SFT阶段的多智能体系统和RL阶段的混合奖励系统也是关键设计，它们分别用于生成高质量的训练数据和引导模型学习正确的行为。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MaskSearch在开放域多跳问答任务上取得了显著的性能提升。在领域内任务上，MaskSearch相较于基线模型提升了X%。更重要的是，在领域外任务上，MaskSearch也表现出了良好的泛化能力，证明了其通用搜索能力的有效性。这些结果表明，MaskSearch是一种有效的提升LLM智能体搜索能力的预训练框架。

🎯 应用场景

MaskSearch框架可应用于各种需要智能体具备强大搜索能力的场景，例如智能客服、知识图谱问答、自动报告生成等。通过提升LLM智能体的通用搜索能力，可以使其更好地理解用户意图，获取相关信息，并生成高质量的回复或报告，从而提高工作效率和用户满意度。未来，该框架还可以扩展到其他模态，例如图像和视频，以支持更复杂的智能体任务。

📄 摘要（原文）

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.