Разреженное и линейное внимание, варианты голов, FlashAttention и эволюция активаций

Комбинирует три типа связей вместо полного графа внимания: локальное окно (каждый токен видит соседей), случайные связи (несколько случайных токенов на каждый запрос) и глобальные токены (небольшой набор токенов, которые видят всех и которых видят все — аналог CLS). Число связей на токен константно, поэтому сложность падает до O(n).

Авторы показали теоретически, что при достаточном числе глобальных токенов такая разреженная схема остаётся универсальным аппроксиматором последовательностей — то есть не теряет выразительности полного внимания, несмотря на линейную сложность. Хорошо работает на задачах с длинными документами: QA, суммаризация, геномные последовательности.

Строится на LSH-attention (locality-sensitive hashing): вместо сравнения каждого запроса со всеми ключами, запросы и ключи хэшируются в бакеты так, что похожие векторы с высокой вероятностью попадают в один бакет. Внимание считается только внутри бакета — токены сортируются по хэшу и разбиваются на чанки, что даёт O(n log n).

Второе ключевое решение — обратимые (reversible) остаточные слои, заимствованные из RevNet: активации каждого слоя восстанавливаются из следующего слоя при обратном проходе вместо хранения в памяти. Это убирает необходимость хранить активации всех слоёв для backprop — память перестаёт расти с глубиной сети. В сумме с чанкованным FFN это позволяло обрабатывать последовательности до 64k токенов на одном GPU.

Отталкивается от эмпирического наблюдения: матрица внимания в реальных моделях почти всегда низкого ранга. Вместо того чтобы проецировать K и V по размерности признаков, Linformer проецирует их по оси длины последовательности: обучаемые матрицы E, F размера n×k сжимают K и V с n строк до k ≪ n строк перед вычислением внимания.

