Parallel Diffusion Solver via Residual Dirichlet Policy Optimization

Deze paper introduceert de EPD-Solver, een nieuwe ODE-oplosser die de hoge latentie van diffusiemodellen verlaagt door parallelle gradientberekeningen en een tweestaps-optimatieframework met reinforcement learning te gebruiken om de beeldkwaliteit bij snelle generatie te behouden.

De kern

Stel je voor dat je een meesterkunstenaar bent die een prachtig schilderij moet maken, maar je hebt een heel lastige opdracht: je mag alleen één kwaststreek per seconde zetten, en elke streek moet perfect zijn. Als je te snel gaat (weinig streeken), wordt het schilderij wazig en onherkenbaar. Als je te langzaam gaat (veel streeken), duurt het uren voordat het af is.

Dit is precies het probleem met Diffusiemodellen (de AI's die foto's maken zoals Midjourney of DALL-E). Ze werken door een ruisig beeld heel langzaam "op te poetsen" tot het een duidelijke foto is. Dit proces is echter traag omdat het stap-voor-stap moet gebeuren.

Deze paper introduceert een nieuwe oplossing: de EPD-Solver. Laten we kijken hoe dit werkt met een paar simpele vergelijkingen.

1. Het oude probleem: De "Gokker" vs. De "Wiskundige"

Stel je voor dat je een berg afdaalt.

• De oude methoden (zoals DDIM) kijken alleen naar de grond direct onder hun voeten en zeggen: "Ik ga hier een stap zetten." Als de berg een scherpe bocht maakt, raken ze de weg kwijt en vallen ze in een kuil (dit noemen ze fouten of truncation errors).
• Andere methoden (zoals EDM) kijken naar de grond onder hun voeten én naar de plek waar ze over 10 seconden zouden zijn, en nemen het gemiddelde. Dit is beter, maar nog steeds niet perfect bij scherpe bochten.

2. De EPD-oplossing: De "Parallellische Verkenner"

De EPD-Solver heeft een slimme truc bedacht. In plaats van één of twee punten te checken, doet hij alsof hij meerdere vrienden heeft die tegelijkertijd de berg aflopen.

• De Parallelle Kracht: Stel je voor dat je in één seconde vijf vrienden de berg afstuurt naar verschillende plekken in de buurt. Ze kijken allemaal naar de helling (de "gradiënt") op hun plek.
• De Samenvatting: Vervolgens komen ze terug en zeggen: "Oké, als we al onze waarnemingen samenvoegen, weten we precies welke kant we op moeten, zelfs als de weg heel krom is."
• Het Magische: Omdat al deze vrienden tegelijkertijd werken (parallel), duurt het niet langer dan het kijken van één vriend. De AI maakt dus dezelfde tijd, maar krijgt veel meer informatie. Het is alsof je met één oogopslag vijf verschillende hoeken van een object ziet in plaats van maar één.

3. Twee fases van leren: Van "Kopieer" naar "Meester"

De auteurs laten de AI in twee stappen leren, net zoals een leerling die eerst een meester kopieert en dan zijn eigen stijl vindt.

Fase 1: De Distillatie (Het kopiëren)
De AI kijkt naar een zeer langzame, perfecte "meester-AI" die 50 stappen maakt. De EPD-Solver probeert deze perfecte route na te bootsen, maar dan in slechts 5 stappen. Hij leert de "bochten" van de weg te begrijpen door de parallelle vrienden te gebruiken. Dit is als een student die de oplossingen van een wiskundig genie bestudeert om zelf snel te kunnen rekenen.

Fase 2: Reinforcement Learning (Het menselijke gevoel)
Soms is een wiskundig perfecte route niet wat mensen mooi vinden. Misschien is de foto technisch correct, maar ziet de hond er een beetje raar uit.

• Hier komt RL (Versterkend Leren) om de hoek kijken.
• De auteurs gebruiken een slimme truc: ze laten de AI niet de hele zware "meester" herschrijven (dat kost te veel tijd en energie). In plaats daarvan laten ze alleen de stuurknoppen van de EPD-Solver aanpassen.
• Ze gebruiken een Dirichlet-beleid: Denk hierbij aan een magische kompasnaald. De AI leert hoe ze deze naald moet draaien om de beste resultaten te krijgen, gebaseerd op wat mensen mooi vinden (bijvoorbeeld: "deze foto ziet er realistischer uit").
• Ze gebruiken een methode genaamd RLOO: De AI maakt 5 versies van een foto, laat een "jury" (een reward model) ze beoordelen, en past de stuurknoppen aan op basis van welke versie het beste scoorde.

