Квадратичная сложность и линейно растущий KV cache в классическом механизме внимания, используемом в современных LLM, — серьёзная преграда для обработки экстремально длинных контекстов, особенно в условиях ограниченной видеопамяти на GPU. В этой статье мы рассмотрим различные варианты Linear Attention (линейного внимания) и техники дистилляции LLM в субквадратичные варианты механизма внимания, а также расскажем о наших экспериментах с GigaChat.

Введение

Традиционный механизм self-attention, известный со времён статьи «Attention Is All You Need» и активно применяемый в современных LLM, имеет существенные недостатки: квадратичную сложность вычисления относительно длины последовательности и линейно растущее потребление памяти. Это создаёт серьёзные препятствия для обработки экстремально длинных последовательностей, что может быть важно в некоторых задачах. Сложности возникают не только на этапе обучения модели на больших контекстах, но и на этапе инференса: мы упираемся в растущий KV cache и ограничения VRAM. За время существования оригинальной архитектуры трансформера было предложено много модификаций механизма внимания, так или иначе направленных на то, чтобы сделать какой-то субквадратичный аналог и снизить потребление памяти. Часть из этих модификаций успешно применяется на практике в ключевых проприетарных и open-source моделях:

• GQA, MQA, MLA — эти модификации значительно снижают размер KV cache, но он всё ещё линейно растёт с увеличением длины последовательности.

• NSA, DSA — модификации от DeepSeek, хотя и не сокращающие потребление памяти, но значительно ускоряющие вычисления. Причём DSA можно обучить в режиме post-train (после дорогой основной процедуры обучения).

• SWA (Sliding Window Attention) — классический механизм внимания, но с окном фиксированного размера. Использование окна позволяет зафиксировать размер KV cache и ускорить вычисления. За счёт использования нескольких слоёв SWA теоретически может иметь достаточно большой receptive field, но на практике применяется вкупе с чередованием с softmax-вниманием на всю длину последовательности (cм. gpt-oss, Gemma3, Mistral).

В этой же статье мы расскажем о Linear Attention — более эффективном варианте self-attention, — а также о способах получения моделей с линейным вниманием из уже предобученных чекпоинтов.

Это может быть важно в случаях, требующих обработки с помощью LLM очень длинных контекстов — скажем, более 256 тыс. токенов — при ограниченных ресурсах. Более эффективные варианты механизма внимания релевантны также в случае использования рассуждающих (reasoning) моделей в парадигме масштабирования за счёт увеличения длины контекста.

Рассмотрим субквадратичное внимание детальнее. Для начала вспомним вид классического self-attention:

где:

• — длина последовательности;

• — внутренняя размерность головы;

• — матрицы запросов, ключей, значений и выходов внимания;

• — каузальная маска.

Замечание. Для простоты повествования мы опускаем часть с масштабированием под softmax (деление на ), а также ничего не говорим о способах позиционного кодирования. По умолчанию будем считать, что используется RoPE.

В итеративной форме выход внимания на позиции можно записать следующим образом:

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

Для softmax-внимания мы можем принять, что — моделирование взаимодействия векторов и через скалярное произведение и последующее применение экспоненты.

Немного поговорим о Kernel Attention. Это модификация, в которой — то есть моделирование взаимодействия — осуществляется через скалярное произведение векторов после применения некоторого преобразования, называемого kernel’ом. Что даёт такая форма? На самом деле мы избавились от softmax и теперь можем пользоваться линейностью!

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

По сути, мы лишь избавились от softmax. Тут снова, подобно Kernel Attention’у, легко прийти к рекуррентному инференсу:

Вообще, есть множество модификаций линейного внимания. Можно добавлять константы затухания, использовать механизмы гейтинга и разные параметризации. Существует даже некоторая форма дуализма с SSM (Mamba, Mamba-2). Мы попытаемся избежать глубокого погружения в эту тему и утилитарно, исходя из чисто прикладных соображений, отметим GatedDeltaNet — модификацию, отличающуюся достаточно высокими результатами на общепринятых бенчмарках.

Обновление состояния в процессе генерации у разных вариантов линейного внимания

Можем заключить, что Kernel Attention и различные варианты Linear Attention позволяют вывести для механизма внимания рекуррентный вид, что даёт нам следующие преимущества:

• Константные затраты по памяти на этапе инференса — требуется лишь хранить матрицу и/или вектор состояний.

