Variance-Tilted Diffusion Models for Diverse Sampling

Technische Samenvatting: Variantie-gekantelde Diffusiemodellen voor Diverse Sampling

Probleemstelling
Diffusiemodellen zijn de standaard geworden voor generatieve modellering binnen diverse domeinen, variërend van hoogresolutie beeldsynthese tot wetenschappelijke ontdekkingen (bijv. eiwitvouwing en moleculaire docking). De standaard sampling-paradigma steunt echter op het trekken van onafhankelijke en identiek verdeelde (i.i.d.) monsters. Deze aanpak is vaak suboptimaal voor downstream-taken die een diverse set kandidaten vereisen die verschillende regio's van de datamanifold beslaan, zoals het genereren van structureel diverse liganden in medicijnontwerp of robuuste trajectvoorstellen in robotische bewegingsplanning.

Bestaande methoden om diversiteit af te dwingen, zoals particle guidance, introduceren repulsieve potentialen in de sampling-dynamiek. Hoewel effectief, modificeren deze methoden de drift heuristisch, wat resulteert in een impliciete doelverdeling die computationeel onhandelbaar is om exact te evalueren of te controleren.

Methodologie
De auteurs stellen een "target-first" formulering voor waarbij de batch-doelverdeling expliciet wordt gedefinieerd voordat de exacte sampler wordt afgeleid. De kerninnovatie is een variantie-gewogen batchverdeling die collecties monsters bevoordeelt met een grote empirische spreiding na een voorgeschreven lineaire feature-map.

Voor een batch van $n$ monsters $\mathbf{x} = (x^{(1)}, \dots, x^{(n)})$ wordt de doelverdeling $\pi^A_0(\mathbf{x})$ gedefinieerd als:
$$\pi^A_0(\mathbf{x}) \propto \text{Var}^A_n(\mathbf{x}) \prod_{i=1}^n p_0(x^{(i)})$$
waarbij $p_0$ de oorspronkelijke dataverdeling is, $A: \mathbb{R}^d \to \mathbb{R}^k$ een lineaire operator is die monsters naar een taakspecifieke feature-ruimte mapt, en $\text{Var}^A_n(\mathbf{x})$ de empirische variantie van de features $A(x^{(i)})$ is.

Om uit deze verdeling te samplen, leiden de auteurs de exacte Doob $h$-transformatie van de onafhankelijke diffusiedynamiek af. Dit houdt in dat de conditionele verwachting van de batch-variantie wordt berekend gegeven de ruisige staat op tijdstip $t$:
$$h^A_t(\mathbf{x}_t) := \mathbb{E}[\text{Var}^A_n(\mathbf{X}_0) \mid \mathbf{X}_t = \mathbf{x}_t]$$

Met behulp van Tweedie's identiteiten en een bias-variantie-decompositie leiden de auteurs een gesloten vorm af voor de score-correctie $\nabla_{\mathbf{x}_t} \log h^A_t(\mathbf{x}_t)$. Deze correctie bestaat uit twee afzonderlijke termen:

1. Interactieterm (Repulsie): Een term die de posterior gedenoised gemiddelden ($\mu_t$) van de deeltjes van elkaar weg duwt in de feature-ruimte. Dit wordt efficiënt berekend met behulp van vector-Jacobian producten (VJP) via de denoiser.
2. Curvatuurterm: Een term die deeltjes naar regio's met een hogere feature-variantie beweegt, waarbij de divergentie van de scorefunctie (specifiek $\text{div}(A^\top A \nabla s_t)$) betrokken is. Deze wordt geschat via eindige verschillen of stochastische trace-estimators (bijv. Hutchinson-estimators).

De resulterende sampler is een interagerend-deeltjesproces waarbij de score voor elk deeltje de som is van de standaard score en de afgeleide Doob-correctie.

Belangrijkste Bijdragen

1. Target-First Formulering: Het artikel kadert diverse generatie als een exact sampling-probleem vanuit een expliciete batch-doelverdeling, waardoor de ongecontroleerde repulsieve krachten van heuristische methoden worden vermeden.
2. Closed-Form Doob Correctie: De auteurs bieden een analytische afleiding van de score-correctie voor een algemene lineaire operator $A$, waarbij de correctie wordt gedecomponeerd in een repulsieterm voor posterior gemiddelden en een curvatuurterm voor feature-variantie.
3. Tractable Sampling Algoritme: Een gediscretiseerde interacting-particle sampler wordt geïntroduceerd. De leidende interactieterm maakt gebruik van efficiënte VJP-berekeningen, terwijl de curvatuurterm wordt afgehandeld via numerieke schatting, wat de methode praktisch maakt voor bestaande diffusiemodellen.

Resultaten
De auteurs hebben de aanpak geëvalueerd op beeldgeneratie met behulp van Stable Diffusion v1.5 (512x512) met diverse prompts. Ze vergeleken drie configuraties:

• Standaard Classifier-Free Guidance (CFG).
• Variance-Tilted (VT) zonder de divergentieterm.
• Variance-Tilted + Divergence (VT + Divergence).

Kwalitatieve resultaten (Figuren 1–5) laten zien dat de voorgestelde methoden een diversere set monsters produceren vergeleken met standaard CFG. De gegenereerde afbeeldingen beslaan een breder scala aan poses, achtergronden en visuele stijlen, in plaats van slechts minimale variaties van hoog-waarschijnlijkheidsmodi te produceren. De inclusie van de divergentieterm verfijnt de diversiteit verder, hoewel de auteurs opmerken dat de volledige methode de beeldkwaliteit licht kan negatief beïnvloeden en de computationele kosten verhoogt door de divergentie-schatting.

Betekenis en Claims
Het artikel claimt een probabilistisch alternatief te bieden voor heuristische deeltjesrepulsie. Door de exacte sampler voor een variantie-gewogen doelverdeling af te leiden, biedt de methode een transparant probabilistisch objectief in plaats van een impliciet objectief. De auteurs positioneren dit als een rigoureuze aanpak voor diverse sampling die praktische implementatieroutes behoudt.

De auteurs blijven bescheiden over beperkingen en erkennen dat, hoewel de methode de diversiteit kwalitatief verbetert, het de beeldkwaliteit licht kan verslechteren en hogere computationele kosten met zich meebrengt door de schatting van tweede-orde afgeleiden (divergentie). Ze suggereren dat toekomstig werk zich kan richten op het afvlakken (annealing) van de Doob-correctie om kwaliteitsproblemen te mitigeren en het ontwikkelen van efficiëntere estimators voor de curvatuurterm.