Technische Samenvatting: Diffusion Fine-tuning met Rewarded Moment Matching Distillation

Probleemstelling

Diffusiemodellen zijn de standaard geworden voor hoogwaardige beeld-synthese en wetenschappelijke voorspellingen vanwege hun stabiele trainingsdoelstellingen. De praktische implementatie wordt echter gehinderd door de hoge computationele kosten van inferentie, die tientallen tot honderden denoising-stappen vereist. Hoewel distillatie-technieken deze modellen comprimeren tot generatoren met weinig stappen, en Reinforcement Learning (RL) fine-tuning ze afstemt op specifieke beloningsfuncties (bijv. menselijke voorkeuren of wetenschappelijke metrieken), blijft het combineren van deze twee fasen uitdagend.

Bestaande benaderingen om distillatie en beloningsoptimalisatie te versmelten, kampen met significante beperkingen:

1. Fragiliteit in gezamenlijke training: Het samenvoegen van distillatie en beloningsfine-tuning in één fase zorgt er vaak voor dat de gegenereerde monsters afwijken van de inputdistributie van het docentmodel (teacher model), waardoor het distillatiesignaal ongeldig wordt.
2. Geheugen- en biasproblemen: Het fine-tunen van een reeds gedistilleerd model via multi-step backpropagation (bijv. DRaFT) is geheugenintensief. Het afkappen van backpropagation (bijv. ReFL) introduceert bias door beloningen te evalueren op ruisige latents.
3. Structurele beperkingen: Methoden zoals HyperNoise vermijden keten-backpropagation door initiële ruis te verstoren, maar zijn structureel beperkt tot laagfrequente beeldveranderingen, wat hun vermogen beperkt om complexe, hoogfrequente structurele beloningen te optimaliseren.

Methodologie: Rewarded Moment Matching Distillation (RMMD)

De auteurs stellen RMMD voor, een tweefasig framework dat distillatie en beloningsfine-tuning principieel verbindt terwijl de hoge getrouwheid ("naturalness") van geavanceerde distillatie behouden blijft.

Fase 1: Moment-Matching Distillation (MMD)

Het proces begint met het distilleren van een basis docentmodel met behulp van Moment-Matching Distillation (MMD). Deze methode stemt de tussenliggende denoising-distributies langs het sampling-traject af, wat een multi-step studentmodel oplevert dat de marginalen van de docent nauwgezet volgt. Deze student wordt vervolgens bevroren om als een stabiele distributionele referentie te dienen.

Fase 2: On-Policy Reward Fine-tuning

In de tweede fase wordt het studentmodel getuned om een beloningsfunctie te maximaliseren, terwijl de distributionele getrouwheid behouden blijft. Belangrijke innovaties zijn:

1. On-Policy Sampling: In plaats van off-policy datapunten te gebruiken, samplet de methode ruisige toestanden $x^{pol}_t$ uit de huidige policy van de student (door de student $K_{student}$ stappen te laten draaien, de output opnieuw te verruisen en de gradiënt te stoppen). Dit zorgt ervoor dat de beloningsdoelstelling de verwachte beloning van de student direct benadert en voorkomt dat distributionele drift het trainingssignaal ongeldig maakt.
2. Recycled Regularization: De kernbijdrage is het hergebruiken van de Moment-Matching loss als regularisator tijdens de fine-tuning. In plaats van een heuristische straf, straft deze loss expliciet de discrepantie af tussen de tussenliggende denoising-distributies van de student en die van het bevroren referentiemodel. Dit biedt een interpreteerbare knop (het regularisatiegewicht) om de afweging te maken tussen beloningsmaximalisatie en distributionele getrouwheid.
3. Single-Step Gradiënt: De methode legt het volledige frequentiebereik vast zonder multi-step unrolling door een enkele denoising-stap uit te voeren op gecorrumpeerde on-policy samples en de beloning te evalueren.
4. L2 Regularisatie: Een aanvullende L2-term straft significante afwijkingen van de initiële MMD-gedistilleerde modelvoorspellingen af om de training verder te stabiliseren.

