TRACE: как распределять rollout-бюджет в агентном RL и тратить вычисления умнее
Обучение LLM-агентов через reinforcement learning с верифицируемыми наградами (RLVR) стало стандартом для задач, где конечный результат можно автоматически проверить: математика, многошаговый поиск, вызов функций. Но этот стандарт обходится дорого. Каждый rollout требует генерации длинных цепочек рассуждений и взаимодействия со средой, а значит — GPU-часов, электроэнергии и времени. Вопрос не в том, сколько rollout'ов сгенерировать, а в том, куда их направить, чтобы каждый давал максимум обучающего сигнала.
Новый фреймворк TRACE (Tree Rollout Allocation for Contrastive Exploration) предлагает ответ: вместо равномерного распределения бюджета по всем промптам и траекториям, он строит дерево rollout'ов и перенаправляет вычисления к «якорям» — корням и префиксам, где потомки с наибольшей вероятностью дадут противоположные исходы. Результат: при том же бюджете точность растет, а доля информативных rollout'ов увеличивается вдвое.
Что такое RLVR и почему rollout'ы дороги
RLVR — reinforcement learning with verifiable rewards — подход, где модель учится на задачах с бинарной проверкой результата: правильно или неправильно. В отличие от human feedback (RLHF), здесь не нужны аннотаторы — верификатор может быть программой, тестом или компилятором. Это делает масштабирование дешевле, но не бесплатно: каждый rollout требует полной генерации цепочки рассуждений, часто с инструментами — Python-интерпретатором, поисковиком, API.
Типичный RLVR-пайплайн, такой как GRPO (Group Relative Policy Optimization), берет батч промптов, генерирует для каждого фиксированное число rollout'ов, сравнивает их между собой и обновляет политику. Проблема в том, что многие из этих rollout'ов — мертвый груз. Слишком простые задачи решаются с вероятностью близкой к 1, слишком сложные — к 0. В обоих случаях группа rollout'ов дает одинаковые исходы, и сравнивать нечего. Обучающий сигнал нулевой, а GPU-часы потрачены.
Что такое TRACE
TRACE — это единый фреймворк распределения rollout-бюджета для многоходового агентного RLVR. Он формулирует три задачи, которые раньше решались отдельно, как одну: фильтрацию промптов, распределение количества rollout'ов по промптам и ветвление на уровне префиксов внутри одного промпта. Все три — это бюджетные решения над деревом rollout'ов.
Ключевая идея — mixed-reward contrast allocation. Фиксированный бюджет полезен ровно настолько, насколько он порождает pairwise-сравнения: rollout'ы с общим корнем или префиксом, но разными исходами, дают неявные предпочтения между продолжениями. Если все потомки якоря заканчиваются одинаково — успех или провал — обучающий сигнал нулевой. TRACE направляет бюджет к якорям с промежуточной условной вероятностью успеха, где контраст максимален.
Почему равномерное распределение — это трата денег
В типовых пайплайнах RLVR каждый промпт получает фиксированное число rollout'ов, независимо от сложности. Но эмпирически многие промпты дают исходы, скученные у 0 или 1: слишком простые задачи решаются всегда, слишком сложные — никогда. На этих якорях контраст исчерпан, и дополнительные rollout'ы не добавляют сигнала.
На уровне траектории та же проблема глубже. Когда длинный многоходовой rollout получает одну терминальную награду, кредит за отдельные решения распределяется слабо. Модель не знает, какой именно шаг привел к успеху или провалу. Process-supervision решает это, но требует ручных промежуточных наград, уязвимых к reward hacking — явлению, когда модель находит способ получить высокую промежуточную награду, не решая задачу. TRACE предлагает обойтись без промежуточных наград, извлекая локальный кредит из структуры дерева rollout'ов.
Почему дерево лучше, чем плоский список
Ключевой инсайт TRACE — переход от плоских rollout'ов к дереву. В традиционном RLVR каждый rollout — атомарная траектория от начала до конца. Но в многоходовом агентном взаимодействии каждый ход — thought, action, observation — создает естественную точку ветвления. Если модель после трех ходов оказалась в состоянии, откуда разные продолжения ведут к разным исходам, стоит сгенерировать несколько альтернативных продолжений именно отсюда, а не начинать новые rollout'ы с нуля.
Это аналогично тому, как шахматист анализирует не все партии от начала, а ветвит расчет на критических позициях. TRACE автоматизирует этот принцип: предиктор оценивает, какие префиксы «горячие» — то есть, какие позиции обещают контраст — и направляет туда дополнительные вычисления.
Как работает TRACE: два уровня распределения
TRACE работает в два этапа, оба управляются общим предиктором условной вероятности успеха.
Глобальное распределение на корнях. На входе — пул кандидатов из B промптов. Предиктор оценивает для каждого корня вероятность успеха v_i. Значение якоря при выделении m rollout'ов — это вероятность того, что среди них окажутся и успех, и провал: V_root = 1 - v_i^m - (1-v_i)^m. Задача распределения — максимизировать сумму значений при фиксированном общем бюджете M. Промпт с m_i = 0 пропускается, с m_i >= 2 активируется. Динамическое программирование решает эту задачу за время, пренебрежимо малое по сравнению с генерацией rollout'ов.
Локальное расширение префиксов. После генерации bare rollout'ов для активного промпта каждый посещенный префикс становится кандидатом на ветвление. Предиктор оценивает условную вероятность успеха от этого префикса. Значение выделения k продолжений — вероятность, что хотя бы одно продолжение даст исход, противоположный уже наблюденному. Локальный бюджет фиксирован: сумма K_i,j,t по всем префиксам промпта равна m_i * N. Это означает, что Stage 2 не ждет других промптов: как только bare rollout'ы для промпта завершены, его локальное расширение запускается немедленно, минуя межпромптовое ожидание.
Рекурсивная оценка предиктора. После сбора дерева rollout'ов терминальные листья дают бинарные награды. Внутренние узлы обучаются bottom-up: эмпирическая частота успеха среди потомков. Предиктор Ṽ_ψ обучается на среднеквадратичную регрессию к этим рекурсивным целям по множеству корней и информативных внутренних узлов. Один предиктор служит и для корней, и для префиксов — именно это делает его обобщаемым.
Теория: почему префиксы лучше корней
Авторы доказывают три предложения, которые обосновывают архитектуру.
Предложение 1: информация префикса улучшает предсказание сложности. Оптимальная среднеквадратичная ошибка предсказания средней терминальной награды m продолжений не ухудшается с углублением префикса: E_{t+1,m}^* <= E_{t,m}^*. Другими словами, чем больше взаимодействия наблюдено, тем точнее прогноз — и значит, префиксное распределение не хуже корневого, а при наличии промежуточных ходов строго лучше.
Предложение 2: неопределенность префикса как потенциал контраста. Ожидаемая квадратичная вариация прогноза условного успеха вдоль суффикса равна V_t^π(1-V_t^π). Это не статическая метрика неопределенности, а ожидаемая накопленная движение прогноза ниже префикса — мера того, сколько контраста еще осталось.
Предложение 3: распределение по активации доминирует равномерное. При нормализованной условной энергии градиента ожидаемая квадратичная норма локального градиента при TRACE не меньше, чем при равномерном распределении, для каждого этапа в отдельности и для их комбинации. Неравенства строгие, когда равномерное распределение не оптимально.
Результаты: конкретные цифры
Эксперименты проведены на трех задачах: математические рассуждения (DeepScaler), многошаговый QA (HotpotQA) и вызов функций (BFCL v4). Бэкбоны — Qwen3-8B и Qwen3-14B. Сравнение с GRPO, PCL и TreePO при одинаковом бюджете.
Точность. На математике TRACE поднимает среднюю точность с 70.0 до 71.1 для Qwen3-8B и с 73.5 до 74.9 для Qwen3-14B. Аналогичные приросты — на QA и функциях. Сравнение с TreePO особенно показательно: оба метода используют деревья и tree-aware обновления, но TRACE направляет ветви к предсказанно информативным якорям, а не случайно.
Эффективность rollout'ов. Метрика effective ratio — доля промптов в батче, чьи деревья содержат и успешные, и провальные терминальные листья. На математике TRACE поднимает ее с 26.8% до 60.6% для Qwen3-8B и с 34.7% до 59.7% для Qwen3-14B. Это означает, что при том же бюджете более чем вдвое больше обновлений несут невырожденный контрастный сигнал.
Абляция. На HotpotQA с Qwen3-8B активное корневое распределение дает +0.3 пункта точности и +6.3 пункта effective ratio. Активное префиксное — +0.5 и +4.5. Вместе — +1.1 и +9.5. Эффекты складываются: корневое распределение выбирает промпты, склонные к невырожденности, а префиксное тратит продолжения на точки, где новые потомки еще могут дать контраст.
Форма бюджета важнее размера. При общем бюджете 2048 конфигурация (1024 корня, N=2) сильнее, чем (512 корней, N=6), хотя ожидаемое число rollout'ов одинаково. Узкое место — не количество rollout'ов, а достижение состояний, где награды могут образовать контраст. Это контринтуитивно: казалось бы, больше бюджета — всегда лучше, но TRACE показывает, что бюджет без контраста — просто дорогой шум.
Предиктор обобщается на префиксы. Диагностика показывает, что предиктор, обученный в основном на корневых исходах, надежно ранжирует сложность внутренних префиксов. Spearman-корреляция между предсказанной и фактической частотой успеха остается значимой даже на глубоких префиксах, что подтверждает: префиксные истории несут обучаемую неопределенность, а не просто шум.
Практические выводы
Для инженеров, строящих агентные системы, TRACE дает три урока.
Во-первых, не все промпты равны. Фильтрация по предсказанной сложности — не оптимизация «по желанию», а прямое следствие теории контраста. Промпты с предсказуемым исходом можно смело пропускать, экономя бюджет для пограничных случаев.
Во-вторых, внутрипромптовая структура matters. ReAct-взаимодействие естественно пакетирует каждое решение в узел thought-action-observation, и эти узлы — осмысленные точки ветвления. TRACE показывает, что префиксы несут обучаемую неопределенность, а не просто шумовые фрагменты траекторий.
В-третьих, предиктор условного успеха обобщается. Один легковесный предиктор, обученный на корневых исходах, переносит сигнал на внутренние префиксы. Это означает, что сложность префикса — не эвристика, а измеримая величина, которую можно предсказывать и использовать для распределения.
Часто задаваемые вопросы
Работает ли TRACE только для математики? Нет. Авторы оценивают на трех типах задач: математика, многошаговый QA и вызов функций. Фреймворк требует верифицируемых терминальных наград, но не привязан к конкретной предметной области.
Нужен ли process reward model? Нет. TRACE работает с outcome-only наградами, извлекая локальный кредит из структуры дерева. Это избегает уязвимости reward hacking, присущей ручным промежуточным наградам.
Насколько дорого решать задачи распределения? Динамическое программирование для корневого и префиксного распределения решается за время, пренебрежимо малое по сравнению с генерацией rollout'ов. Вычислительные накладные расходы незначительны.
Можно ли использовать TRACE с PPO или другими оптимизаторами? Да. TRACE отделяет сбор rollout'ов от оптимизации политики. Собранные деревья можно передавать любому tree-aware оптимизатору — с group-relative backups, process-reward backups или другими правилами кредита.
Итог
TRACE переформулирует агентное RLVR с вопроса «сколько rollout'ов сгенерировать» на вопрос «где в дереве ветвиться». При фиксированном бюджете он направляет вычисления к якорям с максимальным контрастным потенциалом, удваивая долю информативных обновлений и поднимая точность на 1–2 пункта. Для команд, которые тренируют агентов на верифицируемых задачах, это прямой путь к экономии GPU-часов без потери качества — или к лучшему качеству при том же бюджете.