Test-Time Compositional Generalization in Diffusion Models via Concept Discovery

Technische Samenvatting: Compositional Generalization op het Moment van Testen in Diffusiemodellen via Conceptontdekking

Probleemstelling

Compositional generalization—het vermogen om bekende onderdelen te hercombineren tot nieuwe configuraties—is een kenmerk van menselijke intelligentie, maar blijft een aanzienlijke uitdaging voor deep learning-modellen. Waar neurale netwerken vaak goed interpoleren binnen trainingsverdelingen, hebben ze moeite wanneer bekende primitieven voorkomen in onbekende combinaties. Bestaande benaderingen voor compositional generation in diffusiemodellen vertrouwen doorgaans op een vooraf gedefinieerde bibliotheek van concepten, conditionele signalen of externe labels om de generatie van nieuwe composities te sturen. Dit creëert een kloof tussen conceptontdekking en compositional generalization: eerdere methoden gaan ervan uit dat de relevante primitieven al bekend en beschikbaar zijn, en slagen er niet in scenario's aan te pakken waarin een model query-specifieke concepten vanaf nul moet ontdekken op het moment van testen, zonder een vaste vocabulaire of externe supervisie.

Methodologie

Het artikel stelt een raamwerk voor dat een vooraf getraind diffusiemodel hergebruikt als een impliciete hiërarchie van op tijd geïndexeerde dichtheidsmodi om concepten op het moment van testen te ontdekken en te combineren. De methode werkt in drie hoofdfasen:

1. Conceptontdekking via Mode-Ascent:
   In plaats van het diffusiemodel uitsluitend te gebruiken voor reverse sampling, maken de auteurs gebruik van de geleerde op tijd geïndexeerde scorefunctie $s_\theta(x_t, t) \approx \nabla_{x_t} \log p_t(x_t)$ om lokale dichtheidsmodi van de verstoord marginaalverdelingen $p_t(x_t)$ te lokaliseren. Uitgaande van een enkele out-of-distribution (OOD) query $x_q$, voert de methode gradiëntascent uit op de onvoorwaardelijke score op meerdere ruis-tijdstappen $t$ om lokale modi te herstellen.

   • Theoretische Basis: Door gebruik te maken van de formule van Tweedie tonen de auteurs aan dat deze modi overeenkomen met ruisvrije "conceptprototypen" in de schone data-ruimte ($x_0$), die effectief fungeren als discrete centroiden op verschillende niveaus van abstractie, afhankelijk van het ruisniveau.
   • Representatie: Elke ontdekte modus wordt teruggemapt naar de $x_0$-ruimte en benaderd als een lokaal Gaussisch expert $q_j(x_0) = \mathcal{N}(x_0; m_j, \Sigma_j)$, waarbij het gemiddelde wordt afgeleid uit de modus en de covariantie wordt geschat via Hutchinson-probes.

2. Product-of-Experts (PoE) Compositie:
   De ontdekte Gaussische experts worden gecombineerd tot een enkel analytisch leermodel met behulp van een Product-of-Experts (PoE)-formulering.

   • Selectie: Een gulzige submodulaire optimalisatiealgoritme selecteert het subset van prototypen dat het meest relevant is voor de query, door een log-likelihood-dekkingsdoelstelling per dimensie te maximaliseren.
   • Compositie: De geselecteerde experts worden gecombineerd tot een gewogen PoE-verdeling $q_T$. Omdat de experts Gaussisch zijn, is de resulterende compositie ook Gaussisch met een gesloten-vorm scorefunctie. Dit maakt het mogelijk om een analytische ruisvoorspelling $\epsilon_{q_T}(x_t, t)$ af te leiden zonder dat een getraind conditioneel token nodig is voor de nieuwe compositie.

3. Compositional Sampling en Distillatie:

   • Analytische Sampling: De analytische score van de PoE-leerder wordt geïntegreerd in een classifier-free guidance (CFG)-raamwerk. Het model contrasteert de PoE-voorspelling met de onvoorwaardelijke voorspelling om de sampling te sturen naar de ontdekte compositie.
   • LoRA-distillatie: Om generatie binnen de variëteit te verbeteren en een herbruikbare conditie te creëren, distilleert de methode optioneel de PoE-leerder in het basismodel. Dit omvat het trainen van een nieuwe klasse-embedding en een Low-Rank Adaptation (LoRA)-adapter op een pool van samples gegenereerd door de PoE-leerder. Deze stap absorbeert de ontdekte compositie terug in de scorevariëteit van het model.

