ParetoSlider: Diffusion Models Post-Training for Continuous Reward Control

1. Het Probleem

Huidige methoden voor het afstemmen van generatieve modellen (zoals Diffusion en Flow-matching modellen) op menselijke voorkeuren maken vaak gebruik van Reinforcement Learning (RL). De standaardpraktijk is echter gebaseerd op single-objective RL, waarbij meerdere criteria worden samengevoegd tot één vaste, gewogen som van beloningen (rewards) vóór het trainen begint. Dit proces staat bekend als early scalarization.

Deze aanpak heeft twee fundamentele beperkingen:

1. Vastzitten aan één compromis: Het model leert slechts één specifiek trade-off punt (bijvoorbeeld een vaste balans tussen "prompt-adherentie" en "bronbehoud").
2. Geen controle tijdens inferentie: Gebruikers kunnen tijdens het genereren niet dynamisch schuiven tussen conflicterende doelen. Als ze een ander compromis willen, moeten ze een nieuw model trainen of meerdere checkpoints beheren, wat inefficiënt en kostbaar is.

Bovendien kunnen directe aggregaties van ruwe beloningen leiden tot "reward hijacking", waarbij beloningen met een hoge schaal of variatie de leerprocessen domineren en andere doelen verdringen.

2. Methodologie: ParetoSlider

ParetoSlider introduceert een Multi-Objective Reinforcement Learning (MORL) framework dat een enkel diffusion model traint om de volledige Pareto-front (de set van optimale trade-offs) te benaderen, in plaats van één statisch punt.

De kerncomponenten van de methode zijn:

• Preference-Conditioned Policy: In plaats van een statisch model, wordt het model getraind met een voorkeursvector $\omega$ als conditioneringssignaal. Deze vector (waarbij $\sum \omega_m = 1$) specificeert de gewenste weging tussen de verschillende beloningsfuncties. Tijdens de inferentie kan de gebruiker $\omega$ continu variëren om het gewenste compromis te kiezen zonder het model opnieuw te hoeven trainen.
• Late Scalarization Strategie: Om het probleem van "reward hijacking" en schaalverschillen op te lossen, gebruikt ParetoSlider late scalarization:
  1. Voor elke beloningsfunctie $m$ wordt een apart voordeel (advantage) berekend en genormaliseerd binnen een groep van gegenereerde samples.
  2. Er wordt een aparte DiffusionNFT-verliesfunctie ($L_m$) berekend voor elke beloning.
  3. Pas op het allerlaatste moment worden deze verliezen samengevoegd volgens de voorkeursvector $\omega$: $L_{total} = \sum \omega_m L_m$.
     Dit zorgt ervoor dat de optimalisatie wordt gestuurd door de voorkeur van de gebruiker en niet door de arbiträre schaal van de onderliggende beloningsfuncties.
• Architecturale Integratie: De voorkeursvector $\omega$ wordt op een lichte manier in de transformer-architectuur ingebracht (bijv. via LoRA-adapters en modulatieve lagen zoals AdaLN). Dit gebeurt op verschillende niveaus, afhankelijk van het model (SD3.5, FluxKontext, LTX-2), zodat het model de voorkeursinformatie effectief kan verwerken zonder de pre-train kennis te verliezen.
• Training Pipeline: Het framework bouwt voort op DiffusionNFT. Het proces omvat het genereren van een groep samples voor een gegeven prompt en voorkeursvector, het scoren met meerdere reward-modellen, het berekenen van genormaliseerde voordelen per doel, en het updaten van het beleid via een gewogen som van de verliezen.

3. Belangrijkste Bijdragen

• Eén Model voor Alle Trade-offs: ParetoSlider traint één enkel model dat de volledige continuüm van optimale oplossingen (Pareto-front) omvat. Dit elimineert de noodzaak om meerdere modellen te trainen voor verschillende wegingen.
• Continuous Inference-Time Control: Gebruikers kunnen tijdens het genereren soepel schuiven tussen conflicterende doelen (bijv. van fotorealisme naar schets, of van bronbehoud naar instructie-opvolging) door simpelweg de parameter $\omega$ aan te passen.
• Late Scalarization: De paper introduceert en valideert een nieuwe trainingsstrategie die de stabiliteit van multi-objective learning verbetert door beloningsnormering te scheiden van de voorkeursgewichten.
• Brede Toepasbaarheid: De methode is succesvol getoetst op drie verschillende state-of-the-art flow-matching backbones:
  • SD3.5 voor tekst-naar-beeld (T2I).
  • FluxKontext voor instructie-gebaseerde beeldbewerking (I2I).
  • LTX-2 voor tekst-naar-video (T2V).

4. Resultaten

De auteurs evalueren ParetoSlider tegenover bestaande baselines zoals Fixed-Weights (separate modellen voor vaste wegingen), Flow-Multi (Pareto-selectie zonder conditionering), en Prompt Rewriting.

• Kwalitatieve Resultaten: De visualisaties tonen soepele, continue overgangen tussen extremen. Bijvoorbeeld, bij beeldbewerking kan een gebruiker soepel schuiven van "volledig behoud van het origineel" naar "volledige uitvoering van de instructie", met een evenwichtig middenpunt dat beide aspecten respecteert.
• Kwantitatieve Resultaten:
  • Pareto-front Benadering: ParetoSlider bereikt een hogere Hypervolume (HV) score dan alle baselines, wat aangeeft dat het een bredere en betere set van oplossingen dekt.
  • Prestatie: Het enkele voorkeur-geconditioneerde model presteert gelijk aan of beter dan meerdere apart getrainde modellen op hun respectievelijke vaste punten.
  • Ablatie Studies: Vergelijkingen tonen aan dat late scalarization superieur is aan early scalarization en Smooth Tchebycheff methoden, omdat het soepelere overgangen mogelijk maakt en minder geneigd is om in te storten op één enkel doel.

5. Betekenis en Impact

ParetoSlider stelt een schaalbaar paradigma neer voor multi-objective alignment in visuele generatieve modellen. Het lost het fundamentele probleem op van het vastzitten aan statische compromissen in RL-finetuning.

De belangrijkste implicaties zijn:

• Gebruiksvriendelijkheid: Het opent de deur voor intuïtieve gebruikersinterfaces (zoals sliders) waarmee niet-experts complexe trade-offs kunnen navigeren zonder technische kennis van modeltraining.
• Efficiëntie: Het elimineert de kosten en het opslagvermogen die nodig zijn voor het beheren van tientallen modellen voor verschillende voorkeuren.
• Flexibiliteit: Het maakt dynamische aanpassing mogelijk aan veranderende gebruikersbehoeften tijdens de inferentie, wat essentieel is voor real-world toepassingen waar de definitie van "kwaliteit" contextafhankelijk is.

Kortom, ParetoSlider transformeert multi-objective optimalisatie van een statisch trainingsprobleem naar een dynamisch, controleerbaar inferentieproces.