Difusión de atención

La atención acumulada es un método para interpretar modelos basados en Transformer (architecture)s, principalmente modelos de lenguaje grandes, mediante el cálculo de la atención acumulativa que fluye desde cada token de entrada hasta cada otro token a través de todas las capas de la red. Introducido como respuesta a la observación de que los pesos de atención en capas individuales suelen ser ruidosos y no reflejan directamente la decisión final del modelo, la atención acumulada agrega las matrices de atención a través de las capas para producir un mapa de importancia único y coherente. Esta técnica se utiliza ampliamente en el campo de la IA explicable para visualizar qué partes de una entrada contribuyen más a la predicción de un modelo, lo que ayuda en la depuración, la detección de sesgos y la comprensión del comportamiento del modelo.

La idea central se basa en que, en un transformer, el mecanismo de atención de cada capa calcula una suma ponderada de las representaciones de la capa anterior. Al multiplicar las matrices de atención de capas sucesivas, se puede rastrear cómo la información de los tokens iniciales se propaga a través de la red hasta la salida final. Esto es análogo a seguir un camino en un grafo dirigido, donde los nodos son los tokens y las aristas son los pesos de atención. La matriz resultante de la atención acumulada proporciona una visión global de la influencia de cada token, que suele ser más interpretable que examinar el mapa de atención de una sola capa.

Antecedentes y motivación

Los transformers, introducidos en el artículo de 2017 "Attention Is All You Need" por investigadores de Google DeepMind y University of Toronto, dependen de mecanismos de atención para capturar dependencias entre tokens. Cada capa calcula puntuaciones de atención que indican cuánto atiende cada token a los demás. Sin embargo, estas puntuaciones no son directamente interpretables como importancia. Por ejemplo, un token podría atender a un signo de puntuación que tiene poco peso semántico, o la atención podría distribuirse de manera difusa entre muchos tokens. Además, la atención en capas posteriores se calcula sobre representaciones que ya han sido transformadas por capas anteriores, por lo que los pesos brutos no tienen en cuenta este efecto acumulativo.

Los primeros intentos de explicabilidad se centraron en visualizar mapas de atención de cabezas o capas individuales, pero a menudo producían resultados contradictorios o confusos. Investigadores como Anima Anandkumar y Jakob Uszkoreit señalaron que los patrones de atención pueden variar mucho entre capas y cabezas, lo que dificulta extraer conclusiones. Esto motivó el desarrollo de métodos que agregaran información a través de toda la red.

El algoritmo

La atención acumulada opera bajo el supuesto de que los pesos de atención representan una forma de enrutamiento de información. El algoritmo procede de la siguiente manera:

1. Para cada capa \( l \), se obtiene la matriz de atención \( A_l \) de forma (longitud_de_secuencia, longitud_de_secuencia), donde \( A_l[i][j] \) es el peso de atención desde el token \( i \) al token \( j \). Esta matriz se promedia típicamente sobre todas las cabezas de atención de esa capa.
2. Se añade la matriz identidad a cada matriz de atención para tener en cuenta la conexión residual, que permite que cada token conserve su propia información. La matriz resultante es \( \tilde{A}_l = A_l + I \).
3. Se normaliza cada fila de \( \tilde{A}_l \) para que sume 1, asegurando que las matrices sigan siendo estocásticas.
4. Se calcula la matriz de atención acumulada \( R \) multiplicando las matrices normalizadas a través de todas las capas: \( R = \tilde{A}_1 \cdot \tilde{A}_2 \cdot ... \cdot \tilde{A}_L \), donde \( L \) es el número de capas.

La matriz resultante \( R \) proporciona la atención total que fluye desde cada token hasta cualquier otro, considerando todos los caminos posibles a través de la red. La entrada \( R[i][j] \) puede interpretarse como la proporción de información del token \( j \) que influye en la representación del token \( i \) en la salida.

Este método fue propuesto por primera vez en un artículo de 2020 por investigadores de OpenAI y Stanford AI Lab, quienes demostraron su eficacia en GPT-2 y otros modelos. Mostraron que la atención acumulada podía identificar tokens relevantes en tareas como análisis de sentimiento o resolución de pronombres, superando a menudo a la atención de una sola capa en términos de alineación con juicios humanos.

Aplicaciones en la interpretación de modelos

