Sliding Window Attention

La atención de ventana deslizante (SWA) limita la atención de cada token a un número fijo de tokens cercanos —los w tokens más recientes en un modelo de lenguaje causal— en lugar de abarcar toda la secuencia, como hace la autoatención completa. El cálculo de la atención crece entonces de forma lineal con la longitud de la secuencia, en lugar de cuadrática, y un modelo generativo solo tiene que conservar en su caché KV las claves y los valores de los últimos w tokens. Longformer convirtió la ventana deslizante en la base de su atención en 2020, y Mistral 7B de Mistral AI la incorporó en 2023 a un modelo de lenguaje de gran tamaño (LLM) abierto y muy utilizado, con una ventana de 4.096 tokens.

Cómo crece la ventana con la profundidad#

Una capa de SWA solo ve w tokens, pero las capas de Transformer se apilan. En la capa 2, un token lee estados ocultos que ya resumen w tokens de la capa 1, por lo que la información puede avanzar unas w − 1 posiciones más en cada capa, igual que las convoluciones apiladas amplían el campo receptivo de una red neuronal convolucional (CNN). En PyTorch, la ventana es una máscara de atención booleana, y multiplicar esa máscara por sí misma permite seguir hasta dónde llega la información:

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

La ventana calcula 58 puntuaciones por capa para 16 tokens, en lugar de 136, y el alcance del último token crece 3 tokens por capa: de 4 a 7, 10 y 13. Se puede pasar la misma máscara como attn_mask a scaled_dot_product_attention de PyTorch. En Mistral 7B, 32 capas con una ventana de 4.096 tokens dan un alcance de atención teórico de unos 131K tokens, aunque la información más lejana tiene que sobrevivir a muchos pasos para llegar.

Mistral 7B también guarda las claves y los valores en un búfer circular de w posiciones, escribe la posición i en la posición i mod w y sobrescribe la entrada más antigua. En una secuencia de 32K tokens, esto redujo ocho veces la memoria de la caché sin afectar a la calidad, y el tamaño fijo de la caché mantiene estable el uso de memoria, por larga que sea la generación.

Intercalado de capas locales y globales#

Las ventanas puras no pueden copiar un dato de mucho antes en un solo paso, así que los modelos recientes combinan capas SWA con capas globales completas. Gemma 2 de Google las alternaba una a una con una ventana de 4.096 tokens, y Gemma 3 utiliza cinco capas locales por cada capa global con una ventana de solo 1.024 tokens, lo que mantiene pequeña la caché KV en una ventana de contexto de hasta 128K tokens. Solo las capas globales almacenan las claves y los valores de todo el contexto, por lo que cinco de cada seis capas contienen solo 1.024 tokens cada una.

Los transformers de visión utilizan la misma combinación para imágenes de alta resolución, donde cada parche es un token. SAM 3 de Meta, que se ejecuta en el paquete Ultralytics para la segmentación de imágenes, aplica atención dentro de ventanas no superpuestas de parches de 24 × 24 en 28 de los 32 bloques de su codificador de imágenes, y atención global en toda la cuadrícula de parches de 72 × 72 de una imagen de 1.008 píxeles en los bloques 7, 15, 23 y 31.

SWA frente a términos relacionados#

• Swin Transformer: Divide la imagen en ventanas fijas y no superpuestas que comparten todos los tokens de cada ventana, y desplaza la cuadrícula entre capas para que las ventanas vecinas intercambien información.
• Atención dispersa: La familia más amplia de patrones fijos. SWA es su miembro más sencillo; otros añaden ventanas dilatadas, tokens globales o conexiones aleatorias.
• Sumideros de atención: Los primeros tokens de una secuencia, que absorben el peso de atención sobrante. Si se aplica una ventana durante la inferencia a un modelo entrenado con atención completa, expulsar esos tokens interrumpe la generación; StreamingLLM los mantiene en la caché junto a la ventana.
• Atención deformable: Cada consulta atiende a unos pocos puntos de muestreo cuyos desplazamientos aprende, de modo que el vecindario depende del contenido.
• Atención Flash: Un kernel de GPU para la atención exacta. Sus implementaciones aceptan un tamaño de ventana, por lo que los modelos SWA omiten por completo los bloques que quedan fuera de la ventana.