Attention Rollout/es

Article
Topic area	Deep Learning
Prerequisites	Transformer, Attention Mechanism, Self-Attention

Resumen

El rollout de atención es una técnica para cuantificar cuánta información fluye desde cada token de entrada hasta cada posición de salida a lo largo de todas las capas de un Transformer, multiplicando recursivamente las matrices de pesos de atención de cada capa. Fue introducido por Abnar y Zuidema en 2020 como remedio al conocido problema de que los pesos de atención sin procesar de una sola capa son una aproximación pobre de la contribución de cada token de entrada a una predicción dada, ya que los tokens se mezclan repetidamente con sus vecinos por la Self-Attention en cada capa.

El rollout trata la pila de capas como un grafo y calcula la accesibilidad ponderada por trayectorias desde cualquier nodo de salida hacia atrás hasta las entradas. La matriz resultante se utiliza ampliamente como un mapa de saliencia para modelos basados en transformer, en particular Vision Transformers (ViTs) y codificadores al estilo de BERT, donde es una de las herramientas estándar para la interpretabilidad y para producir localización discriminativa por clase sin reentrenamiento.

Motivación

Los pesos de atención de una sola capa tienen un atractivo intuitivo como explicaciones: la posición $ i $ "mira" la posición $ j $ con peso $ a_{ij} $. Sin embargo, en un transformer profundo, la representación en la capa $ \ell $ ya es una mezcla de todos los tokens de entrada, por lo que un valor alto de $ a_{ij} $ en una capa tardía no significa que la predicción en la posición $ i $ dependa principalmente del token de entrada original $ j $. Varios estudios han argumentado que "la atención no es explicación" precisamente debido a esta contaminación a través de las capas.

El rollout de atención aborda esto formulando una pregunta diferente: si trazamos la información hacia atrás a través de la red, ¿cuánta de la representación en una posición de salida dada se originó en cada token de entrada? La respuesta requiere componer las matrices de atención por capa, no inspeccionarlas de forma aislada.

Formulación

Configuración

Considere un Transformer con $ L $ capas, longitud de secuencia $ n $ y $ H $ cabezas de atención por capa. Sea $ A^{(\ell)} \in \mathbb{R}^{n \times n} $ la matriz de pesos de atención en la capa $ \ell $, donde las filas corresponden a posiciones de consulta y las columnas a posiciones de clave, y cada fila suma uno. Cuando la capa tiene varias cabezas, las matrices por cabeza $ A^{(\ell, h)} $ se promedian primero:

$ {\displaystyle A^{(\ell)} = \frac{1}{H} \sum_{h=1}^{H} A^{(\ell, h)}.} $

A veces se utilizan otras reducciones, como tomar el máximo entre cabezas o una suma ponderada aprendida.

Corrección residual

Una capa de pura autoatención no es el único camino a través de un bloque transformer: una conexión residual suma la entrada de la capa a su salida, por lo que una fracción no trivial de la señal en cada posición es simplemente la representación previa de esa misma posición. Para tener esto en cuenta, el rollout reemplaza la matriz de atención sin procesar por

$ {\displaystyle \tilde{A}^{(\ell)} = \frac{1}{2}\bigl(A^{(\ell)} + I\bigr),} $

donde $ I $ es la matriz identidad. El factor de un medio preserva la propiedad estocástica por filas: cada fila de $ \tilde{A}^{(\ell)} $ sigue sumando uno, por lo que puede interpretarse como una distribución de probabilidad sobre las posiciones de entrada. El término identidad codifica la contribución del flujo residual.

Producto recursivo

La matriz de rollout en la capa $ \ell $ es el producto acumulado

$ {\displaystyle R^{(\ell)} = \tilde{A}^{(\ell)} \, \tilde{A}^{(\ell-1)} \cdots \tilde{A}^{(1)} = \prod_{k=\ell}^{1} \tilde{A}^{(k)}.} $

La entrada $ R^{(L)}_{ij} $ se interpreta como la proporción de la representación en la posición de salida $ i $ que puede atribuirse, a través de trayectorias de atención compuestas, a la posición de entrada $ j $. Como cada $ \tilde{A}^{(\ell)} $ es estocástica por filas, también lo es $ R^{(L)} $, lo que hace que el resultado pueda visualizarse directamente como un mapa de calor sobre los tokens de entrada.

Para un ViT de clasificación o un modelo al estilo de BERT que utiliza un token CLS especial, la fila de $ R^{(L)} $ correspondiente al índice CLS proporciona un mapa de saliencia sobre los parches o tokens de entrada para la clase predicha.

Algoritmo

El procedimiento es sencillo de implementar sobre cualquier transformer que exponga sus pesos de atención:

1. Realice un pase hacia adelante y almacene $ A^{(\ell, h)} $ para cada capa y cada cabeza.
2. Para cada capa, promedie sobre las cabezas para obtener $ A^{(\ell)} $.
3. Sume la identidad y renormalice para obtener $ \tilde{A}^{(\ell)} $.
4. Inicialice $ R \leftarrow I $ y multiplique $ R \leftarrow \tilde{A}^{(\ell)} R $ desde la primera capa hasta la última.
5. Extraiga la fila de interés (típicamente CLS) y reorganícela en la cuadrícula de entrada.

El coste total es de $ O(L n^2) $ productos de matriz-matriz, dominado por la multiplicación por capa. Para las longitudes de secuencia modestas típicas en clasificación de imágenes, todo el rollout es barato en comparación con un pase hacia adelante.