La atención acumulada se ha aplicado a una variedad de tareas en el procesamiento del lenguaje natural y más allá. En la investigación de aprendizaje automático, se utiliza para:

• Identificar tokens clave: Examinando la matriz de atención acumulada, los profesionales pueden ver qué palabras o subpalabras de entrada influyen más en la salida del modelo. Por ejemplo, en una tarea de clasificación de sentimientos, la atención acumulada puede resaltar palabras como "terrible" o "increíble" como altamente influyentes.
• Detectar sesgos: Si un modelo atiende consistentemente a términos demográficos de una manera que conduce a predicciones sesgadas, la atención acumulada puede revelar estos patrones, ayudando a los investigadores a mitigar comportamientos injustos.
• Depurar errores del modelo: Cuando un modelo produce una predicción incorrecta, la atención acumulada puede mostrar si se centró en partes irrelevantes de la entrada, guiando mejoras en los datos de entrenamiento o en la arquitectura.
• Comparar modelos: Calculando matrices de atención acumulada para diferentes modelos, se puede comparar su flujo interno de información, lo que es útil para la selección de modelos y para entender diferencias arquitectónicas.

Más allá del texto, la atención acumulada se ha adaptado para transformers de visión (ViTs) utilizados en tareas de visión por computadora. Por ejemplo, investigadores de Google Cloud y Amazon Web Services la han aplicado a la clasificación de imágenes para visualizar qué regiones de una imagen contribuyen a una predicción, lo que ayuda en el análisis de imágenes médicas y en sistemas de conducción autónoma.

Limitaciones y críticas

A pesar de su popularidad, la atención acumulada tiene varias limitaciones. En primer lugar, el supuesto de que los pesos de atención pueden tratarse como distribuciones de probabilidad del flujo de información no siempre es válido. Los pesos de atención se calculan sobre representaciones aprendidas y no reflejan necesariamente una relación causal. Críticos como Melanie Mitchell y Aleksander Madry han argumentado que la atención no es explicación y que métodos como la atención acumulada pueden producir interpretaciones engañosas.

En segundo lugar, la multiplicación de matrices de atención a través de capas puede llevar a problemas numéricos, especialmente en redes profundas con muchas capas. La matriz de atención acumulada puede volverse demasiado difusa, con todos los tokens teniendo importancia similar, o puede concentrarse en unos pocos tokens debido a la multiplicación repetida. Esto puede hacer que los resultados sean menos útiles en la práctica.

En tercer lugar, la atención acumulada no tiene en cuenta las transformaciones no lineales (como las redes feed-forward o la normalización de capas) que ocurren entre las capas de atención. Estas transformaciones pueden alterar significativamente el flujo de información, por lo que el modelo multiplicativo simple puede pasar por alto efectos importantes.

En cuarto lugar, el método promedia la atención sobre todas las cabezas, lo que puede ocultar los roles distintos de diferentes cabezas. Algunas cabezas pueden capturar relaciones sintácticas, mientras que otras capturan relaciones semánticas; promediarlas pierde esta matiz.

Variantes y extensiones

Para abordar algunas de estas limitaciones, los investigadores han propuesto varias variantes:

• Flujo de atención: Este método, introducido por investigadores de MIT CSAIL, trata la atención como un problema de flujo en un grafo y utiliza algoritmos de flujo máximo para calcular la importancia de los tokens, evitando el supuesto multiplicativo.
• Atención acumulada con pesos residuales: Algunas implementaciones ponderan la matriz identidad por un factor que refleja la fuerza de la conexión residual, en lugar de añadirla uniformemente.
• Propagación de relevancia por capas (LRP): Esta técnica, proveniente de la comunidad de aprendizaje profundo, propaga puntuaciones de relevancia hacia atrás a través de la red, proporcionando una alternativa a los métodos basados en atención.
• Gradientes integrados y SHAP: Estos son métodos de atribución de características que no dependen de los pesos de atención, sino de información de gradientes o conceptos de teoría de juegos. A menudo se utilizan como verificación de cordura frente a los métodos basados en atención.

A pesar de estas alternativas, la atención acumulada sigue siendo una línea base popular debido a su simplicidad y eficiencia computacional. No requiere entrenamiento adicional ni cálculo de gradientes, lo que la hace fácil de aplicar a cualquier transformer preentrenado.

