Technische Samenvatting: DiffusionOPD

Probleemstelling

Versterkingsleren (RL) heeft zich bewezen als effectief voor het optimaliseren van op diffusie gebaseerde tekst-naar-beeldmodellen tegen enkele-objectief beloningssignalen (bijvoorbeeld esthetiek of tekstgetrouwheid). Echter, bestaande methoden kampen met moeilijkheden in multi-task omgevingen waar een enkel model meerdere, potentieel conflicterende doelstellingen tegelijkertijd moet vervullen. Huidige benaderingen staan voor twee primaire beperkingen:

1. Gecombineerde optimalisatie: Het gelijktijdig trainen van alle taken leidt vaak tot interferentie tussen taken en een ongelijkgewicht in taakmoeilijkheid, waarbij eenvoudigere taken de leerdynamiek domineren en signalen van uitdagendere taken onderdrukken.
2. Cascaderend RL: Het sequentieel optimaliseren van taken vermijdt directe gradiëntconflicten, maar is omstandig, vereist meerdere trainingsfasen en zorgvuldige planning. Bovendien is het vatbaar voor catastrofaal vergeten, waarbij aanpassing aan latere taken de prestaties op eerder geleerde taken verslechtert.

Het artikel betoogt dat multi-task RL ontkoppeld moet worden in twee distincte processen: single-task on-policy exploratie en multi-task integratie van capaciteiten.

Methodologie: DiffusionOPD

De auteurs stellen DiffusionOPD voor, een unificerend multi-task trainingsparadigma gebaseerd op Online Policy Distillatie (OPD). Het kader werkt in twee fasen:

1. Trainingsfase: Onafhankelijke Specialisatie van Leraren

In plaats van gecombineerde optimalisatie, traint de methode eerst een set taakspecifieke leraarmodellen onafhankelijk van elkaar. Elke leraar wordt geoptimaliseerd voor een enkele taak (bijvoorbeeld OCR, esthetiek, compositie-uitlijning) met behulp van standaard kant-en-klare RL-algoritmen. Dit zorgt ervoor dat elke leraar een hoog prestatiedak bereikt op zijn specifieke doelstelling zonder interferentie van andere taken.

2. Integratiefase: On-Policy Distillatie

Vervolgens wordt een unificerend studentmodel getraind om de capaciteiten van deze gespecialiseerde leraren te distilleren. Cruciaal is dat de student leert langs zijn eigen rollout-trajecten (on-policy).

• Proces: De student genereert een denoisend traject. Voor elke stap in dit traject levert de corresponderende taakspecifieke leraar supervisie.
• Doelstelling: De student wordt getraind om het transitiegedrag van de leraar op elke denoisende stap na te bootsen.

Theoretische Afleiding: Van Discreet naar Continu

