Priming: преобразование Transformer в более быстрые гибридные SSM-модели при затратах на обучение менее 0,5%

ИИ-агенты становятся более долгоживущими, автономными и требовательными к вычислительным ресурсам.

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

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

К ним относятся:

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

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

Для этого могут потребоваться триллионы токенов и сотни тысяч GPU-часов.

Исследователи Amazon Web Services предлагают иной подход в статье «Priming: Hybrid State Space Models From Pre-trained Transformers».

Вместо повторного обучения всему с нуля Priming переносит знания из существующего предобученного Transformer в новую гибридную архитектуру.

Полный конвейер использует менее 0,5% исходного бюджета токенов на предобучение модели, сокращая стоимость исследования архитектур более чем на 99%.

Исследователи применяют метод к:

• Qwen;
• Llama;
• Mistral;
• плотным моделям;
• моделям mixture-of-experts;
• системам масштаба 8B и 32B.

В наиболее сильных результатах обслуживания длинного контекста замена части Transformer на слои пространства состояний Gated KalmaNet существенно сокращает объем KV-кеша и повышает пропускную способность декодирования.

Примированная гибридная 32B-модель быстрее достигает порогов точности AIME, чем ее исходный Transformer Qwen3.

Результат — не просто более быстрая модель.

Priming превращает проектирование архитектуры в задачу переноса знаний, а не в задачу полного предобучения.

Ключевые выводы

• Priming повторно использует веса Transformer для инициализации гибридных моделей attention–SSM, требуя менее 0,5% исходного бюджета токенов на предобучение.
• Гибридная 32B-модель Gated KalmaNet может соответствовать качеству сопоставимого дообученного Transformer, используя примерно вдвое меньше памяти KV-кеша.
• При контексте 128K гибрид 32B GKA достигает примерно 2,06× пропускной способности декодирования при меньшем числе итераций решателя; 8B-модели достигают до примерно 2,29×.
• В контролируемых сравнениях наблюдаемая иерархия производительности на длинном контексте в целом выглядит так: GKA > GDN > Mamba-2.
• Гибриды на основе SSM превосходят гибриды с вниманием со скользящим окном на тестах длинного контекста, даже когда SWA получает существенно больше памяти.
• GKA предоставляет во время выполнения регулятор соотношения скорости и точности, поскольку его итеративный решатель может использовать больше или меньше итераций без переобучения.
• Композиция гибридных состояний может расширить нативный контекст 128K до 256K без дополнительного обучения.

Механизм внимания

Чтобы понять, почему работает Priming, полезно начать с обычного внимания в Transformer.

Внимание использует три обучаемые проекции:

• запрос — что ищет текущий токен;
• ключ — что предлагает другой токен;
• значение — информацию, которую этот токен может внести.

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

Токены с более сильным совпадением вносят большую часть своих значений в выходной результат.

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

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

Этот кеш чрезвычайно полезен.

Но он же является источником существенных затрат памяти.

По мере увеличения входной последовательности растет и кеш.

Для коротких диалогов это может быть приемлемо.

Для рассуждений на 128K токенов, крупных кодовых баз или постоянных агентов он становится одним из главных ограничений для размера батча и эффективности обслуживания.

Два типа слоев памяти: эйдетическая и затухающая

Команда Amazon описывает полное внимание как форму эйдетической памяти.

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

Это свойство дословного доступа делает возможными многие полезные способности:

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

Но не каждому слою Transformer нужен полный доступ ко всей истории.

Некоторые слои преимущественно выполняют локальные операции:

• отслеживают текущую фразу;
• копируют близлежащую информацию;
• обновляют локальный синтаксис;
• смешивают недавний контекст.

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

Может быть достаточно сжатого состояния фиксированного размера.

Авторы называют это затухающей памятью.

Таким образом, гибридная модель может разделить работу:

Цель не в том, чтобы полностью убрать внимание.

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

Priming: знания не привязаны к архитектуре

Центральная идея Priming состоит в том, что предобученные знания не обязаны оставаться привязанными к точной архитектуре, которая их изначально усвоила.

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

Обе модели можно понимать как применяющие операцию смешивания к последовательности значений.

Слой внимания строит матрицу смешивания из проекций запросов и ключей.

Слой пространства состояний создает связанную форму смешивания через рекуррентное состояние фиксированной размерности.

Если SSM может достаточно близко аппроксимировать поведение конкретного слоя внимания, значительная часть уже изученной структуры Transformer может быть повторно использована.

