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

1. Il Problema: Bias Contestuali e Correlazioni Spurie

I modelli di deep learning, in particolare i Vision Transformer (ViT), tendono a fare affidamento su indizi contestuali (sfondi, oggetti co-occorrenti) per imparare le rappresentazioni degli oggetti, piuttosto che sulle proprietà intrinseche degli oggetti stessi.

• Conseguenza: Questo porta a correlazioni spurie. Ad esempio, un modello potrebbe imparare che i cani sono sempre all'aperto e i gatti sempre al chiuso. Quando il modello incontra dati fuori distribuzione (OOD), come un cane in un interno, fallisce perché la correlazione appresa non è più valida.
• Limitazione delle soluzioni attuali: Le tecniche esistenti di "masking" (mascheramento) applicano solitamente una maschera di attenzione tardiva (late masking) su mappe di feature di alto livello. Tuttavia, poiché le architetture moderne hanno campi ricettivi vasti, le feature di alto livello sono già "contaminate" dalle informazioni dello sfondo prima che la maschera venga applicata. Inoltre, l'uso di maschere di attenzione "soft" (non binarie) permette una perdita residua di informazioni dalle regioni spurie.

2. Metodologia: iFAM (Inherently Faithful Attention Maps)

Gli autori propongono iFAM, un framework a due stadi che utilizza un masking precoce (early masking) e binario per garantire che il modello si basi esclusivamente sulle regioni rilevanti dell'immagine.

Architettura a Due Stadi

1. Stadio 1: Il Selettore (Selector)

   • Processa l'immagine intera.
   • Utilizza un metodo di scoperta delle parti (basato su PDiscoFormer) per identificare le regioni di primo piano (foreground) rilevanti per il compito.
   • Genera una maschera binaria discreta (0 o 1) che seleziona i token corrispondenti alle parti dell'oggetto, ignorando lo sfondo.
   • È addestrato con perdite di "shape delle parti" per garantire coerenza semantica.

2. Stadio 2: Il Predittore (Predictor)

   • Riceve in input solo i token selezionati dallo Stadio 1.
   • È un ViT il cui meccanismo di auto-attenzione è modificato per essere architettonicamente cieco alle regioni mascherate.
   • Viene applicata una maschera di attenzione rigida: i token mascherati non possono scambiare informazioni con gli altri token (i valori di attenzione sono impostati a $-\infty$).
   • Questo garantisce che il campo ricettivo del predittore sia limitato fisicamente alle regioni selezionate, eliminando qualsiasi possibilità di leakage informativo dallo sfondo.

Addestramento e Interventi

• Addestramento congiunto: I due stadi sono addestrati insieme, permettendo allo Stadio 2 di affinare la selezione dello Stadio 1 tramite gradienti.
• Interventi a tempo di test (Test-Time Interventions): Grazie alla natura esplicita e auditabile delle maschere semantiche, è possibile intervenire manualmente o automaticamente durante l'inferenza:
  • Rimozione di parti: Se una parte scoperta cattura sistematicamente un elemento spurio (es. uno sfondo specifico), può essere rimossa manualmente o tramite analisi "leave-one-out".
  • Rimozione di token: I token assegnati a parti con bassa confidenza (lontani dal prototipo della parte nello spazio delle feature) possono essere rimossi per gestire oggetti OOD.

3. Contributi Chiave

1. Architettura a due stadi: Sposta l'elaborazione da un'analisi dell'immagine completa a un'elaborazione selettiva, mitigando le correlazioni spurie fin dall'inizio.
2. Masking Precoce e Binario: A differenza del masking tardivo su feature map, questa approccio applica la maschera a livello di input (token), garantendo che il campo ricettivo del modello sia fisicamente limitato alle regioni rilevanti.
3. Auditabilità e Interventi: Il modello offre trasparenza intrinseca. Le maschere semantiche permettono agli umani di verificare il ragionamento del modello e di applicare strategie di intervento (come la rimozione di parti spurie) senza riaddestramento, migliorando ulteriormente la robustezza.

4. Risultati Sperimentali

Il metodo è stato valutato su una serie di benchmark progettati per testare la robustezza contro bias di sfondo:

• Dataset: MetaShift (gatti/cani), Waterbirds (uccelli acquatici/terrestri), CUB/Waterbird200 (classificazione fine-grained con sfondi OOD), SIIM-ACR (medico, rilevamento pneumotorace), e ImageNet-9.
• Metriche: Accuratezza Media (AA), Accuratezza del Gruppo Peggiore (WGA) e BG-GAP (differenza di accuratezza tra sfondi corretti e casuali).
• Performance:
  • iFAM supera significativamente i baselines, inclusi metodi di masking tardivo (come PDiscoFormer) e tecniche di debiasing avanzate (GroupDRO, MaskTune).
  • Su Waterbird, iFAM raggiunge un WGA del 97.0% contro il 95.8% del miglior baseline self-supervised.
  • Su MetaShift, ottiene un WGA del 88.6% contro l'83.0% di metodi simili.
  • Su ImageNet-9, riduce il BG-GAP a 2.4, dimostrando una dipendenza minima dagli indizi di sfondo rispetto a modelli molto più grandi (ViT-L, ResNet-152).
  • Gli interventi a tempo di test hanno ulteriormente migliorato le prestazioni (es. +3 punti WGA su MetaShift combinando rimozione di parti e token).

5. Significato e Conclusioni

Questo lavoro dimostra che la robustezza nei modelli di visione artificiale non può essere ottenuta solo tramite tecniche di regolarizzazione o addestramento su dati aumentati, ma richiede un'architettura che garantisca fedeltà intrinseca (inherent faithfulness).

• Limitare fisicamente il campo ricettivo del modello alle sole regioni rilevanti è superiore all'uso di meccanismi di attenzione soft su feature contaminate.
• La capacità di rendere il processo decisionale del modello auditabile e modificabile a runtime apre nuove strade per la collaborazione uomo-AI, permettendo di correggere errori sistematici senza riaddestrare l'intero modello.
• Il costo computazionale aggiuntivo (due passaggi in avanti) è compensato dalla maggiore efficienza dello stadio 2 (che elabora meno token) e dalla significativa riduzione del bias, rendendo iFAM una soluzione scalabile e robusta per applicazioni critiche.