Unsupervised Decomposition and Recombination with Discriminator-Driven Diffusion Models

1. Il Problema

L'apprendimento di rappresentazioni che decompongono i dati complessi in fattori latenti indipendenti (disentanglement) è fondamentale per la generalizzazione e la generazione di nuovi campioni attraverso la ricombinazione di componenti. Tuttavia, il disentanglement completamente non supervisionato è teoricamente impossibile senza assunzioni aggiuntive.
Le sfide principali includono:

• Mancanza di supervisione: Senza etichette sui fattori, è difficile garantire che i fattori appresi corrispondano a attributi semantici o fisici reali (es. illuminazione, forma, posizione).
• Incoerenza nella ricombinazione: I modelli generativi esistenti (come VAE o Diffusion Models standard) spesso producono artefatti o campioni "fuori distribuzione" quando si tentano di ricombinare fattori latenti da diverse immagini o video. Questo accade perché la struttura composizionale dei dati non è esplicitamente modellata o regolarizzata.
• Limitazioni nella robotica: Nei dati robotici, le dimostrazioni sono spesso limitate a un sottoinsieme ristretto di stati. La capacità di generare nuove traiettorie fisicamente plausibili ricombinando azioni apprese è cruciale per l'esplorazione dello spazio degli stati, ma i metodi attuali faticano a garantire la coerenza fisica.

2. Metodologia

Gli autori propongono un framework che combina modelli di diffusione con un segnale di addestramento avversariale guidato da un discriminatore per migliorare sia la scoperta dei fattori che la qualità della generazione composizionale.

Architettura del Modello

• Decomposizione Fattoriale: Un encoder mappa un'immagine o un video in un insieme di $K$ fattori latenti indipendenti $z = (z_1, ..., z_K)$.
• Modello Generativo: Viene utilizzato un modello di diffusione condizionato. La previsione del denoiser $\epsilon_\theta$ è decomposta come la media dei contributi di ciascun fattore latente:
  $$\epsilon_\theta(x_t, t, z) = \frac{1}{K} \sum_{k=1}^K \epsilon_\theta(x_t, t, z_k)$$
  Questo approccio si basa sul lavoro precedente Decomp Diffusion, che collega la decomposizione energetica (EBM) ai processi di diffusione.

Il Ruolo del Discriminatore

Il cuore dell'innovazione è l'introduzione di un discriminatore $D_\psi$ addestrato per distinguere tra:

1. Campioni a singola fonte: Generati ricombinando fattori estratti da una singola immagine di input (o mantenendo tutti i fattori originali).
2. Campioni ricombinati: Generati mescolando fattori latenti provenienti da due diverse immagini di input ($x_A$ e $x_B$).

Funzione di Perdita:

• Il discriminatore minimizza la perdita di classificazione per distinguere i due tipi di campioni.
• Il generatore (encoder e denoiser) viene addestrato in modo avversariale per "ingannare" il discriminatore, minimizzando la probabilità che i campioni ricombinati vengano classificati come falsi.
• L'obiettivo totale per il generatore è: $L_{gen} = L_{MSE} + \lambda L_{adv}$, dove $L_{MSE}$ è la perdita di ricostruzione standard della diffusione e $L_{adv}$ è la perdita avversariale.

Meccanismo di Feedback:
Il discriminatore agisce come un proxy scalabile per la "plausibilità percettiva". Se una ricombinazione produce un'immagine incoerente (es. texture che non corrispondono, illuminazione sbagliata), il discriminatore la penalizza. Questo forza il modello a imparare una struttura latente in cui i fattori sono indipendenti e ricombinabili senza violare la varietà dei dati reali.

3. Contributi Chiave

1. Segnale Adversariale per la Diffusione: Introduzione di un segnale di feedback basato su discriminatore specifico per modelli di diffusione fattorizzati, che migliora sia la scoperta dei fattori che la generazione composizionale senza supervisione esplicita.
2. Miglioramento del Disentanglement e della Qualità: Dimostrazione che l'addestramento con questo segnale porta a spazi latenti più disaccoppiati e a ricombinazioni visivamente coerenti, superando i baseline esistenti.
3. Applicazione alla Robotica: Applicazione innovativa a video robotici. Ricombinando componenti di azione apprese, il metodo genera nuove traiettorie video che aumentano significativamente la copertura dello spazio degli stati per l'esplorazione, superando i limiti delle dimostrazioni originali.

4. Risultati Sperimentali

Dataset Immagine (CelebA-HQ, Virtual KITTI, CLEVR, Falcor3D)

• Qualità Percettiva (FID): Il metodo proposto ottiene punteggi FID (Fréchet Inception Distance) significativamente più bassi rispetto al baseline Decomp Diffusion su tutti i dataset. Ad esempio, su CelebA-HQ il FID scende da 82.70 a 43.98.
• Disentanglement: Miglioramento nei metrici MIG (Mutual Information Gap) e MCC (Mean Correlation Coefficient) sul dataset Falcor3D, indicando una migliore separazione dei fattori generativi indipendenti.
• Qualità Visiva: Le immagini ricombinate mostrano coerenza strutturale e semantica, evitando gli artefatti comuni nei metodi precedenti (es. sovrapposizioni di texture incoerenti).

Applicazione Robotica (LIBERO Benchmark)

• Copertura dello Spazio degli Stati: Nel benchmark LIBERO (manipolazione di oggetti su tavolo), la ricombinazione di componenti latenti ha portato a un aumento sostanziale della copertura dello spazio degli stati esplorati.
  • Scene 5: +35.8% di copertura rispetto al baseline.
  • Scene 6: +73.8% di copertura rispetto al baseline.
• Coerenza Fisica: Le traiettorie generate, sebbene nuove, sono fisicamente plausibili e possono essere eseguite da un policy pre-addestrata, dimostrando che il modello ha appreso primitivi d'azione riutilizzabili.

5. Significato e Impatto

Questo lavoro rappresenta un passo avanti significativo nell'apprendimento non supervisionato di rappresentazioni composizionali.

• Superamento delle limitazioni teoriche: Offre una soluzione pratica al problema del disentanglement non supervisionato, utilizzando il feedback del discriminatore come vincolo induttivo per guidare l'apprendimento verso strutture fattoriali.
• Generalizzazione Composizionale: Dimostra che i modelli possono imparare a ricombinare concetti (oggetti, attributi, azioni) in modi mai visti durante l'addestramento, mantenendo la coerenza fisica e semantica.
• Impatto sulla Robotica: La capacità di generare nuove traiettorie di esplorazione attraverso la ricombinazione di azioni apprese apre nuove possibilità per l'addestramento di agenti robotici in ambienti complessi, riducendo la dipendenza da grandi quantità di dati di dimostrazione specifici.

In sintesi, l'approccio proposto trasforma il processo di generazione da una semplice ricostruzione a un meccanismo di "ricombinazione intelligente", dove il discriminatore funge da garante della coerenza, permettendo ai modelli di scoprire e manipolare fattori causali sottostanti in modo non supervisionato.