Technische Samenvatting: Discriminator-Guided RL voor Flow Matching

Probleemstelling

Flow-gebaseerde en score-matching modellen zijn het dominante paradigma geworden voor generatieve modellering vanwege hun simulatievrije, op regressie gebaseerde trainingsdoelstellingen. Deze modellen worden echter doorgaans getraind in een meerfasige pijplijn: eerst wordt een basismodel gefit aan een datadistributie $q$ met behulp van flow of score matching (FSM); ten tweede wordt reinforcement learning (RL) toegepast om het model te sturen naar gewenste regio's gedefinieerd door een beloning $r(x)$.

Het artikel identificeert een structurele mismatch in dit paradigma. Hoewel FSM-verliezen de $\ell_2$-regressiefout van snelheid of scorevelden onder de marginalen tijdens de training minimaliseren, slagen ze er vaak niet in om intrinsieke eigenschappen van de datadistributie te herstellen—zoals visuele realisme, coherente objectstructuur of fysiek correcte beweging—die theoretisch wel aanwezig zijn in de data. Gevolgelijk vertrouwen beoefenaars op voorkeursgebaseerde RL om deze eigenschappen te herstellen. Deze aanpak is problematisch om twee redenen:

1. Structureel Falen: Het is raadselachtig dat modellen deze werkelijke eigenschappen niet direct uit data kunnen leren via matching-verliezen, maar RL nodig hebben om ze te herstellen.
2. Kosten en Conflatie: Het trainen van beloningen vereist doorgaans dure menselijke voorkeursdata, wat de objectieve realiteit van de data conflueert met de subjectieve neigingen van annotatoren.

De auteurs stellen dat het falen van matching-doelstellingen om deze eigenschappen te vangen voortkomt uit een "train-test mismatch" (waarbij fouten accumuleren langs de rollout-distributie van het model in plaats van langs de marginalen tijdens de training) en een "geometrische mismatch" (waarbij kleine fouten in de snelheidruimte kunnen leiden tot grote fouten in de samplekwaliteit/beloningsruimte).

Methodologie: Discriminator-Guided RL (DRL)

Het artikel stelt Discriminator-Guided RL (DRL) voor, een methode om flow-matching modellen te corrigeren zonder afhankelijk te zijn van menselijke voorkeursdata. De kerninzicht is dat de optimale beloning voor het afstemmen van een basismodel $p_{base}$ op een doelverdeling $q$ de log-densiteitsratio $\log(q(x)/p_{base}(x))$ is. Omdat deze ratio onhandelbaar is, schat DRL deze in met behulp van een discriminator.

De DRL-pijplijn bestaat uit twee fasen:

Fase 1: Beloningsschatting in de Representatieruimte

In plaats van een discriminator direct in de pixelruimte te trainen (wat statistisch moeilijk en semantisch onbetrouwbaar is), traint DRL een discriminator in een pre-trained self-supervised learning (SSL) representatieruimte.

• Een bevroren encoder $\phi$ (bijv. DINOv2) brengt afbeeldingen naar een feature-ruimte $Z$.
• Een discriminator $D$ wordt getraind om onderscheid te maken tussen echte datasamples $x \sim q$ en basismodel-samples $x_{fake} \sim p_{base}$ met behulp van hun gecodeerde representaties $\phi(x)$.
• De beloningsfunctie wordt gedefinieerd als de logit van deze discriminator: $\hat{r}(x) = \log(D(\phi(x)) / (1 - D(\phi(x))))$.
• Theoretische Rechtvaardiging: Door de beloning te beperken tot de SSL-ruimte, lost DRL een geconstrueerd optimalisatieprobleem op: het vindt de distributie $p$ die het dichtst bij $p_{base}$ ligt (in KL-divergentie) zodanig dat de pushforward van $p$ door $\phi$ overeenkomt met de pushforward van $q$ door $\phi$. Dit zorgt ervoor dat correcties worden uitgevoerd langs semantisch betekenisvolle assen die zichtbaar zijn voor de encoder.

Fase 2: KL-geregulariseerde RL via Adjoint Matching

De geschatte beloning $\hat{r}$ wordt gebruikt om het basismodel te verfijnen via KL-geregulariseerde RL.

• Het doel is om de verwachte beloning te maximaliseren terwijl de KL-divergentie tussen de traject-distributies van het verfijnde model en het basismodel wordt geminimaliseerd.
• De auteurs maken gebruik van Adjoint Matching, een state-of-the-art algoritme voor flow-modellen dat de traject-niveau doelstelling optimaliseert via gradiëntpropagatie door het samplingsproces.
• Om stabiele training te garanderen met de geheugenloze SDE die vereist is voor ongebiaste endpoint tilting, introduceren de auteurs een Local Linear Integrator. Deze integrator gaat effectiever om met de numerieke stijfheid nabij $t=0$ (waar drift- en diffusiecoëfficiënten divergeren) dan standaard Euler-Maruyama discretisaties met offsets.