Сюда входят обученные проекции исходного слоя.

Новой модели нужно изучить главным образом параметры, специфичные для динамики пространства состояний.

Какие слои внимания можно заменить?

Не каждый слой внимания одинаково пригоден для преобразования.

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

Слой, который в основном ведет себя локально, — гораздо более подходящий кандидат для SSM.

Исследователи используют недорогой эмпирический прокси-метод:

1. временно заменить рассматриваемый слой внимания на внимание со скользящим окном;
2. сохранить только часть исходного контекста;
3. измерить, насколько ухудшается качество модели.

Если производительность почти не меняется, слой, вероятно, не сильно зависит от точного дальнего воспроизведения.

Это делает его перспективным кандидатом для преобразования в SSM-слой с фиксированным состоянием.

Прокси-метод связан с величиной из теории динамических систем, называемой рангом Ханкеля матрицы смешивания слоя.

Неформально ранг Ханкеля отражает, какой объем состояния необходим для точного воспроизведения входно-выходного поведения слоя.

Слой с более простой временной динамикой часто можно представить через меньший узкий канал состояния.

Три этапа Priming

Полный конвейер Priming состоит из трех основных этапов.

1. Выбор слоев

Определить слои внимания, которые можно сжать в представления состояния фиксированного размера без значительной потери качества.

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

2. Перенос весов и выравнивание

Инициализировать каждую заменяющую SSM весами того слоя Transformer, который она заменяет.

Затем выполнить короткую фазу дистилляции, чтобы SSM научилась воспроизводить поведение исходного слоя.

3. Сквозная адаптация

Дообучить полную гибридную модель на предсказании следующего токена или на целевом downstream-рецепте.

На этом этапе слои внимания и SSM могут совместно адаптироваться как единая модель.

Priming выбирает локальные слои внимания, переносит их проекции в SSM-слои, выравнивает поведение через дистилляцию, а затем дообучает полный гибрид.

Весь процесс потребляет менее 0,5% исходного бюджета токенов на предобучение модели.

Именно это делает практичными контролируемые эксперименты с архитектурами.

Вместо того чтобы тратить полный запуск предобучения на каждый вариант дизайна, исследователи могут запускать каждый вариант с одной и той же сильной предобученной инициализации.

Контролируемое сравнение слоев затухающей памяти

Когда преобразование архитектуры становится недорогим, проще исследовать еще один вопрос:

Если данные, бюджет обучения и доля гибридной архитектуры неизменны, какой SSM-слой работает лучше?

В статье сравниваются три дизайна пространства состояний.

Mamba-2

Mamba-2 использует относительно однородный механизм затухания своего состояния.

При поступлении новой информации модель обновляет состояние, тогда как более старая информация затухает.

Механизм прост и эффективен, но процесс забывания сравнительно неизбирателен.

Gated DeltaNet

Gated DeltaNet добавляет более целевые обновления.

Вместо равномерного затухания всего содержимого она может избирательно стирать или переписывать значение, связанное с определенным ключом.

Это дает ей более выразительный контроль над тем, что сохраняется.

Gated KalmaNet

Gated KalmaNet идет дальше.

Ее обновление зависит не только от текущего ключа, но и от статистик, накопленных за всю историю.

Слой построен вокруг итеративной онлайн-формулировки гребневой регрессии, вдохновленной фильтрацией Калмана.

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

В контролируемых экспериментах с гибридами 8B GKA в целом опережает GDN, а GDN, в свою очередь, стабильно превосходит Mamba-2.

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

GKA > GDN > Mamba-2

Различие усиливается в задачах рассуждений с длинным контекстом, где качество сжатой памяти становится важнее.

В масштабе 32B гибридная модель GKA остается в пределах примерно 1% от сопоставимого Transformer, дообученного на тех же данных, при этом используя примерно вдвое меньший KV-кеш.

Затухающая память превосходит дословную оконную память

Один из наиболее неочевидных выводов получен при сравнении SSM с вниманием со скользящим окном.

Внимание со скользящим окном сохраняет точные представления только самых последних токенов.

На первый взгляд это кажется привлекательным.

Ничто внутри окна не сжимается и не теряется.

Но окно отбрасывает все, что старше его фиксированного диапазона.

SSM делает противоположное.

Она сжимает всю историю в с потерями состояние фиксированного размера.

Авторы обнаружили, что сжатая сводка всей истории часто работает лучше.