De kern theoretische bijdrage is het tillen van het OPD-kader van discrete token-overgangen (LLM's) naar continu-staat Markov-processen (Diffusiemodellen).

• Markov-kettingformulering: Het diffusie-denoisingsproces wordt gemodelleerd als een discrete-tijd Markov-ketting die wordt geïnduceerd door een tijd-omgekeerde Stochastische Differentiaalvergelijking (SDE).
• Gaussische overgangen: Zowel de student als de leraar definiëren één-staps Gaussische overgangskernen. Aangezien ze dezelfde covariantie delen (bepaald door de planner en het ruisniveau) maar verschillen in hun gemiddelden, staat de omgekeerde Kullback-Leibler (KL)-divergentie toe tot een gesloten vorm uitdrukking:
  $$\text{KL}(p_S \parallel p_T) = \frac{\|\mu_S - \mu_T\|^2}{2\sigma^2}$$
  waarbij $\mu_S$ en $\mu_T$ respectievelijk de gemiddelden van de student- en leraarstransitie zijn.
• Optimalisatie: De student wordt geoptimaliseerd door het minimaliseren van deze gesloten vorm KL-doelstelling via directe backpropagatie.

Vergelijking met PPO-achtige Gradiënten

Het artikel contrasteert deze gesloten vorm-benadering met een PPO-achtige policy-gradiëntformulering (waarbij de leraar wordt behandeld als een procesbeloningsmodel).

• Variance-reductie: Hoewel PPO-achtige gradiënten onbevooroordeeld zijn in verwachting, introduceren ze een score-functie term evenredig aan Gaussisch ruis, wat leidt tot een hogere gradiëntvariatie.
• Deterministische compatibiliteit: De gesloten vorm KL-doelstelling breidt zich natuurlijk uit naar deterministische ODE-samplers (waarbij de overgang een directe $L_2$-match is), terwijl PPO-achtige doelstellingen afhankelijk zijn van stochastische policy-dichtheden.
• Conclusie: De gesloten vorm KL biedt een optimalisatieprincipe met lagere variatie en grotere generaliteit dat van toepassing is op zowel SDE- als ODE-samplers.

Belangrijkste Bijdragen

1. DiffusionOPD-kader: Een nieuw paradigma dat single-task exploratie ontkoppelt van multi-task integratie, gebruikmakend van taakspecifieke leraren om een unificerend studentmodel te superviseren langs zijn eigen trajecten.
2. Theoretische unificatie: Afleiding van een principiële, gesloten vorm per-stap KL-doelstelling voor diffusiemodellen die stochastische SDE en deterministische ODE-verfijning unificeert, en optimalisatie met lagere variatie biedt dan PPO-achtige gradiënten.
3. Empirische validatie: Uitgebreide experimenten die aantonen dat DiffusionOPD zowel multi-task RL-baselines (gecombineerde optimalisatie) als cascaderende RL-baselines overtreft in trainings-efficiëntie en uiteindelijke prestaties.

Experimentele Resultaten

De auteurs evalueerden DiffusionOPD op diverse benchmarks, waaronder GenEval (compositie-generatie), OCR (tekstweergave) en Esthetiek (visuele kwaliteit), met SD3.5-Medium als basismodel.

• Prestaties: DiffusionOPD behaalde State-of-the-Art (SOTA) resultaten op alle geëvalueerde benchmarks, en overtrof Multi-Task GRPO-Guard, Multi-Task NFT en Cascaderend NFT.
  • Gemiddelde Score: DiffusionOPD (0,929) versus Cascaderend NFT (0,851) versus Multi-Task NFT (0,715).
• Efficiëntie: De methode toonde aanzienlijk snellere convergentie en een hoger prestatiedak in vergelijking met multi-task RL-baselines.
• Traintijd: DiffusionOPD vereiste minder totale wandkloktijd om doelscores te bereiken in vergelijking met gecombineerde RL-methoden, terwijl het de sequentiële overhead van cascaderende methoden vermijdt.
• Ablatiestudies:
  • Verliesformulering: De gesloten vorm KL-doelstelling overtrof consistent PPO-achtige policy-gradiënten in zowel convergentiesnelheid als uiteindelijke beloning.
  • Ruisniveau: Lagere ruisniveaus (naderend bij ODE-sampling) leidden tot snellere convergentie en hogere prestaties, waarbij de ODE-sampler tot vijf keer efficiënter was dan de SDE-sampler met hoge ruis.

Betekenis en Claims

Het artikel beweert dat DiffusionOPD de fundamentele beperkingen van huidig multi-task RL voor diffusiemodellen aanpakt door exploratie te ontkoppelen van integratie. Door gebruik te maken van de analytische eigenschappen van diffusie-overgangen, biedt het een unificerend, laag-variatie trainingsprincipe dat van toepassing is op verschillende types samplers (SDE en ODE).

De auteurs positioneren DiffusionOPD niet louter als een specifieke distillatiemethode, maar als een algemeen kader voor on-policy distillatie in diffusiemodellen. Zij stellen dat deze benadering effectief de afweging oplost tussen optimalisatieconflict (bij gecombineerde training) en catastrofaal vergeten/inefficiëntie (bij cascaderende training), en zo een nieuwe basis legt voor multi-task en voorkeurs-uitgelijnde diffusiemodellering.