Если вы тренируете большую языковую модель через reinforcement learning, до 80% времени уходит не на градиентные обновления, а на rollout — генерацию ответов, которые потом оценивает reward model. Казалось бы, Multi-Token Prediction (MTP) должен спасать: draft-модель предсказывает несколько токенов за раз, target-модель проверяет их параллельно, и inference ускоряется в 2–3 раза. Но на практике, когда MTP встраивают в RL-пайплайн, acceleration rate обычно падает до жалких 10–15%. Почему?

Новая работа Bebop (Breaking Entropy Bounds for Optimal Prediction) даёт точный диагноз: энтропия политики. Когда модель обучается через RL, она исследует пространство ответов, её распределение становится более размытым — и draft-модель, обученная на более «уверенных» данных, перестаёт попадать в цель. Авторы показывают, что acceptance rate MTP линейно падает с ростом энтропии, и предлагают решение, которое не требует дорогого онлайн-переобучения MTP во время RL: probabilistic rejection sampling + end-to-end TV loss. Результат — ускорение до 1.8× на реальных задачах reasoning и агентного взаимодействия.

Что такое Multi-Token Prediction и почему он ломается в RL

Multi-Token Prediction — это разновидность speculative decoding. Основная модель (target) делает дорогой forward pass, но вместо того чтобы генерировать по одному токену, она проверяет сразу несколько кандидатов, предложенных маленькой draft-моделью. Если draft угадал — сэкономили время, если нет — откатываемся и генерируем заново.

Ключевой показатель — acceptance rate: какая доля предложенных токенов проходит проверку. В supervised fine-tuning (SFT) draft-модель обучается на том же распределении, что и target, и acceptance rate держится на уровне 70–90%. Но в RL ситуация меняется кардинально.

Во-первых, RL поощряет исследование: политика намеренно поддерживает высокую энтропию, чтобы пробовать разные стратегии ответа. Во-вторых, веса target-модели обновляются на каждом шаге RL, а draft-модель обычно заморожена — возникает mismatch между распределениями. Казалось бы, оба фактора должны влиять, но авторы Bebop провели декомпозицию и выяснили: энтропия отвечает за 95% падения acceptance rate, а mismatch от policy updates — менее 5%.

Энтропийный предел: почему greedy sampling обречён

Авторы формализуют проблему через информационную теорию. Пусть p — распределение target-модели, q — распределение draft-модели, H(p) — энтропия target. При target-only sampling draft выбирает токен жадно (argmax q), и acceptance rate равен вероятности этого токена по мнению target: α_TO = p(argmax q).

Когда энтропия низкая, target «уверен» в ответе — распределение пиковое, p(argmax q) близок к 1, и даже приблизительный draft попадает. Но когда энтропия растёт, распределение размывается, максимальная вероятность падает, и α_TO падает вместе с ней. Авторы показывают чёткую отрицательную линейную зависимость с коэффициентом наклона около −1.68: рост энтропии на единицу снижает acceptance rate примерно на 1.7%.

Проблема в том, что target-only sampling жёстко привязан к max_y p(y), а этот максимум неизбежно падает при росте энтропии. Даже идеально обученный draft не может превзойти этот предел — он определяется не качеством draft, а формой target-распределения.

Rejection sampling: как разорвать связь с энтропией

Альтернатива target-only — probabilistic rejection sampling. Здесь acceptance rate определяется не одним токеном, а полным перекрытием распределений: α_RS = Σ_v min(p(v), q(v)). Это совсем другая математика: вместо того чтобы спрашивать «угадал ли draft топ-1 токен?», мы спрашиваем «насколько близки распределения draft и target?».

И здесь происходит ключевое отличие. Rejection sampling зависит от энтропии гораздо слабее, потому что даже если target «неуверен» и распределение размыто, draft может повторить эту же размытость — и перекрытие останется высоким. Авторы экспериментально показывают, что при rejection sampling наклон зависимости acceptance rate от энтропии падает с −1.68 до −0.06 — то есть связь практически исчезает. Draft перестаёт бороться с энтропией и начинает работать вместе с ней.

