本文提出 MASKSEARCH 框架，通过 Retrieval-Augmented Mask Prediction (RAMP) 预训练任务结合监督微调和强化学习，显著提升了大型语言模型在开放域多跳问答任务中的代理搜索能力。
Large Language Model, Retrieval-Augmented Generation, Reasoning, Reinforcement Learning, Supervised Learning, Multi-Agent
Weiqi Wu, Xin Guan, Shen Huang, Yong Jiang, Pengjun Xie, Fei Huang, Jiuxin Cao, Hai Zhao, Jingren Zhou
Tongyi Lab, Alibaba Group
Generated by grok-3
Background Problem
大型语言模型（LLMs）虽然在多种任务中表现出色，但常因幻觉问题和对领域特定或实时任务的处理能力不足而受到限制。检索增强语言模型（RALMs）通过引入外部知识检索机制来增强生成能力，但其检索与生成的分离限制了模型在多步任务中的适应性。近年来，自主 AI 代理的兴起为 LLMs 提供了利用工具进行检索、规划和推理的能力，然而现有基于训练的方法受限于任务特定数据的特性，泛化能力不足。本文提出了一种通用预训练框架 MASKSEARCH，旨在通过 Retrieval-Augmented Mask Prediction (RAMP) 任务增强 LLMs 的代理搜索能力，解决多步检索和推理任务中的泛化问题。
Method
MASKSEARCH 是一个两阶段训练框架，核心在于预训练任务 Retrieval-Augmented Mask Prediction (RAMP)，其目标是通过多步搜索和推理预测输入文本中的掩码内容。
- 核心思想：在预训练阶段，模型学习利用搜索工具从外部知识库中获取信息以填补掩码内容，从而培养通用的检索和推理能力；在下游任务阶段，通过进一步训练提升性能。
- 具体步骤：
- RAMP 任务设计：基于维基百科数据，提取显著跨度（如命名实体、术语等）并随机掩码，模型需通过搜索工具预测掩码内容。
- 监督微调（SFT）：结合代理方法和蒸馏方法生成 Chain-of-Thought (CoT) 数据，初期通过多代理系统（Planner, Rewriter, Observer）合成推理轨迹，后期采用自进化蒸馏策略迭代更新数据集。
- 强化学习（RL）：采用 Dynamic Sampling Policy Optimization (DAPO) 算法，设计混合奖励系统（格式奖励和答案奖励），优化搜索和推理过程。
- 课程学习：根据掩码数量逐步增加任务难度，从简单到复杂训练模型。
- 关键问题与批评：RAMP 任务虽然创新，但其复杂性可能导致学习目标模糊，尤其在掩码数量多时模型可能难以收敛。此外，多代理系统的设计虽然有趣，但其生成数据的质量和多样性可能存在瓶颈，自进化蒸馏也可能因缺乏外部监督而陷入局部最优。
Experiment
实验在多个开放域多跳问答数据集（如 HotpotQA, FanoutQA, Musique 等）上评估 MASKSEARCH 的效果，使用 Qwen2.5 和 LLaMA 系列模型（1B 到 8B 参数规模）进行测试。
- 实验设置：预训练阶段构建了 10M 规模的 RAMP 数据集，SFT 阶段使用 58K CoT 轨迹进行下游任务训练，RL 阶段采用 DAPO 算法。基线包括 RAG-PE、Agent-PE、Distilled Search-R1 和 Search-R1，评估指标为 token-level Recall。
- 结果分析：MASKSEARCH 在所有数据集上均显著优于基线，尤其在小模型（如 Qwen2.5-1.5B）上提升明显，例如在 Bamboogle 数据集上提升了 11.78 分。RL 训练在 HotpotQA 上比 SFT 提升了 3-5 分，显示出更高的性能上限。课程学习在部分模型（如 Qwen2.5-7B）上优于混合训练策略。
- 实验设计评价：实验设置较为全面，覆盖了不同规模模型和多种数据集，证明了框架的可扩展性。然而，基线设置可能不够公平，例如 Search-R1 直接在下游数据上训练，可能低估了其潜力。此外，RL 奖励设计中模型判别奖励（Model-Based Reward）虽有效，但依赖大型模型（如 Qwen2.5-72B）作为判别器，增加了计算成本且可能引入偏差。实验对长期训练效果和过拟合风险的分析不足，结果是否可持续存疑。
Further Thoughts
MASKSEARCH 框架在增强代理搜索能力方面展现了潜力，但其依赖于复杂的数据生成和奖励设计可能限制了实际应用场景。我认为其多代理系统生成 CoT 数据的思路可以进一步扩展到其他领域，如教育领域的个性化学习路径生成，或医疗领域的多步诊断推理支持系统。此外，RL 奖励设计中模型判别奖励（Model-Based Reward）有效避免了奖励欺骗问题，这或许可以与最近的一些对齐技术（如 RLHF）结合，进一步提升模型在复杂任务中的一致性和可信度。然而，论文未充分探讨 RAMP 任务在非问答任务（如创意写作或代码生成）中的适用性，未来可以尝试将这一框架迁移到更广泛的生成任务中，测试其通用性。同时，课程学习的实现细节和难度递增策略需要更明确的定义，以确保其在不同数据集和模型架构上的稳定性。