Flujo de atención

El artículo original de Abnar y Zuidema también propone el flujo de atención (attention flow), una variante relacionada pero más costosa. En lugar de multiplicar las matrices por capa, el flujo de atención trata el grafo de atención como un grafo dirigido capacitado y calcula un flujo máximo desde cada nodo de entrada hasta cada nodo de salida. El flujo es más fiel a los cuellos de botella de la red, ya que la multiplicación puede sobrecontar trayectorias que comparten una arista, pero su coste es superlineal en $ n $ y rara vez se utiliza a gran escala. El rollout es la variante dominante en la práctica.

Aplicaciones

El rollout de atención se ha convertido en una herramienta de visualización por defecto para varias familias de modelos:

• Vision Transformers: los mapas de calor de rollout, restringidos a la fila CLS de $ R^{(L)} $ y reorganizados en la cuadrícula de parches, producen localizaciones discriminativas por clase competitivas con métodos basados en gradientes como Grad-CAM en ImageNet y en bancos de pruebas de segmentación débilmente supervisada.
• Modelos de lenguaje: el rollout se utiliza para inspeccionar en qué tokens de entrada se apoya un codificador BERT o RoBERTa para una predicción dada, complementando los estudios de sondeo por cabeza.
• Transformers multimodales: en modelos al estilo de CLIP y modelos imagen-texto, las matrices de rollout restringidas a las capas de atención cruzada revelan qué regiones de la imagen sustentan un determinado token de texto, lo que apoya la segmentación y el anclaje de vocabulario abierto.
• Auditoría de modelos: comparando los mapas de rollout antes y después de un paso de ajuste fino o de un cambio de dominio, los profesionales pueden detectar si un modelo ha cambiado su dependencia de una región de entrada a otra.

Variantes y extensiones

Varias extensiones del rollout básico mejoran su fidelidad o su especificidad por clase:

• Rollout ponderado por gradiente: combina la matriz de atención con los gradientes de la puntuación de la clase predicha respecto a los valores de atención, reemplazando $ A^{(\ell)} $ en el producto de rollout por $ (\nabla_A y)\odot A^{(\ell)} $ evaluado para la clase relevante. Chefer et al. muestran que esto produce mapas más nítidos y más discriminativos por clase que el rollout convencional en ViTs.
• Rollout específico por cabeza: en lugar de promediar las cabezas, el producto de rollout se calcula por cabeza y los resultados se agregan, exponiendo el papel de cada cabeza individual.
• Rollout codificador-decodificador: para modelos de traducción y seq2seq, el rollout se calcula por separado a lo largo de la cadena de autoatención del codificador, las capas de atención cruzada y la cadena de autoatención del decodificador, y las matrices se componen de extremo a extremo.
• Rollout disperso y podado: se descartan los ceros por debajo de un umbral antes de la multiplicación para enfocar la visualización en las trayectorias dominantes.
• Rollout de atención top-k: conserva solo las entradas top-k por fila en cada $ \tilde{A}^{(\ell)} $ antes de la composición, motivado por la observación de que las distribuciones de atención suelen tener colas pesadas.

Comparación con otros métodos de saliencia

El rollout es una de varias familias de herramientas de explicación aplicadas a transformers:

• En comparación con la atención sin procesar, el rollout corrige la mezcla multicapa que hace que la atención de una sola capa sea engañosa.
• En comparación con Grad-CAM y otros métodos basados en gradientes, el rollout es puramente forward y no requiere derivar a través de la predicción; esto lo hace barato y agnóstico a la arquitectura, pero también menos específico por clase a menos que se incorporen gradientes.
• En comparación con gradientes integrados y métodos al estilo de SHAP, el rollout es mucho más rápido (un pequeño número de productos matriciales frente a muchos pases hacia adelante), pero no satisface las mismas garantías axiomáticas.
• En comparación con el flujo de atención, el rollout es una aproximación tratable que produce mapas de saliencia cualitativamente similares en la mayoría de los casos.

Limitaciones

Al usar el rollout de atención como explicación se aplican varias salvedades:

• No utiliza la ruta de los valores: el rollout solo inspecciona los pesos de atención e ignora las proyecciones de valor, que son las que transportan el contenido real. Dos cabezas con patrones de atención idénticos pero matrices de valor muy distintas producirían rollouts indistinguibles.
• La corrección por identidad es heurística: el peso igual entre $ A $ e $ I $ es un convenio en lugar de una cantidad derivada. La verdadera masa residual depende de las normas relativas de las contribuciones residual y de atención, que varían entre capas y tokens.
• Agnosticismo de clase: el rollout convencional produce el mismo mapa independientemente de la clase predicha; se necesitan variantes ponderadas por gradiente para obtener explicaciones discriminativas por clase.
• Debates sobre fidelidad: las evaluaciones formales de las explicaciones basadas en atención han producido resultados mixtos, y el rollout hereda gran parte de esta crítica. Los usuarios deben tratar los mapas de rollout como resúmenes informativos en lugar de como atribuciones de verdad fundamental.
• Suposiciones arquitectónicas: el rollout, en su forma estándar, supone una pila de bloques de atención idénticos con conexiones residuales. Adaptarlo a mezcla de expertos, atención dispersa o transformers basados en enrutamiento requiere cuidado.