Self-Attention

La autoatención es el tipo de atención en el que las consultas, las claves y los valores proceden de la misma secuencia, de modo que cada token construye su nueva representación a partir de los demás tokens de esa secuencia. Una capa de autoatención puede conectar en un solo paso la primera y la última palabra de un párrafo, o esquinas opuestas de una imagen, mientras que una red neuronal recurrente (RNN) necesita un paso por posición y una red neuronal convolucional (CNN) necesita apilar muchas capas para ampliar su campo de visión. Apilada con capas prealimentadas, forma el Transformer presentado en el artículo de 2017 La atención es todo lo que necesitas, que constituye la base de los modelos de lenguaje grandes (LLM), los Transformers de visión (ViT) y los detectores basados en Transformer.

Cómo funciona la autoatención#

Cada token de entrada, representado como un embedding, se multiplica por tres matrices aprendidas para obtener una consulta, una clave y un valor. La salida del token i es una media ponderada del vector de valor de cada token, con pesos definidos mediante un softmax aplicado a los productos escalares escalados de la consulta del token i con todas las claves; este es el mecanismo de atención de producto escalar escalado descrito en la página sobre el mecanismo de atención. Como las consultas, las claves y los valores son proyecciones de los mismos tokens, una secuencia de n tokens produce una matriz de pesos n × n. En «El animal no cruzó la calle porque estaba demasiado cansado», la fila de «estaba» puede asignar la mayor parte de su peso a «animal» e incorporar el significado de esa palabra a su propia representación.

La autoatención multicabezal ejecuta en paralelo varias de estas operaciones con proyecciones más pequeñas y concatena los resultados, de modo que las distintas cabezas pueden seguir relaciones diferentes, como la sintaxis en una cabeza y la correferencia en otra. El Transformer original utilizaba 8 cabezas de 64 dimensiones cada una en un modelo de 512 dimensiones.

Ceguera al orden y enmascaramiento causal#

La autoatención trata la entrada como un conjunto. Si se desordenan los tokens, las salidas se desordenan del mismo modo y no cambia nada más, por lo que los modelos añaden información posicional, como los embeddings sinusoidales del Transformer original o el embedding posicional rotatorio (RoPE) de la mayoría de los LLM actuales. Los codificadores como BERT permiten que cada token atienda en ambas direcciones. Los LLM solo con decodificador aplican una máscara causal para que cada token atienda únicamente a sí mismo y a los tokens anteriores, lo que permite al modelo generar texto de izquierda a derecha y reutilizar las claves y los valores de los tokens anteriores desde una caché KV. MultiheadAttention de PyTorch muestra ambas propiedades cuando se pasa el mismo tensor como consulta, clave y valor:

import torch

torch.manual_seed(0)
tokens = torch.randn(1, 5, 16)  # una secuencia de 5 tokens con 16 características cada uno
attn = torch.nn.MultiheadAttention(embed_dim=16, num_heads=2, batch_first=True)

# Autoatención: las consultas, las claves y los valores son proyecciones de los mismos tokens
out, weights = attn(tokens, tokens, tokens)
print("Attention weights shape:", tuple(weights.shape))

# Al desordenar los tokens, solo se desordenan las salidas, por lo que el orden de las palabras es invisible sin información posicional
perm = torch.tensor([3, 0, 4, 1, 2])
out_shuffled, _ = attn(tokens[:, perm], tokens[:, perm], tokens[:, perm])
print("Same outputs after shuffling:", torch.allclose(out_shuffled, out[:, perm], atol=1e-6))

# Una máscara causal bloquea la atención a los tokens posteriores, como en un LLM solo con decodificador
mask = torch.triu(torch.ones(5, 5, dtype=torch.bool), diagonal=1)
_, causal_weights = attn(tokens, tokens, tokens, attn_mask=mask)
print(causal_weights[0].detach().round(decimals=2))

Imprime la forma de los pesos (1, 5, 5), True para la prueba de desorden y una matriz triangular inferior: el primer token asigna un peso de 1.00 a sí mismo, el segundo reparte 0.31 y 0.69 entre los dos primeros tokens, y todas las entradas por encima de la diagonal son 0.00.

Autoatención y términos relacionados#

• Mecanismo de atención: Operación general que pondera los valores según la similitud entre consultas y claves. La autoatención es el caso en el que los tres proceden de una misma secuencia.
• Atención cruzada: Las consultas proceden de una secuencia, y las claves y los valores, de otra. Un decodificador de traducción atiende a la frase de origen codificada, las consultas de objetos de un decodificador DETR atienden a las características de la imagen y un modelo de difusión de texto a imagen atiende a la indicación.
• Convolución: Aplica el mismo filtro pequeño en cada posición, de modo que cada salida ve una ventana local fija cuyos pesos no dependen del contenido. La autoatención calcula sus pesos a partir del contenido y abarca toda la entrada.

El Transformer original de codificador-decodificador combina autoatención bidireccional, autoatención causal y atención cruzada:

Coste y variantes eficientes#

Calcular la puntuación de cada par de tokens hace que el tiempo y la memoria crezcan con el cuadrado de la longitud de la secuencia: 1.000 tokens generan un millón de puntuaciones por cabeza y 100.000 tokens generan diez mil millones. FlashAttention calcula la autoatención exacta en bloques sin almacenar la matriz completa, la atención de consultas agrupadas (GQA) comparte claves y valores entre cabezas para reducir la caché KV, y la atención de ventana deslizante, dispersa y lineal reduce el número de pares. En visión, el número de tokens depende de la resolución de la imagen, por lo que muchos detectores aplican autoatención solo donde los mapas de características son más pequeños. YOLO26 ejecuta autoatención multicabezal solo en el bloque C2PSA al final de su backbone y en un bloque C3k2 con atención del neck, ambos sobre el mapa de características de 20×20 y stride 32 de una imagen de 640 píxeles, es decir, 400 tokens. Su módulo de atención calcula las consultas, las claves y los valores con una convolución 1×1 sobre ese mapa y añade información posicional mediante una convolución depthwise 3×3 aplicada a los valores.