Selective Aggregation of Attention Maps Improves Diffusion-Based Visual Interpretation

1. Il Problema

I modelli generativi da testo a immagine (T2I), come Stable Diffusion e Imagen, hanno raggiunto risultati eccezionali nella generazione di immagini di alta qualità. Per interpretare il comportamento di questi modelli e migliorarne le applicazioni (es. editing, segmentazione), la ricerca si affida ampiamente alle mappe di attenzione incrociata (cross-attention maps).

Tuttavia, esiste un limite significativo nell'approccio attuale:

• Aggregazione Indiscriminata: Metodi consolidati come DAAM (Diffusion-based Attention Map) calcolano la mappa di attenzione aggregando le mappe di tutti i testine di attenzione (attention heads) e di tutti i timestep di generazione.
• Mancanza di Specificità: Si presume che la media di tutte le testine fornisca una rappresentazione accurata del concetto, ma le caratteristiche distinte delle diverse testine rimangono poco esplorate. Non è chiaro se tutte le testine contribuiscano ugualmente alla rappresentazione di un concetto specifico o se alcune siano "rumorose" o irrilevanti.
• Diagnosi Limitata: L'aggregazione totale può mascherare errori di interpretazione del prompt o focalizzazioni errate, rendendo difficile diagnosticare perché un modello abbia generato un oggetto indesiderato o ambiguo.

2. Metodologia

Gli autori propongono un metodo chiamato Selezione e Aggregazione Selettiva delle Mappe di Attenzione. L'obiettivo è utilizzare solo un sottoinsieme di testine di attenzione che sono statisticamente più rilevanti per un concetto target specifico.

Il processo si articola in tre fasi principali:

1. Valutazione della Rilevanza (HRV):

   • Si utilizza il metodo HRV (Head Relevance Vector) proposto in uno studio precedente [7].
   • Per ogni concetto visivo (es. "Animali"), vengono generate parole chiave rappresentative.
   • Durante la generazione casuale di immagini, si calcola quanto ogni testina di attenzione risponde attivamente a queste parole chiave.
   • Si ottiene un vettore di rilevanza per ogni concetto, che indica quali testine sono più sensibili a quel concetto specifico.

2. Selezione delle Testine:

   • Invece di usare tutte le $H$ testine, il metodo seleziona solo il 20-25% delle testine con i punteggi di rilevanza più alti per il concetto target.
   • Nel paper, per il benchmark "Animali", vengono selezionate le 30 testine più rilevanti su un totale di 128 (modello Stable Diffusion v1.4).

3. Aggregazione Selettiva:

   • Si calcola la mappa di attenzione aggregata ($\hat{M}^{Ours}$) mediando le mappe delle sole testine selezionate ($H_n$) e di tutti i timestep, invece di includere tutte le testine come fa DAAM.
   • La formula modificata è:
     $$\hat{M}^{Ours} = \frac{1}{|H_n|} \sum_{h \in H_n} F_h \left( \frac{1}{T} \sum_{t=1}^{T} M^{(h)}_t \right)$$
     dove $F_h$ è l'interpolazione bicubica per adattare la risoluzione spaziale.

3. Contributi Chiave

• Superiorità dell'Aggregazione Selettiva: Dimostrazione empirica che l'uso di un sottoinsieme di testine rilevanti produce mappe di attenzione più accurate e semanticamente coerenti rispetto all'aggregazione totale.
• Diagnosi Visiva degli Errori: Il metodo permette di visualizzare come il modello interpreta parole ambigue. Se un prompt genera due concetti diversi (es. "mouse" come animale e come dispositivo), le testine rilevanti per "Animali" si focalizzeranno solo sull'animale, mentre quelle per "Elettronica" solo sul dispositivo. DAAM, aggregando tutto, mostrerebbe una sovrapposizione confusa.
• Miglioramento della Segmentazione: L'approccio offre una nuova via per la segmentazione semantica basata sulla diffusione, superando i metodi attuali senza bisogno di riaddestrare il modello.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su Stable Diffusion v1.4 con un benchmark focalizzato sulla categoria "Animali" (100 immagini generate per 10 animali diversi). La segmentazione è stata confrontata con ground-truth ottenuti tramite Grounded-SAM.

• Metrica: Mean Intersection over Union (mIoU) calcolata su diverse soglie di binarizzazione (0.3, 0.4, 0.5).
• Confronto DAAM vs. Ours:
  • Il metodo proposto ha ottenuto costantemente punteggi mIoU più alti rispetto a DAAM.
  • Esempio (Soglia 0.4): DAAM = 0.7540 vs. Ours = 0.7765.
  • Esempio (Soglia 0.5): DAAM = 0.6261 vs. Ours = 0.6785.
• Analisi delle Testine Rilevanti vs. Irrilevanti:
  • Utilizzando le 30 testine meno rilevanti per "Animali", le prestazioni crollano drasticamente (mIoU scende a ~0.46-0.61), confermando che le testine irrilevanti catturano caratteristiche non pertinenti o rumore.
  • Le testine più rilevanti catturano i concetti specifici con molta più precisione.
• Studio Ablativo: È stato testato il numero di testine da selezionare (20, 30, 40). Il risultato ottimale è stato ottenuto con 30 testine, suggerendo che un numero troppo basso perde informazioni e un numero troppo alto include testine meno pertinenti.

5. Significato e Implicazioni

Questo lavoro offre un cambio di paradigma nell'interpretazione dei modelli T2I:

1. Interpretabilità Migliorata: Sposta il focus dall'analisi "media" a un'analisi "selettiva", rivelando che la struttura interna del modello è altamente specializzata per concetto.
2. Controllo e Diagnosi: Fornisce uno strumento potente per diagnosticare perché un modello fallisce o interpreta male un prompt, permettendo agli utenti di isolare visivamente la fonte dell'errore (es. distinguere tra significati multipli di una parola).
3. Applicabilità Futura: Sebbene testato su Stable Diffusion v1.4, i risultati suggeriscono che la selezione delle testine è un principio generale applicabile a diverse architetture T2I per migliorare il controllo generativo, l'editing e la sicurezza (es. rilevamento di contenuti indesiderati).

In sintesi, il paper dimostra che non tutte le testine di attenzione sono uguali e che la loro selezione intelligente è fondamentale per ottenere una comprensione visiva più profonda e accurata dei modelli generativi.