Sparse-трансформеры: как урезать O(n²) внимания

Почему квадрат в attention упирается в стену

Стандартный self-attention считает попарные веса между всеми токенами. Для последовательности длиной n это n × n матрица: при n = 1024 это около миллиона элементов на одну голову, при n = 16384 — уже 268 миллионов. Память и время растут как O(n²), и именно этот квадрат мешает подавать в модель длинные тексты, аудиодорожки или развёрнутые в последовательность изображения.

Разреженное внимание (sparse attention) решает проблему в лоб: каждый токен смотрит не на все предыдущие, а на заранее выбранное подмножество. Если это подмножество имеет размер порядка √n или log n, сложность падает до O(n·√n) или O(n·log n). На длинных последовательностях разница между n² и n·√n — это разница между «не помещается в память GPU» и «обучается за разумное время».

Что именно разрежают

В плотном attention для позиции i вычисляется softmax по ключам всех позиций j ≤ i (в авторегрессионном режиме). Разреженная схема заменяет полный набор j на фиксированный шаблон связей. Ключевая идея: если шаблоны нескольких слоёв или голов подобрать так, чтобы их композиция покрывала всю последовательность, каждый токен всё равно сможет косвенно получить информацию из любой точки — просто за несколько шагов, а не за один.

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

Основные шаблоны разреженности

Исторически подход популяризировала работа OpenAI «Generating Long Sequences with Sparse Transformers» (2019), где для генерации изображений и аудио использовались две фиксированные схемы. С тех пор появилось семейство методов, различающихся тем, какие связи считать.

• Strided (шаговое) внимание. Токен смотрит на несколько ближайших соседей плюс на позиции с фиксированным шагом. Хорошо ложится на данные с периодической структурой — например, изображение, развёрнутое построчно, где шаг равен ширине строки.
• Fixed (блочное фиксированное). Последовательность делится на блоки; внутри блока внимание локальное, а между блоками работают выделенные «summary»-позиции. Подходит там, где нет очевидной периодичности.
• Sliding window (скользящее окно). Каждый токен видит w соседей слева и справа. Стек таких слоёв даёт рецептивное поле, растущее линейно с глубиной — идея из Longformer.
• Global tokens (глобальные токены). Несколько специальных позиций (например, служебный токен последовательности) видят всех и видны всем. Их добавляют поверх локальных схем, чтобы сохранить канал для «дальнобойной» информации.

Комбинация — норма, а не исключение

На практике одну схему почти не используют. Longformer сочетает скользящее окно с глобальными токенами, BigBird добавляет к окну ещё и случайные связи, доказывая при этом, что такой шаблон сохраняет выразительность полного трансформера. Разные головы в одном слое могут нести разные шаблоны: часть — локальные, часть — strided.

Сравнение подходов

Отдельно стоит FlashAttention: это не разреженный, а точный алгоритм, который переписывает вычисление attention так, чтобы не материализовать всю n × n матрицу в медленной памяти. FLOPs остаются квадратичными, но потребление памяти становится линейным, а на практике он часто оказывается быстрее наивной разреженности. Прежде чем городить кастомный sparse-паттерн, проверьте, не хватит ли вам плотного attention на FlashAttention — код проще, а качество не страдает.

Как это выглядит в коде

Разреженность на уровне модели — это чаще всего маска, определяющая, какие пары (i, j) участвуют в softmax. Простейший вариант скользящего окна с причинностью:

import torch

def sliding_window_mask(n, window):
    # True = разрешённая связь
    idx = torch.arange(n)
    i = idx.unsqueeze(1)
    j = idx.unsqueeze(0)
    causal = j <= i
    local = (i - j) < window
    return causal & local

def masked_attention(q, k, v, mask):
    scores = q @ k.transpose(-2, -1) / q.shape[-1] ** 0.5
    scores = scores.masked_fill(~mask, float("-inf"))
    weights = torch.softmax(scores, dim=-1)
    return weights @ v

n, d, window = 8, 4, 3
q = torch.randn(n, d)
k = torch.randn(n, d)
v = torch.randn(n, d)
out = masked_attention(q, k, v, sliding_window_mask(n, window))
print(out.shape)  # torch.Size([8, 4])

Важная оговорка: маскирование через -inf не даёт выигрыша по памяти — матрица scores всё равно n × n. Так удобно прототипировать и проверять корректность. Реальная экономия начинается, когда разреженность реализована на уровне CUDA-ядра и матрица никогда не собирается целиком: примерами служат block-sparse ядра из оригинального Sparse Transformer и специализированные реализации в библиотеках вроде xformers или triton.

Порядок работы на практике

1. Оцените реальную длину последовательностей. Если n стабильно меньше нескольких тысяч, плотный attention с FlashAttention обычно проще и не медленнее.
2. Проверьте, есть ли в данных структура (периодичность, блоки, локальность). Она подсказывает, какой шаблон разреженности осмыслен.
3. Начните с готовой реализации (Longformer, BigBird, block-sparse), а не с собственного ядра — так вы отделите вопрос «работает ли идея» от вопроса «оптимально ли ядро».
4. Сравнивайте не только скорость, но и качество на вашей задаче: разреженность иногда роняет метрики на данных, где важны редкие дальние зависимости.

Когда разреженность оправдана, а когда нет

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

Ещё одна ловушка — реализация. Разреженные матричные операции хуже утилизируют тензорные ядра GPU, которые заточены под плотные регулярные вычисления. Поэтому теоретический выигрыш в FLOPs не всегда превращается в выигрыш по времени на реальном железе. Отсюда популярность block-sparse схем: они разрежают крупными блоками, сохраняя внутри блока плотную структуру, удобную для GPU.