• Удобный рекуррентный пересчёт, линейная сложность подсчёта.

Обзор

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

Линеаризация LLM — это преобразование предобученного трансформера в линейные или субквадратичные архитектуры с константным потреблением памяти и сохранением исходного качества.

Рассмотрим некоторые из работ, вышедших в конце 2024-2025 гг., которые касаются линеаризации и важны для общего контекста.

Статья «LoLCATs: On Low-Rank Linearizing of Large Language Models»

LoLCATs — это двухэтапный метод для замены квадратичного softmax-внимания на линейный аналог. В статье используется вариант Kernel Attention поверх уже обученных матриц Q, K, V.

Рассмотрим предложенный процесс обучения подробнее.

1. Attention Transfer. Оригинальные параметры модели замораживаются. Обучаются лишь добавленные kernel'ы (feature map) поверх , минимизируя MSE между выходами оригинального внимания и линейного, подсчитанного через модифицированные векторы и . По сути, это шаг послойной дистилляции исходного внимания в линейный вариант.

2. Файнтюн через LoRA поверх . Использование LoRA-адаптеров позволяет обучать лишь 0,2 % параметров модели, что сильно экономит вычислительные ресурсы.

Вариант замены исходного механизма внимания только на Kernel Attention оказывается неконкурентоспособным. Авторы рекомендуют использовать Linear Attention в комбинации со Sliding Window Attention с небольшим окном (в экспериментах — 64 токена). В статье заявляется о довольно небольшом performance gap'e между получившейся моделью и исходным вариантом. Однако остаётся существенный разрыв на фактологическом бенчмарке MMLU и задачах на retrieval-способности модели типа Needle-in-a-Haystack. Авторы устраняют этот разрыв файнтюном на синтетических датасетах, что вызывает некоторые сомнения относительно обобщающей способности модели, особенно на длинных контекстах.

Статья «Liger: Gated Linear Recurrent Conversion without New Parameters»

Здесь раскрываются две основные идеи: механизмы гейтинга, которые особенно успешны в RNN/SSM, и переиспользование параметров уже обученной модели. Обучение линейной модели в масштабах, сопоставимых с размерами современных LLM, — это рискованно и дорого, поэтому авторы предлагают конвертировать исходный трансформер в вариант Gated Recurrent без добавления новых параметров. В работе сделан выбор в пользу Gated Linear Attention:

Для подсчёта GLA используют уже имеющиеся query, key, value. Механизм гейтинга реализован поверх key-векторов. Как и в предыдущей статье, авторы не применяют Linear Attention в чистом виде, а комбинируют его с SWA с небольшим окном. Это позволяет добиться лучшего качества. После преобразования модели применили быстрый файнтюн на небольшом количестве токенов через LoRA-адаптеры. После дообучения адаптеры можно интегрировать в исходные слои, поэтому новых параметров у нас не появляется. К сожалению, в работе не была полноценно исследована эффективность метода на длинных контекстах, хотя именно на длинных контекстах раскрываются все преимущества линейного внимания.

Статья «Lizard: An Efficient Linearization Framework for Large Language Models»

Ещё один подход к линеаризации, очень напоминающий предыдущие две статьи. В основе — всё те же GLA с kernel'ами, которые обучают поверх исходных . GLA используют вкупе со Sliding Window Attention. Из особенностей хочется отметить, что в реализации задействована механика RoPE-free, при которой позиционное кодирование через RoPE применяется в исходной модели (в учителе), а вот приближение через GLA & SWA вычисляется поверх исходных , без использования RoPE.

В статье также раскрывается идея добавления некоторых обучаемых meta-memory-токенов (логитов) — attention sink'ов. Здесь под attention sink'ом подразумевается некоторый специальный токен, стабилизирующий распределение внутри голов внимания после взятия softmax. Оказывается, довольно часто вероятностная масса внутри механизма внимания распределена неравномерно, и очень большие значения получают первые токены, вне зависимости от их семантики. Чтобы стабилизировать подсчёт SWA через softmax, и добавляют такой специальный логит, который может «взять на себя» большую вероятностную массу, при этом никак не влияя на итоговый результат через value-векторы. Для любителей разобраться в деталях оставляем несколько статей про attention sink'и и причины возникновения больших активаций: один, два, три, четыре.

Хочется отметить, что такие обучаемые attention sink'и использовались в недавнем релизе OpenAI — gpt-oss, — поскольку там тоже применяют SWA.