Гибрид на основе SSM стабильно превосходит гибрид со скользящим окном во всех представленных задачах длинного контекста.

Каждый гибрид на основе SSM в исследовании превосходил свой вариант SWA на бенчмарках длинного контекста.

Это оставалось верным, даже когда внимание со скользящим окном получало примерно в два-четыре раза больше памяти, чем слой пространства состояний.

Результат предполагает полезный принцип:

Компактная сводка всего прошлого может быть полезнее безупречной записи только недавнего прошлого.

Это важно для постоянных агентов, где старый контекст может по-прежнему содержать важную информацию, даже если он уже не является недавним.

Регулятор соотношения скорости и точности во время выполнения

GKA обладает еще одним свойством, которого Mamba-2 и GDN не предлагают аналогичным образом: переменными вычислениями во время инференса.

Mamba-2 и GDN применяют по существу один и тот же вычислительный бюджет к каждому токену.

GKA итеративно решает внутреннюю задачу регрессии.

Число итераций можно менять во время обслуживания.

Переобучение не требуется.

Изменение числа итераций GKA перемещает модель вдоль фронтира скорости и точности во время обслуживания.

В экспериментах статьи на AIME 2025 изменение GKA с обычных 30 итераций до 10 дает примерно:

~10% более быстрое декодирование
~1% потери точности

Это означает, что одна обученная модель может поддерживать разные уровни задержки.

Система обслуживания может использовать:

• меньше итераций, когда важнее всего скорость;
• больше итераций, когда важнее максимальное качество.

Это превращает усилие инференса в параметр времени выполнения, а не в свойство, зафиксированное при обучении.

Более быстрые рассуждения, а не только более быстрые токены

Сырая пропускная способность токенов может вводить в заблуждение.

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

Поэтому исследователи также измеряют время до целевой точности по настенным часам.

Они сравнивают примированную гибридную 32B GKA с ее исходной моделью Qwen3-32B, используя один и тот же рецепт рассуждений.

Обе системы генерируют цепочки рассуждений сходной длины:

• около 14 000 медианных токенов на полном наборе AIME 2025;
• около 28 000 на наиболее сложном подмножестве.

Гибрид выигрывает не за счет меньшей глубины размышлений.

Он выигрывает потому, что может обрабатывать сопоставимые цепочки рассуждений с более высокой пропускной способностью.

При реалистичных уровнях конкурентной нагрузки гибрид 32B достигает порогов целевой точности:

• примерно в 1,2×–1,4× быстрее на полном бенчмарке;
• до примерно в 1,6× быстрее на самых сложных задачах.

Главная причина — память.

Замена половины слоев внимания SSM-слоями с фиксированным состоянием примерно вдвое снижает нагрузку KV-кеша на последовательность.

Это позволяет системе обслуживания запускать больше последовательностей одновременно на том же оборудовании.

Преимущество в пропускной способности увеличивается по мере роста контекста и давления со стороны KV-кеша.

Для 32B-модели с 10 итерациями GKA заявленный выигрыш пропускной способности декодирования составляет приблизительно:

8B-модели демонстрируют еще больший выигрыш: примерно от 1,94× на 16K до 2,29× на 128K.

Поэтому аннотация статьи на arXiv суммирует более широкий результат как до примерно 2,3× более высокой пропускной способности декодирования.

Последствия для RL-дообучения

Выигрыш в пропускной способности может быть важен и во время дообучения с подкреплением.

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

Amazon сообщает о росте скорости получения данных для RL-обучения на 20%–100% в протестированных условиях.

Это не меняет сам алгоритм RL.

Это меняет скорость, с которой модель может производить данные, необходимые процессу RL.

Качество на бенчмарках рассуждений

Выпущенные гибридные модели оптимизированы не только для скорости.

В статье также сравнивается качество их рассуждений с исходными Transformer Qwen3.

Примированные гибриды GKA остаются конкурентоспособными с исходными Transformer или превосходят их на представленных бенчмарках рассуждений.

В аннотации статьи на arXiv сообщается, что гибрид 32B GKA улучшает средний результат рассуждений на 3,8 балла по сравнению со своим исходным чекпоинтом Qwen3-32B, оставаясь в пределах 1% от Transformer, дообученного на тех же данных.

Важный методологический момент состоит в том, что гибридная архитектура оценивается не после полного независимого запуска предобучения.

Она наследует большую часть знаний от существующей модели и получает лишь относительно небольшой объем обучения на выравнивание и адаптацию.