K' = E·K   # (k×d), было (n×d)
V' = F·V   # (k×d), было (n×d)
attn = softmax(Q·K'ᵀ / √d) · V'   # O(n·k) вместо O(n²)

Использует механизм FAVOR+ (Fast Attention Via positive Orthogonal Random features): softmax-ядро exp(q·k) приближается через случайные признаковые отображения φ(·), такие что softmax(Q·Kᵀ) ≈ φ(Q)·φ(K)ᵀ. Это превращает внимание в обычное линейное произведение, которое можно переставить по ассоциативности:

# вместо: (φ(Q) φ(K)ᵀ) V  — снова матрица n×n
# считаем в другом порядке:
KV = φ(K)ᵀ @ V        # (r×d), не зависит от n
out = φ(Q) @ KV       # O(n·d·r), матрица n×n не строится никогда

В отличие от Linformer, аппроксимация ядра не завязана на фиксированную длину последовательности — она обобщается на произвольные n, а оценка несмещённая (unbiased) за счёт положительных ортогональных случайных признаков, что даёт низкую дисперсию приближения softmax.

Исходит из двух эмпирических свойств softmax-внимания, которые стоит сохранить при линеаризации: неотрицательность весов и локальность (внимание концентрируется на близких токенах). cosFormer заменяет softmax на ReLU (гарантирует неотрицательность весов линейного внимания) и добавляет косинусное перевзвешивание по относительной позиции токенов, имитируя локальность без явной нормализации softmax.

Q' = ReLU(Q), K' = ReLU(K)
weight(i,j) = cos((i−j)·π / 2M)     # локальность через относительную позицию
attn(i) = Σⱼ weight(i,j) · Q'ᵢ·K'ⱼ · Vⱼ   # считается за O(n), как и Performer

Результат — строго линейная сложность O(n) и рекуррентная форма вычисления, удобная для авторегрессивной генерации (можно писать как RNN-подобное накопление состояния), при качестве, сопоставимом или превосходящем Linformer/Performer на ряде задач.

Каждая из h голов имеет собственные проекции Q, K, V в подпространство размерности d_model/h, считает внимание независимо, результаты конкатенируются и проецируются обратно. Разные головы обучаются захватывать разные типы зависимостей (синтаксис, кореференции, позиционные паттерны).

Все query-головы используют одну общую KV-голову. Число различных проекций K/V падает с h до 1 — KV-кэш при генерации сжимается во столько же раз, во сколько раз h больше 1, что резко снижает нагрузку на память при инкрементальном декодировании.

Цена — некоторая потеря выразительности и качества по сравнению с MHA, из-за чего MQA использовался ограниченно (PaLM, ранние версии Falcon) до появления промежуточного варианта.

Query-головы делятся на g групп, внутри каждой группы головы делят одну общую KV-голову (g между 1 — это MQA, и h — это MHA). При g=8 для модели с h=64 головами KV-кэш сжимается в 8 раз при качестве, близком к полному MHA.

Сегодня это стандарт для моделей с длинным контекстом: Llama 2/3, Mistral, большинство современных открытых LLM используют GQA как золотую середину между качеством MHA и эффективностью MQA.

Ключевое наблюдение: обычная реализация внимания упирается не в число FLOPs, а в скорость обмена данными между быстрой SRAM (на кристалле, мало памяти) и медленной HBM (много памяти, но на порядок медленнее). Материализация полной матрицы n×n внимания в HBM и есть главный источник задержки.

Тайлинг: Q, K, V разбиваются на блоки, которые целиком помещаются в SRAM. Внимание считается блок за блоком, результат накапливается прямо на кристалле — полная матрица n×n никогда не материализуется в HBM.

Online softmax: softmax по строке обычно требует видеть всю строку сразу (для нормировки). FlashAttention считает softmax инкрементально, блок за блоком, постоянно корректируя частичную сумму и максимум — математически эквивалентно обычному softmax, но без необходимости держать всю строку в памяти одновременно.

Пересчёт вместо хранения: при обратном проходе матрица внимания не хранится, а пересчитывается заново из Q, K, V — это дороже по FLOPs, но дешевле по IO, а именно IO — узкое место.

# схематично: один блок тайлинга
for block_k, block_v in tiles(K, V):        # блоки помещаются в SRAM
    s_block = Q_block @ block_k.T          # маленькая матрица, не n×n
    m_new = max(m_running, s_block.max())  # online softmax: обновляем максимум
    # переcчитываем накопленную сумму под новый максимум и добавляем блок
    acc = acc * exp(m_running - m_new) + exp(s_block - m_new) @ block_v
    m_running = m_new

FlashAttention-2 (2023) улучшает распределение работы между потоками GPU: лучше параллелит вычисления по оси длины последовательности (а не только по батчу и головам), сокращает долю немatmul-операций — почти двукратное ускорение над первой версией. FlashAttention-3 (2024) нацелен на архитектуру Hopper (H100): использует асинхронное выполнение и специализацию варпов между matmul и softmax, а также низкую точность FP8 — дополнительный прирост на новом железе.

Простейшая нелинейность: обнуляет отрицательные значения, пропускает положительные без изменений. Дешева в вычислении, не страдает от насыщения градиента на положительной части (в отличие от sigmoid/tanh), поэтому долго была стандартом. Проблема — «умирающие» нейроны: если вход стабильно отрицателен, градиент через нейрон равен нулю, и он перестаёт обучаться.

Взвешивает вход его перцентилем по стандартной нормальной функции распределения Φ(x) — по сути, стохастическое обоснование: вход «пропускается» с вероятностью, зависящей от его величины. Гладкая и немонотонная (небольшой провал в отрицательной области вместо жёсткого нуля), что даёт более мягкий сигнал градиента, чем ReLU. Стала стандартом в BERT, GPT-2, GPT-3.

GELU(x) ≈ 0.5x · (1 + tanh[√(2/π)·(x + 0.044715x³)])   # практическая аппроксимация

Тот же принцип самогейтинга, что и у GELU, но вентиль — обычная сигмоида, а не нормальная CDF: SiLU(x) = x · sigmoid(x). Кривая почти неотличима от GELU визуально (см. график выше), но дешевле вычислять. Используется в EfficientNet, а главное — стал строительным блоком для гейтинга в современных FFN-слоях LLM.

Это не отдельная функция активации, а вариант Gated Linear Unit, применённый ко всему FFN-блоку: вход проецируется двумя независимыми линейными слоями, один из результатов пропускается через SiLU и служит «воротами», которые поэлементно умножаются на второй, линейный, результат.

FFN_SwiGLU(x) = (SiLU(x·W) ⊙ (x·V)) · W2   # W, V, W2 — три отдельные матрицы

Эмпирически SwiGLU-FFN стабильно превосходит варианты на чистом ReLU или GELU при равном вычислительном бюджете, поэтому стал де-факто стандартом FFN в PaLM, LLaMA/Llama 2/3, Mistral и большинстве современных открытых LLM.