Процесс обучения модели Lizard снова состоит из двух шагов: аппроксимации через дистилляцию и последующего файнтюна. В результате удаётся получить в целом линеаризованные версии моделей с результатами, сопоставимыми по ключевым метрикам с бейзлайном. В качестве бейзлайнов использовались — Mistral-7B и LLaMa-8B. В таблице ниже представлены результаты различных методов, включая вышеупомянутые Liger и LoLCATs.

Авторы Lizard проверяют полученные модели на retrieval-задаче с помощью классического бенчмарка Needle-in-a-Haystack. Для этого снова используют предварительное обучение на синтетическом датасете, состоящем из примеров на 2 тыс. токенов. Качество оценивали на контекстах вплоть до 8 тыс.

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

Статья «RAttention: Minimal-Window Local-Global Attention»

Ещё одна работа на тему использования линейного внимания вместе с SWA. Однако здесь, в отличие от предыдущих статей, модель обучили с нуля. Авторы предлагают уже известную схему с Linear Attention и SWA с маленьким окном и использованием общих весов для матриц . Обучаются модели небольшого размера 3 млрд и 12 млрд на 4 тыс. контекста. В результате удаётся получить неплохих кандидатов, сопоставимых по качеству и даже где-то превосходящих бейзлайн с полным вниманием. Особо пристально авторы изучают оценку качества на контекстах свыше 4 тыс. Эксперименты показывают, что предложенная схема экстраполируется на длинные контексты лучше, чем бейзлайн, в каком-то виде наблюдается zero-shot length generalization. Статья демонстрирует перспективу обучения моделей с линейным вниманием с нуля.

Статья «Jet-Nemotron: Efficient Language Model with Post Neural Architecture Search»

В работе описан подход к линеаризации на основе Post Neural Architecture Search. Авторы отмечают, что для лучшего восстановления исходного качества модели требуется оставлять несколько слоёв softmax-внимания. При этом авторы показывают, что простое чередование «каждый k-й» даёт неоптимальный результат. Поэтому до полноценной линеаризации они предлагают провести поиск слоёв, которые следует оставить с softmax-вниманием. Для этого на основе предобученного чекпоинта обучают модель, в которой случайным образом в разных слоях переключаются между полноценным и линейным вниманием; так получили модель, которая до какой-то степени устойчива к изменениям текущей конфигурации слоёв внимания. Затем с помощью beam search авторы ищут конфигурации для этой модели, максимизирующие конкретные метрики при фиксированном вычислительном бюджете. Таким образом, они получают информацию о том, в каких слоях следует оставлять softmax-внимание при обучении.

После получения финальной конфигурации слоёв внимания авторы линеаризуют модели подобно другим работам через послойную дистилляцию и файнтюн. В качестве базовых моделей используют популярные Qwen2.5-1.5B и Qwen2.5-3B. Для дистилляции применяют 50 млрд токенов из открытого датасета и 350 млрд токенов из закрытого для финального файнтюна. В качестве линейного внимания используют модифицированный Gated DeltaNet. Также авторы приходят к выводу, что ключевой фактор, определяющий скорость генерации токенов, — размер KV cache. Поэтому после фиксации таким образом вычислительного бюджета они перебирают различные гиперпараметры: размеры представлений, количество голов и так далее, чтобы дополнительно улучшить финальное качество без изменения скорости инференса.

Общие выводы:

Изучив вышеперечисленные работы, можно сделать несколько важных промежуточных выводов, на которые мы опирались в собственных экспериментах и к которым независимо пришли в своих запусках:

• Получить линеаризованную модель в режиме post-train через дистилляцию более чем реально. При этом можно обучать лишь kernel'ы поверх или полноценный слой через дистилляцию с нуля.

• Почти всегда линеаризация делается в два этапа: послойная дистилляция и последующий файнтюн. На этапе файнтюна для вычислительной эффективности можно обойтись только LoRA-адаптерами.

• Какие-то слои могут аппроксимироваться лучше. Самый большой послойный loss обычно наблюдается на последних слоях.

• Linear Attention очень слаб в задачах на in-context retrieval. Возникают проблемы и на MMLU-like бенчмарках, требующих выбора одного из предложенных вариантов ответа.

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

• Недавний релиз Qwen3-Next иллюстрирует потенциал гибридных моделей, совмещающих в себе полноценный механизм внимания и более эффективные аналоги.