За пределами окна обучения: расширение контекста через композицию состояний

Примированные гибриды нативно поддерживают 128K токенов.

Но реальные сессии агентов, большие репозитории и длинные юридические или технические документы могут превышать этот объем.

Обычно расширение контекста может требовать еще одного цикла обучения на длинном контексте.

Вместо этого исследователи адаптируют идею из своей более ранней работы PICASO.

Почему расширение гибридного контекста сложнее

Чистые Transformer могут использовать методы масштабирования позиций, такие как YaRN.

Чистые SSM могут компонировать рекуррентные состояния.

Гибридная модель содержит оба компонента:

• KV-кеши от внимания;
• рекуррентные состояния от SSM-слоев.

Обе системы памяти нужно расширять совместно.

Процедура без обучения

Предложенный метод состоит в следующем:

1. Разделить длинный вход на фрагменты, короче нативного окна контекста модели.
2. Добавить один и тот же общий префикс к каждому фрагменту.
3. Обработать фрагменты независимо и параллельно.
4. Объединить KV-кеши Transformer из фрагментов конкатенацией.
5. Объединить состояния SSM с помощью метода композиции состояний на основе взвешенного усреднения, представленного в PICASO.
6. Продолжить обработку финального суффикса, используя объединенную гибридную память.

Длинные входы обрабатываются как независимые фрагменты; кеши внимания конкатенируются, а состояния SSM объединяются.

Это расширяет примированные гибридные модели с:

Нативный контекст 128K
        ↓
Скомпонованный контекст 256K

без дополнительного обучения.

В статье сообщается о существенном сохранении точности на:

• задачах извлечения RULER NIAH;
• задачах рассуждений BABILong.

Тот же порядок выразительности остается заметным после расширения:

GKA > GDN > Mamba-2

Пространство гибридных архитектур теперь открыто

Главная ценность Priming может быть шире любого отдельного бенчмарка.

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

Это делало чрезвычайно дорогими такие вопросы, как:

• Какой тип слоя должен заменять внимание?
• Сколько слоев внимания можно заменить?
• Где следует размещать SSM-слои?
• Какое гибридное соотношение обеспечивает лучший компромисс?
• Какой дизайн SSM дает лучшее качество на байт памяти?
• Как модели следует балансировать FLOPs инференса и задержку?

Priming делает эти эксперименты значительно дешевле, потому что каждый кандидат может начинаться с одной и той же предобученной основы.

Amazon выпустила реализацию под названием Hybrid Model Factory по лицензии Apache 2.0.

Релиз включает:

• примированные гибридные модели 8B и 32B;
• варианты, настроенные на инструкции и рассуждения;
• код обучения Priming;
• поддержку обучения с параллелизмом последовательностей для длинного контекста;
• оптимизированные ядра Triton;
• реализации GKA, GDN, Mamba-2 и дополнительных гибридных слоев;
• плагин обслуживания vLLM;
• непрерывную обработку батчей;
• фрагментированный prefill;
• параллелизм тензоров;
• чекпоинты моделей Hugging Face.

Репозиторий в настоящее время поддерживает несколько семейств предобученных Transformer и классов моделей, включая Qwen, модели семейства Llama, модели семейства Mistral, плотные архитектуры и варианты MoE.

Для обслуживания длинного контекста практическая идея проста:

Сохранять внимание там, где важно точное воспроизведение
        +
Использовать память фиксированного состояния там, где оно не требуется
        ↓
Меньший KV-кеш
        ↓
Больше одновременных последовательностей
        ↓
Более высокая пропускная способность

Для исследований архитектур выгода иная:

Сильный предобученный Transformer
        ↓
Недорого примировать множество гибридных вариантов
        ↓
Сравнивать архитектуры в контролируемых условиях

Общая тема заключается в том, что знания можно повторно использовать даже при изменении архитектуры памяти.

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

Что такое Priming в гибридных языковых моделях?

Priming — это метод инициализации новой гибридной архитектуры attention–SSM на основе предобученного Transformer. Вместо повторного изучения всех знаний модели с нуля он переносит существующие веса, выравнивает заменяющие SSM-слои, а затем выполняет короткую сквозную адаптацию.

Сколько обучения требуется для Priming?

В статье сообщается, что полный конвейер Priming использует менее 0,5% исходного бюджета токенов на предобучение модели. Это соответствует более чем 99-процентному сокращению потребностей в данных масштаба предобучения по сравнению с созданием новой архитектуры с нуля.

