Single Task Algorithmic Reasoning Models: как с помощью маленькой модели обойти большую LLM?

Результаты коллег по цеху и наши исследования позволяют сделать однозначный вывод: в 2026-м для рецепта отличной reasoning-модели уже недостаточно лишь текстовых рассуждений. Для качественных ответов на большинство вопросов пользователей нужна работа с инструментами, мультимодальные режимы рассуждений, и нередко даже такая экзотика, как рассуждения в латентных представлениях сети.

Мы рады поделиться с сообществом фреймворком для обучения рекурсивных reasoning-моделей STARM, существенно опережающих по качеству предыдущие open source-аналоги. 

О чём расскажем?

STARM (Single Task Algorithmic Reasoning Models) — это результат наших экспериментов с обучением рекурсивных reasoning-моделей, в ходе которых мы:

• Нашли лучшую архитектуру и получили SOTA-результат. Провели эксперименты по оценке влияния всех архитектурных деталей из работ HRM, TRM и URM и подобрали оптимальный рецепт, дающий максимальный прирост.

• Стабилизировали рекурсивные модели. Адресовали основную проблему таких архитектур и доработали фреймворк для воспроизводимого и стабильного обучения.

• Расширили набор доменов. Добавили новые, более сложные классы алгоритмических задач, позволяющие проверить обобщаемость.

• Провели честное сравнение. Сравнили полученные модели со специализированными трансформерами, обученными под задачу, open source LLM и моделями аналогичного семейства.

• Проверили пределы масштабируемости. Протестировали способности моделей к test-time scaling и обучили единую архитектуру для всех задач.

Как появились рекурсивные reasoning-модели?

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

Одним из перспективных направлений стали рекурсивные reasoner’ы — небольшие модели, которые вместо однократного ответа могут многократно перерабатывать внутреннее представление решения. На каждом шаге модель уточняет текущую гипотезу, исправляет ошибки и постепенно приближается к правильному результату. Главная идея этих архитектур — возможность адаптивно тратить вычислительные ресурсы: простые задачи решать быстро, а сложные — «обдумывать» дольше, выполняя больше итераций улучшения. Таким образом, рекурсивные reasoner’ы пытаются приблизить процесс решения задач к тому, как человек сначала формирует черновой вариант ответа, а затем постепенно его уточняет.

Первой громкой работой стала Hierarchical Reasoning Model (HRM). Авторы, вдохновившись концепцией мышления человека, предложили двухуровневую архитектуру: «быстрый» модуль (модуль ) занимается локальными обновлениями, а «медленный» (модуль ) хранит более абстрактное состояние задачи. Главный вывод работы: даже модель всего с миллионами параметров способна конкурировать с в разы большими LLM в алгоритмических задачах.

Однако сразу возник вопрос: действительно ли для этого нужна сложная иерархия?

Ответ попытались дать в статье Tiny Recursive Model (TRM). Авторы заменили -голову, которая ранее училась с помощью RL-алгоритма Q-learning, на классическое бинарное предсказание необходимости остановки. Также сделали общие веса для – и -модулей и добавили — усреднение весов за последние шагов для итогового чекпойнта. И модель, у которой в два раза меньше параметров, из-за общих весов – и -модулей смогла достичь лучшего качества.

Следующий этап развития — это Universal Reasoning Model (URM) и похожие модели. Здесь исследователи искали более удачные компоненты внутри рекурсивного блока, которые позволяют лучше переносить алгоритмы между различными задачами, как, например, свёртки, spectral encoding или группировка токенов. Схема обучения URM практически ничем не отличалась от схемы TRM, кроме добавления ConvSwiGLU и оптимизатора Muon для большей стабилизации.

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

Проверка способностей к обобщению

В большинстве работ по рекурсивным reasoning-моделям оценивают качество только на решении судоку, прохождении лабиринтов и ARC-AGI-1. Рекурсивные reasoner’ы малого размера достаточно сильно переобучаются, поэтому нам нужны были наборы с тестами, выходящими за пределы обучающего распределения. Мы решили тестировать модели на шести доменах:

• ARC-AGI и ARC-AGI-2. Эти бенчмарки напоминают те самые тесты IQ, которые каждый из нас хотя бы раз проходил. По нескольким примерам из контекста нужно понять правило и заполнить следующий шаблон в соответствии с ним.

