Sliding Window Attention

Sliding Window Attention — это механизм внимания (attention), используемый в архитектурах нейронных сетей, в частности в трансформерах, который ограничивает область вычисления внимания для каждого токена фиксированным окном из соседних токенов. В отличие от полного (глобального) внимания, где каждый токен взаимодействует со всеми остальными токенами последовательности, скользящее окно учитывает только локальный контекст, что значительно снижает вычислительную сложность и потребление памяти, особенно при работе с длинными последовательностями.

История

Механизм внимания был впервые предложен в 2014 году для задач машинного перевода, а в 2017 году архитектура трансформера, основанная на полном внимании, произвела революцию в обработке естественного языка. Однако квадратичная сложность полного внимания (O(n²), где n — длина последовательности) стала серьёзным ограничением для обработки длинных текстов, геномов, временных рядов и других данных. Для решения этой проблемы были разработаны различные методы аппроксимации внимания, одним из которых стало скользящее окно.

Идея локального внимания существовала и ранее в сверточных нейронных сетях, но её адаптация к трансформерам произошла в конце 2010-х годов. В 2019 году модель Longformer (Beltagy et al.) представила комбинацию скользящего окна и глобального внимания для обработки длинных документов. В 2020 году модель BigBird (Zaheer et al.) предложила похожий гибридный подход, добавив случайное внимание. В 2021 году архитектура Mistral (Jiang et al.) полностью отказалась от глобального внимания в пользу скользящего окна, продемонстрировав, что для многих задач достаточно локального контекста. В 2023 году модель Gemma от Google также использовала модифицированное скользящее окно.

Принцип работы

В стандартном трансформере с полным вниманием для каждого токена вычисляется взвешенная сумма всех токенов последовательности. В скользящем окне для токена с индексом i вычисляется внимание только к токенам, находящимся в пределах окна: от i — w/2 до i + w/2, где w — размер окна. На практике часто используется одностороннее окно (каузальное), когда токен видит только предыдущие токены, что необходимо для авторегрессивной генерации.

Математическая формализация

Пусть дана последовательность токенов x₁, x₂, ..., xₙ. Для токена xᵢ вычисляется вектор внимания:

Attention(Qᵢ, Kⱼ, Vⱼ) = softmax( (Qᵢ · Kⱼ) / √d ) · Vⱼ, где j ∈ [max(1, i — w), min(n, i + w)]

Здесь:

• Qᵢ — запрос (query) для токена i,
• Kⱼ — ключ (key) для токена j,
• Vⱼ — значение (value) для токена j,
• d — размерность векторов,
• w — размер окна.

Маскирование

Для реализации скользящего окна используется маска внимания — матрица, в которой разрешённые взаимодействия помечаются единицами, а запрещённые — нулями или отрицательной бесконечностью (после softmax они становятся нулями). В случае каузального скользящего окна маска представляет собой диагональную полосу шириной w, где все элементы выше главной диагонали обнулены.

Виды скользящего окна

Одностороннее (каузальное) окно

Используется в авторегрессивных моделях (например, Mistral, GPT). Каждый токен видит только себя и предыдущие w токенов. Это обеспечивает линейную сложность по длине последовательности и позволяет генерировать текст по одному токену.

Двустороннее окно

Применяется в моделях для понимания текста (например, BERT с локальным вниманием). Токен видит w/2 токенов слева и w/2 справа. Сложность также линейна, но требует доступа ко всей последовательности.

Гибридное окно

Комбинация скользящего окна с глобальным вниманием для отдельных токенов (например, специальные токены [CLS] в Longformer). Глобальные токены имеют доступ ко всей последовательности, а остальные — только к локальному окну.

Динамическое окно

Размер окна может меняться в зависимости от позиции токена или контекста. Например, в модели Gemma используется «сдвиг» окна в разных слоях, чтобы информация могла распространяться дальше.

Преимущества и недостатки

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

• Вычислительная сложность: O(n·w) вместо O(n²). При фиксированном w (например, 4096) сложность линейна по длине последовательности, что позволяет обрабатывать последовательности длиной до 1 миллиона токенов.
• Потребление памяти: Матрица внимания хранит только w·n элементов, а не n². Это критически важно для графических процессоров с ограниченной памятью.
• Скорость обучения и инференса: Меньшее количество операций ускоряет как обучение, так и генерацию текста.
• Параллелизация: Окна независимы, что позволяет эффективно распараллеливать вычисления.

Недостатки

• Ограниченный контекст: Модель не может напрямую учитывать зависимости между токенами, находящимися дальше размера окна. Для преодоления этого используются глубокие сети (многослойные трансформеры), где информация распространяется через слои.
• Необходимость настройки размера окна: Слишком маленькое окно теряет долгосрочные зависимости, слишком большое — снижает преимущества в скорости.
• Сложность реализации: Требуется специальная оптимизация для эффективного вычисления на GPU (например, использование FlashAttention с поддержкой масок).

Применение

Обработка естественного языка

• Генерация текста: Модели Mistral, GPT-4 (частично), Llama (с использованием скользящего окна в некоторых версиях) применяют скользящее окно для генерации длинных текстов без потери производительности.
• Анализ документов: Longformer и BigBird используются для обработки научных статей, юридических документов и книг длиной до 10 000+ токенов.

Биоинформатика

• Анализ геномов: Модель DNABERT с локальным вниманием обрабатывает последовательности ДНК длиной до 100 000 пар оснований, выявляя регуляторные участки.
• Предсказание структуры белков: AlphaFold2 использует локальное внимание для моделирования взаимодействий между аминокислотами.

Компьютерное зрение

• Vision Transformers (ViT): Скользящее окно применяется в моделях Swin Transformer (Liu et al., 2021) для обработки изображений, где окно перемещается по пространственным измерениям, обеспечивая иерархическое представление.

Временные ряды

• Прогнозирование: Модели на основе трансформеров со скользящим окном (например, Informer) обрабатывают длинные временные ряды (до 100 000 шагов) для задач финансового прогнозирования или мониторинга оборудования.

Примеры моделей

• Longformer (2020): Использует скользящее окно размером 512 токенов с дополнительными глобальными токенами. Обрабатывает документы до 4096 токенов.
• BigBird (2020): Комбинирует скользящее окно (размер 3), случайное внимание (3 токена) и глобальное внимание (2 токена). Масштабируется до 4096 токенов.
• Mistral 7B (2023): Полностью основана на скользящем окне размером 4096 токенов. Достигает производительности, сопоставимой с моделями в 3 раза большего размера, благодаря эффективному использованию памяти.
• Gemma (2024): Использует скользящее окно с «сдвигом» в разных слоях, что позволяет модели охватывать контекст до 8192 токенов при окне 4096.

Сравнение с другими механизмами внимания

Критика

Скользящее окно критикуется за невозможность моделировать долгосрочные зависимости без глубоких сетей. Эксперименты показывают, что для задач, требующих понимания структуры всего документа (например, реферирование), гибридные подходы (Longformer, BigBird) работают лучше, чем чистое скользящее окно. Кроме того, размер окна часто выбирается эмпирически, и оптимальное значение может сильно варьироваться в зависимости от задачи.

Источники

• Beltagy, I., Peters, M. E., & Cohan, A. (2020). Longformer: The Long-Document Transformer. arXiv:2004.05150.
• Zaheer, M., et al. (2020). Big Bird: Transformers for Longer Sequences. NeurIPS 2020.
• Jiang, A. Q., et al. (2023). Mistral 7B. arXiv:2310.06825.
• Liu, Z., et al. (2021). Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. ICCV 2021.
• Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017.