Sliding Window Attention

A atenção com janela deslizante (SWA) limita cada token a prestar atenção a um número fixo de tokens vizinhos — os w mais recentes num modelo de linguagem causal — em vez de toda a sequência, como acontece na autoatenção completa. O cálculo da atenção passa então a crescer linearmente com o comprimento da sequência, em vez de quadraticamente, e um modelo generativo só precisa de manter no cache KV as chaves e os valores dos últimos w tokens. O Longformer fez da janela deslizante o elemento central da sua atenção em 2020, e o Mistral 7B da Mistral AI introduziu-a em 2023 num modelo de linguagem de grande escala (LLM) aberto e muito utilizado, com uma janela de 4 096 tokens.

Como a janela cresce com a profundidade#

Uma camada de SWA vê apenas w tokens, mas as camadas do Transformer são empilhadas. Na camada 2, um token lê estados ocultos que já resumem w tokens da camada 1, pelo que a informação pode percorrer cerca de mais w − 1 posições em cada camada, tal como as convoluções empilhadas alargam o campo recetivo de uma rede neural convolucional (CNN). No PyTorch, a janela é uma máscara de atenção booleana, e multiplicar essa máscara por si própria mostra até onde a informação se propaga:

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

Para 16 tokens, a janela calcula 58 pontuações por camada, em vez de 136, e o alcance do último token aumenta 3 tokens por camada: de 4 para 7, 10 e 13. A mesma máscara pode ser transmitida como attn_mask para a função scaled_dot_product_attention do PyTorch. No Mistral 7B, 32 camadas com uma janela de 4 096 tokens proporcionam um alcance teórico de atenção de cerca de 131 mil tokens, embora a informação mais distante tenha de sobreviver a muitas etapas para chegar ao destino.

O Mistral 7B também armazena as chaves e os valores num buffer circular de w posições, escrevendo a posição i na posição i mod w e substituindo a entrada mais antiga. Numa sequência de 32 mil tokens, isto reduziu oito vezes a memória da cache sem prejudicar a qualidade; o tamanho fixo da cache mantém a memória estável, independentemente da duração da geração.

Intercalando camadas locais e globais#

Janelas puras não conseguem copiar de uma só vez um fato distante no contexto, por isso os modelos recentes combinam camadas SWA com camadas globais completas. O Gemma 2, do Google, alternava as camadas na proporção de uma para uma, com uma janela de 4.096 tokens, e o Gemma 3 usa cinco camadas locais para cada camada global, com uma janela de apenas 1.024 tokens, mantendo pequeno o cache KV ao longo de uma janela de contexto de até 128K tokens. Apenas as camadas globais armazenam chaves e valores para todo o contexto; assim, cinco de cada seis camadas mantêm apenas 1.024 tokens cada.

Os transformers de visão usam a mesma combinação para imagens de alta resolução, nas quais cada patch é um token. O SAM 3, da Meta, que é executado no pacote Ultralytics para segmentação de imagens, aplica atenção dentro de janelas não sobrepostas de patches de 24 × 24 em 28 dos 32 blocos do codificador de imagens e atenção global em toda a grade de patches de 72 × 72 de uma imagem de 1.008 pixels nos blocos 7, 15, 23 e 31.

SWA e termos relacionados#

• Swin Transformer: Divide a imagem em janelas fixas e não sobrepostas, compartilhadas por todos os tokens da janela, e desloca a grade entre camadas para que janelas vizinhas troquem informações.
• Atenção esparsa: A família mais ampla de padrões fixos. A SWA é o seu membro mais simples; outras abordagens acrescentam janelas dilatadas, tokens globais ou conexões aleatórias.
• Sumidouros de atenção: Os primeiros tokens de uma sequência, que absorvem o peso de atenção excedente. Quando uma janela é aplicada durante a inferência a um modelo treinado com atenção completa, removê-los interrompe a geração, e o StreamingLLM os mantém no cache junto à janela.
• Atenção deformável: Cada consulta presta atenção a alguns pontos de amostragem cujos deslocamentos aprende, de modo que a vizinhança depende do conteúdo.
• Atenção Flash: Um kernel de GPU para atenção exata. As implementações aceitam um tamanho de janela; assim, os modelos SWA ignoram completamente os blocos fora da janela.