Attention Sinks

Un sumidero de atención es un token que recibe una gran proporción del peso de atención en muchas cabezas y capas, aunque su contenido no sea pertinente para la consulta; suele ser el primer token de la secuencia. Guangxuan Xiao y sus colegas del MIT, Meta AI y Carnegie Mellon University dieron nombre a este efecto en Modelos de lenguaje de transmisión eficientes con sumideros de atención, publicado en ICLR 2024. Demostraron que los modelos de lenguaje grandes (LLMs), como Llama 2, MPT, Falcon y Pythia, fallan cuando una caché KV deslizante expulsa los primeros tokens, y que conservar cuatro de ellos restablece una generación estable.

Por qué se forman los sumideros#

La autoatención convierte las puntuaciones de cada consulta en pesos mediante una softmax, y esos pesos siempre suman uno. Cuando ningún token del contexto es pertinente para una cabeza, el peso tiene que asignarse igualmente a algún sitio. En un modelo de lenguaje causal, todos los tokens posteriores pueden ver los primeros, así que el entrenamiento los convierte en el destino predeterminado del peso que una cabeza no necesita. En el análisis de StreamingLLM sobre Llama 2, todas las capas salvo las dos inferiores concentraban la atención en los tokens iniciales en todas las cabezas. Un logit de sumidero añadido a la softmax muestra el efecto:

import torch

torch.manual_seed(0)
values = torch.randn(6, 4)
# A query that matches none of the 6 tokens: every score is close to zero
scores = torch.tensor([0.1, -0.2, 0.0, 0.2, -0.1, 0.0])

# Softmax must hand out a total weight of 1, so the output mixes all six values
weights = scores.softmax(dim=-1)
print(f"No sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

# A sink logit joins the softmax but carries no value, so it absorbs the spare weight
sink = torch.tensor([3.0])
weights = torch.cat([sink, scores]).softmax(dim=-1)[1:]
print(f"With sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

Sin un sumidero, la cabeza reparte todo su peso entre seis valores irrelevantes y devuelve una salida con norma 0,83. Con un logit de sumidero de 3, el sumidero recibe el 77 % del peso y la salida se reduce a 0,19, por lo que la cabeza apenas contribuye a esta consulta. Los modelos de pesos abiertos gpt-oss de OpenAI incorporan esta función con un logit de sumidero aprendido para cada cabeza de atención, de modo que cada cabeza dispone de un lugar para el peso sobrante que no aporta nada a la salida.

Generación en streaming con StreamingLLM#

Un modelo que genera texto indefinidamente supera cualquier ventana de contexto y no puede conservar todos los tokens en su caché KV, y la atención de ventana convencional, que solo conserva los tokens más recientes, falla en cuanto desaparecen los primeros tokens. StreamingLLM conserva permanentemente las claves y los valores de cuatro tokens iniciales junto a una ventana móvil de los más recientes, y asigna las posiciones según el lugar que ocupan en la caché en lugar de en el texto original, volviendo a aplicar las incrustaciones posicionales rotatorias (RoPE) en cada paso:

Esto mantuvo estable el modelado del lenguaje durante más de 4 millones de tokens y alcanzó una velocidad hasta 22,2 veces superior a la de una ventana deslizante que recalcula la caché con cada token nuevo. La ventana sigue limitando lo que puede recordar el modelo: los tokens expulsados de la parte intermedia desaparecen para siempre, así que un chatbot o un agente de IA en streaming solo recuerda sus interacciones más recientes. Los autores también preentrenaron modelos con un token de marcador de posición dedicado y entrenable al principio de cada muestra, que asumía por sí solo la función de sumidero.

Sumideros de atención y términos relacionados#

• Atención de ventana deslizante: Limita cada token a los vecinos recientes. Aplicada durante la inferencia a un modelo entrenado con atención completa, falla cuando los sumideros salen de la ventana.
• Deterioro del contexto: Descenso de la precisión a medida que crece un prompt, aunque siga cabiendo en la ventana de contexto.

Sumideros en modelos de visión#

Los transformadores de visión (ViT) muestran un comportamiento relacionado. En Los transformadores de visión necesitan registros, investigadores de Meta e Inria detectaron tokens con normas inusualmente altas en parches de fondo con poca información de ViT supervisados y autosupervisados, reutilizados para el cálculo interno. Añadir unos cuantos tokens de registro a la entrada absorbía esa función, eliminaba los artefactos de los mapas de características y de los mapas de atención, y mejoraba la predicción densa.

SAM 2 de Meta, disponible en el paquete Ultralytics, gestiona la memoria de vídeo con un plan similar que conserva el ancla y una ventana: en cada fotograma nuevo, su atención de memoria lee las entradas del banco de memoria de los fotogramas indicados por el usuario, que se conservan durante todo el vídeo, además de las de los 6 fotogramas más recientes, de modo que el coste de atención por fotograma se mantiene fijo, por larga que sea la duración del vídeo.