• Арифметика. Один из добавленных нами тестов, в котором модель должна расставить знаки между числами, чтобы получить заданное итоговое число. Мы разделили выборку на части in distribution и out of distribution, чтобы проверить способность к обобщению.

• «Игра в жизнь» Джона Конвея. Второй из добавленных нами тестов. Задача состоит в предсказании состояния клеточного автомата через несколько этапов эволюции в соответствии с правилами игры. Здесь тоже есть часть out of distribution: модель должна обобщиться на новые паттерны и этапы эволюции.

• Лабиринты. Нужно найти кратчайший путь от случайно расположенной стартовой точки до отмеченной конечной точки в лабиринте на . Модель должна выучить эффективный алгоритм обхода лабиринта, сделав это за наименьшее количество рекурсивных шагов.

• Судоку. По правилам одноимённой игры, нужно расположить числа в сетке на .

Рекурсивные модели обучаются совсем не так, как LLM

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

Рекурсивные модели постоянно пытаются развалиться

Первое, что объединяет все рекуррентные архитектуры, — это нестабильность обучения, которая возникает из-за циклического переиспользования общих весов, поэтому неточность на ранней итерации может привести к провалу метрик или взрыву градиентов.

Предыдущие работы использовали достаточно минималистичный набор нормализаций: ограничение глубины прохода градиентов (HRM), warmup, weight decay и оптимизатор Adam-atan2 в TRM или Muon в URM. При воспроизведении конфигураций из опубликованных работ было заметно, что кривые обучения продолжали выглядеть прерывисто, а результаты при обучении одной и той же модели могли отличаться друг от друга на процентные пункты. Поэтому мы приступили к поиску дополнительных способов стабилизации обучения.

Прежде всего мы стали логировать нормы градиентов, чтобы отслеживать нестабильность в обучении, и закономерно увидели, что при взрыве градиентов эксперименты часто расходятся. Для предотвращения этого мы решили выставить grad_clip = 1.

Далее мы пошли в сторону регуляризации модели, потому что из-за большого количества эпох и малого числа параметров рекурсивные reasoner’ы сильно склонны к переобучению. Чтобы это починить, мы использовали dropout с фиксированной маской внутри всей -траектории, и зашумление входных эмбеддингов для того, чтобы представления большого размера не страдали от малого разнообразия токенов. Также мы добавили взвешенное обновление, которое вычислялось по следующей схеме:

1) Регион доверия (Trust Region):

где — текущий, — предыдущий выходы-модуля, — евклидова норма, — коэффициент региона доверия.

2) Адаптивное обновление (Gate):

Пусть — линейный слой с матрицей весов и смещением :

c начальной инициализацией:

где — исходная случайная матрица.

Тогда адаптивное обновление можно записать в следующем виде:

где – предыдущий выход -модуля, — коэффициент адаптивного обновления.

Помимо этого, как и в моделях URM и TRM, мы использовали –копию модели для сглаживания эффекта от неудачных шагов модели.

Использование этих методов и фиксация начального seed’а помогли нам стабилизировать запуски и снизить дельту между экспериментами с одинаковыми параметрами практически до .

Больше батч — выше качество

Большой вклад в стабилизацию обучения и рост качества внесло увеличение размера батча.

Ключевая идея была вдохновлена работой «Simplifying Deep Temporal Difference Learning», авторы которой показали, что стабильность Q-learning’а можно обеспечить с помощью регуляризации и синхронного сбора данных с большего количества параллельных окружений (в нашем случае большего количества примеров в батче).

Протестировав гипотезу, мы увидели, что этот эффект проявляется и при обучении рекурсивных reasoning-моделей. 

Увеличение размера батча сделало обучение более предсказуемым и позволяло получать приросты от процентных пунктов при тех же настройках обучения. Этот эффект объясним чувствительностью рекурсивных архитектур к шуму и стабилизацией обучения -головы. 

Градиентная аккумуляция не заменяет большой батч

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

Например, в «игре в жизнь» конфигурация с размером батча и двумя шагами аккумуляции оказалась значительно хуже эксперимента с реальным размером батча . Это показывает, что модель из-за своего размера сильно чувствительна к динамике оптимизации.

Длинный путь градиента важнее экономии памяти