Удаляет ли Priming все слои внимания Transformer?

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

Какой слой пространства состояний показал лучший результат?

В контролируемых экспериментах Gated KalmaNet в целом показывает лучший результат, за ним следует Gated DeltaNet, а затем Mamba-2. Разрыв становится существеннее в задачах рассуждений с длинным контекстом.

Насколько быстрее примированные гибридные модели?

Заявленный выигрыш зависит от размера модели, длины контекста и числа итераций GKA. Для 32B-модели пропускная способность декодирования достигает примерно 2,06× от исходного Transformer при контексте 128K и 10 итерациях GKA; 8B-модели достигают до примерно 2,29×.

Может ли Priming расширять длину контекста без обучения?

Priming сам по себе создает гибридную модель, тогда как отдельный метод композиции гибридных состояний расширяет ее контекст без дополнительного обучения. В статье продемонстрировано расширение с нативного контекста 128K до 256K за счет композиции KV-кешей внимания и состояний SSM.

Являются ли модели Priming и код открытыми?

Да. Amazon выпустила Hybrid Model Factory по лицензии Apache 2.0 вместе с кодом обучения, чекпоинтами моделей, оптимизированными ядрами, инструментами параллелизма последовательностей и плагином обслуживания vLLM.

Могу ли я обслуживать выпущенные модели с помощью vLLM?

Да. Hybrid Model Factory включает плагин vLLM, предназначенный для промышленного обслуживания ее пользовательских гибридных архитектур. Репозиторий документирует непрерывную обработку батчей, параллелизм тензоров, кеширование префиксов и развертывание на основе Docker.

Связанные инструменты

• Hybrid Model Factory: открытый инструментарий Amazon для Priming, обучения и обслуживания гибридных моделей attention–SSM.
• Примированные гибридные модели на Hugging Face: выпущенная Amazon коллекция примированных моделей 8B и 32B.
• vLLM: фреймворк высокопроизводительного обслуживания LLM, используемый плагином инференса Hybrid Model Factory.
• LlamaFactory: фреймворк обучения, расширенный Hybrid Model Factory для рабочих процессов обучения гибридных моделей.
• Triton: язык программирования GPU, используемый для оптимизированных ядер гибридных моделей.
• Qwen: семейство моделей, используемое в качестве основной предобученной основы для выпущенного Amazon набора моделей Primed Hybrid.

Связанные ссылки

• Исходная статья Amazon Science: исходная статья Луки Занкато.
• Статья о Priming на arXiv: полная исследовательская статья, описывающая метод Priming и контролируемые эксперименты с гибридными архитектурами.
• Репозиторий Hybrid Model Factory на GitHub: код обучения, плагин инференса, документация, ядра и конфигурации выпущенных моделей.
• Документация конвейера Priming: пошаговая документация по преобразованию Transformer в гибрид, дистилляции и адаптации.
• Руководство по инференсу гибридных моделей: официальные инструкции по обслуживанию через vLLM и Hugging Face.
• Статья PICASO: более ранняя работа о композиции состояний, на которой основан метод расширения контекста для гибридных моделей.
• Статья о Gated KalmaNet: исследовательская статья, описывающая механизм затухающей памяти и переменных вычислений GKA.
• Статья Mamba-2 / State Space Duality: статья ICML 2024, связывающая Transformer и SSM и представляющая Mamba-2.

Итог

Priming превращает исследование гибридных моделей из задачи полного предобучения в задачу переноса знаний. Существующие веса Transformer становятся отправной точкой, выбранные слои внимания преобразуются в SSM-слои, а короткий процесс выравнивания и адаптации восстанавливает качество модели, используя менее 0,5% исходного бюджета токенов на предобучение.

Полученные гибриды снижают нагрузку KV-кеша, сохраняя внимание там, где полезно точное воспроизведение. В представленных экспериментах Gated KalmaNet обеспечивает наиболее сильные результаты затухающей памяти, предоставляет управление соотношением скорости и точности во время выполнения и дает значительный выигрыш пропускной способности на длинном контексте.

Ту же архитектуру можно вывести за пределы ее нативного окна обучения: композиция состояний расширяет гибриды 128K до 256K без дополнительного обучения, комбинируя кеши внимания и рекуррентные состояния SSM.

Ключевая идея проста: усвоенные моделью знания не обязаны быть привязаны к архитектуре, которая первоначально их усвоила.