Selective Aggregation of Attention Maps Improves Diffusion-Based Visual Interpretation

Resumen Técnico

1. Planteamiento del Problema

Los modelos generativos de texto a imagen (T2I), como Stable Diffusion e Imagen, han logrado resultados impresionantes alineando imágenes de alta calidad con prompts de texto. Estos modelos utilizan capas de atención cruzada (cross-attention) para guiar la generación. Aunque métodos existentes como DAAM (Diffusion-based Attention Map Analysis) han demostrado que promediar los mapas de atención de todos los cabezales (heads) y pasos de tiempo proporciona una interpretabilidad visual útil, existe una brecha en la comprensión de las características individuales de los diferentes cabezales de atención.

El problema central identificado es que no todos los cabezales de atención contribuyen por igual a conceptos visuales específicos. La agregación indiscriminada de todos los cabezales (como hace DAAM) puede diluir la señal relevante con ruido procedente de cabezales no relacionados, limitando la precisión en tareas de segmentación y el diagnóstico de errores de interpretación del modelo.

2. Metodología

Los autores proponen un enfoque de agregación selectiva que no utiliza todos los cabezales de atención, sino solo un subconjunto de los más relevantes para un concepto dado. El proceso se divide en los siguientes pasos:

• Identificación de Cabezales Relevantes (HRV): Se utiliza el método propuesto en el trabajo previo HRV (Head Relevance Vector) para cuantificar la relevancia de cada cabezal de atención hacia un conjunto de conceptos visuales.
  • Se generan palabras representativas para un concepto (ej. "Animales").
  • Se codifican y proyectan en el espacio de claves (key embeddings) de cada cabezal.
  • Se calcula un vector de relevancia que indica qué cabezales responden más fuertemente a cada concepto.
• Selección del Subconjunto: Se selecciona el 20-25% superior de los cabezales con las puntuaciones de relevancia más altas para el concepto objetivo (en los experimentos, se seleccionaron los 30 cabezales más relevantes de un total de 128).
• Agregación Selectiva: En lugar de promediar todos los cabezales (ecuación de DAAM), se promedian únicamente los mapas de atención de los cabezales seleccionados ($H_n$).
  • La fórmula modificada es: $\hat{M}^{Ours} = \frac{1}{|H_n|} \sum_{h \in H_n} F_h(\frac{1}{T} \sum_{t=1}^{T} M_t^{(h)})$.
• Evaluación: Los mapas resultantes se comparan con mapas de segmentación semántica de "verdad terreno" (ground-truth) generados mediante Grounded-SAM. Se calcula la Intersección sobre Unión (IoU) media tras aplicar umbrales binarios (0.3, 0.4, 0.5).

3. Contribuciones Clave

• Mejora en la Segmentación: Demostraron que la agregación selectiva supera al método estándar DAAM en tareas de segmentación de objetos basada en difusión.
• Características Específicas del Concepto: Se evidenció que los cabezales más relevantes capturan características específicas del concepto con mayor precisión que los menos relevantes o el promedio global.
• Diagnóstico de Errores (Misinterpretación): Se propuso una nueva herramienta visual para diagnosticar errores en la interpretación de prompts. Si un modelo malinterpreta una palabra ambigua (ej. "mouse" como animal vs. dispositivo electrónico), los cabezales relevantes para "Animales" se enfocarán solo en el animal, mientras que los de "Electrónica" se enfocarán en el ratón, permitiendo aislar la fuente del error.
• Análisis de Ablación: Se identificó que seleccionar un número específico de cabezales (30 en este caso) optimiza el rendimiento, evitando el ruido de los cabezales irrelevantes.

4. Resultados Experimentales

Los experimentos se realizaron utilizando Stable Diffusion v1.4 con un conjunto de datos centrado en la categoría "Animales" (100 imágenes generadas).

• Rendimiento Cuantitativo (IoU):
  • El método propuesto ("Ours") superó consistentemente a DAAM en todos los umbrales de prueba.
  • Ejemplo (Umbral 0.4): DAAM obtuvo un IoU de 0.7540, mientras que el método propuesto alcanzó 0.7765.
  • Comparación de Cabezales: Al usar los 30 cabezales menos relevantes para "Animales", el IoU cayó drásticamente a 0.6172 (umbral 0.4), demostrando que la selección es crítica.
• Resultados Cualitativos:
  • Las visualizaciones mostraron que DAAM a menudo tenía focos de atención no deseados o difusos (marcados con círculos rojos en las figuras), mientras que el método selectivo delimitaba con mayor precisión los objetos objetivo.
  • En el caso de la palabra "mouse" (ratón/animal y ratón/electrónico), la agregación selectiva permitió visualizar claramente cómo diferentes cabezales se asignaban a diferentes interpretaciones dentro de la misma imagen generada.

5. Significado e Impacto

Este trabajo sugiere que la selección de cabezales de atención es un paso fundamental para mejorar la interpretabilidad y el control en la generación de imágenes T2I.

• Interpretabilidad: Proporciona una visión más clara y precisa de cómo el modelo asocia tokens de texto con regiones espaciales, eliminando el ruido de cabezales irrelevantes.
• Control y Depuración: Ofrece una vía para diagnosticar por qué un modelo genera conceptos incorrectos o ambiguos, permitiendo a los investigadores entender la "caja negra" de la generación difusa.
• Futuro: Aunque el estudio se centró en Stable Diffusion v1.4 y la categoría de animales, los autores proponen que esta metodología es generalizable a otras arquitecturas T2I y conceptos visuales más amplios, abriendo nuevas direcciones para el control fino en la generación de contenido.

En conclusión, la agregación indiscriminada de mapas de atención no es óptima; una agregación selectiva basada en la relevancia del cabezal conduce a una interpretación visual superior y a una mejor comprensión del comportamiento del modelo.