Expanding the Role of Diffusion Models for Robust Classifier Training

Questo lavoro dimostra che integrare le rappresentazioni interne dei modelli di diffusione, oltre ai dati sintetici generati, nell'addestramento avversariale migliora significativamente la robustezza e la disentanglement delle caratteristiche dei classificatori di immagini su diversi dataset.

L'essenziale

Immagina di voler insegnare a un bambino a riconoscere le differenze tra un gatto e un cane. Normalmente, gli mostri mille foto di gatti e cani. Ma se un "bullo" (un attaccante informatico) prende quelle foto e ci aggiunge un po' di "polvere invisibile" (rumore) che l'occhio umano non vede, il bambino potrebbe confondersi e dire che un gatto è un cane.

Nel mondo dell'intelligenza artificiale, questo è il problema della robustezza: come fare in modo che l'IA non si confonda quando qualcuno le nasconde piccoli trucchi?

Questo articolo parla di un nuovo modo per addestrare queste intelligenze artificiali, usando una tecnologia chiamata Modelli Diffusivi. Ecco come funziona, spiegato in modo semplice:

1. Il vecchio trucco: "Facciamo finta di tutto"

Fino a poco tempo fa, per rendere l'IA più forte contro i "bulli", gli scienziati usavano i modelli diffusivi come una fotocopiatrice magica.

• Come funzionava: Il modello generava milioni di foto finte (sintetiche) di gatti e cani, perfette e variegate.
• Il risultato: Addestrando l'IA su queste foto finte, diventava molto brava a riconoscere i veri animali, anche se qualcuno le nascondeva un po' di polvere.
• Il limite: Si usava il modello solo per creare le foto, ignorando tutto il resto.

2. La nuova scoperta: "Ascolta anche il suo pensiero"

Gli autori di questo studio si sono chiesti: "E se non usassimo solo le foto che crea il modello, ma anche il modo in cui il modello 'pensa' mentre le crea?"

Immagina il modello diffusivo come un artista esperto che sta dipingendo un quadro.

