Beyond Pairwise Preferences: Listwise Reward-Aware Alignment for Diffusion Models

Technische Samenvatting: Diffusion LAIR

Probleemstelling
Huidige methoden voor voorkeursoptimalisatie om tekst-naar-beeld diffusiemodellen te aligneren, zoals Direct Preference Optimization (DPO) en diens diffusiespecifieke varianten (bijv. Diffusion-DPO, DSPO), vertrouwen grotendeels op paarsgewijze supervisie. Deze methoden reduceren trainingsdata tot binaire "winnaar-verliezer"-vergelijkingen voor een gegeven prompt. De auteurs betogen dat deze reductie om twee hoofdredenen beperkend is:

1. Datastructuur: Praktische voorkeursdatasets (bijv. Pick-a-Pic) bevatten vaak meerdere kandidaatbeelden voor één prompt, wat een natuurlijke lijstgewijze structuur creëert die paarsgewijze methoden negeren.
2. Informatieverlies: Moderne beloningsmodellen leveren continue scores. Het reduceren hiervan tot één enkel binair label gooit de relatieve kwaliteit van niet-geselecteerde kandidaten weg en negeert de omvang van de beloningsgaten tussen samples.
   Bovendien zijn bestaande online reinforcement learning (RL)-benaderingen voor diffusie-alignment computatiekostenintensief, moeilijk te paralleliseren en gevoelig voor distributieveranderingen vanwege hun on-policy aard. Hoewel offline methoden efficiënter zijn, hebben ze het rijkere signaal dat beschikbaar is in lijstgewijze, met beloning gescoorde data niet volledig benut.

Methodologie: Diffusion LAIR
Het paper stelt Diffusion LAIR (Listwise Advantage-weighted Implicit Reward) voor, een offline voorkeursoptimalisatiemethode die alignment behandelt als een lijstgewijs, beloningsbewust probleem.

• Datarepresentatie: In plaats van onafhankelijke paren, verwerkt de methode een prompt $c$ met een set van $N_c$ kandidaatbeelden $\{x^{(i)}_0\}$ en hun bijbehorende continue beloningsscores $\{r_i\}$.
• Advantage-gewichting: De beloningsscores worden omgezet in gecentreerde advantage-weights ($w_i$). Eerst wordt een softmax-kans $p_i$ berekend met een temperatuur $\tau$. Vervolgens worden de weights gecentreerd tegen een uniforme baseline:
  $$w_i = p_i - \frac{1}{N_c}$$
  Dit zorgt ervoor dat $\sum w_i = 0$, waarbij hoge beloningsamples positieve weights krijgen en lage beloningsamples negatieve weights, relatief ten opzichte van het groepsgemiddelde.
• Optimalisatiedoelstelling: De methode optimaliseert een advantage-gewogen regressiedoelstelling op de impliciete beloning. De impliciete beloning wordt gedefinieerd als de verbetering in het ontdoofingsverlies van het huidige model ($\theta$) ten opzichte van een vast referentiemodel ($ref$):
  $$s_\theta(x_0, x_t, t, c, \epsilon) = \omega(\lambda_t) (l_{ref} - l_\theta)$$
  De verliesfunctie is:
  $$L_{LAIR} = \mathbb{E} \left[ -\sum_{i=1}^{N_c} w_i s^{(i)}_\theta + \frac{\lambda}{N_c} \sum_{i=1}^{N_c} (s^{(i)}_\theta)^2 \right]$$
  De eerste term moedigt het model aan om de impliciete beloning te verhogen voor hoge beloningsamples en te verlagen voor lage beloningsamples. De tweede term is een kwadratische straf die de grootte van de impliciete beloning regulariseert, waardoor agressieve afwijkingen van het referentiemodel worden voorkomen.

Theoretische Analyse
De auteurs bieden een theoretische karakterisering van de doelstelling:

1. Gesloten-vorm optimum: De doelstelling staat een unieke, begrende gesloten-vorm optimum toe in de ruimte van de impliciete beloning: $s^*_i = \frac{N_c}{2\lambda} w_i$. Dit verduidelijkt dat de regularisatiesterkte $\lambda$ direct de grootte van de voorkeursupdate controleert.
2. Surrogaat KL-grens: Onder specifieke aannames (log-ratio benadering en begrenste volledige ondersteuningsextensie) leiden de auteurs een grens af voor de Kullback-Leibler (KL)-divergentie tussen de geïnduceerde verdeling en de referentieverdeling: $D_{KL} \le \frac{\Delta_c}{\eta}$. Dit suggereert dat $\lambda$ de scherpte van de voorkeurshelling controleert, waardoor de update conservatief blijft.
3. Voordeel ten opzichte van paarsgewijs: In tegenstelling tot paarsgewijze doelstellingen (bijv. Bradley-Terry), die de impliciete beloning naar oneindig kunnen drijven om de marge tussen een winnaar en een verliezer te maximaliseren, wijst de lijstgewijze formulering van LAIR eindige, relatieve doelen toe gebaseerd op de volledige set kandidaten, waardoor de rangschikkingstructuur en de beloningsafstanden behouden blijven.

Experimentele Resultaten
De methode werd geëvalueerd op Stable Diffusion 1.5 (SD1.5) en Stable Diffusion XL (SDXL) met behulp van de Pick-a-Pic v2 dataset.

• Opzet: Beelden werden offline gescoord met het PickScore beloningsmodel. De modellen werden fine-tuned en geëvalueerd op tekst-naar-beeld generatie (HPD, Parti-prompts), compositiegeneratie (GenEval) en instructiegebaseerde beeldbewerking (InstructPix2Pix).
• Prestaties: Diffusion LAIR presteerde beter dan sterke baselines (Diffusion-DPO, DSPO, Diffusion-KTO, MaPO, InPO) op de meeste metrieken.
  • Op SD1.5 behaalde het de hoogste scores op PickScore, HPS, Aesthetics, CLIP en ImageReward voor zowel de Parti-prompt als de HPD datasets.
  • Op SDXL presteerde het eveneens beter dan baselines, waaronder MaPO en Diffusion-DPO.
  • Generalisatie: Opmerkelijk is dat, ondanks dat het model uitsluitend met PickScore werd getraind, het aanzienlijke verbeteringen liet zien in andere metrieken (bijv. ImageReward), wat suggereert dat de lijstgewijze doelstelling de algemene samplekwaliteit verbetert, verder dan de specifieke beloningsmodel die werd gebruikt.
  • Compositie & Bewerking: De methode toonde consistente winsten in GenEval (compositie-redenering) en InstructPix2Pix (beeldbewerking) winpercentages in vergelijking met paarsgewijze baselines.
• Efficiëntie: Het trainingsproces was aanzienlijk efficiënter dan online RL-baselines en vergelijkbaar of efficiënter dan andere offline methoden, waarbij ongeveer 24 GPU-uren nodig waren voor SD1.5 en 140 A100 GPU-uren voor SDXL.

Betekenis en Claims
Het paper claimt dat Diffusion LAIR een eenvoudige en effectieve richting biedt voor het aligneren van diffusiemodellen door verder te gaan dan de beperkingen van paarsgewijze reductie. De belangrijkste bijdragen zijn:

1. Lijstgewijs leren: Het is de eerste methode die direct gebruikmaakt van groepen met beloning gescoorde beelden voor diffusie-voorkeursoptimalisatie, waarbij de volledige structuur van relatieve samplekwaliteit wordt benut.
2. Theoretische helderheid: Het biedt een begrensd gesloten-vorm optimum, wat inzicht geeft in hoe regularisatie distributieverandering en updategrootte controleert.
3. Empirische superioriteit: Het demonstreert sterke prestaties op generatie-, compositie- en beweringstaken, wat bewijst dat lijstgewijze, beloningsbewuste optimalisatie bredere verbeteringen oplevert dan paarsgewijze methoden.

De auteurs merken beperkingen op, specifiek de afhankelijkheid van de kwaliteit van het beloningsmodel dat voor training wordt gebruikt en de geïdealiseerde aannames in de theoretische KL-grens. Ze stellen echter dat de methode robuust is ten opzichte van ruisachtige voorkeurslabels die veel voorkomen in door mensen geannoteerde datasets.