Как работает attention с линейной сложностью (Linformer, Performer, Longformer)?

Краткий тезис

Стандартный механизм self-attention в трансформерах имеет квадратичную сложность O(n²) по длине последовательности n, что делает его непрактичным для длинных контекстов (100k+ токенов). Linformer, Performer и Longformer предлагают разные подходы к снижению сложности до линейной или почти линейной: Linformer проецирует ключи (K) и значения (V) на низкоранговое пространство, Performer аппроксимирует softmax через случайные признаки (FAVOR+), а Longformer комбинирует sliding window с глобальным attention. Все три метода жертвуют частью качества ради скорости, но позволяют обрабатывать последовательности, недоступные для полного attention.

1. Проблема квадратичной сложности стандартного attention

В классическом трансформере (Vaswani et al., 2017) attention вычисляется как:

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

• Q, K, V — матрицы размером n × d (n — длина последовательности, d — размерность).
• Матрица QK^T имеет размер n × n, её вычисление и хранение требуют O(n²) времени и памяти.

Для n = 100k это ~10 млрд операций — неприемлемо для GPU. Поэтому для работы с длинными документами, геномами, видео или многотурными диалогами необходимы альтернативы.

Термин «квадратичная сложность» означает, что время выполнения растёт пропорционально квадрату длины входа. Даже при n = 10k полный attention уже становится узким местом.

2. Обзор подходов к линейному attention

Все методы снижения сложности можно разделить на три категории:

Мы детально разберём три наиболее известных метода: Linformer, Performer и Longformer.

3. Linformer: проекция K и V на низкоранговое пространство

Linformer (Wang et al., 2020) основан на наблюдении, что матрица attention (softmax(QK^T)) часто имеет низкий эффективный ранг. Авторы предлагают сжать K и V с помощью обучаемых проекционных матриц E и F размером n × k, где k — гиперпараметр (обычно 128–512, k << n).

Формула:

K' = E K   (размер k × d)
V' = F V   (размер k × d)
Attention = softmax(Q K'^T / sqrt(d)) V'

• Q остаётся размером n × d.
• Q K'^T имеет размер n × k → сложность O(n·k·d) вместо O(n²·d).
• Память: O(n·k) вместо O(n²).

Термин «низкоранговое пространство» означает, что мы представляем исходные матрицы K и V в подпространстве меньшей размерности, теряя часть информации, но сохраняя основные паттерны.

Плюсы

• Простота реализации.
• Хорошее качество при k ≥ 128 (близко к full attention).
• Сложность O(n) при фиксированном k.

Минусы

• Гиперпараметр k нужно подбирать; слишком малое k ухудшает качество.
• Проекционные матрицы E и F — дополнительные параметры (но их число невелико: 2·n·k).

Пример кода (упрощённо):

import torch
import torch.nn as nn

class LinformerAttention(nn.Module):
    def __init__(self, d_model, n_heads, k=256):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.k = k
        self.head_dim = d_model // n_heads

        self.E = nn.Linear(d_model, k, bias=False)  # проекция K
        self.F = nn.Linear(d_model, k, bias=False)  # проекция V

    def forward(self, Q, K, V, mask=None):
        # Q, K, V: (batch, n, d_model)
        batch, n, _ = Q.shape
        # проецируем K и V
        K_proj = self.E(K)  # (batch, n, k)
        V_proj = self.F(V)  # (batch, n, k)
        # attention с пониженной размерностью
        attn = torch.matmul(Q, K_proj.transpose(-2, -1)) / (self.head_dim ** 0.5)
        if mask is not None:
            attn = attn.masked_fill(mask == 0, float('-inf'))
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, V_proj)
        return out

4. Performer (FAVOR+): аппроксимация softmax через случайные признаки

Performer (Choromanski et al., 2020) использует математический трюк: softmax можно представить как скалярное произведение в пространстве случайных признаков. Для этого применяется положительное случайное отображение (FAVOR+).

Идея: exp(q·k) ≈ φ(q)·φ(k), где φ — случайное отображение размерности r (r << n). Тогда:

Attention(Q, K, V) ≈ D^{-1} (φ(Q) · φ(K)^T) V

где D — диагональная матрица нормализации. Сложность O(n·d·r) — линейна по n.

Термин «случайные признаки» — это векторы, полученные путём применения случайной проекции и нелинейности (например, exp(ω·x) для ω ~ N(0,1)). Performer использует специальную конструкцию для гарантии положительности и низкой дисперсии.

Плюсы

• Теоретически обоснованная аппроксимация с контролируемой ошибкой.
• Не требует дополнительных обучаемых параметров (кроме случайных весов, которые фиксируются после инициализации).
• Сложность O(n) при фиксированном r.

Минусы

• Для высокой точности нужно r порядка 256–1024, что увеличивает константу.
• На практике может быть медленнее, чем Linformer, из-за вычисления φ.

Пример (упрощённо):

import torch
import torch.nn as nn
import math

