Sliding Window Attention

Внимание в скользящем окне (SWA) ограничивает множество токенов, которые может учитывать каждый токен: в причинно-следственной языковой модели это последние w токенов, а не вся последовательность, как при полном самовнимании. В результате объём вычислений для внимания растёт линейно, а не квадратично относительно длины последовательности, и генерирующей модели достаточно хранить ключи и значения последних w токенов в KV-кэше. В 2020 году Longformer сделала скользящее окно основой своего механизма внимания. В 2023 году Mistral AI перенесла этот подход в широко используемую открытую большую языковую модель (LLM) Mistral 7B с окном в 4 096 токенов.

Как окно расширяется с глубиной#

Один слой SWA видит только w токенов, но слои Transformer выстраиваются друг за другом. На втором слое токен считывает скрытые состояния, каждое из которых уже суммирует w токенов первого слоя, поэтому на каждом последующем слое информация может проходить примерно на w − 1 позиций дальше. Похожим образом расширяется рецептивное поле свёрточной нейронной сети (CNN) при добавлении свёрточных слоёв. В PyTorch окно задаётся булевой маской внимания, а умножение этой маски на саму себя показывает, как далеко распространяется информация:

import torch

seq_len, window = 16, 4
pos = torch.arange(seq_len)
offset = pos[:, None] - pos[None, :]

# Causal sliding window: each token attends to itself and the window - 1 tokens before it
mask = (offset >= 0) & (offset < window)
print(f"Scores per layer: {mask.sum().item()} with the window, {(offset >= 0).sum().item()} with full causal attention")

# Stacking layers widens the reach: each layer adds window - 1 more tokens of history
reach = mask
for layer in range(1, 5):
    print(f"Layer {layer}: last token draws on {reach[-1].sum().item()} tokens")
    reach = (mask.int() @ reach.int()) > 0

Для 16 токенов в каждом слое окно вычисляет 58 оценок вместо 136, а область охвата последнего токена увеличивается на 3 токена за слой: с 4 до 7, 10 и 13. Ту же маску можно передать в PyTorch как attn_mask функции scaled_dot_product_attention. Для Mistral 7B 32 слоя с окном в 4 096 токенов дают теоретический охват внимания примерно в 131K токенов, хотя информации из далёких слоёв нужно пройти много шагов, чтобы дойти до текущего.

Mistral 7B также хранит ключи и значения в циклическом буфере из w слотов: позиция i записывается в слот i mod w, заменяя самую старую запись. Для последовательности из 32K токенов это уменьшило объём памяти под кэш в восемь раз без потери качества, а фиксированный размер кэша не меняется при увеличении длины генерации.

Чередование локальных и глобальных слоёв#

Чисто оконные слои не могут за один шаг скопировать факт из далёкого прошлого, поэтому в современных моделях слои SWA чередуются с полностью глобальными слоями. В Gemma 2 от Google они чередовались один к одному с окном в 4 096 токенов, а в Gemma 3 на каждый глобальный слой приходится пять локальных с окном всего в 1 024 токена, что позволяет поддерживать небольшой кэш KV при контекстном окне длиной до 128K токенов. Только глобальные слои хранят ключи и значения для всего контекста, поэтому в пяти из каждых шести слоёв содержится всего по 1 024 токена.

Визуальные трансформеры используют такое же сочетание для изображений высокого разрешения, где каждый патч — это токен. SAM 3 от Meta, работающая в пакете Ultralytics для сегментации изображений, использует внимание в неперекрывающихся окнах размером 24 × 24 патча в 28 из 32 блоков энкодера изображений и глобальное внимание по всей сетке 72 × 72 патча изображения размером 1 008 пикселей в блоках 7, 15, 23 и 31.

SWA и связанные термины#

• Swin Transformer: Разбивает изображение на фиксированные неперекрывающиеся окна, общие для всех токенов внутри окна, и сдвигает сетку между слоями, чтобы соседние окна обменивались информацией.
• Разреженное внимание: Более широкое семейство фиксированных шаблонов. SWA — его простейший вариант; другие методы добавляют расширенные окна, глобальные токены или случайные связи.
• Стоки внимания: Первые токены последовательности, которые поглощают избыточный вес внимания. Если при инференсе применить окно к модели, обученной с полным вниманием, удаление этих токенов нарушит генерацию; StreamingLLM сохраняет их в кэше рядом с окном.
• Деформируемое внимание: Каждый запрос обращается к нескольким точкам выборки, смещения которых модель обучается предсказывать, поэтому окрестность зависит от содержимого.
• Flash attention: Ядро GPU для точного вычисления внимания. Его реализации поддерживают размер окна, поэтому модели SWA полностью пропускают блоки за пределами окна.