За пределами self-attention: что придёт после трансформеров

Архитектура трансформеров уже восемь лет движет ИИ. Все крупные языковые модели, большинство систем генерации изображений и всё больше моделей для аудио и видео построены на механизме self-attention из статьи «Attention Is All You Need». Но у self-attention есть фундаментальная проблема: вычислительные затраты и потребление памяти растут квадратично с длиной последовательности. Удвоили длину входа — стоимость выросла вчетверо.

Для коротких последовательностей это не важно. Но для контекстных окон в 128K токенов, к которым мы движемся, и миллионных окон, которых все хотят, это серьёзное узкое место. Волна исследований ищет альтернативы: attention residuals, которые переиспользуют вычисления между слоями, линейные варианты внимания, которые убирают квадратичную стоимость, и гибридные архитектуры, сочетающие внимание с более дешёвыми механизмами. Трансформер никуда не уходит, но его переделывают.

Почему self-attention дорогой

Чтобы разобраться в альтернативах, нужно понимать, что именно вычисляет self-attention. Для последовательности из N токенов self-attention считает оценку релевантности для каждой пары токенов. Токен 1 и токен 2, токен 1 и токен 3, ..., токен 1 и токен N, затем токен 2 со всеми остальными и так далее. Это N² пар.

import torch
import torch.nn.functional as F
def self_attention(Q, K, V):
"""
Standard self-attention.
Q, K, V: (batch, seq_len, d_model)
The attention matrix is seq_len × seq_len.
For seq_len = 1024:   ~1M entries   (manageable)
For seq_len = 32768:  ~1B entries   (expensive)
For seq_len = 131072: ~17B entries  (very expensive)
"""
d_k = Q.size(-1)
# This matmul creates the N×N attention matrix
scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
weights = F.softmax(scores, dim=-1)
return torch.matmul(weights, V)

При 4K токенов матрица внимания содержит 16 миллионов элементов — для современных GPU это не проблема. При 128K токенов — уже 16 миллиардов элементов. При миллионе токенов — больше триллиона. Даже Flash Attention (который не уменьшает объём вычислений, но заметно улучшает паттерны доступа к памяти) в конце концов проигрывает квадратичному росту.

Поэтому ранние трансформеры ограничивались 512 или 1024 токенами. Каждое новое поколение железа и оптимизаций поднимает планку, но мы упираемся в математическую стену. Линейное масштабирование (O(N)) принципиально лучше квадратичного (O(N²)), и именно к нему стремится большинство альтернативных архитектур.

Attention Residuals: переиспользуем то, что уже посчитано

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

Наблюдение такое: в глубокой модели (скажем, из 32 слоёв) паттерны внимания соседних слоёв часто удивительно похожи. Слой 15 и слой 16 обычно смотрят на похожие позиции с небольшими поправками. Считать полную матрицу N² с нуля на каждом слое избыточно — значительная часть работы уже сделана слоем назад.

Attention residuals используют это, вычисляя «остаточный» паттерн внимания: разницу между тем, на что этот слой хочет смотреть, и тем, что посчитал предыдущий слой. Если разница небольшая (а в средних слоях обычно так и есть), вычисления дешевле. Итоговый паттерн внимания — это паттерн предыдущего слоя плюс остаток текущего.

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

На практике attention residuals снижают вычислительную стоимость внимания на 30–50% в средних слоях глубоких моделей с минимальным влиянием на качество. Первые и последние несколько слоёв всё равно требуют полного расчёта внимания (их паттерны более специфичны), а средние слои, которых большинство, получают заметное ускорение.

Линейное внимание: избавляемся от квадратичной стоимости

Варианты линейного внимания пытаются переформулировать attention так, чтобы он масштабировался как O(N), а не O(N²). Общий подход: вместо явного вычисления матрицы внимания N×N найти способ получить тот же (или примерно тот же) результат линейными операциями.

Математический трюк опирается на разложение softmax в ядро. Стандартное внимание считает softmax(QK^T)V. Если заменить softmax на другую ядерную функцию, которую можно разложить как φ(Q) · φ(K)^T, порядок вычислений можно переставить: вместо (φ(Q) · φ(K)^T) · V (с промежуточной матрицей N×N) считаем φ(Q) · (φ(K)^T · V) (с промежуточной матрицей d×d, где d — размерность модели). Поскольку для длинных последовательностей d << N, это заметно дешевле.