Линеаризация

В этой части статьи нам бы хотелось рассказать о собственных экспериментах по линеаризации. Мы начали с воспроизведения подхода, предложенного в статье LoLCATs, слегка модифицировав его: взяли один из промежуточных чекпоинтов нашей базовой модели GigaChat-Lite. Добавили обучаемый Hedgehog Kernel: , а также использовали SWA и attention sink. Затем сделали послойную дистилляцию с последующим LoRA-файнтюном. Потратив не так много вычислений, мы смогли получить линеаризованный аналог, немного потеряв в ключевых метриках.

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

• классический MMLU, переведённый на русский язык ruMMLU и усложнённую версию MMLU-Pro;

• HellaSwag, BoolQ, Winogrande.

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

Для предварительной оценки модели на длинных контекстах мы использовали бенчмарк LIBRA (arxiv, github) от коллег из соседней команды. Замеры на простейших задачах типа Passkey на длинном контексте показали упомянутую выше проблему in-context retrieval: мы получили там околонулевые значения даже для контекстов, на которых модель обучалась.

Сравнение метрик

Мы использовали чекпоинт, обученный на контексте в 8 тыс. токенов. Эксперименты с дистилляцией и файнтюном мы начали на небольших контекстах в 2–4 тыс. токенов, но на более длинные последовательности модель не обобщилась. После повторения экспериментов на контексте в 8 тыс. токенов мы столкнулись со всё ещё околонулевыми результатами на простейшей retrieval-задаче. Проблема отчасти решалась возвращением части полноценных softmax-слоёв. Идея, в целом, естественная и встречалась в упомянутых ранее работах. Здесь мы пробовали наивно чередовать слои и оставлять, например, каждый шестой слой модели с исходным softmax-вниманием.

Passkey LIBRA на бейзлайне, репликации LoLCATs и LoLCATs-модель с чередованием исходных слоёв

В ходе дальнейших экспериментов мы попробовали ряд идей. Во-первых, мы отказались от файнтюна через адаптеры в пользу полноценного дообучения, поскольку наши ресурсы позволяли это сделать. Во-вторых, мы понимали, что качество приближения слоёв внимания разное: где-то аппроксимация лучше, где-то хуже. Поэтому мы использовали различные группы параметров в оптимизаторе для разных слоёв, а также с помощью специального обучения предварительно снизили энтропию в слоях внимания перед последующей дистилляцией для лучшего приближения. Однако эти серии экспериментов оказались, скорее, неудачными: сходу восстановить retrieval и сложные бенчмарки на общие знания (MMLU-Pro) нам не удалось.

В результате мы пришли к гипотезе, что ряд softmax-слоёв оказывается очень важным, и заменять их невыгодно с точки зрения конечных целевых метрик. Замена таких слоёв приводит к значительному падению метрик на ключевых датасетах, включая те, что проверяют общие знания и способности на retrieval. Восстановление происходит крайне медленно, не всегда достигая исходных значений. В связи с этим мы договорились использовать чередование видов внимания в модели, но выбирать слои под замену на линейный вариант более разумным способом. Часто в разных публикациях предлагают выбирать каждый n-й слой, например каждый пятый, но в действительности такой подход не оптимален. Мы решили зафиксировать прокси-метрики для общих знаний — MMLU, а также для базовой способности к retrieval’у — Passkey. В MMLU модель должна выбрать правильный ответ на вопрос среди четырёх вариантов, а в Passkey — извлечь ключ из нескольких цифр, спрятанный в произвольном тексте. Основываясь на этих прокси-метриках, мы можем оценить важность исходного слоя и затем заменить на более эффективный вариант.

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

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

Пайплайн

Общий пайплайн для линеаризации состоит из нескольких шагов:

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

• Подбор конфигурации гибридной модели. Используя различные стратегии, пытаемся подобрать наиболее удачную конфигурацию, не просадив общие знания и способности к базовому retrieval’у. Оставляем важные softmax-слои, остальные заменяем на линейное внимание.

• Дообучение. После выбора слоёв и замены их обратно на softmax мы дообучаем нашу модель для восстановления исходных метрик. Для этого мы используем высококачественный healing set, который позволяет модели быстрее восстанавливать своё качество.

Выбор важных слоёв