Waarom is dit zo cool?

1. Snelheid: Het is net zo snel als de oude methoden, maar de foto's zijn veel scherper.
2. Kwaliteit: Op tests (zoals het maken van foto's van mensen of dieren) scoort deze methode beter dan alles wat er voorheen bestond, zelfs als je maar heel weinig stappen (5 of 20) gebruikt.
3. Flexibiliteit: Het werkt als een "plug-in". Je kunt het op bijna elke bestaande foto-AI plakken om die direct sneller en beter te maken.

Kort samengevat:
De EPD-Solver is als een team van verkenners dat samenwerkt om een berg af te dalen. In plaats van één persoon die langzaam en voorzichtig loopt, sturen ze een groepje dat tegelijkertijd de weg verkent. Hierdoor vinden ze de perfecte route veel sneller, zonder dat het langer duurt. En door een beetje "menselijke feedback" toe te voegen, zorgen ze ervoor dat het eindresultaat er niet alleen technisch goed uitziet, maar ook echt mooi is voor het menselijk oog.

Technische samenvatting

Titel: Parallel Diffusion Solver via Residual Dirichlet Policy Optimization (EPD-Solver)

1. Het Probleem

Diffusiemodellen (DMs) hebben state-of-the-art prestaties bereikt in generatieve taken zoals beeldsynthese en videoproduktie. Echter, een groot nadeel is de hoge sampling-latentie. Dit komt doordat deze modellen data genereren via een sequentieel proces van ruisreductie (denoising), wat veel iteratieve stappen vereist.

• Bestaande oplossingen:
  • Solver-gebaseerde methoden: Proberen het aantal stappen te verminderen door snellere numerieke solvers te gebruiken. Dit leidt echter vaak tot grote truncatiefouten en kwaliteitsverlies bij een laag aantal functiebeoordelingen (NFE).
  • Distillatie-gebaseerde methoden: Trainen een studentmodel om direct van ruis naar data te gaan. Dit is extreem snel maar vereist hoge trainingskosten en mist flexibiliteit.
  • Parallelle methoden: Proberen rekentijd te ruilen voor snelheid, maar zijn vaak nog niet optimaal voor kwaliteitsverbetering.
• De kernuitdaging: Hoe kan men de sampling-snelheid verhogen (laag NFE) zonder de beeldkwaliteit te laten verslechteren door opgehoopte numerieke fouten, en dit zonder de inferentie-tijd (latency) significant te verhogen?

2. Methodologie

De auteurs stellen de Ensemble Parallel Direction solver (EPD-Solver) voor, een nieuwe ODE-solver die gebruikmaakt van parallelle gradient-berekeningen om truncatiefouten te minimaliseren. Het werk rust op twee pijlers: een theoretisch inzicht en een tweestaps optimalisatieframework.

A. Theoretisch Inzicht: Parallelle Gradienten en het Gemiddelde Waarde Theorema

• Geometrische Observatie: De auteurs analyseren dat sampling-trajecten van diffusiemodellen grotendeels beperkt zijn tot een laag-dimensionale variëteit (manifold), specifiek een 2D-vlak, zelfs in hoge dimensies.
• Theorema: Ze gebruiken het Gemiddelde Waarde Theorema voor vector-waardige functies. Dit theorema stelt dat de exacte integraal van een vectorfunctie over een interval kan worden uitgedrukt als een gewogen som (simpliciale combinatie) van gradienten die op meerdere punten binnen dat interval worden geëvalueerd.
• EPD-Solver: In plaats van één gradient (zoals bij Euler/DIM) of twee (zoals bij Heun/EDM) te gebruiken, evalueert EPD-Solver $K$ parallelle gradienten op geleerde tussenliggende tijdstippen binnen één integratiestap. Omdat deze berekeningen onafhankelijk zijn, kunnen ze volledig parallel worden uitgevoerd op moderne hardware, waardoor de wandkloktijd (latency) niet toeneemt ondanks de extra berekeningen.

B. Tweestaps Optimalisatie Framework
Om de parameters van deze solver (tijdstippen en gewichten) te leren, gebruiken de auteurs een tweestapsaanpak:

1. Stap 1: Distillatie-gebaseerde Parameteroptimalisatie

   • Het doel is een robuuste initialisatie te vinden die de trajecten van een hoog-NFE "leraar"-solver (bijv. DPM-Solver-2) nabootst.
   • Ze optimaliseren een kleine set leerbare parameters ($\Theta_n$) die de schaal van de netwerkoutput en de tijdstippen aanpassen om blootstellingsbias (exposure bias) te verminderen.
   • Dit resulteert in een stabiele basis-solver die de kromming van het traject goed benadert.

2. Stap 2: Residual Dirichlet Policy Optimization (RDPO)

   • Distillatie alleen is vaak onvoldoende voor zeer lage stap-aantallen en aligneert niet perfect met menselijke perceptie.
   • De solver wordt herformuleerd als een stochastisch beleid (policy). De parameters (tijdstippen en gewichten) worden geparametriseerd via Dirichlet-verdelingen.
   • In plaats van de hele backbone van het diffusiemodel te finetunen, wordt alleen een "residueel" beleid geleerd rondom de gedistilleerde basis.
   • Ze gebruiken Reinforcement Learning (RL) met een PPO-variant (Proximal Policy Optimization) en een Reward Leave-One-Out (RLOO) baseline.
   • De beloning (reward) wordt bepaald door menselijke voorkeurmodellen (zoals HPSv2.1), waardoor de solver wordt getraind om semantisch en esthetisch betere beelden te genereren zonder de onderliggende DM te wijzigen.

3. Belangrijkste Bijdragen

• EPD-Solver: Een nieuwe ODE-solver die parallelle gradient-evaluaties gebruikt om truncatiefouten te reduceren zonder extra latency.
• EPD-Plugin: Een flexibele plugin die bestaande ODE-samplers kan verbeteren door hun gradiënt-schatting te vervangen door de parallelle EPD-methode.
• Parameter-efficiënt RL: Een innovatief finetuning-schema (RDPO) dat de solver als een Dirichlet-beleid behandelt. Dit maakt het mogelijk om menselijke voorkeuren te aligneren met minimale trainingskosten (slechts een paar parameters) en zonder het risico van "reward hacking" op het grote model.
• Theoretische onderbouwing: Een bewijs dat de exacte integraal van een vector-veld kan worden benaderd door een convexe combinatie van gradienten binnen het interval, wat de basis vormt voor de parallelle aanpak.

4. Resultaten

De auteurs testen hun methode op diverse benchmarks, variërend van CIFAR-10 tot grote Text-to-Image (T2I) modellen zoals Stable Diffusion v1.5 en SD3-Medium.

• Unconditional Generatie (Benchmarks):
  • Bij een laag latency-budget (5 NFE) behaalt EPD-Solver state-of-the-art FID-scores:
    • CIFAR-10: 4.47
    • FFHQ: 7.97
    • ImageNet: 8.17
    • LSUN Bedroom: 8.26
  • Dit is een aanzienlijke verbetering ten opzichte van bestaande solvers zoals AMED-Solver (bijv. 13.20 op LSUN).
• Text-to-Image (T2I):
  • Op Stable Diffusion v1.5 en SD3-Medium verbetert de RL-gefinetuned solver de menselijke voorkeurscores (HPSv2.1, ImageReward) aanzienlijk.
  • Op SD3-Medium overtreft de EPD-Solver met slechts 20 stappen de officiële 28-staps baseline (DDIM) op menselijke voorkeur, terwijl de inferentiekosten met 60% worden verlaagd.
  • De methode sluit de kloof tussen inferentie-efficiëntie en hoge-fideliteit generatie.
• Efficiëntie:
  • De parallelle berekeningen veroorzaken verwaarloosbare vertraging. Bij $K=2$ (twee parallelle gradienten) blijft de latency bijna gelijk aan de baseline, terwijl de kwaliteit sterk verbetert.

5. Betekenis en Impact

Dit paper biedt een doorbraak in het versnellen van diffusiemodellen zonder kwaliteitsverlies.

• Efficiëntie vs. Kwaliteit: Het lost het fundamentele compromis op tussen het aantal stappen (snelheid) en de beeldkwaliteit door slimme parallelisatie en theoretisch onderbouwde numerieke integratie.
• Menselijke Alignement: Door Reinforcement Learning toe te passen op de solver in plaats van het model, biedt het een nieuwe, kosteneffectieve route om generatieve modellen beter af te stemmen op menselijke perceptie en voorkeuren.
• Toepasbaarheid: De methode is plug-and-play en kan worden toegepast op bestaande solvers (via EPD-Plugin) en werkt op schaal van kleine datasets tot grote T2I-modellen.

Kortom, EPD-Solver stelt een nieuwe standaard voor snelle, hoogwaardige generatie door de wiskundige beperkingen van numerieke integratie te omzeilen via parallelisme en menselijke feedback.