def linear_attention(Q, K, V, feature_map=None):
"""
Linear attention via kernel feature maps.
Cost: O(N * d^2) instead of O(N^2 * d)
"""
if feature_map is None:
# ELU+1 is a common choice (from Katharopoulos et al.)
feature_map = lambda x: F.elu(x) + 1
Q = feature_map(Q)  # (batch, seq_len, d)
K = feature_map(K)  # (batch, seq_len, d)
# Key insight: compute K^T @ V first (d × d matrix)
# instead of Q @ K^T first (N × N matrix)
KV = torch.einsum('bnd,bnm->bdm', K, V)  # (batch, d, d)
# Then multiply by Q
output = torch.einsum('bnd,bdm->bnm', Q, KV)  # (batch, N, d)
# Normalize
Z = torch.einsum('bnd,bd->bn', Q, K.sum(dim=1))  # normalization
output = output / Z.unsqueeze(-1)
return output

Подвох в том, что замена softmax на другую ядерную функцию меняет распределение внимания, и модели, обученные с softmax-вниманием, не обязательно хорошо переносятся на линейное. Разрыв в качестве заметно сократился — современные линейные варианты достигают 95–98% качества softmax-внимания, — но он сохраняется, особенно в задачах, требующих точного поиска на больших дистанциях.

State Space Models: другая парадигма

State space models (SSM), такие как Mamba, идут принципиально другим путём. Вместо попарных связей между токенами они обрабатывают последовательность через рекуррентность, поддерживая скрытое состояние фиксированного размера, которое обновляется на каждом токене. Это по своей природе O(N): обработка вдвое большего числа токенов занимает вдвое больше времени, а не в четыре раза.

Изюминка современных SSM в том, что параметры рекуррентности зависят от входных данных (selective state spaces). Это даёт модели форму внимания, основанного на содержимом: она может «выбирать», что запомнить, а что забыть, без квадратичной стоимости. Модели в стиле Mamba на многих бенчмарках не уступают трансформерам, при этом заметно быстрее на длинных последовательностях.

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

Гибридные архитектуры: прагматичный путь

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

• Полное внимание для глобальных рассуждений. Некоторым слоям нужно смотреть на всю последовательность — находить релевантный контекст в тысячах токенов от текущей позиции. Эти слои получают стандартное (возможно, оптимизированное через Flash) self-attention.
• Локальное внимание для ближайшего контекста. Многие слои в основном смотрят на соседние токены (скользящее окно). Фиксированное окно в 256–1024 токена снижает стоимость до O(N·W), где W — размер окна.
• Линейное внимание для широкого контекста. Некоторым слоям нужно агрегировать информацию по всей последовательности, но не нужны точные веса внимания. Линейное внимание даёт это за O(N).
• Слои SSM для последовательной обработки. Слои в стиле Mamba могут эффективно обрабатывать последовательные зависимости без какого-либо вычисления внимания.

Такие модели, как Jamba (AI21), и различные исследовательские архитектуры чередуют эти механизмы в зависимости от роли слоя. Ранние слои используют локальное внимание (обрабатывая синтаксис и локальные паттерны). Средние слои используют линейное внимание или SSM (строя более широкие представления). Несколько стратегически расположенных слоёв используют полное внимание (глобальные рассуждения и поиск). Это даёт почти линейное масштабирование в целом, сохраняя качество модели, которому нужно некоторое полное внимание.

На что стоит обратить внимание разработчикам

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

• Контекстные окна будут расти. По мере снижения стоимости внимания окна расширяются. Это меняет архитектуру приложений: вместо сложных RAG-пайплайнов, чтобы уместить релевантный контекст в окно 4K, можно просто запихнуть всё в промпт на 1M токенов. Простота привлекает, но задержки и стоимость зависят от архитектуры по-разному.
• Профили задержек меняются. У трансформеров задержка до определённой точки примерно постоянна, а потом растёт квадратично. У моделей с линейным вниманием и SSM рост задержки более плавный, линейный. Для приложений, где важно время отклика, нужно понимать, как модель масштабируется.
• Разница в качестве зависит от задачи. Модели с линейным вниманием могут немного уступать в задачах, требующих точного извлечения из конкретных позиций в длинном контексте («что было третьим пунктом в списке на странице 47?»). В задачах на общее понимание они работают не хуже. Знайте свой сценарий.
• Оптимизация инференса важнее. По мере усложнения архитектур (смешивания разных типов внимания) движкам инференса нужно эффективно обрабатывать гетерогенные вычисления. vLLM, TensorRT-LLM и похожие фреймворки адаптируются, но кастомные архитектуры могут поначалу не поддерживаться.

Трансформер не заменяют — его развивают. Self-attention остаётся самым выразительным механизмом для моделирования связей между токенами. Но его не нужно использовать везде, в каждом слое, с полной стоимостью N². Модели ближайших лет будут применять внимание точечно: с полной точностью там, где это важнее всего, и более дешёвыми альтернативами во всём остальном. В результате модели станут быстрее, будут работать с более длинным контекстом и дешевле в эксплуатации, при этом не уступая текущему качеству или превосходя его. Это стоит внимания.