Learn to Rank: Visual Attribution by Learning Importance Ranking

Resumen Técnico: AHA (Amortized Hybrid Attribution)

1. El Problema

La interpretabilidad de los modelos de visión por computadora complejos (especialmente los Transformers de Visión - ViT) es crucial para establecer confianza y responsabilidad en dominios críticos. Las técnicas existentes de generación de mapas de atribución visual enfrentan un compromiso de tres vías difícil de resolver:

• Métodos basados en propagación (ej. Grad-CAM, Attention Rollout): Son eficientes computacionalmente, pero pueden estar sesgados hacia características de bajo nivel (bordes) y son específicos de la arquitectura. A menudo no capturan bien las interacciones no lineales.
• Métodos basados en perturbación (ej. RISE, SHAP, MDA): Son causalmente sólidos (miden el cambio real en la salida al modificar la entrada), pero son extremadamente costosos computacionalmente. Además, en los ViT, suelen operar a nivel de "parches" (patches), lo que genera mapas de atribución gruesos y poco precisos espacialmente que requieren un upsampling tosco.
• Explicadores basados en aprendizaje (Amortizados): Son rápidos en la inferencia, pero suelen optimizar objetivos sustitutos heurísticos o distilar de maestros heurísticos, lo que puede heredar sus sesgos o basarse en suposiciones restrictivas (como una descomposición binaria señal/ruido).

El objetivo es desarrollar un método que combine la eficiencia de los explicadores aprendidos, la solidez causal de los métodos de perturbación y la precisión a nivel de píxel, optimizando directamente las métricas de calidad de la atribución.

2. Metodología: AHA (Amortized Hybrid Attribution)

Los autores proponen AHA, un marco de atribución basado en aprendizaje que optimiza directamente las métricas de Deletion (Borrado) e Insertion (Inserción) mediante un proceso de aprendizaje de clasificación (Learning to Rank).

A. Formulación del Problema como Aprendizaje de Permutación
Las métricas de Deletion e Insertion dependen de ordenar los píxeles por importancia y eliminarlos/insertarlos secuencialmente. Dado que la operación de ordenamiento (sorting) es no diferenciable, no se puede optimizar directamente mediante retropropagación.

• Solución: Los autores reformulan el problema como un problema de aprendizaje de permutaciones. En lugar de ordenar valores fijos, aprenden una matriz de permutación suave que asigna los valores de atribución a posiciones de rango ideales.

B. Relajación Diferenciable (Gumbel-Sinkhorn)
Para hacer diferenciables las operaciones de ordenamiento y selección de los $k$ mejores píxeles:

1. Se utiliza el algoritmo Gumbel-Sinkhorn.
2. Se calcula una matriz de similitud entre las puntuaciones de atribución y una distribución objetivo de rangos.
3. Se añade ruido Gumbel para introducir estocasticidad y explorar diferentes permutaciones durante el entrenamiento.
4. Se aplica el algoritmo Sinkhorn-Knopp para obtener una matriz de permutación doblemente estocástica suave ($P^{soft}$).
5. A partir de esta matriz, se generan máscaras top-$k$ suaves (valores continuos entre 0 y 1) en lugar de máscaras binarias duras.

Esto permite calcular aproximaciones diferenciables de las métricas AUC (Área Bajo la Curva) de Deletion e Insertion, permitiendo la optimización end-to-end del modelo explicador.

C. Estrategias de Entrenamiento Eficiente y Robusto
Para evitar el costo computacional de evaluar todas las permutaciones de píxeles y prevenir sobreajuste a artefactos de píxeles individuales:

• Permutación Basada en Regiones: En lugar de ordenar píxeles individuales, la imagen se divide en una cuadrícula de regiones. Se promedian las atribuciones por región, se ordenan las regiones y luego se expanden a la resolución original.
• Muestreo de Pasos de Perturbación: No se evalúan todos los pasos de borrado/inserción, sino un subconjunto muestreado uniformemente para aproximar la AUC.
• Aumento de Cuadrícula: Se varía aleatoriamente el tamaño y la posición de la cuadrícula de regiones durante el entrenamiento para evitar sesgos espaciales.
• Regularización: Se penalizan picos de atribución aislados para fomentar mapas suaves y coherentes.