На практике это означает, что при переходе от target-only к rejection sampling acceptance rate на агентных задачах вырастает с 71% до 90%, а на задачах reasoning — с 75% до 87%. Это не инкрементальное улучшение, а качественный скачок, который делает MTP жизнеспособным в RL-контексте.

Почему CE и KL — неправильные лосс-функции для MTP

Традиционно draft-модель обучают минимизировать cross-entropy (CE) или KL-дивергенцию между q и p. Эти лоссы широко распространены, но авторы Bebop показывают, что они косвенно улучшают rejection sampling: по неравенству Пинскера, TV-расстояние ограничено сверху через KL, но минимизация KL не эквивалентна минимизации TV.

Проблема в том, что CE/KL распределяют оптимизационный бюджет равномерно по всему словарю. Draft учится приближать полное распределение p, включая хвосты — редкие токены, которые почти никогда не влияют на rejection sampling. В результате draft получается «гладким», с низкой энтропией, и плохо перекрывается с размытым target-распределением в RL.

Авторы предлагают end-to-end TV loss — прямую минимизацию total variation distance между draft и target:

L_TV = d_TV(p, q) = 1 − Σ_v min(p(v), q(v))

Градиент TV loss концентрируется на токенах, близких к decision boundary — там, где p(v) ≈ q(v), и небольшое смещение draft может перекинуть токен из «отвергнутых» в «принятые». Токены в хвостах распределения игнорируются, что экономит оптимизационный бюджет и делает draft «острее» — его энтропия приближается к энтропии target, а не занижается искусственно.

Что показывают эксперименты

Авторы тестируют Bebop на семействе моделей Qwen (3.5, 3.6, 3.7) в трёх режимах RL: reasoning (математика, код, инструкции, до 64K токенов), SWE (многоходовое редактирование кода, до 128K токенов и 200 ходов) и агентные задачи. Все эксперименты проводятся в асинхронном RL-фреймворке на базе veRL и SGLang.

Сравнение лосс-функций при SFT. На Qwen3.5-35A3B e2e TV loss улучшает rejection sampling acceptance rate на 3–8% по сравнению с CE baseline в зависимости от задачи. Наибольший прирост — на агентных задачах: с 90.3% (CE) до 97.0% (TV). На out-of-distribution MT-Bench прирост скромнее — 2.3%, что подтверждает хорошую генерализацию метода.

Стабильность в RL. Главное преимущество TV loss проявляется именно во время RL-обучения. На графиках accept length по ходу тренировки видно, что target-only sampling (TO) и rejection sampling с CE (RS w/ CE) показывают устойчивое снижение acceptance rate по мере роста энтропии политики. Rejection sampling с TV loss (RS w/ TV) остаётся практически плоским — наклон зависимости от энтропии снижается на 95%.

Throughput. Улучшение acceptance rate линейно транслируется в ускорение обучения. MTP с rejection sampling и TV loss даёт 1.5–1.8× ускорение per-step latency по сравнению с обучением без MTP. На агентных задачах rollout phase ускоряется до 2.4×. При масштабе сотен тысяч GPU-часов это экономия в десятки тысяч часов вычислений.

Масштабирование с размером модели. Acceptance rate растёт с размером модели: Qwen3.7-Max достигает 95% на агентных задачах при γ=3 (target проверяет 4 токена за раз). Это означает, что draft почти слился с target — разница в размере компенсируется качеством обучения. Для меньших моделей (Qwen3.6-27B) acceptance rate ниже — 79.9% на математике, 67.5% на MT-Bench — но всё равно значительно выше CE baseline.

Нужно ли обновлять MTP веса во время RL?

Важный практический вопрос: если target-модель меняется во время RL, не нужно ли обновлять и draft-модель онлайн? Несколько недавних работ предлагали именно это — online MTP training во время RL. Но авторы Bebop показывают, что это не нужно.

Во-первых, mismatch от policy updates оказывается пренебрежимо малым — менее 5% влияния на acceptance rate. Во-вторых, когда MTP веса всё-таки обновляют во время RL (даже с TV loss), acceptance rate постепенно сходится к тому же уровню, что и без обновлений. А если обновлять CE loss — draft становится гладче, и acceptance rate деградирует до уровня RS w/ CE baseline.