class PerformerAttention(nn.Module):
    def __init__(self, d_model, n_heads, r=256):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.r = r
        self.head_dim = d_model // n_heads
        # случайная матрица для проекции (фиксируется)
        self.register_buffer('omega', torch.randn(self.head_dim, r))

    def _feature_map(self, x):
        # x: (batch, n, head_dim)
        # φ(x) = exp( -||x||^[[2 Как вы решаете проблему lost in the middle при работе с длинными контекстами|2]]/[[2 Как вы решаете проблему lost in the middle при работе с длинными контекстами|2]] ) * exp(ω·x)
        # упрощённо: используем exp(ω·x)
        proj = torch.matmul(x, self.omega)  # (batch, n, r)
        return torch.exp(proj)  # положительные признаки

    def forward(self, Q, K, V, mask=None):
        batch, n, _ = Q.shape
        Q_feat = self._feature_map(Q)
        K_feat = self._feature_map(K)
        # attention через скалярное произведение признаков
        attn = torch.matmul(Q_feat, K_feat.transpose(-2, -1))  # (batch, n, n)
        # нормализация (упрощённо, без D)
        attn = attn / (n ** 0.5)
        if mask is not None:
            attn = attn.masked_fill(mask == 0, float('-inf'))
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, V)
        return out

Примечание: реальная реализация Performer использует более сложную нормализацию и ортогонализацию случайных матриц.

5. Longformer: sliding window + global attention

Longformer (Beltagy et al., 2020) не пытается аппроксимировать полный attention, а заменяет его разреженной схемой:

• Sliding window attention: каждый токен «видит» только w соседей слева и справа (w — размер окна, обычно 512). Сложность O(n·w).
• Dilated sliding window: как в CNN, окно с шагом (dilation) для увеличения рецептивного поля без роста w.
• Global attention: для некоторых токенов (например, [CLS] или специальных маркеров) attention считается полным (глобальным). Число таких токенов g фиксировано (обычно 1–128). Сложность O(g·n).

Итоговая сложность: O(n·w + g·n) ≈ O(n) при w, g = const.

Термин «sliding window» — это окно фиксированной ширины, которое скользит по последовательности, обрабатывая локальные контексты.

Плюсы

• Интуитивно понятен, легко реализовать.
• Хорошо работает на задачах, где важна локальная информация (текст, геномы).
• Поддерживает очень длинные последовательности (до 32k+).

Минусы

• Не может моделировать дальние зависимости между произвольными токенами (только через глобальные).
• Требует выбора w и g, а также разметки глобальных токенов.

Пример (концептуально):

import torch
import torch.nn as nn

class LongformerAttention(nn.Module):
    def __init__(self, d_model, n_heads, window=512, global_tokens=1):
        super().__init__()
        self.window = window
        self.global_tokens = global_tokens
        # ... стандартные линейные проекции

    def forward(self, Q, K, V, global_mask=None):
        batch, n, _ = Q.shape
        # 1. sliding window attention (локальное)
        # создаём диагональную маску окна
        # 2. global attention (полное)
        # 3. комбинируем результаты
        # на практике используется специальный слой Longformer из библиотеки transformers
        pass

6. Сравнение методов

Термин «качество относительно full attention» означает, насколько близки результаты (loss, accuracy) к модели с полным attention при одинаковом количестве параметров.

7. Компромисс: скорость vs качество

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

• Если в данных преобладают локальные зависимости (текст, код) — Longformer даёт наилучшее качество при минимальных потерях.
• Если нужна теоретическая гарантия аппроксимации и нет желания подбирать окна — Performer.
• Если важна простота и низкая задержка — Linformer.

На практике для RAG с длинными документами (100k+ токенов) часто используют комбинацию: сначала retrieval находит релевантные чанки, а затем Longformer или Performer обрабатывает их в одном контексте.

8. Применение в RAG и Agentic RAG

В Agentic RAG агент может получать на вход несколько документов (каждый длиной до 10k токенов) и должен обработать их совместно. Полный attention здесь невозможен. Эффективные attention механизмы позволяют:

• Обрабатывать контекст из нескольких документов как одну последовательность.
• Использовать global attention для специальных токенов (например, [CLS] или маркеров запроса).
• Снижать latency при многотурных диалогах.

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

9. Пет-проект для закрепления

Задача Реализовать упрощённую версию Linformer и сравнить её с полным attention на задаче классификации длинных текстов.

Инструменты PyTorch, Hugging Face Transformers (для baseline), датасет IMDB (рецензии, можно обрезать до 2048 токенов).

Шаги:

1. Загрузить датасет IMDB и токенизировать с максимальной длиной 2048.
2. Реализовать класс LinformerAttention (как в разделе 3).
3. Создать простую модель классификации (embedding + attention + pool + linear).
4. Обучить две версии: с полным attention и с Linformer (k=128).
5. Сравнить:
   • Время обучения на эпоху.
   • Accuracy на тесте.
   • Потребление памяти (GPU).
6. Построить графики зависимости accuracy от длины последовательности (например, обрезать до 512, 1024, 2048).

Ожидаемый результат Вы увидите, что Linformer обучается в 2–4 раза быстрее и использует меньше памяти, но accuracy может упасть на 1–3%. Для длинных последовательностей (2048+) полный attention может не поместиться в память, а Linformer — да.

10. Связь с другими вопросами

Эти вопросы связаны с эффективностью обработки длинных контекстов, оценкой и оптимизацией RAG-систем.

11. Навигация

• Предыдущий: 645
• Следующий: 647
• Индекс: 00. Индекс разборов

Навигация

• Предыдущий: 645
• Следующий: 647
• Индекс: 00. Индекс разборов