De totale doelstelling combineert de on-policy beloningsgradiënt met de moment-matching regularisator en de L2-term:
$$\nabla_\theta \mathcal{L}_{RMMD} = \mathbb{E} \left[ \left(\frac{\partial \hat{x}_0}{\partial \theta}\right)^\top \left( m_{student} - m_{teacher} - \lambda \nabla_{\hat{x}_0} R(\hat{x}_0) \right) \right] + \lambda_{reg} \mathcal{L}_{L2}$$

Belangrijkste Bijdragen

1. RMMD-procedure: Een nieuwe methode voor diffusiedistillatie die tegelijkertijd een beloningsfunctie maximaliseert, gebruikmakend van een on-policy versie van de moment-matching loss voor regularisatie.
2. Empirische Superieuriteit: Heeft aangetoond dat RMMD consequent beter presteert dan single-step baselines (DI++) en multi-step concurrenten (DRaFT, HyperNoise) op FID-Reward Pareto-fronten over diverse beloningsfuncties, variërend van pixelniveau-metrieken tot semantische uitlijning (CLIP).
3. Wetenschappelijke Toepassing: Succesvol toegepast RMMD op GenCast, een state-of-the-art diffusie-gebaseerd weervoorspellingsmodel, waarbij de Continuous Ranked Probability Score (CRPS) werd geoptimaliseerd.

Resultaten

ImageNet Experimenten

• Multi-Step Voordeel: Het gebruik van een 8-staps MMD-gedistilleerde student als startpunt presteerde aanzienlijk beter dan 1-staps baselines (DI++). De 8-staps RMMD bereikte superieure afwegingen op FID-Reward Pareto-fronten.
• Vergelijking: RMMD genereerde betere Pareto-fronten dan DRaFT en HyperNoise over verschillende resoluties (64x64 en 512x512) en beloningsvormen (CLIP-uitlijning, Inception Score, gladheid).
• Kwalitatieve Analyse: Visualisaties toonden aan dat terwijl DRaFT-1 de inhoud behield maar adversariële artefacten introduceerde, en DRaFT-2 leed onder reward hacking (drastische distributieverschuivingen), RMMD erin slaagde subtiele modificaties te integreren om beloningen te verhogen zonder significant af te wijken van de oorspronkelijke datadistributie.

GenCast Weervoorspelling

• Versnelling: Het gedistilleerde model bereikte een 7,5× versnelling (8 stappen versus 59 functie-evaluaties) vergeleken met de docent.
• Nauwkeurigheid: Het RMMD-model, geoptimaliseerd met CRPS, presteerde beter dan het docentmodel op 93% van de doelweervariabelen (vergeleken met 75% voor de beste MMD-alleen baseline).
• Kalibratie: RMMD pakte het probleem van onder-dispersie aan dat gebruikelijk is bij MMD-distillatie. Hoewel de on-policy variant iets slechtere dispersiemetrieken vertoonde dan de off-policy variant, bood het de beste algehele CRPS-verbeteringen.
• Generalisatie: Optimalisatie van de CRPS bij een voorspeltijd van 12 uur resulteerde in consistente verbeteringen voor voorspellingen tot 7 dagen, wat suggereert dat de methode de onderliggende distributionele modellering verbetert in plaats van alleen te overfitten op de directe beloning.

Betekenis en Claims

Het artikel claimt dat RMMD een principiële oplossing biedt voor het "drift"-probleem dat inherent is aan gezamenlijke distillatie en fine-tuning. Door de moment-matching loss te hergebruiken als regularisator en on-policy sampling te gebruiken, stelt de methode het mogelijk om complexe, differentieerbare doelstellingen (zoals CRPS in weervoorspelling) te optimaliseren zonder de generatieve kwaliteit of stabiliteit van het gedistilleerde model op te offeren.

De auteurs benadrukken dat deze aanpak schaalbaar is naar complexe, hoogdimensionale wetenschappelijke domeinen. De succesvolle toepassing op GenCast laat zien dat diffusiemodellen zowel aanzienlijk versneld kunnen worden als hun probabilistische voorspellingen kunnen verbeteren, wat de weg vrijmaakt voor efficiëntere en nauwkeurigere diffusie-gebaseerde modellen in wetenschappelijke berekeningen waar differentieerbare doelstellingen bestaan. Het werk suggereert dat het handhaven van een stabiele distributionele referentie via moment matching cruciaal is voor effectieve beloningsfine-tuning in diffusiemodellen.