«Скорость маленькой модели при надёжности фронтирной»: AWS дообучила поискового агента многотуровым RL на SageMaker AI

Какую задачу решали

Поисковые агенты на основе больших языковых моделей меняют то, как компании добывают информацию: вместо идеально сформулированного запроса агент сам решает, что искать, какую стратегию извлечения применить и когда остановиться, уточняя подход по результатам предыдущих раундов.

Сложность в том, что базовая модель не знает ваших инструментов и вашей среды: маленькую модель трудно заставить надёжно вести себя в многошаговом режиме, а фронтирная работает, но за это приходится платить задержкой и стоимостью.

Дообучение — третий путь: модель учат вашим инструментам напрямую, чтобы получить скорость и цену маленькой модели при надёжности, для которой иначе нужна фронтирная.

Классические подходы не подходят. Supervised fine-tuning (SFT) опирается на экспертные демонстрации идеальных многошаговых траекторий, которые дорого собирать и которых обычно просто нет для вашего сценария, а однотуровое reinforcement learning (RL) с проверяемыми наградами (RLVR) оценивает по одному ответу за раз и теряет взаимозависимость решений между шагами.

Multi-turn reinforcement learning (MTRL) оптимизирует агента сразу по всей траектории: сигналу награды достаточно отражать, был ли хорошим конечный результат.

Что даёт SageMaker MTRL

Amazon SageMaker AI MTRL позволяет дообучать модели с подкреплением в многошаговых сценариях: агентная задача представляется как последовательность решений, обучающие данные генерируются многотуровыми роллаутами, а модель оптимизируется алгоритмами policy gradient.

Сервис предлагает модульный интерфейс «агент — среда», в котором настраиваются свои функции награды, циклы инструментов и формы многошаговой беседы; исполнение при этом бессерверное и оплачивается по токенам, без провижининга GPU-кластеров.

Генерация и обновление градиентов идут параллельно с ограниченной устареваемостью вне политики (bounded off-policy staleness), поэтому обучение остаётся быстрым и не уходит слишком далеко от текущей политики.

В библиотеку алгоритмов входят Proximal Policy Optimization (PPO), Clipped Importance Sampling Policy Optimization (CISPO) и importance-sampling (IS) потери в паре с групповыми оценщиками преимущества: GRPO, GRPO pass@k, RLOO и другими.

Есть поддержка возобновления обучения — длинные прогоны можно разбивать между заданиями, обходя лимит времени одного задания, — наблюдаемость траекторий и награды по шагам в MLflow под управлением SageMaker AI, а также задания оценки, где до развёртывания отчитываются reward, pass@k и метрики траекторий.

Для поискового агента это естественное совпадение: есть ясный сигнал награды (качество извлечения), многошаговый цикл взаимодействия, в котором агент шлёт запросы и получает результаты, и хорошо определённая среда — инструменты поиска.

Как поставили эксперимент

В разборе на Amazon SageMaker AI MTRL дообучали модель Qwen3.6-27B в регионе US West (Oregon).

Агент работал в сценарии корпоративного поиска с двумя инструментами: лексический поиск BM25 для точных совпадений по словам и векторный поиск для семантических или концептуальных запросов.

Число туров ограничивали, чтобы агент не генерировал чрезмерно длинные ответы и искал эффективнее.

Обучение шло на открытых наборах FRAMES, BRIGHT, Enterprise RAG (более 500 тыс. синтетических корпоративных документов и 500 вопросов), ESCI, Musique и MLQA, а проверялось на FreshStack, WixQA, BrowseComp-Plus (около 100 тыс. проверенных людьми открытых документов) и Wands (более 42 тыс. оценок релевантности «запрос — товар»).

Пять процентов обучающих примеров внутри каждого набора отвели под валидацию.

Основной метрикой и одновременно функцией награды выбрали nDCG@10 — стандартную метрику информационного поиска, которая измеряет, насколько хорошо первые десять найденных документов совпадают с идеальной выдачей; 1,0 означает идеальный порядок, 0,0 — что релевантных документов не нашлось.

Если агент достигал максимума туров или лимита токенов на один тур, ему назначали награду −1, чтобы прямо отучать модель от таких сбоев без ручной сборки сложных промежуточных награждений.

Конфигурация оказалась минимальной: менялись всего три гиперпараметра — max_epochs=1, global_batch_size=128 и rollout_max_concurrency=32, остальное оставалось по умолчанию, а запуск шёл через SDK MultiTurnRLTrainer.

Что показали результаты

Дообучение улучшило агента на трёх из четырёх отложенных наборов и заметно повысило надёжность на всех.

nDCG@10 вырос с 0,5136 до 0,6354 на BrowseComp-Plus (+23,7%), с 0,5725 до 0,6781 на WixQA (+18,4%) и с 0,5762 до 0,6112 на Wands (+6%).

На FreshStack качество немного просело — с 0,4112 до 0,4089.

Самым сильным результатом стала надёжность: на BrowseComp-Plus доля отказов упала с 22,89% до 0,68%, то есть агент научился не только лучше искать, но и укладываться в лимит туров и токенов.

Среднее число туров на вопрос снизилось с 7,0 до 6,3 на BrowseComp-Plus и с 3,1 до 2,8 на FreshStack, но выросло с 4,3 до 4,5 на WixQA и с 2,2 до 2,9 на Wands.

Кривые награды на обучении и валидации растут и выходят на плато к концу прогона, что говорит о том, что дальнейшее обучение вряд ли дало бы выигрыш.

Контекст и ограничения

MTRL-обучение может длиться несколько дней, а у сервиса есть лимит задания по умолчанию в 24 часа, который настраивается через JSON-схему CreateJob; остановленное или упавшее по таймауту задание можно продолжить с предыдущего чекпойнта.

Заявленные улучшения относятся к конкретному сценарию корпоративного поиска и выбранной функции награды: AWS описывает собственные прогоны оценки, а внешнего воспроизведения этих чисел Mantis не нашла.

Авторы отдельно напоминают про уборку ресурсов: остановить или удалить обучающее задание в консоли SageMaker AI, удалить артефакты модели в Amazon S3 и развёрнутые для оценки эндпоинты, чтобы не платить за них дальше.

Материал подписан Хубином Шэнем (Huibin Shen), Пуджей Карадги (Pooja Karadgi), Сифэй Ли (Sifei Li) и Лоренцо Стеллой (Lorenzo Stella) и отнесён к уровню Advanced (300).