Поскольку нам не избежать чередования обычного внимания с линейным, хотелось бы получить наилучшие метрики при наименьшем числе softmax-слоёв. Для этого мы реализовали два алгоритма поиска важных слоёв.

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

Мы предлагаем два алгоритма для выбора важных слоёв внимания: randomized и beam search. С их реализациями вы можете ознакомиться в нашем GitHub-репозитории.

Алгоритм randomized search

• Рандомизированный алгоритм устроен следующим образом:

• Фиксируем случайных конфигураций с softmax-слоями.

• Замеряем валидации MMLU и Passkey для получившихся конфигураций.

• Для каждого слоя считаем среднее значение метрики в запусках, где он был softmax'ом.

• Оставляем по пять самых важных слоёв для каждой метрики.

• Полученный набор дополнительно уменьшаем с помощью перебора с отсечением менее важных слоёв.

Алгоритм beam search

В алгоритме beam search мы заранее фиксируем необходимое нам количество softmax-слоёв n и количество поддерживаемых лучей , через обозначим общее количество слоёв в модели. Для каждого шага алгоритма мы поддерживаем промежуточных конфигураций: например, при после трёх шагов у нас могут получиться множества или . На каждом шаге мы будем пытаться получить новый набор множеств так, чтобы максимизировать необходимую нам метрику. Для этого к каждому множеству добавим ещё один новый слой. Так, для шага мы получим кандидатов (на самом деле меньше, поскольку можно выбирать слои с номерами строго больше, чем в выбранном множестве). Для всех кандидатов мы посчитаем выбранную метрику. Из этих кандидатов мы оставим только с наилучшими результатами. Алгоритм продолжается, пока мы не получим конфигурации с слоями, после чего мы можем взять конфигурацию с максимальной метрикой или попытаться её дополнительно уменьшить с помощью перебора.

Результат поиска слоёв

В результате для MMLU и Passkey действительно выделяются несколько слоёв. Можно заметить, что для этих метрик важны разные слои. Также интересно, что для Passkey эти слои расположены ближе к началу сети, а для MMLU — чуть дальше середины.

Для нашей модели мы выделили наиболее удачную конфигурацию из 7 слоёв, 3 из которых важны для retrieval-способностей, а 4 — для общих знаний. Такая конфигурация позволила нам уменьшить KV cache в 5 раз!

Далее приведены метрики для двух запусков с разным количеством токенов на этапе файнтюна. Более длительное обучение позволяет ещё немного улучшить интересующие нас метрики на бенчмарках на общие знания, но из-за маленького контекста обучения (4096) начинают падать retrieval-способности. Мы видим в этом ограничение текущей реализации и планируем поддержать контекстный параллелизм для улучшения результата.

Сравнение метрик

Мы использовали чекпоинт, обученный на контексте в 8 тыс. токенов. Эксперименты с дистилляцией и файнтюном мы начали на небольших контекстах в 2–4 тыс. токенов, но на более длинные последовательности модель не обобщилась. После повторения экспериментов на контексте в 8 тыс. токенов мы столкнулись со всё ещё околонулевыми результатами на простейшей retrieval-задаче. Проблема отчасти решалась возвращением части полноценных softmax-слоёв. Идея, в целом, естественная и встречалась в упомянутых ранее работах. Здесь мы пробовали наивно чередовать слои и оставлять, например, каждый шестой слой модели с исходным softmax-вниманием.

Passkey LIBRA на бейзлайне и двух вариантах линеаризации с разным числом токенов на файнтюне

Итоги

Проведённые эксперименты открывают дорогу для дальнейших постоптимизаций наших моделей. Использование более эффективного линейного варианта значительно уменьшает объём KV cache и сокращает время обработки длинных последовательностей во время prefill-стадии и генерации; гибридная архитектура помогает не терять в метриках на retrieval-задачах. Сейчас мы сконцентрированы в том числе на оценке производительности перед возможным внедрением. Помимо этого, результаты гибридных моделей свидетельствуют о возможности обучения таких моделей с нуля, что тоже может быть интересно. Мы публикуем на GitHub код для нашего инструмента выбора слоёв.

Авторы: Денис Титов, Илья Щукин.

Отдельная благодарность:

• Глебу Мезенцеву за помощь с воспроизведением LoLCATs и другими экспериментами;

• Валерию Терновскому и Ивану Оселедцу за ценные идеи;

• Евгению Косареву и команде pretrain за полезный фидбек и модели для экспериментов.