Cola2 обходит Llama-3-8B на HellaSwag, потратив в 43 раза меньше токенов

Команда исследователей показала Cola2, фреймворк иерархического языкового моделирования, в котором смысл живёт в непрерывном латентном пространстве, а текст собирается дискретным декодером. На 346,03 млрд исходных токенов ещё не сошедшийся чекпоинт набрал в среднем 60,10 по восьми задачам и обогнал Llama-3-8B на HellaSwag и StoryCloze. Заявленный бюджет предобучения Llama-3 при этом больше более чем в 43 раза.

Что именно сжимает латентное пространство

Идея иерархической генерации давно витает в воздухе: сначала модель планирует смысл крупными мазками, потом реализует его в токенах. Слабое место всегда одно. Если латенты сжимают не то или не так, декодер либо игнорирует их, либо превращается в простой копировальный аппарат.

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

Диффузионный декодер и баланс между копированием и самостоятельностью

Декодер здесь блочно-каузальный, маскированный диффузионный. Его обучают на трёх типах условий: с чистыми латентами, с зашумлёнными и с полностью выброшенными. Это даёт возможность применить classifier-free guidance. Меняя силу guidance, можно сдвигать поведение модели между точным восстановлением смысла из латента и автономной генерацией.

Отдельная инженерная деталь касается потоков градиентов. Градиенты предсказания токенов идут через латенты в энкодер, и представления подстраиваются под нужды декодера. А вот обучение приора на Flow Matching отвязано от обучения представлений через stop-gradient. Приор учится предсказывать следующие латенты, но не тянет за собой энкодер и не разрушает пространство, которое тот строит.

Есть и тонкость на инференсе. Сгенерированный текст заново кодируется при фиксированной разметке спанов, чтобы история латентов совпадала с реальным текстовым контекстом, а не с тем, что модель «думала» раньше. Без этого шага ошибки между непрерывной и дискретной частями копились бы от блока к блоку.

Цифры и что с ними не так

• Обучающий бюджет в основных экспериментах: около 210 млрд токенов, пять downstream-задач.
• Масштабирование: авторы описывают его как конкурентное и ненасыщенное.
• Крупный запуск: 346,03 млрд исходных токенов, чекпоинт ещё не сошёлся.
• Среднее по восьми задачам: 60,10.
• Обгоняет Llama-3-8B на HellaSwag и StoryCloze при бюджете предобучения, меньшем более чем в 43 раза.

Сравнение с Llama-3-8B стоит читать аккуратно. HellaSwag и StoryCloze измеряют правдоподобие продолжения и здравый смысл, то есть как раз то, что хорошо ловит модель с явным слоем смысловой структуры. На знаниевых и рассуждающих бенчмарках картина может оказаться другой, и по тексту работы победа заявлена только на двух задачах из восьми. Зато кривая масштабирования не выходит на плато, а значит, потенциал на больших бюджетах не исчерпан.

Где это может пригодиться

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

Для исследователей диффузионных языковых моделей Cola2 даёт готовый рецепт стыковки непрерывной и дискретной частей: отрицательная обратная связь для сжатия, три режима условий для декодера, stop-gradient для приора, пересчёт латентов при генерации. Каждый из этих элементов можно проверять и переносить отдельно.

Для индустрии это сигнал, что альтернативы чисто авторегрессионной схеме набирают вес. Пока речь о модели в районе 8B-класса и незавершённом обучении, но если масштабирование подтвердится на десятках и сотнях миллиардов параметров, иерархические латентные модели станут серьёзным кандидатом на следующий шаг после классического next-token prediction.