Belangrijkste Bijdragen

1. Theoretische Analyse van Matching-beperkingen: Het artikel biedt een formele analyse die aantoont dat flow/score matching-verliezen geen controle garanderen over beloningskloven (eigenschappen van belang) vanwege de verschuivingen tussen train- en testdistributies en de misuitlijning tussen de geometrie van de snelheidruimte en de sample-ruimte beloningsgeometrie.
2. Discriminator-Guided RL (DRL): Een nieuw framework dat een discriminator gebruikt, getraind in een bevroren SSL-representatieruimte, om een beloningssignaal te genereren. Dit stelt het model in staat om dataproperties te herstellen zonder menselijke voorkeurslabels.
3. Local Linear Integrator: Een nieuwe numerieke integrator voor de geheugenloze SDE die gebruikt wordt bij RL-training, welke vroege norm-explosie voorkomt en zowel de samplingskwaliteit als de trainingsstabiliteit verbetert.
4. Empirische Validatie: Demonstratie dat DRL de distributionele uitlijning en beeldkwaliteit aanzienlijk verbetert over diverse architecturen (SiT, JiT, REPA, RAE) en feature-ruimtes (DINOv2, DINOv3, SigLIP, Inception).

Resultaten

De auteurs evalueerden DRL op vier ImageNet-pre-trained flow modellen. Belangrijkste bevindingen zijn:

• Distributionele Uitlijning: DRL vermindert de Fréchet Distance (FD) aanzienlijk. Bijvoorbeeld, op SiT daalde de guidance-vrije FD in de DINOv3-ruimte van 88,2 naar 19,3. Over alle modellen en feature-ruimtes heen verbeterde DRL consequent de uitlijning zonder CFG (Classifier-Free Guidance) te gebruiken.
• Beeldkwaliteit zonder Voorkeursdata: DRL verbeterde de scores op onafhankelijke menselijke voorkeurs-beloningsmodellen (ImageReward, HPSv2, PickScore, Aesthetics), ondanks dat het nooit op deze getraind was. Dit suggereert dat DRL succesvol de "realisme"-eigenschappen herstelt die standaard matching-verliezen missen.
• Verbeterde Fundament voor Voorkeursuitlijning: Wanneer DRL wordt gebruikt als startpunt voor daaropvolgende voorkeursgebaseerde RL (PRL), levert het een superieur Pareto-front tussen beloning en beeldgetrouwheid op. DRL absorbeert de structurele correcties die uit de data te herstellen zijn, waardoor PRL zich kan richten op genuante subjectieve voorkeuren in plaats van op het oplossen van laag-niveau artefacten (bijv. oververzadiging, extreme helderheid).
• Distillatie-gat: Experimenten toonden aan dat een student-model getraind via standaard flow matching op samples van een DRL-teacher de prestaties van de teacher niet kon repliceren, zelfs niet na 50 miljoen samples. Dit bevestigt dat de RL-doelstelling zelf, en niet alleen de samples, noodzakelijk is om deze eigenschappen te herstellen.

Betekenis en Claims

Het artikel claimt dat RL post-training voor flow-gebaseerde modellen niet louter moet worden gezien als een hulpmiddel voor het optimaliseren van externe voorkeuren, maar als een complementair mechanisme voor het herstellen van structuur in de trainingsdata die imperfect wordt gevangen door standaard matching-gebaseerde doelstellingen.

Door een discriminator in een pre-trained representatieruimte te gebruiken, biedt DRL een beloningssignaal dat:

1. Uitgelijnd is met Dataproperties: Het richt zich direct op de datadistributie in plaats van op subjectieve menselijke voorkeuren.
2. Geometrisch Sound is: Het omzeilt de obstructies van de snelheidruimte-geometrie door direct te optimaliseren op het beloningslandschap.
3. Efficiënt is: Het vereist aanzienlijk minder rekenkracht dan pretraining en vermijdt de kosten van het verzamelen van voorkeursdata.

De auteurs concluderen dat matching-doelstellingen en RL geen methoden met conflicterende doelen zijn, maar eerder twee complementaire instrumenten om het nut van trainingsdata te maximaliseren. Het werk suggereert dat de "beloning" voor hoogwaardige generatie vaak "altijd al in de data zat", maar een on-policy, beloningsgestuurd mechanisme vereist om volledig gerealiseerd te worden.