Two-stage Vision Transformers and Hard Masking offer Robust Object Representations

1. El Problema: Correlaciones Espurias y Fugas de Información

Los modelos de aprendizaje profundo, especialmente los basados en Visión por Computadora (como los Transformers de Visión o ViT), tienden a depender excesivamente de pistas contextuales (fondos, objetos co-ocurrentes) en lugar de las propiedades intrínsecas de los objetos. Esto genera correlaciones espurias que funcionan bien en datos de entrenamiento pero fallan catastróficamente en escenarios de distribución fuera de entrenamiento (OOD - Out-of-Distribution), donde el contexto cambia.

El artículo identifica dos fallos críticos en los métodos de atención existentes (enmascaramiento tardío o late masking):

1. Contaminación de características: Los métodos actuales aplican máscaras de atención sobre mapas de características de alto nivel. Dado que las arquitecturas modernas tienen campos receptivos grandes, la información del fondo ya ha "contaminado" las características antes de que se aplique la máscara.
2. Fugas residuales: Las máscaras de atención suelen ser "suaves" (soft), asignando pesos no nulos a todas las ubicaciones. Esto permite que información espuria del fondo se filtre residualmente hacia la predicción, comprometiendo la fiabilidad (faithfulness) del modelo.

2. Metodología: iFAM (Inherently Faithful Attention Maps)

Los autores proponen iFAM, un marco de trabajo de dos etapas diseñado para garantizar que el modelo solo utilice información de las regiones relevantes del objeto, eliminando físicamente la posibilidad de que el fondo influya en la decisión.

Arquitectura de Dos Etapas

1. Etapa 1: Selector (Descubrimiento de Partes)

   • Utiliza un método de descubrimiento de partes (basado en PDiscoFormer) que procesa la imagen completa.
   • Su objetivo es identificar regiones de primer plano (foreground) relevantes para la tarea, generando mapas de atención semánticos.
   • Utiliza pérdidas de "moldeado de partes" (part-shaping losses) para asegurar que las partes descubiertas sean consistentes y semánticamente significativas, no solo predictivas.

2. Etapa 2: Predictor (Clasificación con Enmascaramiento Temprano)

   • Recibe exclusivamente los tokens seleccionados por la Etapa 1.
   • Mecanismo Clave (Enmascaramiento Temprano / Early Masking): En lugar de aplicar una máscara a las características finales, se aplica una máscara binaria discreta directamente a la entrada del Transformer.
   • Se modifica el mecanismo de auto-atención del ViT: se establece un valor de $-\infty$ en la matriz de atención para cualquier token que haya sido enmascarado (fondo). Esto impide físicamente que los tokens de fondo intercambien información con los tokens del objeto.
   • Entrenamiento Joint: Ambas etapas se entrenan conjuntamente, permitiendo que la Etapa 2 refine la selección de la Etapa 1 mediante retropropagación.

Intervenciones Auditable en Tiempo de Prueba

Una ventaja única de iFAM es su transparencia. Dado que las máscaras son explícitas y semánticas, permiten intervenciones humanas o automáticas sin reentrenar el modelo:

• Eliminación de Partes (Part-Removal): Si una parte descubierta captura sistemáticamente un elemento espurio (ej. un tubo de drenaje en una radiografía), se puede excluir manualmente o automáticamente (análisis leave-one-out) de la entrada del clasificador.
• Eliminación de Tokens (Token-Removal): Se pueden eliminar tokens asignados a partes con baja confianza (distancia grande al prototipo en el espacio de características) para manejar objetos OOD que se detectan erróneamente como primer plano.

3. Contribuciones Clave

1. Arquitectura de Dos Etapas: Un cambio de paradigma del procesamiento de imagen completa al procesamiento selectivo para mitigar correlaciones espurias.
2. Mecanismo de Enmascaramiento Temprano: Una solución arquitectónica que restringe el campo receptivo del predictor a las regiones seleccionadas, garantizando la fiabilidad frente a la fuga de información del fondo.
3. Robustez y Auditabilidad: El modelo logra un estado del arte (SOTA) en robustez y permite intervenciones en tiempo de prueba para mejorar aún más el rendimiento, algo ausente en modelos de caja negra.

4. Resultados Experimentales

El método se evaluó en diversos benchmarks diseñados con sesgos de fondo conocidos:

• MetaShift (Gato vs. Perro): iFAM alcanzó un 88.6% de precisión en el grupo más difícil (WGA), superando significativamente a métodos de enmascaramiento tardío (73.5% - 83.0%) y a modelos ERM estándar.
• Waterbirds: Logró un 97.0% de WGA, superando a todos los baselines, incluidos los que usan enmascaramiento tardío o pre-entrenamiento auto-supervisado avanzado.
• ImageNet-9 (Desafío de Fondos): iFAM obtuvo el menor BG-GAP (diferencia de precisión entre fondos correctos y aleatorios) de 2.4, demostrando una dependencia mínima del fondo en comparación con ViTs grandes y ResNets.
• SIIM-ACR (Radiografías): En detección de neumotórax, iFAM mejoró el AUC del grupo más difícil (WG AUC) de 48.1% (PDiscoFormer) a 65.9%, acercándose al rendimiento de modelos con máscaras de ground-truth, sin usar anotaciones de segmentación.
• Intervenciones: La combinación de eliminación de partes y tokens mejoró aún más los resultados (ej. +4 puntos en WGA en MetaShift), validando la utilidad de la auditabilidad.

5. Significado y Conclusión

El trabajo demuestra que la fiabilidad inherente (inherent faithfulness) es crucial para la robustez en visión por computadora. A diferencia de los métodos que intentan "aprender" a ignorar el fondo a través de regularización suave, iFAM elimina físicamente la posibilidad de que el fondo influya en la decisión mediante un enmascaramiento binario temprano y arquitectónico.

Esto resuelve el dilema de que el contexto es necesario para la detección pero perjudicial para la generalización: el sistema descubre el contexto relevante en la primera etapa y luego lo aísla estrictamente en la segunda. Aunque introduce un costo computacional adicional (dos pasadas), el segundo paso es altamente eficiente al procesar solo una fracción de los tokens, y la ganancia en robustez y capacidad de intervención humana justifica el enfoque.

Código: Disponible en GitHub: https://github.com/ananthu-aniraj/ifam