• Il vecchio metodo: Prendi solo il quadro finito (l'immagine sintetica) e mostralo all'allievo.
• Il nuovo metodo: Mentre l'artista dipinge, lo allievo guarda anche le bozze intermedie e i pensieri dell'artista. L'artista, mentre toglie il "rumore" dalla tela per rivelare l'immagine, impara a vedere le forme vere e a ignorare i dettagli inutili.

Gli scienziati hanno scoperto che questi "pensieri" (chiamati rappresentazioni) contengono informazioni molto preziose: sono resistenti (non si confondono facilmente con la polvere) e diversi (vedono il mondo da molte angolazioni).

3. La soluzione: "Due maestri per uno studente"

Hanno creato un nuovo metodo di addestramento (chiamato DRA) che funziona così:

1. Hanno uno studente (l'IA che deve imparare a classificare).
2. Hanno un maestro esperto (il modello diffusivo congelato, che non cambia).
3. Mentre lo studente studia le foto (vere o finte), il maestro gli sussurra: "Ehi, guarda qui, questa è la forma vera del gatto, ignora quel puntino strano!".

In pratica, l'IA impara non solo guardando le immagini, ma allineando il suo modo di pensare a quello del maestro esperto.

Perché è importante? (Le metafore finali)

• Il puzzle: Immagina che l'IA debba risolvere un puzzle.

  • I vecchi metodi le davano solo pezzi di puzzle aggiuntivi (le foto finte).
  • Questo nuovo metodo le dà anche la immagine sulla scatola (le rappresentazioni del modello diffusivo) che le dice come dovrebbero combaciare i pezzi, rendendo il puzzle molto più facile da risolvere anche se qualcuno ha rotto alcuni pezzi (attacco avversario).

• La pulizia: I modelli precedenti a volte imparavano a riconoscere i "rumori" ad alta frequenza (come il fruscio di un foglio di carta) invece delle forme vere. Il modello diffusivo, invece, è come un pulitore esperto: sa distinguere la polvere dal volto vero. Insegnando all'IA a pensare come questo pulitore, l'IA diventa molto più brava a ignorare i trucchi dei "bulli".

I risultati

Hanno provato questo metodo su tre "palestre" diverse (insiemi di dati chiamati CIFAR-10, CIFAR-100 e ImageNet) e i risultati sono stati eccellenti:

• L'IA è diventata più precisa nel riconoscere le immagini.
• È diventata molto più difficile da ingannare per gli attaccanti.
• Ha imparato a "separare" meglio i concetti (come distinguere chiaramente un gatto da un cane senza confondersi), rendendo il suo cervello più ordinato e meno soggetto a errori.

In sintesi: Non si tratta più solo di creare più immagini finte per allenare l'IA, ma di farle imparare a pensare come un esperto che sa già come rimuovere il rumore dal mondo. È come passare dal far copiare a un bambino un libro di testo, a fargli ascoltare le spiegazioni di un professore mentre legge.

Sommario tecnico

1. Il Problema

I modelli di apprendimento automatico sono notoriamente vulnerabili agli esempi avversariali: input perturbati da rumore semanticamente impercettibile che possono alterare drasticamente le previsioni del modello.

• Adversarial Training (AT): Rimane l'approccio più efficace per ottenere robustezza, ma soffre di un fenomeno noto come robust overfitting, dove la robustezza sul set di test diminuisce durante l'addestramento nonostante l'accuratezza sulle immagini pulite rimanga stabile.
• Soluzioni Esistenti: L'uso di dati sintetici generati da modelli di diffusione (DM-AT) ha mostrato risultati promettenti, trattando i modelli di diffusione principalmente come generatori di dati.
• Il Gap: Tuttavia, i modelli di diffusione producono anche rappresentazioni intermedie (attivazioni interne) che codificano caratteristiche semantiche significative. La domanda centrale della ricerca è: queste rappresentazioni interne possono essere sfruttate oltre alla semplice generazione di dati sintetici per migliorare ulteriormente l'addestramento di classificatori robusti?

2. Metodologia

Gli autori propongono di estendere il ruolo dei modelli di diffusione integrando le loro rappresentazioni interne come segnale di apprendimento ausiliario durante l'Adversarial Training.

A. Analisi Preliminare delle Rappresentazioni

Prima di proporre il metodo, gli autori analizzano le proprietà delle rappresentazioni estratte dai modelli di diffusione (denoising intermediate activations):

• Robustezza Parziale: Le rappresentazioni mostrano una correlazione positiva con la robustezza e sono meno sensibili al rumore ad alta frequenza rispetto ai modelli supervisionati standard.
• Diversità: A differenza dell'addestramento avversariale puro che tende a ridurre la diversità delle feature, le rappresentazioni di diffusione mantengono un'alta uniformità (diversità delle feature).
• Frequenza: L'analisi di frequenza mostra che le rappresentazioni di diffusione si concentrano su componenti a bassa frequenza, evitando la dipendenza da segnali ad alta frequenza tipica dei modelli ricostruttivi (come MAE), che sono più vulnerabili agli attacchi.

B. Diffusion Representation Alignment (DRA)

Per sfruttare queste proprietà, gli autori modificano la ricetta DM-AT introducendo un modulo di allineamento:

1. Architettura: Si utilizza un classificatore $f_{CLS}$ (composto da un encoder e un head di classificazione) e un modello di diffusione congelato $g_{\phi}$.
2. Allineamento: Durante l'addestramento avversariale, le rappresentazioni del classificatore ($h_{CLS}$) ottenute da un esempio avversariale $\hat{x}$ vengono allineate alle rappresentazioni estratte dal modello di diffusione congelato ($h_{DR}$) per lo stesso input (o condizione).
3. Funzione di Perdita: Viene introdotta una perdita di regolarizzazione $L_{DRA}$ basata sulla similarità (cosine similarity) tra le due rappresentazioni, proiettate tramite un piccolo head di proiezione (MLP) per allinearne gli spazi.
   $$L_{AT-DRA} = L_{AT} + \lambda L_{DRA}$$
   Dove $\lambda$ controlla la forza della regolarizzazione.

3. Contributi Chiave

1. Nuovo Paradigma di Integrazione: Dimostrano che le rappresentazioni dei modelli di diffusione non servono solo per generare dati, ma possono agire come priori di feature robusti e diversificati per guidare l'addestramento del classificatore.
2. Miglioramento della Robustezza: L'integrazione delle rappresentazioni di diffusione (DRA) nell'AT porta a miglioramenti consistenti sia nell'accuratezza pulita che in quella avversariale su diversi dataset e architetture.
3. Analisi Meccanicistica:
   • Disentanglement: L'uso di modelli di diffusione incoraggia l'apprendimento di rappresentazioni più facilmente disaccoppiate (disentangled), riducendo la sovrapposizione di feature (feature superposition) che può essere sfruttata dagli attacchi avversariali.
   • Ruoli Complementari:
     • I dati sintetici generati favoriscono rappresentazioni a basso rango (low-rank) con buone proprietà di generalizzazione.
     • L'allineamento delle rappresentazioni incoraggia il modello a sfruttare efficacemente le dimensioni rappresentazionali per codificare feature robuste (non necessariamente a basso rango).
     • La combinazione dei due approcci offre benefici complementari.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su CIFAR-10, CIFAR-100 e ImageNet utilizzando diverse architetture (WideResNet, ViT, ConvNeXt) e configurazioni di dati sintetici.

• Performance: L'approccio DM-AT + DRA supera sistematicamente lo stato dell'arte (SOTA) basato solo su DM-AT.
  • Su CIFAR-10 (ViT-B/2 con 50M dati sintetici): L'accuratezza robusta (AutoAttack) sale dal 71.31% (DM-AT) al 71.77% (DM-AT + DRA), con un aumento parallelo dell'accuratezza pulita.
  • Su ImageNet (ConvNeXt-B): L'accuratezza robusta sale dal 54.44% al 56.07%.
• Analisi delle Rappresentazioni:
  • Le metriche di uniformità e allineamento mostrano che DRA sposta il compromesso (frontier) verso rappresentazioni più diversificate e robuste.
  • L'uso di Sparse AutoEncoders (SAE) rivela che i modelli addestrati con DRA hanno una perdita di ricostruzione normalizzata inferiore, indicando rappresentazioni più facili da disaccoppiare.
  • L'analisi della dimensione di classificazione mostra che l'allineamento con le rappresentazioni di diffusione porta a un uso più efficace delle dimensioni del vettore di feature per la robustezza.

5. Significato e Implicazioni

Questo lavoro rappresenta un passo avanti significativo nella ricerca sulla robustezza avversariale:

• Oltre la Generazione di Dati: Sposta il focus dall'uso dei modelli di diffusione come semplici generatori di dati sintetici al loro utilizzo come fonti di conoscenza semantica robusta attraverso le loro rappresentazioni interne.
• Efficienza: A differenza di metodi di purificazione o classificatori generativi bayesiani che richiedono costi di inferenza elevati, DRA è un metodo di addestramento che non introduce overhead durante l'inferenza.
• Ispirazione Futura: Suggerisce che l'integrazione di rappresentazioni pre-addestrate (non solo dati) può essere una strategia potente per migliorare la generalizzazione e la robustezza dei modelli di visione artificiale, aprendo nuove direzioni per l'uso dei modelli generativi nell'addestramento discriminatoro.

In sintesi, il paper dimostra che allineare le rappresentazioni del classificatore con quelle di un modello di diffusione congelato è una strategia efficace per costruire classificatori più robusti, sfruttando la capacità intrinseca dei modelli di diffusione di catturare feature semantiche robuste e diversificate.