MaskSearch: Универсальная предобучающая структура для повышения поисковых возможностей агентов

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

May 26, 2025

Авторы: Weiqi Wu, Xin Guan, Shen Huang, Yong Jiang, Pengjun Xie, Fei Huang, Jiuxin Cao, Hai Zhao, Jingren Zhou

cs.AI

Аннотация

Модели языка с усилением поиска (Retrieval-Augmented Language Models, RALMs) представляют собой классическую парадигму, в которой модели улучшают свои генеративные способности за счет использования внешних знаний, извлекаемых с помощью специализированного модуля. Последние достижения в области агентных технологий позволяют крупным языковым моделям (Large Language Models, LLMs) автономно использовать инструменты для поиска, планирования и рассуждений. Хотя существующие методы, основанные на обучении, демонстрируют перспективные результаты, их агентные возможности ограничены внутренними характеристиками данных, специфичных для задач, используемых в процессе обучения. Для дальнейшего улучшения универсальной поисковой способности агентов мы предлагаем новую структуру предварительного обучения, MaskSearch. На этапе предварительного обучения мы вводим задачу предсказания маски с усилением поиска (Retrieval Augmented Mask Prediction, RAMP), в которой модель учится использовать поисковые инструменты для заполнения замаскированных фрагментов на большом объеме данных предварительного обучения, тем самым приобретая универсальные способности к поиску и рассуждению для LLMs. После этого модель обучается на последующих задачах для достижения дополнительного улучшения. Мы применяем как контролируемую тонкую настройку (Supervised Fine-tuning, SFT), так и обучение с подкреплением (Reinforcement Learning, RL). Для SFT мы комбинируем методы, основанные на агентах и дистилляции, для генерации обучающих данных, начиная с мультиагентной системы, состоящей из планировщика, переписчика, наблюдателя, и заканчивая саморазвивающейся учительской моделью. Для RL мы используем DAPO в качестве структуры обучения и применяем гибридную систему вознаграждений, состоящую из вознаграждений за ответы и вознаграждений за формат. Дополнительно мы вводим подход поэтапного обучения, который позволяет модели постепенно обучаться от более простых к более сложным примерам в зависимости от количества замаскированных фрагментов. Мы оцениваем эффективность нашей структуры в сценарии открытого домена с многошаговым ответом на вопросы. В ходе обширных экспериментов мы демонстрируем, что MaskSearch значительно улучшает производительность поисковых агентов на основе LLM как на внутридоменных, так и на внедоменных последующих задачах.

English

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.