В процессе экспериментов мы нашли второй параметр, помимо dropout’а, от которого зависело качество обобщения на новые примеры. Речь о количестве циклов, через которые мы пропускаем градиенты.

В исходной HRM градиент проходит только через последнюю итерацию – и -циклов. В более поздних работах — TRM, URM и SE-RRM — градиент по-прежнему ограничен последним -циклом, но уже распространяется через все -итерации внутри него. Такое решение выглядит естественным — оно уменьшает объём вычислений и экономит память. Однако никто не пытался проверить, насколько именно эта схема влияет на способность модели обобщаться.

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

Вариант (,), соответствующий исходному HRM, и вариант (,_cycles) с обобщением на новые примеры справлялись значительно хуже. Предыдущие работы по рекурсивным reasoning-моделям не исследовали влияние схемы распространения градиента на качество модели в задачах out of distribution.

Второй модуль был лишним

Одним из главных архитектурных вопросов было: «Действительно ли необходимо разделение на два рекурсивных модуля?».

В нашей работе преимущество оказалось на стороне одномодульной TRM-подобной архитектуры. Практически на всех исследованных доменах она обучалась стабильнее и показывала более высокое качество, чем HRM-подобная схема.

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

Свёртки не стали серебряной пулей

В URM рекурсивный блок дополняется свёрточными слоями. На первый взгляд это выглядит естественным решением для задач, представленных в виде двумерных сеток: свёртки находят локальные паттерны, что должно помогать в лабиринтах, «игре в жизнь» и ARC-AGI-задачах.

Однако в наших экспериментах этот эффект оказался гораздо слабее, чем ожидалось. В большинстве доменов использование свёрток не только не улучшило качество, но и сделало обучение менее стабильным. Исключениями стали ARC-AGI-1 и ARC-AGI-2, где свёрточные слои давали устойчивый прирост.

Итоговый вид архитектуры и результаты

Следуя описанным выше законам, мы сформировали финальную версию архитектуры STARM, объединив наиболее эффективные компоненты из предыдущих подходов и собственные улучшения.

В первую очередь нам удалось решить ключевую проблему рекурсивных моделей — нестабильность обучения. За счёт комбинации dropout’ов, зашумления входных эмбеддингов и адаптивного взвешивания выходов мы получили воспроизводимый процесс обучения без существенного разброса между экспериментами.

Вторым важным результатом стала оптимизация самой архитектуры. Итоговая модель содержит около миллионов параметров, использует четыре слоя в рекурсивном блоке, общие веса для – и -циклов, скрытое состояние размером и применяет свёрточные компоненты только там, где они действительно дают прирост качества, — в задачах семейства ARC-AGI.

Помимо улучшения архитектуры, нам удалось повысить способность модели к обобщению. Увеличение длины пути распространения градиента, использование больших реальных размеров батча и добавление dropout’ов позволили STARM лучше работать с примерами, выходящими за пределы обучающего распределения. В итоге мы получили финализированную архитектуру модели.

Эта архитектура показала следующие результаты в рассмотренных доменах:

Для объективного сравнения с прошлыми работами мы обучили сопоставимый с архитектурой по параметрам трансформер без рекурсивных шагов, воспроизвели оригинальные HRM, TRM и URM и проверили LLM на тех же тестовых наборах в режиме few shot. Для всех датасетов мы замеряли -версии моделей.

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

Масштабируемость: потенциальные возможности моделей?

Одна модель не смогла выучить всё сразу

Те, кто не погружается глубоко в статью и чтение работ, скорее всего, не знают, что модели рекурсивного reasoning’a обычно обучаются только на одном домене, при том что отдельной мотивации этого решения в статьях не было описано.

Поэтому логичный вопрос, который нас интересовал: «Смогут ли рекурсивные модели одновременно освоить несколько алгоритмов, или каждая задача требует специализированного solver’а?»

Для этого мы объединили все домены в один многозадачный датасет и обучили единую модель. Эксперимент оказался значительно дороже обычного обучения: даже после тысяч итераций, что равно нескольким экспериментам с однозадачными моделями, обучение не успело полностью сойтись.

Несмотря на то, что loss на обучении продолжал снижаться, качество на отдельных задачах практически не росло. Особенно сильно пострадали судоку и лабиринты, где точность на тестовой выборке оставалась близкой к нулю. Аналогичная ситуация наблюдалась и в «игре в жизнь». Стабильно обучалась только арифметика.

