Technische Samenvatting: Reinforce Adjoint Matching (RAM)

Probleemstelling

Diffusie- en flow-matching-modellen bereiken generatie van hoge fideliteit omdat hun pretraining berust op supervised regressie: een schone steekproef wordt analytisch verstoord, en het model regresseert tegen een gesloten-vorm doel. Het afstemmen van deze modellen op beloningsfuncties (bijvoorbeeld voor compositiecorrectheid, tekstweergave of menselijke voorkeur) via Reinforcement Learning (RL) na-training verstoort echter deze efficiënte structuur.

Bestaande RL na-trainingsmethoden staan voor aanzienlijke computationele en structurele hindernissen:

1. Policy Gradient-methoden: Behandelen het ontdoemen van ruis als een Markov-beslissingsproces, vereisen stochastische SDE-rollouts en optimaliseren variatiele ondergrenzen. Dit introduceert hoge variantie en vereist veel bemonsteringsstappen.
2. Stochastische Optimal Control (Adjoint)-methoden: Naderen het KL-geregulariseerde beloningsmaximalisatie-doel exact, maar vereisen achterwaartse adjoint-sweeps langs opgeslagen SDE-trajecten en beloningsgradienten die via het bemonsteringsproces worden teruggepropageerd. Dit is numeriek instabiel en computationeel duur, vooral voor flow-matching-modellen waar roosters van ruis divergeren nabij het verstoorde einde.
3. Surrogaatverliezen: Recente benaderingen vervangen onberekenbare waarschijnlijkheidsverdelingen door surrogaten (bijvoorbeeld ELBO's of contrastieve eindpunten), ten koste van de schone regressiestructuur die pretraining schaalbaar maakt.

De kernuitdaging is het afleiden van een RL na-trainingsdoel dat de eenvoud en schaalbaarheid van het pretraining-regressiekader behoudt, terwijl het effectief optimaliseert voor beloningen.

Methodologie: Reinforce Adjoint Matching (RAM)

De auteurs stellen Reinforce Adjoint Matching (RAM) voor, een methode die de pretraining-regressiestructuur uitbreidt naar RL na-training zonder SDE-rollouts, achterwaartse adjoint-sweeps of beloningsgradienten te vereisen.

Theoretische Grondslag

De methode is gebaseerd op KL-geregulariseerde beloningsmaximalisatie, geformuleerd als een stochastisch optimal control-probleem.

1. Structuur van het Optimale Proces: Onder KL-regularisatie kantelt het optimale generatieve proces de verdeling van schone eindpunten ($X_0$) naar hoog-beloonde steekproeven ($p^*_0 \propto p_{ref}^0 \exp(r)$), maar laat de conditionele wet van verstoord toestanden gegeven een schoon eindpunt onveranderd. Specifiek blijft het traject van een vast $X_0$ naar een verstoord toestand $X_t$ hetzelfde analytische Gaussische kernel dat wordt gebruikt in pretraining.
2. Adjoint Matching-voorwaarde: De optimale besturing $u_t$ voldoet aan een vast-puntvoorwaarde waarbij de besturing gelijk is aan de negatief geschaalde gradient van de waardenfunctie (de adjoint): $u_t = -\sigma_t \nabla_x V_t(x)$.
3. REINFORCE-decompositie: Om het berekenen van $\nabla r$ te vermijden (wat niet-differentieel kan zijn), passen de auteurs de log-derivaat (REINFORCE)-identiteit toe op de waardenfunctie. Dit decomponeert de adjoint in een beloningsterm gewogen door de achterwaartse brugscore en een pad-kostcorrectieterm.

Het RAM-doel

RAM benadert de adjoint door de pad-kostcorrectie weg te laten (gerechtvaardigd door lage variantie en stabiliteit op schaal) en gebruik te maken van het structurele inzicht dat het verstorende kernel analytisch blijft.

1. Eindpunt-bemonstering: In plaats van een gecontroleerde SDE te simuleren, trekt RAM een schoon eindpunt $X_0$ uit de verdeling van het huidige model met behulp van elke beschikbare ODE-sampler.
2. Analytisch verstoren: Het bemonsterde $X_0$ wordt analytisch verstoord om trainingsstaten te genereren $X_t = (1-t)X_0 + t\epsilon$, exact zoals in pretraining.
3. Gesloten-vorm doel: Het regressiedoel voor het snelheidsveld $v_\theta$ wordt afgeleid door de beloningsterm en de Bayes-brugscore (die een gesloten-vorm uitdrukking heeft voor lineair verstoren) in te vullen in de adjoint-voorwaarde.

De resulterende verliesfunctie is:
$$L_{RAM}(\theta) = \mathbb{E}_t \left[ \left\| v_\theta^t(X_t) - \text{sg}\left( v_{ref}^t(X_t) + r(X_0) \left( (\epsilon - X_0) - v_\theta^t(X_t) \right) \right) \right\|^2 \right]$$
waarbij $X_0 \sim p_\theta^0$, $\epsilon \sim \mathcal{N}(0, I)$, en $\text{sg}(\cdot)$ de stop-gradient-operator aanduidt.

Belangrijke Efficiëntiemechanismen

• Geen SDE-rollouts: Trainingsstaten worden gegenereerd door een enkel eindpunt $K$ keer te verstoren met onafhankelijke ruisvectoren. Dit levert $K$ conditioneel onafhankelijke trainingssteekproeven per model-forward pass op, wat meer gradiëntsignaal per stap biedt dan methoden die vertrouwen op enkele SDE-trajecten.
• Geen beloningsgradienten: De methode werkt met black-box beloningen.
• Geen achterwaartse sweeps: De analytische aard van het verstorende kernel elimineert de noodzaak voor achterwaartse adjoint-integratie.

Belangrijkste Bijdragen

1. Structurele Uitbreiding: Het artikel demonstreert dat de regressiestructuur van pretraining uitbreidt naar RL na-training onder KL-regularisatie. Het optimale proces verandert alleen de verdeling van eindpunten, niet de conditionele verstoringswet.
2. RAM-algoritme: Een nieuw consistentieverlies dat het pretraining-doel corrigeert met beloningsinformatie met behulp van een gesloten-vorm Bayes-brugscore, waardoor de computationele knelpunten van bestaande methoden worden vermeden.
3. Theoretische Rechtvaardiging: Een afleiding die aantoont dat het vaste punt van RAM overeenkomt met het KL-geregulariseerde optimum, met exactheid gegarandeerd voor Gaussische referenties onder lineaire beloningen en over het algemeen overeenkomst op eerste orde.
4. Variantie-analyse: Een analyse van pad-kostschatters die aantoont dat het weglaten van de pad-kostcorrectie (zoals RAM doet) de variantie aanzienlijk vermindert in vergelijking met full-horizon of lokale adjoint-schattingen, waardoor het de enige schatter haalbaar is voor na-training op beeldschaal.

Experimentele Resultaten

De auteurs evalueerden RAM op Stable Diffusion 3.5 Medium (SD3.5M) over drie taken:

1. Compositiese Beeldgeneratie (GenEval): Meten van correctheid van objecten, attributen en ruimtelijke relaties.
2. Visuele Tekstweergave (OCR): Meten van leesbaarheid van tekst in gegenereerde afbeeldingen.
3. Afstemming op Menselijke Voorkeur (PickScore): Afstemmen op menselijke voorkeuren.

Prestaties:

• Beloningsmaximalisatie: RAM behaalde de hoogste trainingsbeloningen op alle drie de taken, presterend beter dan baselines waaronder Flow-GRPO, DiffusionNFT en AWM.
• Trainings-efficiëntie: RAM bereikte de piekbeloning van Flow-GRPO in 50× minder stappen voor GenEval, 48× minder voor OCR en 34× minder voor PickScore.
• Beeldkwaliteit: In tegenstelling tot sommige baselines die leden onder "reward hacking" (verslechterende esthetische kwaliteit voor taakmetrieken), behield RAM of verbeterde het de beeldkwaliteitsmetrieken (Aesthetic, DeQA, ImageReward, HPSv2) terwijl het taakbeloningen maximaliseerde.
• CFG-compatibiliteit: RAM blijft compatibel met Classifier-Free Guidance (CFG) na na-training, terwijl baselines zoals DiffusionNFT en AWM zagen dat hun taakbeloningen aanzienlijk daalden wanneer CFG werd toegepast.

Betekenis en Claims

Het artikel claimt dat RAM een paradigma-verschuiving vertegenwoordigt in RL na-training voor diffusie- en flow-matching-modellen. Door gebruik te maken van de specifieke structurele eigenschappen van het optimal control-probleem, herstelt RAM het "eenvoudige regressie"-paradigma dat pretraining schaalbaar maakte.

De auteurs stellen dat:

• RAM de noodzaak elimineert voor kostbare SDE-rollouts en achterwaartse adjoint-sweeps, die de primaire knelpunten zijn geweest voor het schalen van RL naar beeldgeneratie met hoge resolutie.
• De methode "eenvoudig is en schaalt", state-of-the-art resultaten bereikt met aanzienlijk minder trainingsstappen en vergelijkbare per-stap computekosten als bestaande methoden.
• De aanpak modellen succesvol afstemt op complexe beloningen (compositionaliteit, tekst, voorkeuren) zonder zichtbare kwaliteitsverslechtering of reward hacking, wat aantoont dat de pretraining-structuur robuust genoeg is om na-training afstemming aan te kunnen wanneer aan de juiste theoretische voorwaarden wordt voldaan.