La atención de múltiples cabezas es un componente central de la arquitectura transformador en el aprendizaje profundo. Extiende el mecanismo de atención básico calculando varias operaciones de atención en paralelo, cada una con sus propias proyecciones aprendidas de consultas, claves y valores. Esto permite que el modelo capture diferentes tipos de relaciones y patrones dentro de la secuencia de entrada, como dependencias sintácticas, correferencias o señales posicionales, a través de múltiples subespacios de representación. Las salidas de estas cabezas de atención paralelas se concatenan y se proyectan linealmente para producir el resultado final, lo que permite que el modelo atienda conjuntamente a información de diferentes subespacios de representación en diferentes posiciones.
El concepto se introdujo en el artículo de 2017 "Attention Is All You Need" de investigadores de Google, incluyendo a Jakob Uszkoreit, Lukasz Kaiser y Niki Parmar. El artículo propuso el Transformer, que se basa enteramente en mecanismos de atención y prescinde de la recurrencia y las convoluciones, lo que condujo a mejoras significativas en la velocidad de entrenamiento y el rendimiento en tareas de secuencia a secuencia. Desde entonces, la atención de múltiples cabezas se ha convertido en un elemento fundamental de los modelos de lenguaje grandes como GPT, BERT y T5, y se utiliza ampliamente en el procesamiento del lenguaje natural, la visión por computadora y el reconocimiento de voz.
Mecanismo
En el mecanismo de atención, cada token de entrada se representa como un vector, a menudo llamado incrustación de token. Para una secuencia de tokens, el modelo calcula tres matrices: consultas (Q), claves (K) y valores (V). La puntuación de atención entre una consulta y una clave se calcula típicamente como un producto escalar, escalado por la raíz cuadrada de la dimensión de la clave. Estas puntuaciones se normalizan usando una función softmax para producir pesos de atención, que luego se utilizan para calcular una suma ponderada de los valores.
La atención de múltiples cabezas ejecuta este proceso varias veces en paralelo, cada una con diferentes proyecciones lineales aprendidas de las consultas, claves y valores originales. Por ejemplo, con 8 cabezas, el modelo calcula 8 salidas de atención separadas, cada una enfocándose en diferentes aspectos de la secuencia. Las salidas se concatenan y se pasan a través de una capa lineal final. Este diseño permite que el modelo atienda a información de diferentes subespacios de representación, lo que es particularmente útil para capturar patrones lingüísticos diversos.
A diferencia de los pesos "duros", que se establecen durante el entrenamiento, los pesos de atención son "blandos" y se calculan durante el paso hacia adelante, lo que significa que cambian con cada entrada. Esto permite que el modelo se enfoque dinámicamente en partes relevantes de la entrada, como se ilustra con el ejemplo de alineación en la traducción automática: al traducir "I love you" al francés, el modelo asigna pesos de atención altos a "I" al generar "je", a "you" al generar "t'" y a "love" al generar "aime".
Historia
El mecanismo de atención se desarrolló para abordar las debilidades de las redes neuronales recurrentes (RNN), que tienden a favorecer la información al final de una oración y atenuar la importancia de las palabras anteriores. Los primeros mecanismos de atención, como la atención aditiva estilo Bahdanau (2014) y la atención multiplicativa estilo Luong (2015), se integraron en arquitecturas RNN codificador-decodificador para la traducción automática. Estos permitieron que el decodificador accediera directamente a todos los estados ocultos del codificador, en lugar de depender únicamente del estado oculto final.
El gran avance llegó con la autoatención, donde cada elemento en la secuencia de entrada atiende a todos los demás, lo que permite que el modelo capture dependencias globales. Esta idea fue central para el Transformer, que reemplazó la recurrencia por completo con mecanismos de atención. La arquitectura paralelizable del Transformer permitió aceleraciones significativas en el entrenamiento, lo que llevó a su adopción en modelos como BERT, T5 y transformadores generativos preentrenados (GPT).
Variantes
Se han desarrollado varias variantes de atención, muchas implementando pesos blandos. Estas incluyen:
- Programadores de pesos rápidos (1992): Una red neuronal "lenta" genera los pesos "rápidos" de otra red neuronal mediante productos externos. Esto se renombró más tarde como "autoatención linealizada".
- Atención estilo Bahdanau (atención aditiva): Utiliza una red de avance para calcular puntuaciones de alineación.
- Atención estilo Luong (atención multiplicativa): Utiliza productos escalares entre estados del decodificador y del codificador.
- Atención posicional y atención posicional factorizada: Incorporan información posicional en los cálculos de atención.
- Atención espacial y de canal para redes neuronales convolucionales: Operan en dimensiones espaciales o canales de características.
Estas variantes recombinan las entradas del lado del codificador para redistribuir efectos a cada salida objetivo, a menudo usando una matriz de correlación de productos escalares para re-ponderar coeficientes.
Optimizaciones
El tamaño de la matriz de atención es proporcional al cuadrado del número de tokens de entrada, lo que puede ser intensivo en memoria para secuencias largas. Flash attention es una implementación que reduce las necesidades de memoria y aumenta la eficiencia sin sacrificar precisión. Particiona el cálculo de atención en bloques más pequeños que caben en la memoria en chip más rápida de la GPU, reduciendo la necesidad de almacenar grandes matrices intermedias.
FlexAttention, desarrollado por Meta, es un kernel de atención que permite a los usuarios modificar las puntuaciones de atención antes del softmax y elige dinámicamente el algoritmo de atención óptimo, proporcionando flexibilidad para patrones de atención personalizados.
Aplicaciones
La atención se utiliza ampliamente en el procesamiento del lenguaje natural, la visión por computadora y el reconocimiento de voz. En el PLN, mejora la comprensión del contexto en tareas como la respuesta a preguntas y el resumen. En visión, la atención visual ayuda a los modelos a enfocarse en regiones relevantes de la imagen, mejorando la detección de objetos y la generación de subtítulos de imágenes.
En visión por computadora, los transformadores de visión (ViT) aplican atención de múltiples cabezas a parches de imagen. Los mapas de atención, que visualizan las puntuaciones de atención como mapas de calor, se han convertido en una forma rutinaria de inspeccionar el proceso de toma de decisiones de los modelos ViT. Las capas más profundas tienden a mostrar visualizaciones más semánticamente significativas. Attention rollout es un algoritmo recursivo que combina las puntuaciones de atención a través de todas las capas calculando el producto escalar de mapas de atención sucesivos.
Debido a que los transformadores de visión se entrenan típicamente de manera autosupervisada, los mapas de atención generalmente no son sensibles a la clase. Cuando se adjunta una cabeza de clasificación, los mapas de atención discriminativos de clase (CDAM) combinan mapas de atención y gradientes con respecto al token de clase para proporcionar explicaciones específicas de clase.
Impacto
La atención de múltiples cabezas ha sido instrumental en el auge de la IA generativa y los modelos de lenguaje grandes. Permite que los modelos procesen secuencias largas de manera eficiente y capturen dependencias complejas, convirtiéndose en un componente clave en sistemas desarrollados por organizaciones como OpenAI, Anthropic y Google DeepMind. La arquitectura también ha influido en el diseño de hardware, con empresas como NVIDIA y Cerebras optimizando chips para cálculos basados en atención.
La investigación continúa en la mejora de la eficiencia e interpretabilidad de la atención. Técnicas como la atención dispersa, la atención lineal y los métodos basados en kernels buscan reducir la complejidad cuadrática de la atención, mientras que las herramientas de interpretabilidad ayudan a comprender qué aprenden los modelos. La atención de múltiples cabezas sigue siendo un área vibrante de estudio en el aprendizaje automático y la inteligencia artificial.