Такое поведение может быть вызвано недостаточной ёмкостью модели и существенной разницей задач. Один общий рекурсивный блок вынужден одновременно аппроксимировать все разнородные зависимости, из-за чего начинает терять специализацию, а ёмкости исходных эмбеддингов не хватает для хранения информации обо всех выученных правилах.

Исходя из этого мы сделали вывод, что, по крайней мере в текущем виде и размере, рекурсивные reasoning-модели плохо масштабируются на многозадачное обучение.

Обобщение действительно работает

Как мы писали выше, добавление dropout и пропускание градиентов через большее количество – и -циклов увеличило обобщаемость. Расскажем, как мы это проверяли.

«Игра в жизнь»

Для оценки модели мы использовали два типа тестов:

• in-train-out-range — паттерны, присутствовавшие в обучающей выборке, но с большим количеством предсказываемых шагов;

• out-train — данные, отсутствовавшие в обучающей выборке: новые паттерны; этот набор специально создали для оценки обобщения.

В отличие от TRM и URM, модель действительно обобщилась на данные вне обучающего распределения, и даже показывала сопоставимые с in-train-данными результаты.

Арифметика

Похожий эффект, но меньшей интенсивности мы отметили на арифметике. Здесь для замеров использовали два тестовых набора:

• 1–101 — примеры, в которых после расстановки знаков должно получиться от до , похожие на обучающее распределение, но отсутствующие в обучающей выборке;

• 102–110 — примеры, в которых после расстановки знаков должно получиться от до , что выходит за рамки обучающей выборки.

Здесь совпадения по метрикам с 1–101 достичь не удалось, но за счёт стабилизации обучения добились ровного сонаправленного тренда c in distribution-бенчмарком и получили результат выше TRM и URM baseline’ов.

Test-Time Scaling работает не везде

Отличительной особенностью reasoning-моделей является test-time scaling — способность при больших вычислительных затратах на инференсе показывать прирост качества. Мы решили проверить, насколько рекурсивные reasoner’ы обладают похожими свойствами.

Модель хорошо демонстрирует свойство test-time scaling на переборных доменах, к которым относятся только судоку и арифметика. В остальных задачах, где для решения нужно подобрать оптимальный паттерн, дополнительные шаги практически не улучшают качество.

Как и обычные reasoning-модели, рекурсивные reasoner’ы плохо умеют находить свои ошибки и исправлять их.

Половина популярных идей не сработала

Не все гипотезы подтвердились экспериментально. Мы варьировали weight decay и виды scheduler’ов, однако устойчивого улучшения получить не удалось. Некоторые методы работали по отдельности, например lookahead, но при объединении с остальными компонентами их эффект исчезал или становился вовсе отрицательным.

Архитектурные эксперименты тоже не всегда оправдывали ожидания. Увеличение глубины до восьми слоёв не давало выигрыша относительно более компактной четырёхслойной модели, как и изменение размера скрытого состояния, а увеличение количества attention-голов даже ухудшало качество. Попытка заменить Post-Norm на Pre-Norm приводила к нестабильному обучению, хотя в языковых зацикленных моделях эффект ровно обратный.

На очереди — текстовые reasoner’ы

Рассмотренные малые рекурсивные reasoner’ы можно использовать как примитивы для решения символьных алгоритмических задач, где нужна адаптивная глубина вычислений. Например, в сценариях, где LLM выступает в качестве движка планирования, а рекурсивная модель — инструмента для вычислений.

Отдельно, в роли текстового движка, модели такого типа уже были масштабированы в работах HRM-Text и HRM-MoE, и мы начали исследования в этом направлении, результатами которого поделимся с вами.

Благодарности

В завершение благодарим нашу команду: Валерия Терновского, Дмитрия Бочарова, Инессу Фёдорову, Дмитрия Харчева, Юлиану Шахвалиеву и Альберта Корнилова. Именно ваш энтузиазм, глубокая экспертиза и бесчисленные часы, потраченные на обсуждение гипотез и отладку экспериментов, сделали возможным этот результат. Отдельное спасибо всем, кто помогал с рецензированием и давал ценные замечания по ходу работы, — ваша обратная связь стала важной частью этого исследования.