Прунинг рассуждений: PKTP вместо перебора 128 попыток
Одна задача из AIME, 128 попыток рассуждения, бюджет 32 768 токенов на каждую. Это 4,2 миллиона токенов на один вопрос. Большая часть уходит впустую: черновик, который на третьей тысяче токенов свернул не туда, всё равно доводится до конца, потому что остановить его нечем. Именно здесь работает прунинг рассуждений, и именно здесь свежая работа arXiv 2610.31667 предлагает поменять сам критерий отбора.
Прунинг рассуждений (thought pruning) это отсечение части черновиков на ранней стадии: модель генерирует N попыток только до порога T токенов, оставляет k самых перспективных и доводит до полного бюджета лишь их. Остальные выбрасываются, не потратив времени на вычисления, которые всё равно привели бы к неверному ответу.
Проблема в том, что "самые перспективные" можно понимать по-разному, и оба очевидных понимания работают плохо.
Два способа ошибиться при отборе
Первый способ опирается на качество. Внутренние состояния трансформера скармливаются линейному пробнику, тот выдаёт вероятность правильного ответа, и дальше вы просто берёте k черновиков с наибольшим баллом. Логика железная, но у неё есть побочный эффект. Черновики, которые пробник считает правильными, похожи друг на друга, потому что похожи их представления. Пробник ошибается не на одном черновике, а на всём кластере сразу.
Авторы показывают это на одном конкретном примере. Qwen3-8B, задача из AMO-Bench, 128 черновиков, спроецированных на первые две главные компоненты. Шестнадцать лучших по версии пробника лежат плотной группой, и все шестнадцать неверны. Шестнадцать случайно выбранных содержат семь правильных. Вы просто заплатили за 16 полных прогонов, чтобы получить ноль попаданий.
Второй способ опирается на разнообразие. Метод RepExp жадно набирает черновики, представления которых хуже всего объясняются уже выбранными: остаточная дисперсия, обратная матрица, всё по учебнику. Первый выбор случайный, дальше идёт охота за непохожестью. И здесь зеркальная ошибка: самыми непохожими чаще всего оказываются выбросы, то есть как раз те траектории, которые никуда не придут. На той же модели в Omni-MATH четыре выбора RepExp все четыре неверны, а четыре случайных дают одно попадание.
Между этими крайностями лежит классический конфликт exploration и exploitation, и оба метода сидят ровно на его полюсах. Ни один не оптимизирует то, что вам на самом деле нужно.
Что такое pass@k как цель, а не метрика
Нужная величина называется pass@k. Это вероятность того, что хотя бы один черновик из отобранных k, будучи доведён до конца, даст правильный ответ. Формально задача выглядит так: выбрать подмножество S мощности k, которое максимизирует вероятность события "хотя бы один правильный среди S". Всё, что делает PKTP (Pass@k Thought Pruning), подчинено этой формуле.
Ключевой ход в том, что исходы черновиков не независимы, и это надо учитывать явно. Если два черновика идут по одной колее, они падают вместе. Работа моделирует это латентной гауссовой моделью: балл черновика складывается из его собственного качества, общего случайного эффекта, зависящего от представления, и индивидуального шума. Ковариация двух баллов при этом равна скалярному произведению их представлений. Черновики с сонаправленными векторами проваливаются синхронно, и никакая независимая оценка качества этого не увидит.
Обучение сводится к максимизации правдоподобия наблюдаемых картин "кто прошёл, кто нет" по M задачам. Точное правдоподобие содержит интеграл по гауссову случайному эффекту и в замкнутом виде не берётся, поэтому его раскладывают до второго порядка и получают суррогат.
Три слагаемых суррогата
Суррогат оптимизирует не вероятность успеха, а логарифм вероятности полного провала, что то же самое. В нём три части, и каждая читается буквально.
Первая суммирует логарифмы индивидуальных вероятностей провала и тянет отбор к высокому качеству. Если выключить остальные две, получится ровно линейный пробник.
Вторая, отрицательный логарифм детерминанта, требует, чтобы представления выбранных черновиков покрывали разные направления. Это и есть штраф за кластеризацию, причём в явном виде.
Третья, квадратичная форма по сумме взвешенных векторов, ловит то, чего детерминант не видит. Детерминант одинаково относится к вектору и к его противоположности, а проваливаются они по-разному: два черновика с h₂ = h₁ падают вместе чаще, чем с h₂ = −h₁. Штраф различает эти случаи.
Дальше начинается вычислительная часть. Точная минимизация суррогата NP-трудна, авторы доказывают это сведением от задачи о максимальном независимом множестве. Но целевая функция монотонна и супермодулярна по модулю знака, поэтому жадный выбор даёт гарантию (1 − e⁻¹) от оптимума. На практике это означает: перебирать все сочетания из 128 по 4 не нужно, достаточно 4 жадных шага с матричными обновлениями.
Почему это не третий эвристический трюк
Самое интересное в работе не цифры, а предельный анализ. Авторы показывают, что PKTP не висит где-то посередине между двумя лагерями, а строго их интерполирует, причём параметром служит сила корреляции между черновиками.
Если корреляцию обнулить, суррогат вырождается в сумму по качеству, и жадный выбор совпадает с линейным пробником. Если, наоборот, уравнять качества и пренебречь штрафом, остаётся только логарифм детерминанта, а его жадная маргинальная выгода по лемме о детерминанте матрицы даёт ровно критерий остаточной дисперсии, то есть RepExp.
Между пределами работает механика, которую легко представить. Первый выбор всегда определяется качеством: штрафов за сходство ещё нет. Каждый следующий черновик платит столько штрафов, сколько уже выбрано, поэтому к некоторому k накопленные штрафы перевешивают разницу в качестве, и отбор начинает разводить черновики по сторонам, как RepExp. Чем сильнее корреляция, тем раньше происходит перелом.
Приложение C доводит это до точных условий. При малом k, если разрыв между k-м и (k+1)-м качеством превышает величину порядка k·ρ², PKTP отбирает ровно топ-k по качеству, как пробник, и в том же порядке. При большом k объектив сводится к объективу RepExp. То есть поведение метода меняется с бюджетом само, без переключателей и ручных порогов.
Результаты на восьми моделях
Эксперимент устроен тяжело. Восемь открытых reasoning-моделей в двух размерных классах: Qwen3-0.6B, Qwen3-1.7B, OpenMath-Nemotron-1.5B и DeepSeek-R1-Distill-Qwen-1.5B в младшем, Qwen3-8B, OpenMath-Nemotron-7B, DeepSeek-R1-Distill-Qwen-7B и Skywork-OR1-Math-7B в старшем. Генерация при temperature 0.8 и top-p 0.95, полный бюджет B = 32 768 токенов, пул N = 128, порог отсечки T = 4000. Около 600 задач на модель, собранных из AIME 2024, 2025 и 2026, HMMT тех же лет, RIMO-N и дополнительных наборов под каждый класс. Проверка ответов через Qwen3-32B судью по выражению в \boxed{}. Метрика усредняется по 20 повторениям восьмифолдовой кросс-валидации, чтобы каждая задача оценивалась моделью, которая её не видела.
Числа pass@4, где, по словам авторов, выигрыш максимален:
| Модель | Random | Linear Probe | RepExp | PKTP |
|---|---|---|---|---|
| Qwen3-0.6B | 45,4 | 44,3 | 45,0 | 47,2 |
| Qwen3-1.7B | 62,6 | 63,1 | 63,9 | 64,4 |
| OpenMath-Nemotron-1.5B | 69,3 | 69,4 | 68,9 | 69,4 |
| DeepSeek-R1-Distill-Qwen-1.5B | 49,5 | 49,9 | 49,8 | 52,9 |
| Qwen3-8B | 67,2 | 67,4 | 68,1 | 68,6 |
| OpenMath-Nemotron-7B | 66,0 | 67,5 | 66,4 | 68,2 |
| DeepSeek-R1-Distill-Qwen-7B | 58,5 | 56,4 | 58,6 | 59,7 |
| Skywork-OR1-Math-7B | 62,7 | 62,9 | 64,5 | 65,6 |
Разброс выигрыша показателен. На DeepSeek-R1-Distill-Qwen-1.5B при k = 4 PKTP выше линейного пробника на 3,0 пункта и выше RepExp на 3,2. На Qwen3-0.6B он обходит оба baseline больше чем на 2 пункта при любом k ≥ 2. На DeepSeek-R1-Distill-Qwen-7B разрыв с пробником при k = 4 составляет 3,3 пункта.
Есть и неудобные клетки. OpenMath-Nemotron-1.5B при k = 4 даёт 69,4 у PKTP и 69,4 у пробника, то есть ничью. Qwen3-1.7B при k = 1 показывает 39,3 против 39,5 у случайного выбора, то есть чуть хуже. Это ожидаемо: при k = 1 штрафов за сходство ещё нет, и метод вырождается в оценку качества, а оценка качества на этой модели слабая.
Общая закономерность такая. Пробник силён при k = 1 и по мере роста k сползает к случайному выбору, а на Qwen-моделях местами уходит ниже него. RepExp при k = 1 равен случайному по построению, при больших k обгоняет, но в среднем диапазоне тоже проигрывает случайному на OpenMath-Nemotron-1.5B и Qwen3-8B. PKTP единственный, кто держится выше случайного выбора на всех восьми моделях при k ≥ 2.
Чем позже отсечка, тем больше смысла в отборе
Отдельный сюжет, который легко пропустить: порог T. Если резать черновик на 2000-м токене, сигнала почти нет, и любой отбор превращается в гадание. Авторы прогоняют T от 2000 до 8000 и получают устойчивую картину.
На Qwen3-1.7B pass@4 у PKTP растёт с 63,5% при T = 2000 до 65,9% при T = 8000. На DeepSeek-R1-Distill-Qwen-7B с 57,4% до 61,7%. При любом T ≥ 4000 PKTP лучший на обеих моделях.
А теперь контраст. Линейный пробник на Qwen3-1.7B при удлинении отсечки не улучшается, а падает: с 64,2% до 62,3%, уходя ниже случайного выбора. Причина, которую предлагают авторы, касается конкретной реализации представления: экспоненциальное сглаживание с коэффициентом 10⁻³ фактически усредняет последнюю тысячу токенов, поэтому ранний сигнал размывается, когда префикс становится длинным. Это не приговор методу пробника, а указание на то, что представление и критерий отбора надо согласовывать.
Размер пула: больше не всегда лучше
Второй параметр, N, ведёт себя похоже. На Qwen3-1.7B PKTP лучший при любом N от 32 до 128, и оба baseline тоже выигрывают от роста пула: RepExp с 62,9% до 63,9%, пробник с 62,1% до 63,1%.
На DeepSeek-R1-Distill-Qwen-7B картина ломается. Пробник с ростом пула падает с 58,1% до 56,4% и остаётся ниже случайного выбора при любом N. RepExp падает с 59,5% при N = 48 до 58,6% при N = 128. PKTP растёт с 58,8% до 59,7%.
Объяснение простое и слегка неудобное для обоих baseline: большой пул содержит больше выбросов, а и пробник, и RepExp при умеренном k склонны выбирать именно крайние случаи. Один тянется к плотному кластеру, второй к дальним точкам. Чем больше пул, тем вероятнее, что крайность найдётся.
Сколько это экономит на самом деле
Арифметика прозрачна, но её стоит считать аккуратно, потому что маркетинговые формулировки тут легко завышают эффект.
Полный перебор 128 черновиков по 32 768 токенов это 4 194 304 токена на задачу. С отсечкой на 4000 токенах 128 префиксов стоят 512 000 токенов. Дальше каждый выживший стоит ещё 32 768. При k = 4 итог 643 072 токена, то есть в 6,5 раза меньше, экономия 85%. При k = 16 итог 1 036 288, в 4,05 раза меньше.
Но сравнивать надо не с полным перебором, а с честной альтернативой. Наивный вариант "сгенерировать сразу 4 полных черновика" стоит 131 072 токена и даёт на DeepSeek-R1-Distill-Qwen-1.5B pass@4 = 49,5. PKTP при том же k = 4 даёт 52,9, но за 643 тысячи токенов, то есть в 4,9 раза дороже. За +3,4 пункта.
Настоящая выгода проявляется, когда вы уже заплатили за префиксы. Переход с k = 4 на k = 8 добавляет 131 тысячу токенов, то есть +20% к бюджету, и приносит +9,6 пункта на той же модели: с 52,9 до 62,5. Вот это уже хорошая сделка, и именно её продаёт метод: вы один раз платите за широкий пул префиксов, а потом наращиваете k дёшево.
Что с этим делать в продакшене
Схема внедрения требует доступа к скрытым состояниям модели, что для открытых весов нормально, а для закрытого API означает, что метод недоступен. Дальше нужен слой пулинга: в статье это экспоненциальное среднее с коэффициентом 10⁻³, и авторы честно пишут, что более удачное представление напрямую усилит метод, поскольку весь каркас от конкретного способа кодирования не зависит.
Второе требование тяжелее: обучающую выборку. Нужны задачи с проверяемым ответом, прогон по 128 черновиков на каждую и разметка "прошёл или нет". Для математики это делается автоматически через судью, для кода через юнит-тесты, для открытых задач общего вида придётся строить reward-модель и мириться с её ошибками. Сама процедура отбора при этом дешёвая: жадный шаг требует нескольких матрично-векторных умножений на кандидата, а не нового прогона модели.
Третье ограничение касается профиля нагрузки. Если у вас и так генерируется 4 черновика на запрос, прунинг ничего не даст: отсекать нечего. Метод имеет смысл там, где пул широкий, то есть в офлайн-оценке, в синтетике данных, в задачах с высокими ставками и в RL-пайплайнах, где считается pass@k как награда.
Чего в работе нет
Эксперименты полностью на математических бенчмарках: AIME, HMMT, Omni-MATH, AMO-Bench, RIMO-N. Юнит-тесты упоминаются в постановке задачи как пример верификатора, но кодовых бенчмарков в таблицах нет, и это заметная дыра, потому что распределение черновиков в коде устроено иначе, чем в олимпиадной математике.
Второе: представление фиксировано, а его коэффициент выбран без развёрнутого обоснования. Утверждение "более хорошее представление усилит метод" верно по конструкции, но означает, что половина выигрыша может лежать не в критерии отбора, а в качестве эмбеддинга.
Третье: выигрыш на больших k невелик. При k = 16 разница с random на Qwen3-0.6B это +1,7 пункта, на Qwen3-8B +0,7. Если ваш бюджет уже позволяет гонять 16 полных черновиков, платить за инфраструктуру прунинга особого смысла нет.
Часто задаваемые вопросы
Чем прунинг рассуждений отличается от early exit?
Early exit решает, когда остановить одну конкретную траекторию, ориентируясь на уверенность модели в текущем токене. Прунинг решает, какие траектории из уже сгенерированных продолжать, и смотрит на совместное распределение их исходов. Это разные уровни принятия решения, и они спокойно сочетаются: early exit внутри черновика, прунинг между черновиками.
Можно ли применить метод к генерации кода?
Постановка допускает это напрямую, потому что верификатором может быть прогон юнит-тестов. Ограничение в обучающих данных: нужны задачи с достаточно широким разбросом исходов, иначе пул нечего отсекать. В статье кодовых экспериментов нет, так что цифр под этот сценарий пока не существует.
Почему нельзя просто взять лучший черновик по вероятности ответа?
Потому что оценка вероятности не видит корреляций. Она уверенно ставит высокие баллы группе похожих черновиков, а группа проваливается целиком. Пример из статьи: 16 лучших по пробнику на Qwen3-8B оказались все неверны, тогда как 16 случайных дали 7 правильных. Отбор по индивидуальному баллу без штрафа за сходство систематически обманывает.
Итог
Главная мысль работы не в конкретной гауссовой модели и не в жадном алгоритме с гарантией (1 − e⁻¹). Она в том, что pass@k перестаёт быть метрикой и становится целевой функцией. Как только вы оптимизируете именно вероятность "хотя бы один правильный", качество и разнообразие перестают быть конкурирующими эвристиками: их баланс определяется самим критерием, а не тем, какой из двух лагерей выбрал автор.
Практический вывод для тех, кто считает токены: широкий пул префиксов плюс дешёвое наращивание k даёт лучшую отдачу, чем узкий пул полных прогонов. 128 префиксов по 4 тысячи токенов стоят меньше, чем четыре полных черновика, и открывают возможность выбрать восемь.
Если вы гоняете best-of-N в своём пайплайне, посчитайте, во что вам обходится именно перебор, а не генерация. Часто окажется, что узкое место не в модели, а в отсутствии критерия, который умеет считать совместные провалы.