MaskSearch: A Universal Pre-Training Framework to Enhance Agentic Search
  Capability
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.
683eef14fa20482e30563cf2
Weiqi Wu
683eef14fa20482e30563cf3
Xin Guan
683eef14fa20482e30563cf4
Shen Huang
683eef14fa20482e30563cf5
Yong Jiang
683eef14fa20482e30563cf6
Pengjun Xie
683eef14fa20482e30563cf7
Fei Huang
683eef14fa20482e30563cf8
Jiuxin Cao
683eef14fa20482e30563cf9
Hai Zhao
683eef14fa20482e30563cfa
Jingren Zhou
MaskSearch: A Universal Pre-Training Framework to Enhance Agentic Search
  Capability
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.
A novel pre-training framework, MaskSearch, enhances Large Language Models with universal retrieval and reasoning capabilities through a Retrieval Augmented Mask Prediction task, improving their performance in open-domain multi-hop question answering.
Retrieval-Augmented Language Models Agent techniques Large Language Models MaskSearch Retrieval Augmented Mask Prediction Supervised Fine-tuning Reinforcement Learning DAPO curriculum learning open-domain multi-hop question answering