Technische Samenvatting: Generatieve Modellering met Flux Matching

Probleemstelling

De huidige state-of-the-art generatieve modellering, met name in hoogdimensionale domeinen zoals beeldsynthese, proteïneontwerp en robotica, is sterk afhankelijk van score-gebaseerde modellen. Deze modellen leren de Stein-scorefunctie, $\nabla \log p_{\text{data}}(x)$, doorgaans via score matching. Eenmaal geleerd, wordt deze score gebruikt om diffusieprocessen (bijvoorbeeld Langevin-dynamica) aan te sturen om te bemonsteren uit de doeldistributie $p_{\text{data}}$.

De scorefunctie is echter niet het unieke vectorveld dat in staat is om $p_{\text{data}}$ te genereren. Volgens de Fokker-Planck-vergelijking (FPE) zal elk vectorveld $f_\theta$ dat voldoet aan de voorwaarde $\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$, $p_{\text{data}}$ als stationaire distributie hebben. Dit impliceert dat er een oneindige familie van generatieve vectorvelden (niet-conservatieve velden) bestaat die dezelfde doeldistributie kunnen genereren, maar verschillende dynamica induceren (bijvoorbeeld verschillende mengsnelheden, circulatiepatronen of gerichte afhankelijkheden).

Bestaande score-matchingdoelen dwingen het model om te convergeren naar één enkel punt in deze ruimte ($\nabla \log p_{\text{data}}$), waarbij elke afwijking wordt bestraft, zelfs als die afwijking de doeldistributie behoudt. Dit beperkt het vermogen om nuttige inductieve biases, structurele priors of dynamische eigenschappen (zoals snellere menging of causale afhankelijkheden) direct in het generatieve veld te coderen.

Methodologie: Flux Matching

De auteurs stellen Flux Matching voor, een nieuw paradigma dat score matching generaliseert door elk vectorveld binnen de familie van velden die de doeldistributie genereren te leren, in plaats van het model te dwingen om punt voor punt te matchen met de score.

Kernobjectief

In plaats van het minimaliseren van de Fisher-divergentie tussen het model $f_\theta$ en de score $\nabla \log p_{\text{data}}$, minimaliseert Flux Matching de divergentie van de waarschijnlijkheidsflux. De voorwaarde voor een veld om generatief te zijn, is:
$$\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$$
Deze voorwaarde is zwakker dan puntsgewijze gelijkheid en staat een "nulpunt" toe van oneindig veel geldige vectorvelden.

Theoretische Afleiding

Om deze voorwaarde te optimaliseren, definiëren de auteurs een Geprojecteerde Fisher-divergentie:
$$\tilde{J}(\theta) := \mathbb{E}_{x \sim p_{\text{data}}} \left[ \| \Pi_{\text{flux}} f_\theta(x) - \nabla \log p_{\text{data}}(x) \|^2 \right]$$
waarbij $\Pi_{\text{flux}}$ een vectorveld projecteert op de ruimte van gradiëntvelden die dezelfde fluxdivergentie delen. Hoewel $\tilde{J}(\theta)$ theoretisch onderbouwd is, is het in hoge dimensies onberekenbaar.

De auteurs leiden een hanteerbare Flux Matching-verliesfunctie ($L_{\text{flux}}$) af door:

1. Het definiëren van het residu-veld $u_\theta = f_\theta - \nabla \log p_{\text{data}}$ en de Langevin-Stein-operator $r_\theta = \nabla \cdot u_\theta + u_\theta \cdot \nabla \log p_{\text{data}}$.
2. Het gebruik van diffusiesimulaties $d x_t = \nabla \log p_{\text{data}}(x_t) dt + \sqrt{2} dW_t$ om $r_\theta$ over de data-distributie te propageren.
3. Het toepassen van partiële integratie en een stop-gradient-truc om het doel om te zetten in een vorm die backpropagatie door dure pad-afgeleiden vermijdt.

Het uiteindelijke verlies is:
$$L_{\text{flux}}(\theta) := -\mathbb{E}_{t \sim q, x_0 \sim p_{\text{data}}, x_t | x_0} \left[ \frac{1}{q(t)} u_\theta(x_0)^\top \text{sg}\left( \frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t) \right) \right]$$
waarbij $\text{sg}$ de stop-gradient-operator aanduidt. De auteurs bewijzen dat $\nabla_\theta \tilde{J}(\theta) = 2 \nabla_\theta L_{\text{flux}}(\theta)$.

Noise Annealing en Praktische Implementatie

De methode is uitbreidbaar naar noise-annealed settings (vergelijkbaar met diffusiemodellen) door een continuüm van velden te leren voor ruisniveaus $\sigma$.

• Score Benadering: Omdat $\nabla \log p_{\text{data}}$ onbekend is, gebruiken de auteurs een niet-parametrische Kernel Density Estimation (KDE) benadering op mini-batches om de score te schatten.
• Variance Reductie: Om de pad-term $\frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t)$ te schatten, maken de auteurs gebruik van een cross-chain schatter die alle samples in een mini-batch gebruikt, gewogen door Gaussische overgangskansen, in plaats van te vertrouwen op één enkele keten.

