AWS用多轮强化学习在SageMaker AI上微调搜索智能体
AWS博客称,用SageMaker AI的多轮强化学习(MTRL)微调Qwen3.6-27B搜索智能体,仅调整三个超参数,在BrowseComp-Plus上nDCG@10 提升 23.7%,失败率从 22.89% 降到 0.68%。
AI解读:AWS机器学习博客介绍了一次实验:用Amazon SageMaker AI的多轮强化学习(MTRL)微调搜索智能体,让它自己决定搜什么、用哪种检索策略、什么时候停下来,而不是靠人工把每一步示范出来。
之所以要这样做,是因为小模型提示词写死也很难稳定多轮检索,前沿模型效果不错但延迟和成本高;只优化单轮回答的传统强化学习又看不到多轮之间的依赖。MTRL直接拿最终检索结果的好坏当奖励,训练整个过程。
对想要自建企业搜索的人来说,这次配置门槛很低:模型换了、奖励函数按nDCG@10 写、超参数只改max_epochs、global_batch_size、rollout_max_concurrency三个,其余走默认,训练按token计费、不用自己管GPU集群。
实测收益集中在需要多轮深挖的基准上:BrowseComp-Plus的nDCG@10 从 0.5136 升到 0.6354,失败率从 22.89% 降到 0.68%;但FreshStack略降,说明这不是对所有数据集都稳赚的通用提升。
AWS机器学习博客发布了一篇实践文章,介绍如何用Amazon SageMaker AI的多轮强化学习(MTRL)微调一个LLM搜索智能体,并给出在多个检索基准上的实测数字。文章作者为Huibin Shen。
文章描述的搜索智能体不要求用户写出完美查询,而是自主决定搜什么、用哪种检索策略、何时停止,并根据已检索到的内容在多轮交互中调整。文中称,让这种多步行为表现良好并不容易:基础模型不了解你的工具和环境;小模型提示后很少能给出可靠的多轮行为;前沿模型通常可行,但要付出延迟和成本代价;微调提供了第三条路,用小模型的速度和成本获得原本需要前沿模型才有的可靠性。
训练设置与奖励函数
文章把这次实验定位在企业搜索场景,智能体有两个工具:词法搜索(BM25)按词频找精确关键词匹配,适合具体术语或标识符查询;向量搜索把查询和文档转成嵌入向量再算相似度,推荐用于语义或概念类查询。训练还限制了轮数(一轮指一次用户-助手交互),以防智能体生成过长的响应,并鼓励高效搜索行为。
实验使用Amazon SageMaker AI MTRL微调Qwen3.6-27B模型,该模型在美国西部(俄勒冈)区域(us-west-2)受支持。环境要求包括:可访问us-west-2区域Amazon SageMaker AI的AWS账户;按所需格式上传到Amazon S3的训练和验证数据集;一个已部署的智能体端点,暴露BM25和向量搜索工具供MTRL环境在rollout期间调用;以及熟悉Amazon SageMaker AI和Python。
训练数据集包括FRAMES、BRIGHT、Enterprise RAG、ESCI、Musique、MLQA;测试数据集包括FreshStack、WixQA、BrowseComp-Plus、Wands。每个数据集内预留 5% 的训练实例作为验证实例。
奖励函数直接采用nDCG@10(归一化折损累计增益@10),衡量前 10 篇检索文档与理想排序的匹配程度,1.0 表示完美排序,0.0 表示没有检索到相关文档。这是轨迹级奖励:智能体完成整个多轮搜索后,奖励反映最终检索文档与标准答案的匹配程度。当智能体达到最大轮数或单轮最大采样token数时,奖励被设为 -1,以明确教模型避免这类失败模式。
MTRL的任务配置与能力
文章称,SageMaker AI MTRL把智能体任务框定为一系列决策,使用多轮rollout生成训练数据,并用策略梯度算法优化模型。提供的功能包括:模块化智能体-环境接口(可自定义奖励、工具循环和多轮对话形态,集成低代码);无服务器执行(按token计价,无需配置或管理GPU集群);异步rollout与轨迹收集(生成和梯度更新并行,带有界离策略陈旧度);原生算法库(PPO、CISPO、IS损失,搭配GRPO、GRPO pass@k、RLOO等基于组的优势估计器);可恢复训练(长训练可跨多个作业拆分);轨迹与奖励可观测性(可在Amazon SageMaker AI管理的MLflow中逐轮检查智能体行为);评估作业(部署到SageMaker AI端点或Amazon Bedrock前报告奖励、pass@k和轨迹指标)。
- 超参数只改三个:max_epochs为 1,global_batch_size为 128,rollout_max_concurrency为 32。
- 算法、优势估计器、离策略陈旧度边界等通常需要强化学习专业知识的选项均使用默认值。
- 文中称,从基础模型得到下一节的实验结果,只需要上述配置。
测试结果与限制
文章称,MTRL微调在四个留出基准中的三个上提升了搜索智能体,并让它在所有基准上可靠性显著提高。最大提升出现在BrowseComp-Plus(nDCG@10 提升 23.7%)和WixQA(提升 18.4%),Wands上有较小提升,FreshStack上略有回退。可靠性方面的结果更突出:在BrowseComp-Plus上,失败率从 22.89% 降至 0.68%,意味着智能体不仅搜得更好,还学会了在轮数和token预算内完成任务。
- WixQA:400 个问题,基础Qwen3.6-27B的nDCG@10 为 0.5725、失败率 0.67%、平均轮数 4.3;微调后为 0.6781、0.17%、4.5。
- Wands:147 个问题,基础模型 0.5762、失败率 0.00%、平均轮数 2.2;微调后 0.6112、0.00%、2.9。
- FreshStack:672 个问题,基础模型 0.4112、失败率 0.20%、平均轮数 3.1;微调后 0.4089、0.05%、2.8。
- BrowseComp-Plus:830 个问题,基础模型 0.5136、失败率 22.89%、平均轮数 7.0;微调后 0.6354、0.68%、6.3。
- 训练曲线方面,训练集和验证集的nDCG@10 奖励随训练步数稳步上升并趋于饱和,文章据此判断继续训练不会有更多收益。
- MTRL训练可能持续多天,MTRL服务默认时间限制为 24 小时,可通过CreateJob JSON schema调整;作业因超时或基础设施错误停止或失败后,可从先前检查点继续训练。
- 文章提醒,探索该方案后要清理资源以避免持续计费:停止或删除仍在运行的MTRL训练作业,删除不再需要的S3模型产物,删除为评估创建的任何已部署端点。