Test-Time Compositional Generalization in Diffusion Models via Concept Discovery

Riepilogo Tecnico: Generalizzazione Compositiva al Momento del Test nei Modelli Diffusivi tramite Scoperta di Concetti

Enunciato del Problema

La generalizzazione compositiva—la capacità di ricombinare parti familiari in configurazioni nuove—è un tratto distintivo dell'intelligenza umana, ma rimane una sfida significativa per i modelli di deep learning. Mentre le reti neurali spesso interpolano bene all'interno delle distribuzioni di addestramento, faticano quando i primitivi familiari appaiono in combinazioni inedite. Gli approcci esistenti alla generazione compositiva nei modelli diffusivi si basano tipicamente su una libreria predefinita di concetti, segnali di condizionamento o etichette esterne per guidare la generazione di composizioni nuove. Ciò crea un divario tra la scoperta dei concetti e la generalizzazione compositiva: i metodi precedenti assumono che i primitivi rilevanti siano già noti e disponibili, non affrontando scenari in cui un modello deve scoprire concetti specifici della query da zero al momento del test, senza un vocabolario fisso o supervisione esterna.

Metodologia

Il documento propone un framework che riutilizza un modello diffusivo preaddestrato come una gerarchia implicita di modi di densità indicizzati nel tempo per scoprire e comporre concetti al momento del test. Il metodo opera in tre fasi principali:

1. Scoperta di Concetti tramite Ascesa dei Modi:
   Invece di utilizzare il modello diffusivo esclusivamente per il campionamento inverso, gli autori sfruttano la funzione di punteggio indicizzata nel tempo appresa $s_\theta(x_t, t) \approx \nabla_{x_t} \log p_t(x_t)$ per localizzare i modi locali di densità delle marginali rumorose $p_t(x_t)$. Partendo da una singola query fuori distribuzione (OOD) $x_q$, il metodo esegue un'ascesa del gradiente sul punteggio incondizionato a più passi di rumore $t$ per recuperare i modi locali.

   • Base Teorica: Sfruttando la formula di Tweedie, gli autori dimostrano che questi modi corrispondono a "prototipi concettuali" privi di rumore nello spazio dei dati puliti ($x_0$), agendo efficacemente come centroidi discreti a diversi livelli di astrazione a seconda del livello di rumore.
   • Rappresentazione: Ogni modo scoperto viene mappato nuovamente nello spazio $x_0$ e approssimato come un esperto Gaussiano locale $q_j(x_0) = \mathcal{N}(x_0; m_j, \Sigma_j)$, dove la media è derivata dal modo e la covarianza è stimata tramite sonde di Hutchinson.

2. Composizione Prodotto-di-Esperti (PoE):
   Gli esperti Gaussiani scoperti sono combinati in un unico modello insegnante analitico utilizzando una formulazione Prodotto-di-Esperti (PoE).

   • Selezione: Un algoritmo di ottimizzazione submodulare greedy seleziona il sottoinsieme di prototipi più rilevanti per la query, massimizzando un obiettivo di copertura della verosimiglianza logaritmica per dimensione.
   • Composizione: Gli esperti selezionati sono combinati in una distribuzione PoE pesata $q_T$. Poiché gli esperti sono Gaussiani, anche la composizione risultante è Gaussiana con una funzione di punteggio in forma chiusa. Ciò permette di derivare una previsione di rumore analitica $\epsilon_{q_T}(x_t, t)$ senza richiedere un token condizionato addestrato per la composizione nuova.

3. Campionamento Compositivo e Distillazione:

   • Campionamento Analitico: Il punteggio analitico dell'insegnante PoE è integrato in un framework di guida senza classificatore (CFG). Il modello confronta la previsione PoE con la previsione incondizionata per guidare il campionamento verso la composizione scoperta.
   • Distillazione LoRA: Per migliorare la generazione sul manifold e creare una condizione riutilizzabile, il metodo distilla opzionalmente l'insegnante PoE nel modello base. Ciò comporta l'addestramento di un nuovo embedding di classe e di un adattatore Low-Rank Adaptation (LoRA) su un pool di campioni generati dall'insegnante PoE. Questa fase assorbe la composizione scoperta nuovamente nel manifold di punteggio del modello.