Selecteren van Specifieke Velden

Flux Matching leert een familie van velden. Om een specifiek veld met gewenste eigenschappen te selecteren, stellen de auteurs twee strategieën voor:

1. Toepassings-specifiek Verlies: Het toevoegen van aanvullende verliezen (bijvoorbeeld voor mengsnelheid of structurele beperkingen) aan het Flux Matching-doel.
2. Model Parametrering: Architecturale beperkingen (bijvoorbeeld causale attentiemaskers) die specifieke structuren direct in $f_\theta$ afdwingen.

Belangrijkste Bijdragen

1. Flux Matching Paradigma: Introductie van een generatief modelleringkader dat vectorvelden buiten de score leert door de divergentie van de waarschijnlijkheidsflux te matchen.
2. Efficiënte Verliesafleiding: Afleiding van een schaalbare Flux Matching-verliesfunctie die de Fisher-divergentie-geometrie behoudt terwijl het niet-conservatieve vectorvelden toestaat.
3. Noise-Annealed Uitbreiding: Uitbreiding van de methode naar noise-annealed distributies, waardoor toepassing mogelijk wordt op complexe, hoogdimensionale beelddatasets.
4. Gedemonstreerde Flexibiliteit: Empirisch bewijs dat Flux Matching toepassingen mogelijk maakt die onder standaard score matching onmogelijk zijn, waaronder:
   • Snellere Menging: Het leren van vectorvelden die de convergentie van bemonstering versnellen.
   • Interpreteerbaarheid: Het passen van gestructureerde, biologisch zinvolle modellen (bijvoorbeeld RNA-velocity) zonder black-box neurale benaderingen.
   • Gestructureerde Dynamica: Het inbedden van gerichte afhankelijkheden (bijvoorbeeld temporale autoregressie) direct in het generatieve veld.

Resultaten

Het artikel evalueert Flux Matching in vijf settings:

1. Controleerbare Velden (Speelgoeddata): Op een 2D Gaussische mengeling kent Flux Matching een verlies van nul toe aan velden die de doeldistributie behouden maar verschillende dynamica hebben (bijvoorbeeld rotatie, circulatie), terwijl score matching deze afwijkingen bestraft.
2. Interpreteerbare Velden (RNA Velocity): Toegepast op single-cell genomics, past Flux Matching de scVelo dynamische modelparameters via gradiëntafname aan. Het presteert beter dan de standaard op EM gebaseerde scVelo-fitting in consistentie en correctheid over grenzen heen op vijf datasets, wat aantoont dat de winst voortkomt uit het optimalisatieproces en niet uit een verhoogde modeluitdrukkingskracht.
3. Onbeperkte Generatie (Beelden): Op CIFAR-10 en CelebA dient Flux Matching als een zelfstandig doel. Hoewel het concurrerende FID-scores bereikt, is het momenteel 3–4$\times$ trager en gebruikt het 2–3$\times$ meer geheugen tijdens training in vergelijking met Denoising Score Matching (DSM). Het schaalt echter wel naar hoogdimensionale distributies.
4. Snelle Menging: Door het optimaliseren van een mengproxy, leert Flux Matching velden die met aanzienlijk minder bemonsteringsstappen dan standaard Flux Matching of DSM redelijke FID-scores bereiken.
5. Structuur Inbedden: In een taak voor trajectgeneratie (niet-lineaire veren) slaagt Flux Matching erin een causaal attentiemasker op te nemen om temporale autoregressie af te dwingen. Daarentegen verslechtert het toepassen van hetzelfde masker op DSM de prestaties, omdat de symmetrische Jacobiaan van de scorefunctie in conflict komt met gerichte afhankelijkheden.

Betekenis en Claims

De auteurs beweren dat Flux Matching een "nieuwe dimensie" opent in generatieve modellering door het vectorveld zelf te behandelen als een ontwerpkeuze in plaats van een vast doel.

• Voorbij Bemonstering: De primaire waarde ligt niet noodzakelijk in het vervangen van DSM voor standaard beeldgeneratie (waar DSM sterk geoptimaliseerd is), maar in het mogelijk maken van dynamica met nuttige eigenschappen.
• Flexibiliteit: Het stelt praktici in staat om inductieve biases, structurele priors en mechanistische beperkingen direct in het generatieve proces op te leggen.
• Toepassingen: Het artikel benadrukt specifieke use-cases die door deze flexibiliteit worden ontsloten: snellere bemonstering via niet-reverseerbare dynamica, interpreteerbare wetenschappelijke modellen (zoals RNA-velocity) en generatieve modellen met voorgeschreven causale of temporale structuren.

Het artikel concludeert dat Flux Matching, hoewel het momenteel hogere rekenkosten met zich meebrengt, door het vermogen om niet-conservatieve, gestructureerde generatieve velden te leren, een krachtig hulpmiddel biedt voor toepassingen waarbij de dynamica van generatie even belangrijk is als de distributie van de output.