Belangrijkste Bijdragen

1. Overbrugging van Ontdekking en Compositie: Het werk overbrugt conceptontdekking en compositional generalization door een vooraf getraind diffusiemodel te behandelen als een impliciete hiërarchie van dichtheidsmodi. Het maakt het mogelijk om discrete conceptprototypen op het moment van testen te herstellen zonder een vooraf gedefinieerde pool van primitieven.
2. Mechanisme voor Compositie op het Moment van Testen: De auteurs introduceren een methode die gulzig relevante modi selecteert met behulp van een submodulaire likelihood-doelstelling en deze combineert tot een PoE-leerdermodel met een gesloten-vorm score, waardoor directe sampling mogelijk is zonder het basismodel opnieuw te trainen op de specifieke compositie.
3. Empirische Validatie: Het raamwerk wordt geëvalueerd op twee benchmarks, ColorMNIST en CelebA, en toont aan dat ontdekte prototypen en hun PoE-compositie rijkere informatie bevatten over ongezette composities dan sampling alleen op basis van de query of het ophalen van de dichtstbijzijnde getrainde klasse.

Resultaten

Experimenten op ColorMNIST (cijferidentiteit, kleur, achtergrond) en CelebA (gezichtsattributen) tonen aan dat de voorgestelde methode verschillende baselines overtreft:

• Baselines: De methode wordt vergeleken met "Query-only" (sampling rond de query met vaste variantie), "Top-1/Top-3" (ophalen van de dichtstbijzijnde getrainde klassen) en standaard CFG.
• Prestatie: Zowel de analytische PoE-sampler als de LoRA-gedistilleerde variant bereiken superieure FID-, CLIP- en Precisie/Recall-metingen in vergelijking met de baselines.
  • Op ColorMNIST bereikt de LoRA-variant een F1-score van 87,6% (Betrouwbaarheid) en 76,3% (Generalisatie), wat significant beter is dan de Top-1-baseline (~6,5% F1).
  • Op CelebA bereikt de PoE-methode de beste Generalisatie-FID (140,3) en CLIP-scores, terwijl LoRA de balans tussen dekking en fideliteit verder verbetert, met de hoogste Generalisatie-Recall en F1.
• Kwalitatieve Analyse: Visuele resultaten geven aan dat terwijl baselines vaak falen in het behoud van de specifieke combinatie van attributen (bijvoorbeeld het veranderen van cijferidentiteit of afdrijven naar andere attributecombinaties), de voorgestelde methode de OOD-compositie succesvol reconstrueert en generaliseert naar andere instanties van dezelfde klasse.

Betekenis en Claims

Het artikel claimt dat de op tijd geïndexeerde scoregeometrie van diffusiemodellen herbruikbare dichtheidsmodus-concepten bevat die compositional generation op het moment van testen ondersteunen zonder een vooraf gedefinieerde conceptbibliotheek.

• Theoretisch Inzicht: De resultaten suggereren dat diffusiemodellen impliciet een hiërarchie van concepten coderen waarbij modi van verstoord marginaalverdelingen dienen als prototypen op verschillende niveaus van abstractie.
• Praktische Implicatie: De methode maakt het mogelijk om nieuwe configuraties te genereren vanuit een enkele OOD-query, en biedt een pad voor compositional generalization dat niet afhankelijk is van externe supervisie of vooraf gespecificeerde vocabulaires.
• Beperkingen: De auteurs merken bescheiden op dat het protocol aanneemt dat primitieve factoren bekend zijn (standaard compositional setting) en dat de methode afhankelijk is van de inductieve biases van het vooraf getrainde backbone. Hoewel het kan interpoleren tussen bekende primitieven, kan het geen echt nieuwe primitieven fabriceren die extrapoleren buiten de trainingsondersteuning (bijvoorbeeld een witte achtergrond in een model dat alleen is getraind op donkere achtergronden).

Het werk levert zowel theoretisch als empirisch bewijs dat diffusiemodellen kunnen worden hergebruikt voor conceptontdekking en compositie, wat potentieel uitbreidt naar bredere toepassingen in taalmodelleren en robotbesturing.