Contributi Chiave

1. Colmare il Divario tra Scoperta e Composizione: Il lavoro colma il divario tra la scoperta dei concetti e la generalizzazione compositiva trattando un modello diffusivo preaddestrato come una gerarchia implicita di modi di densità. Permette il recupero di prototipi concettuali discreti al momento del test senza un pool di primitivi predefinito.
2. Meccanismo di Composizione al Momento del Test: Gli autori introducono un metodo che seleziona greedy i modi rilevanti utilizzando un obiettivo di verosimiglianza submodulare e li combina in un modello insegnante PoE con punteggio in forma chiusa, permettendo il campionamento diretto senza riaddestrare il modello base sulla specifica composizione.
3. Validazione Empirica: Il framework è valutato su due benchmark, ColorMNIST e CelebA, dimostrando che i prototipi scoperti e la loro composizione PoE codificano informazioni più ricche sulle composizioni non viste rispetto al campionamento basato solo sulla query o al recupero della classe addestrata più vicina.

Risultati

Gli esperimenti su ColorMNIST (identità della cifra, colore, sfondo) e CelebA (attributi del viso) mostrano che il metodo proposto supera diverse linee di base:

• Linee di Base: Il metodo è confrontato con "Solo Query" (campionamento attorno alla query con varianza fissa), "Top-1/Top-3" (recupero delle classi addestrate più vicine) e CFG standard.
• Prestazioni: Sia il campionatore PoE analitico che la variante distillata con LoRA raggiungono metriche FID, similarità CLIP e Precisione/Recall superiori rispetto alle linee di base.
  • Su ColorMNIST, la variante LoRA raggiunge un punteggio F1 del 87,6% (Fedeltà) e del 76,3% (Generalizzazione), superando significativamente la linea di base Top-1 (~6,5% F1).
  • Su CelebA, il metodo PoE ottiene il miglior FID di Generalizzazione (140,3) e i punteggi CLIP, mentre LoRA migliora ulteriormente l'equilibrio tra copertura e fedeltà, raggiungendo il massimo Recall e F1 di Generalizzazione.
• Analisi Qualitativa: I risultati visivi indicano che, mentre le linee di base spesso falliscono nel preservare la combinazione specifica di attributi (ad esempio, cambiando l'identità della cifra o deviando verso combinazioni di attributi diverse), il metodo proposto ricostruisce con successo la composizione OOD e generalizza ad altre istanze della stessa classe.

Significato e Affermazioni

Il documento afferma che la geometria del punteggio indicizzata nel tempo dei modelli diffusivi contiene concetti riutilizzabili basati su modi di densità che supportano la generazione compositiva al momento del test senza una libreria di concetti predefinita.

• Insight Teorico: I risultati suggeriscono che i modelli diffusivi codificano implicitamente una gerarchia di concetti, dove i modi delle marginali rumorose fungono da prototipi a diversi livelli di astrazione.
• Implicazione Pratica: Il metodo permette la generazione di configurazioni nuove da una singola query OOD, offrendo una via per la generalizzazione compositiva che non dipende da supervisione esterna o vocabolari pre-specificati.
• Limitazioni: Gli autori notano modestamente che il protocollo assume che i fattori primitivi siano noti (setting compositivo standard) e che il metodo si basi sui bias induttivi del backbone preaddestrato. Sebbene possa interpolare tra primitivi noti, non può fabbricare primitivi genuinamente nuovi che estrapolano oltre il supporto di addestramento (ad esempio, uno sfondo bianco in un modello addestrato solo su sfondi scuri).

Il lavoro fornisce sia prove teoriche che empiriche che i modelli diffusivi possono essere riutilizzati per la scoperta e la composizione di concetti, potenzialmente estendendosi a applicazioni più ampie nella modellazione linguistica e nel controllo robotico.