Consideraciones de implementación

Implementar la atención acumulada es sencillo con los marcos de trabajo modernos de aprendizaje profundo. En PyTorch o TensorFlow, se pueden utilizar hooks en el paso hacia adelante para capturar las matrices de atención. Los pasos clave son:

1. Registrar hooks en cada capa de atención para almacenar los pesos de atención.
2. Después del paso hacia adelante, promediar los pesos de atención sobre las cabezas.
3. Añadir la matriz identidad y normalizar las filas.
4. Multiplicar las matrices secuencialmente.

Una consideración práctica es que las matrices de atención pueden ser grandes para secuencias largas, lo que genera una sobrecarga de memoria. Para secuencias de longitud 1024, cada matriz es de 1024x1024, y multiplicar muchas de estas matrices puede ser computacionalmente costoso. Sin embargo, para entradas típicas, esto es manejable.

Otra consideración es que el método asume una arquitectura de transformer estándar. Para modelos con atención cruzada (como los modelos codificador-decodificador), la atención acumulada debe adaptarse para manejar la interacción entre la atención del codificador y la del decodificador. En tales casos, se pueden calcular acumulaciones separadas para el codificador y el decodificador, o combinarlas de manera más compleja.

Relación con otros métodos de explicabilidad

La atención acumulada se sitúa dentro de un panorama más amplio de técnicas de explicabilidad para sistemas de inteligencia artificial. A menudo se compara con:

• Métodos basados en gradientes: Estos calculan el gradiente de la salida con respecto a las representaciones de entrada, proporcionando una medida de sensibilidad. Son más fundamentados teóricamente pero requieren retropropagación.
• Métodos basados en perturbaciones: Estos implican modificar los tokens de entrada (por ejemplo, eliminándolos o enmascarándolos) y observar el cambio en la salida. Son independientes del modelo pero computacionalmente costosos.
• Modelos sustitutos: Técnicas como LIME o SHAP entrenan modelos interpretables localmente para aproximar el modelo de caja negra. Son útiles para datos tabulares pero menos para texto.

La atención acumulada es única en el sentido de que se basa puramente en los pesos de atención internos del modelo, sin requerir cálculos adicionales más allá de un paso hacia adelante. Esto la hace especialmente atractiva para aplicaciones en tiempo real, como herramientas de depuración interactivas.

Direcciones futuras

A medida que los modelos transformer continúan creciendo en tamaño y complejidad, aumenta la necesidad de métodos de explicabilidad robustos. Es probable que la atención acumulada evolucione en varias direcciones:

• Integración con métodos causales: Combinar la atención acumulada con técnicas de inferencia causal podría proporcionar atribuciones más precisas.
• Manejo de modelos multimodales: Extender la atención acumulada a modelos que procesan texto, imágenes y audio simultáneamente, como los desarrollados por OpenAI y Anthropic.
• Análisis automatizado: Utilizar las salidas de la atención acumulada para generar automáticamente explicaciones en lenguaje natural de las decisiones del modelo, lo que podría ser útil para cumplimiento normativo y auditorías.
• Evaluación comparativa: Desarrollar estándares de referencia para evaluar la fidelidad de los métodos de explicabilidad, como proponen investigadores de Carnegie Mellon University y BAIR (Berkeley AI Research).

A pesar de sus limitaciones, la atención acumulada se ha convertido en una herramienta fundamental en el kit de herramientas de explicabilidad. Su simplicidad y eficacia la han convertido en una línea base estándar en muchos artículos de investigación y aplicaciones prácticas. A medida que el campo avanza, es probable que se refine y se combine con otras técnicas para proporcionar una comprensión más profunda del funcionamiento interno de las redes neuronales.

Conclusión

La atención acumulada proporciona una forma práctica e intuitiva de entender cómo los transformers procesan la información. Al propagar los pesos de atención a través de las capas, ofrece una perspectiva global sobre la influencia de los tokens que suele ser más útil que inspeccionar capas individuales. Aunque no está exenta de defectos, su facilidad de uso e interpretabilidad han consolidado su lugar en la literatura de explicabilidad. Para cualquier persona que trabaje con modelos transformer, comprender la atención acumulada es un paso valioso para desmitificar estos sistemas potentes pero opacos.