D. Refinamiento en Tiempo de Inferencia (Opcional)
Aunque el método es eficiente en una sola pasada (forward pass), permite un refinamiento opcional en tiempo de prueba. Se pueden realizar unos pocos pasos de gradiente adicionales sobre la imagen de prueba específica para ajustar la atribución, mejorando la fidelidad a costa de un pequeño aumento en el tiempo de cómputo.

3. Contribuciones Clave

1. Marco Híbrido: Un enfoque que une la base causal de la perturbación con la eficiencia amortizada del aprendizaje, optimizando directamente las métricas de evaluación (Deletion/Insertion).
2. Optimización de Métricas Diferenciables: Formulación del problema de optimización de atribución como aprendizaje de permutación, utilizando la relajación Gumbel-Sinkhorn para permitir el entrenamiento end-to-end de métricas no diferenciables.
3. Atribución Densa y Precisa: Generación de mapas de atribución a nivel de píxel (no solo parches) que se alinean estrechamente con los límites de los objetos, incluso en modelos basados en Transformers.
4. Eficiencia y Flexibilidad: Ofrece un modo de alta velocidad (una sola pasada) y un modo de alta fidelidad (con refinamiento), adaptándose a diferentes necesidades de rendimiento.

4. Resultados Experimentales

Los experimentos se realizaron en el conjunto de datos ImageNet utilizando clasificadores ViT-B/16, ViT-B/32 y ConvNeXt Small.

• Rendimiento Cuantitativo: AHA supera consistentemente a los métodos del estado del arte (SOTA) en métricas de fidelidad como Deletion AUC, Insertion AUC, Positive/Negative AUC y ADP (Average Drop Percentage).
  • En ViT-B/16, AHA (sin refinamiento) logra un Deletion AUC de 0.1381 (menor es mejor) frente a 0.1420 de TIS y 0.1556 de MDA.
  • Con solo 3 pasos de refinamiento, mejora aún más, alcanzando el mejor rendimiento general en la mayoría de las métricas.
• Calidad Visual: Los mapas generados son densos y alineados con los bordes de los objetos. A diferencia de los métodos basados en parches (como TIS o ViT-CX) que producen mapas "bloques" o difusos, AHA captura la forma exacta del objeto (ej. la cabeza de un impala o un destornillador de botellas).
• Eficiencia:
  • En modo inferencia estándar ($T=0$), AHA es tan rápido como los métodos de propagación (Grad-CAM, ~16ms), siendo órdenes de magnitud más rápido que los métodos de optimización por muestra (MDA tarda ~13s).
  • Incluso con refinamiento ($T=3$), sigue siendo más rápido que los métodos de perturbación por parche.
• Generalización: El método funciona bien tanto en arquitecturas Transformer como en CNNs (ConvNeXt), demostrando agnosticismo arquitectónico.

5. Significado e Impacto

Este trabajo representa un avance significativo en la explicabilidad de modelos de visión:

• Rompe el compromiso eficiencia-calidad: Demuestra que es posible obtener explicaciones de alta fidelidad (causalmente fundamentadas y precisas) sin el costo computacional prohibitivo de la optimización por muestra.
• Superación de limitaciones de los ViT: Resuelve el problema de la baja resolución espacial en las explicaciones de los Vision Transformers, proporcionando mapas a nivel de píxel en lugar de parches.
• Nueva dirección de investigación: Introduce la idea de tratar la optimización de métricas de atribución como un problema de aprendizaje de permutación diferenciable, abriendo la puerta a optimizar directamente otras métricas de evaluación complejas en el futuro.

En resumen, AHA establece un nuevo estándar para la atribución visual, ofreciendo explicaciones rápidas, precisas y causalmente sólidas que son esenciales para la confianza en sistemas de IA críticos.