Вывод: достаточно одного лёгкого pre-RL этапа обучения MTP с e2e TV loss. Дальше draft работает в режиме inference-only, не требуя дополнительной памяти или вычислительных ресурсов. Это критически важно для production RL-пайплайнов, где каждый гигабайт GPU-памяти и каждая миллисекунда latency на счету.

Почему TV loss делает draft «острее»

Авторы анализируют, как разные лосс-функции влияют на форму draft-распределения. CE/KL создают гладкие распределения с заниженной энтропией — draft «боится» ошибиться и размывает массу по словарю. TV loss, напротив, концентрирует массу на тех же токенах, что и target, даже если target сам неуверен.

Это можно увидеть на метриках: у моделей, обученных TV loss, энтропийный зазор между draft и target меньше (draft энтропийно ближе к target), но KL-дивергенция при этом больше. Казалось бы, парадокс — большее KL при лучшем acceptance rate? Но KL измеряет разницу по всему словарю, а rejection sampling смотрит только на перекрытие. TV loss жертвует точностью в хвостах ради лучшего совпадения в голове распределения — и это именно то, что нужно для speculative decoding.

Корреляция между энтропийным зазором и acceptance rate составляет −0.54, тогда как корреляция между KL и acceptance rate — всего 0.13. Энтропийный зазор, а не KL, является предиктором качества draft для rejection sampling.

Robustness: почему rejection sampling устойчивее к policy updates

Target-only sampling хрупок к изменениям ранжирования. Если RL-градиент меняет топ-1 токен — пусть даже незначительно, с 0.31 до 0.29 вероятности — draft, всё ещё предпочитающий старый топ-1, мгновенно получает отказ. Это дискретный скачок: принято/не принято.

Rejection sampling деградирует плавно. Если target немного сдвинул распределение, overlap Σ_v min(p, q) уменьшается немного — acceptance rate падает постепенно, а не обрушивается. Это свойство делает rejection sampling гораздо более устойчивым к непрерывным policy updates в RL, где распределение меняется на каждом шаге.

Часто задаваемые вопросы

Почему нельзя просто понизить температуру в RL, чтобы снизить энтропию?

Понижение температуры делает политику более детерминированной, что действительно повышает acceptance rate. Но RL требует исследования: если модель не будет пробовать разные варианты ответов, она застрянет в локальном оптимуме. К тому же многие RL-алгоритмы (например, PPO с entropy bonus) явно поощряют высокую энтропию. Bebop решает проблему не за счёт ограничения исследования, а за счёт адаптации draft к естественной энтропии политики.

Работает ли Bebop только для больших моделей?

Нет, метод масштабируется на модели разного размера — от 27B до Max-вариантов. Да, на больших моделях absolute acceptance rate выше (до 99% на агентных задачах для Qwen3.6-Plus), но относительный прирост TV loss над CE baseline сохраняется на всех масштабах. Для небольших моделей, где вычислительный бюджет особенно ограничен, ускорение RL-пайплайна может быть даже ценнее.

Можно ли применить TV loss к существующим MTP-реализациям?

Да, TV loss — это замена лосс-функции на этапе обучения draft-модели, не требующая изменений в inference-коде. Авторы выпустили реализацию в SGLang, но принципиально метод можно интегрировать в любой speculative decoding engine. Главное требование — поддержка rejection sampling на этапе верификации, что уже реализовано в большинстве современных inference-фреймворков.

Итог

Bebop ставит диагноз и предлагает лечение для главного узкого места современного RL-обучения LLM. Проблема не в том, что MTP плохо работает в RL — проблема в том, что мы неправильно его обучали и неправильно сэмплировали. Target-only sampling жёстко привязан к энтропии и обречён на деградацию. CE/KL loss тренирует draft быть «слишком хорошим» в хвостах распределения, вместо того чтобы максимизировать перекрытие в голове.

Решение — probabilistic rejection sampling + end-to-end TV loss — элегантно и практично. Оно не требует онлайн-переобучения MTP во время RL, не добавляет latency, и даёт ускорение до 1.8× на реальных задачах. Для команд, которые тренируют модели через RL на сотнях GPU, это разница между неделями и месяцами обучения. Если вы используете MTP в RL-пайплайне — стоит проверить, не сидит ли ваш draft в энтропийной ловушке.