Unsupervised Decomposition and Recombination with Discriminator-Driven Diffusion Models

Probleemstelling

Het centrale probleem dat dit artikel aanpakt, is de onopgeloste uitdaging om complexe data (zoals afbeeldingen en video's) te ontleden in factoriserende, herbruikbare representaties zonder gebruik te maken van supervisie op factorniveau (bijv. labels voor objecten, achtergronden of bewegingen). Hoewel generatieve modellen zoals Variational Auto-Encoders (VAE's) en Diffusion Models krachtig zijn, missen ze vaak een gestructureerde, ontkoppelde latente ruimte.

Zonder dergelijke structuur is compositional recombination (het combineren van verschillende factoren uit verschillende bronnen om nieuwe, realistische samples te genereren) problematisch. Naïeve recombinatie leidt vaak tot samples die niet op het "data manifold" liggen, wat resulteert in visuele artefacten, fysiek onrealistische video's of semantische inconsistenties. Bestaande methoden voor ontkoppeling vereisen vaak sterke aannames of externe feedback, wat de schaalbaarheid beperkt.

Methodologie

De auteurs introduceren een nieuw kader dat Diffusion Models combineert met een discriminator-gedreven adversariaal trainingsproces om ontkoppelde latente factoren te leren en de kwaliteit van recombinatie te verbeteren.

1. Factoriserend Generatief Model:

   • Het model encodeert een input (afbeelding of video) in $K$ discrete latente componenten $z = (z_1, ..., z_K)$.
   • Het generatieve proces is gebaseerd op Diffusion Models, waarbij het denoising-netwerk $\epsilon_\theta$ wordt opgeleid om de bijdrage van elke individuele factor te voorspellen. De totale voorspelling is het gemiddelde van de bijdragen van de individuele factoren: $\epsilon_\theta(x_t, t) = \frac{1}{K} \sum \epsilon_\theta(x_t, t, z_k)$.
   • Dit benadert een "Product of Experts" structuur, analoog aan Energy-Based Models (EBM), maar zonder de noodzaak van iteratieve optimalisatie tijdens de generatie.

2. Compositional Recombination:

   • Tijdens inferentie worden nieuwe samples gegenereerd door latente factoren uit verschillende bronafbeeldingen te mixen. Een binaire masker $S$ bepaalt welke factoren uit bron A en welke uit bron B worden gebruikt: $\tilde{z} = S \odot z^{(A)} + (1-S) \odot z^{(B)}$.

3. Discriminator-Gedreven Feedback (Kerninnovatie):

   • Om te voorkomen dat recombinaties onrealistisch worden, wordt een discriminator $D_\psi$ getraind om te onderscheiden tussen:
     • Single-source samples: Genereerd uit latente factoren van één enkele bron.
     • Recombined samples: Genereerd uit gemixte factoren van meerdere bronnen.
   • Adversariaal Trainingsdoel: De generator (encoder + diffusion model) wordt getraind om de discriminator te misleiden, zodat gerecombineerde samples ononderscheidbaar zijn van geldige single-source samples.
   • De totale loss voor de generator is een combinatie van de standaard reconstructie-loss (MSE) en een adversariale loss: $L_{gen} = L_{MSE} + \lambda L_{adv}$.
   • Dit creëert een inductieve bias die het model dwingt om factoren te leren die fysiek en semantisch consistent blijven wanneer ze worden gecombineerd.

Belangrijkste Bijdragen

1. Adversariaal Signaal voor Factorisatie: Het introduceren van een discriminator die specifiek getraind is om de plausibiliteit van gerecombineerde samples te beoordelen. Dit signaal verbetert zowel de ontdekking van factoren als de kwaliteit van de compositional generatie.
2. Verbeterde Ontkoppeling en Kwaliteit: Het bewijzen dat deze methode leidt tot een betere ontkoppeling (disentanglement) en visuele consistentie vergeleken met state-of-the-art baselines zoals Decomp Diffusion, β-VAE en InfoGAN.
3. Toepassing op Robotica: Een innovatieve toepassing op robotische video-trajecto's. Door gerecombineerde actiecomponenten te genereren, kan een agent de state-space coverage aanzienlijk vergroten voor exploratie in omgevingen waar trainingdata beperkt is.

Resultaten

De methode werd geëvalueerd op diverse datasets, waaronder CelebA-HQ (gezichten), Virtual KITTI (rijscènes), CLEVR (objecten) en Falcor3D (synthetische scènes), evenals op het LIBERO-benchmark voor robotica.

• Beeldkwaliteit (FID): De methode behaalde consistent lagere FID-scores (Frechet Inception Distance) op alle beelddatasets vergeleken met Decomp Diffusion. Dit duidt op hogere perceptuele kwaliteit en minder artefacten bij recombinatie.
  • Voorbeeld: Op CelebA-HQ daalde de FID van 82.70 (baseline) naar 43.98 (proposed).
• Ontkoppeling (MIG & MCC): De methode behaalde hogere scores op Mutual Information Gap (MIG) en Mean Correlation Coefficient (MCC), wat aantoont dat de latente factoren beter gescheiden zijn en onafhankelijker van elkaar variëren.
• Robotica (LIBERO):
  • Gerecombineerde video's leidden tot actieplannen die fysiek plausibel waren.
  • De state-space coverage (het aantal bezochte toestanden in de discretisatie) nam aanzienlijk toe: met 35,8% in Scene 5 en 73,8% in Scene 6 ten opzichte van de baseline. Dit betekent dat de robot effectiever nieuwe situaties kon verkennen door bestaande vaardigheden te combineren.
• Synthetisch Experiment: Een experiment met synthetische data toonde aan dat de discriminator-feedback de geometrie van de data herstelt en de latente ruimte transformeert naar een structuur die recombinatie toelaat, zelfs zonder semantische labels.

Betekenis en Impact

Dit werk is significant omdat het een brug slaat tussen representatie learning en generatieve modellering zonder afhankelijk te zijn van dure externe supervisie of menselijke feedback.

• Generalisatie: Het stelt modellen in staat om buiten de trainingsverdeling te generaliseren door bestaande componenten op nieuwe manieren te combineren, wat essentieel is voor robuuste AI-systemen.
• Robotica: Voor robotica biedt het een manier om beperkte demonstratiegegevens uit te breiden naar een breed scala aan plausibele acties, wat cruciaal is voor het oplossen van "long-tail" problemen in de echte wereld.
• Theoretische Inzicht: Het artikel biedt theoretische onderbouwing (via Product-of-Experts en mutual information contractie) die uitlegt waarom het straffen van onrealistische recombinaties leidt tot een meer factoriserende latente ruimte.

Samenvattend biedt deze "Discriminator-Driven Diffusion" een schaalbare en effectieve oplossing voor het leren van compositional structuren, met directe toepassingen in beeldmanipulatie, creatieve contentgeneratie en autonoom robotisch plannen.