Как работает sliding window attention в Mistral и Longformer?

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

Sliding window attention — это механизм, ограничивающий область внимания каждого токена фиксированным окном из W предыдущих токенов, что снижает сложность self-attention с O(n²) до O(n·W). В Mistral используется простое окно шириной 4096 токенов, а за счёт 32 слоёв эффективный контекст достигает 131k токенов (receptive field растёт линейно с глубиной). Longformer комбинирует sliding window с global attention на специальные токены и sliding window|dilated sliding window, что позволяет модели захватывать как локальные, так и глобальные зависимости без квадратичного роста затрат.

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

Стандартный self-attention (Vaswani et al., 2017) вычисляет матрицу внимания размером n×n, где n — длина последовательности. Каждый токен attends ко всем остальным, что даёт сложность O(n²) по времени и памяти. Для длинных контекстов (сотни тысяч токенов) это становится непрактичным.

Формула сложности:

• Время: O(n²·d), где d — размерность скрытого состояния.
• Память: O(n²) для хранения матрицы внимания.

Пример: при n=100k, n² = 10¹⁰ — невозможно уместить даже в современные GPU.

2. Идея sliding window attention

Sliding window attention (также known as local attention) ограничивает область внимания каждого токена окном фиксированного размера W, центрированным на текущем токене (обычно — только предыдущие W токенов, как в Mistral, или симметричное окно). Токен attends только к W соседям, а не ко всей последовательности.

Сложность:

• Время: O(n·W·d) — линейно по n.
• Память: O(n·W) — также линейно.

Receptive field (область, которую токен может «видеть» через один слой) равен W. Однако в многослойной сети receptive field растёт: после L слоёв токен может косвенно учитывать информацию из L·W токенов (при условии, что окно однонаправленное и каждый слой сдвигает окно). Это ключевое свойство, используемое в Mistral.

3. Sliding window attention в Mistral

Mistral (Mistral AI, 2023) использует однонаправленное sliding window с размером окна W = 4096 токенов. Модель имеет L = 32 слоя.

Механизм:

• В каждом слое каждый токен attends только к предыдущим 4096 токенам (включая себя).
• После первого слоя receptive field = 4096.
• После второго слоя — до 8192 (токен может через промежуточные представления «дотянуться» до токенов, отстоящих на 2·4096).
• После 32 слоёв — до 32·4096 = 131072 токенов.

Эффективный контекст (effective context) — максимальная дистанция, на которую токен может повлиять на другой токен через цепочку слоёв. В Mistral он равен 131k, что достаточно для большинства задач с длинными документами.

Преимущества:

• Линейная сложность по длине.
• Возможность обрабатывать последовательности до 131k токенов без дополнительных трюков (например, без пересчёта KV-кэша).

Недостатки:

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

4. Sliding window attention в Longformer

Longformer (Beltagy et al., 2020) предлагает гибридный подход: комбинацию sliding window, dilated sliding window и global attention.

4.1 Sliding window (локальное внимание)

Базовый слой использует окно фиксированного размера (например, W = 512). Каждый токен attends к W/2 соседей слева и W/2 справа (симметричное окно). Сложность O(n·W).

4.2 Dilated sliding window (разреженное окно с пропусками)

Для увеличения receptive field без роста W Longformer применяет dilation (разрежение). Вместо того чтобы брать все W соседей, берутся каждый d-й токен в окне. Например, при W=512 и dilation=2 токен attends к 256 токенам, но с шагом 2, что расширяет покрытие до 1024 позиций. На разных слоях dilation может меняться (например, на нижних слоях dilation=1, на верхних — больше).

Формула receptive field с dilation: после L слоёв с dilation d_i на каждом слое receptive field ≈ W + Σ(d_i·W) (упрощённо).

4.3 Global attention (глобальное внимание)

Для задач, где важны глобальные токены (например, [CLS] для классификации, специальные токены для вопросов), Longformer добавляет global attention: некоторые токены (например, [CLS]) могут attends ко всем токенам, и все токены могут attends к ним. Это даёт O(n) дополнительных вычислений на один глобальный токен.

Конфигурация:

• На каждом слое выбирается небольшое количество глобальных токенов (обычно 1–10).
• Для них вычисляется полное внимание (n токенов), остальные используют sliding window.
• Общая сложность: O(n·W + n·G), где G — число глобальных токенов (обычно G << n).

5. Сравнение Mistral и Longformer

Ключевое различие: Mistral полагается на глубину для расширения контекста, а Longformer — на комбинацию dilation и global attention. Mistral оптимизирован для автогрессивной генерации (однонаправленное окно), Longformer — для понимания (bidirectional).

6. Преимущества и недостатки sliding window attention

Преимущества:

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

Недостатки:

• Ограниченный прямой receptive field (без многослойности или dilation).
• Для задач, требующих глобального контекста (например, поиск информации в конце документа), может потребоваться много слоёв или дополнительный механизм.
• Однонаправленное окно (как в Mistral) не подходит для задач, где нужен bidirectional контекст (например, NER, классификация).

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

В Agentic RAG агент часто обрабатывает длинные контексты: историю диалога, несколько документов, результаты поиска. Sliding window attention позволяет:

• Обрабатывать контексты длиной до сотен тысяч токенов без переполнения памяти.
• Сохранять линейное время инференса, что критично для реального времени.
• Использовать многослойность для косвенного «видения» дальних токенов.

Например, Mistral-7B с sliding window attention может быть использован как базовый LLM в RAG-пайплайне, где контекст формируется из нескольких ретриверных чанков.

8. Альтернативы sliding window attention

Sliding window — компромисс между простотой и эффективностью, широко используемый в современных LLM (Mistral, GPT-4 (предположительно), Llama 2 (GQA + sliding window? нет, Llama 2 использует full attention с ограничением контекста)).

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

Задача: Реализовать sliding window attention с нуля на PyTorch и сравнить его производительность с полным attention.

Инструменты: Python, PyTorch, CUDA (опционально), библиотека для бенчмаркинга (timeit, torch.cuda.Event).

Шаги:

1. Реализовать функцию sliding_window_attention(Q, K, V, window_size):
   • Использовать маску внимания, где для каждого токена разрешены только W соседних позиций (слева или симметрично).
   • Вычислить attention scores, применить маску (заменить запрещённые позиции на -inf), softmax, взвешенная сумма.
2. Реализовать full_attention(Q, K, V) — стандартный scaled dot-product attention.
3. Создать тестовые тензоры: batch=1, heads=1, seq_len=1024..16384, dim=64.
4. Замерить время выполнения и пиковое использование памяти (через torch.cuda.max_memory_allocated()).
5. Построить графики: seq_len vs time, seq_len vs memory для обоих методов.
6. Дополнительно: реализовать dilated sliding window (с параметром dilation) и global attention (добавить один глобальный токен).

Ожидаемый результат:

• При seq_len=8192 sliding window (W=512) работает в ~10 раз быстрее и использует в ~16 раз меньше памяти, чем full attention.
• При seq_len=16384 full attention вызывает OOM (out-of-memory) на GPU с 8GB, а sliding window работает стабильно.
• Dilated sliding window показывает промежуточные результаты по скорости и качеству (можно проверить на простой задаче: предсказание следующего токена).

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

Навигация

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