LoopCD улучшает зацикленные трансформеры без дообучения и экономит вычисления
В статье предложен LoopCD, метод контрастного декодирования для зацикленных трансформеров (looped Transformers). Такие модели экономят параметры тем, что многократно прогоняют один и тот же общий блок через несколько рекуррентных циклов. После каждого цикла получается промежуточное представление, из которого можно декодировать тот же следующий токен, но стандартное декодирование эти ранние состояния отбрасывает.
Идея авторов в том, что ранние циклы выполняют меньше вычислений, поэтому сама рекуррентность даёт готовую пару «слабый и сильный прогноз» без вспомогательных моделей и внешнего обучения. LoopCD не требует обучения: он направляет выбор токена, сопоставляя итоговый прогноз с прогнозом более раннего цикла. Есть два варианта: LoopCD-Logits работает в пространстве логитов и требует одного дополнительного прохода через выходной слой, LoopCD-Hidden работает в пространстве скрытых состояний и не добавляет вычислений на выходе.
Метод проверен на четырёх семействах зацикленных трансформеров, и авторы сообщают о существенных и устойчивых улучшениях при полной рекуррентной глубине. LoopCD-Logits поднимает pass@1 модели Ouro-2.6B-Thinking на AIME 2024 с 61,88% до 73,33%, а LoopCD-Hidden повышает pass@1 модели Huginn на HumanEval с 22,56% до 31,71%.
Главное практическое следствие: за счёт прироста качества число рекуррентных циклов можно сократить вдвое, и результат всё равно будет не хуже базовых запусков без LoopCD на полной глубине. При этом вычисления в прямом проходе (forward FLOPs) снижаются на 22,5, 48,2%. Авторы делают вывод, что LoopCD превращает промежуточные рекуррентные состояния в полезный управляющий сигнал и даёт лучшее качество декодирования при заметно меньших затратах на вывод.
Ключевые факты
- LoopCD, метод контрастного декодирования для зацикленных трансформеров, не требующий обучения: итоговый прогноз сопоставляется с прогнозом более раннего рекуррентного цикла.
- Два варианта: LoopCD-Logits (в пространстве логитов, один дополнительный проход через выходной слой) и LoopCD-Hidden (в пространстве скрытых состояний, без дополнительных затрат на выходе).
- Ouro-2.6B-Thinking: pass@1 на AIME 2024 вырос с 61,88% до 73,33%; Huginn: pass@1 на HumanEval, с 22,56% до 31,71%.
- Вдвое меньше рекуррентных циклов при результате не хуже полноглубинных базовых запусков без LoopCD; forward FLOPs снижаются на 22,5, 48,2%.
- Проверка проведена на четырёх семействах зацикленных трансформеров.
Почему это важно
Зацикленные трансформеры экономят параметры, но ценой многократного прохода по общему блоку, поэтому вычисления на выводе у них растут с числом циклов. LoopCD предлагает использовать то, что при обычном декодировании выбрасывается, промежуточные состояния ранних циклов, как сигнал для выбора токена. Если заявленные в аннотации результаты подтвердятся, качество и экономия вычислений достигаются одновременно и без дообучения.
Кому это важно
Исследователям рекуррентных и зацикленных архитектур, а также тем, кто запускает такие модели на выводе и считает затраты. Для разработчиков, работающих с моделями вроде Ouro и Huginn, метод интересен тем, что не требует ни вспомогательных моделей, ни нового обучения.
Как это применить
Из аннотации следует выбор между двумя вариантами: LoopCD-Logits (один дополнительный выходной проход) и LoopCD-Hidden (без дополнительных затрат на выходе). Практический приём, о котором говорят авторы, сократить число рекуррентных циклов вдвое и компенсировать потерю качества контрастным декодированием. О выпуске кода, лицензии и доступности реализации в аннотации не сказано.
Можно ли доверять
Это данные из аннотации статьи; авторы и организации в ней не названы. Названы результаты только по двум тестам, AIME 2024 и HumanEval, и только для Ouro-2.6B-Thinking и Huginn; остальные из четырёх семейств моделей не перечислены. Неизвестно, усреднены ли цифры по запускам и какие настройки сэмплирования использовались. Диапазон 22,5, 48,2% по FLOPs не привязан к конкретной модели или тесту. Независимой проверки в источнике нет.
Риски и подводные камни
Снижение FLOPs в аннотации приведено только для прямого прохода: данных о задержке или реальном времени работы нет. LoopCD-Logits добавляет один выходной проход, и это надо учитывать в итоговой экономии. Широта эффекта на других моделях и задачах по аннотации оценить нельзя: показаны лишь два теста, а утверждение о «стабильных» улучшениях относится ко всем четырём семействам, цифры по которым не раскрыты.
«Каждый цикл даёт промежуточное представление, декодируемое в тот же следующий токен, однако стандартное декодирование отбрасывает более ранние состояния.»
— из аннотации статьи