MaskSearch: Ein universelles Pre-Training-Framework zur Verbesserung der agentenbasierten Suchfähigkeit

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

May 26, 2025

papers.authors: Weiqi Wu, Xin Guan, Shen Huang, Yong Jiang, Pengjun Xie, Fei Huang, Jiuxin Cao, Hai Zhao, Jingren Zhou

cs.AI

papers.abstract

Retrieval-Augmented Language Models (RALMs) repräsentieren ein klassisches Paradigma, bei dem Modelle ihre generativen Fähigkeiten durch die Nutzung externen Wissens verbessern, das über ein spezialisiertes Modul abgerufen wird. Jüngste Fortschritte in der Agententechnik ermöglichen es Large Language Models (LLMs), Werkzeuge für das Retrieval, die Planung und das logische Denken autonom zu nutzen. Obwohl bestehende, auf Training basierende Methoden vielversprechend sind, werden ihre agentischen Fähigkeiten durch die inhärenten Eigenschaften der aufgabenspezifischen Daten, die während des Trainings verwendet werden, eingeschränkt. Um die universelle Suchfähigkeit von Agenten weiter zu verbessern, schlagen wir ein neuartiges Pre-Training-Framework namens MaskSearch vor. In der Pre-Training-Phase führen wir die Retrieval Augmented Mask Prediction (RAMP)-Aufgabe ein, bei der das Modell lernt, Suchwerkzeuge zu nutzen, um maskierte Abschnitte in einer großen Menge von Pre-Training-Daten zu füllen, wodurch es universelle Retrieval- und Denkfähigkeiten für LLMs erwirbt. Anschließend wird das Modell auf Downstream-Aufgaben trainiert, um weitere Verbesserungen zu erzielen. Wir wenden sowohl Supervised Fine-Tuning (SFT) als auch Reinforcement Learning (RL) für das Training an. Für SFT kombinieren wir agentenbasierte und distilationsbasierte Methoden, um Trainingsdaten zu generieren, beginnend mit einem Multi-Agenten-System, das aus einem Planer, einem Umschreiber, einem Beobachter besteht, gefolgt von einem sich selbst weiterentwickelnden Lehrermodell. Für RL verwenden wir DAPO als Trainingsframework und setzen ein hybrides Belohnungssystem ein, das aus Antwortbelohnungen und Formatbelohnungen besteht. Zusätzlich führen wir einen Curriculum-Learning-Ansatz ein, der es dem Modell ermöglicht, schrittweise von einfacheren zu anspruchsvolleren Instanzen basierend auf der Anzahl der maskierten Abschnitte zu lernen. Wir evaluieren die Effektivität unseres Frameworks im Kontext der Open-Domain-Multi-Hop-Fragebeantwortung. Durch umfangreiche Experimente zeigen wir, dass MaskSearch die Leistung von LLM-basierten Suchagenten sowohl bei in-domain als auch bei out-of-domain Downstream-Aufgaben signifikant verbessert.

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.