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

Riepilogo Tecnico: Diffusion LAIR

Enunciazione del Problema
I metodi attuali di ottimizzazione delle preferenze per l'allineamento dei modelli di diffusione testo-immagine, come l'Ottimizzazione Diretta delle Preferenze (DPO) e le sue varianti specifiche per la diffusione (ad es. Diffusion-DPO, DSPO), fanno largamente affidamento sulla supervisione a coppie. Questi metodi riducono i dati di addestramento a confronti binari "vincitore-perdente" per un dato prompt. Gli autori sostengono che questa riduzione è limitante per due motivi principali:

1. Struttura dei Dati: I dataset pratici di preferenze (ad es. Pick-a-Pic) contengono spesso molteplici immagini candidate per un singolo prompt, creando una struttura naturale di tipo listwise che i metodi a coppie ignorano.
2. Perdita di Informazione: I modelli di reward moderni forniscono punteggi continui. Ridurli a un singolo etichetta binaria scarta la qualità relativa delle candidate non selezionate e ignora l'entità dei gap di reward tra i campioni.
   Inoltre, gli approcci esistenti di apprendimento per rinforzo (RL) online per l'allineamento della diffusione sono computazionalmente costosi, difficili da parallelizzare e sensibili agli spostamenti di distribuzione a causa della loro natura on-policy. Sebbene i metodi offline siano più efficienti, non hanno sfruttato appieno il segnale più ricco disponibile nei dati di tipo listwise e valutati con reward.

Metodologia: Diffusion LAIR
Il documento propone Diffusion LAIR (Listwise Advantage-weighted Implicit Reward), un metodo di ottimizzazione delle preferenze offline che tratta l'allineamento come un problema listwise e consapevole del reward.

• Rappresentazione dei Dati: Invece di coppie indipendenti, il metodo elabora un prompt $c$ con un insieme di $N_c$ immagini candidate $\{x^{(i)}_0\}$ e i loro corrispondenti punteggi di reward continui $\{r_i\}$.
• Ponderazione dell'Advantage: I punteggi di reward sono convertiti in pesi di advantage centrati ($w_i$). Innanzitutto, viene calcolata una probabilità softmax $p_i$ utilizzando una temperatura $\tau$. Quindi, i pesi sono centrati rispetto a una baseline uniforme:
  $$w_i = p_i - \frac{1}{N_c}$$
  Questo garantisce che $\sum w_i = 0$, assegnando pesi positivi ai campioni ad alto reward e pesi negativi ai campioni a basso reward rispetto alla media del gruppo.
• Obiettivo di Ottimizzazione: Il metodo ottimizza un obiettivo di regressione ponderato sull'advantage sul reward implicito. Il reward implicito è definito come il miglioramento della perdita di denoising del modello corrente ($\theta$) rispetto a un modello di riferimento fisso ($ref$):
  $$s_\theta(x_0, x_t, t, c, \epsilon) = \omega(\lambda_t) (l_{ref} - l_\theta)$$
  La funzione di perdita è:
  $$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]$$
  Il primo termine incoraggia il modello ad aumentare il reward implicito per i campioni ad alto reward e a diminuirlo per i campioni a basso reward. Il secondo termine è una penalità quadratica che regolarizza l'entità del reward implicito, prevenendo deviazioni aggressive dal modello di riferimento.

Analisi Teorica
Gli autori forniscono una caratterizzazione teorica dell'obiettivo:

1. Ottimo in Forma Chiusa: L'obiettivo ammette un ottimo unico e limitato in forma chiusa nello spazio del reward implicito: $s^*_i = \frac{N_c}{2\lambda} w_i$. Questo chiarisce che la forza di regolarizzazione $\lambda$ controlla direttamente l'entità dell'aggiornamento delle preferenze.
2. Limite KL Surrogato: Sotto ipotesi specifiche (approssimazione del log-rapporto ed estensione limitata a supporto completo), gli autori derivano un limite sulla divergenza di Kullback-Leibler (KL) tra la distribuzione indotta e la distribuzione di riferimento: $D_{KL} \le \frac{\Delta_c}{\eta}$. Ciò suggerisce che $\lambda$ controlla la nitidezza dell'inclinazione delle preferenze, assicurando che l'aggiornamento rimanga conservativo.
3. Vantaggio rispetto alle Coppie: A differenza degli obiettivi a coppie (ad es. Bradley-Terry) che possono spingere il reward implicito all'infinito per massimizzare il margine tra un vincitore e un perdente, la formulazione listwise di LAIR assegna target finiti e relativi basati sull'intero insieme di candidati, preservando la struttura di ranking e la spaziatura dei reward.

Risultati Sperimentali
Il metodo è stato valutato su Stable Diffusion 1.5 (SD1.5) e Stable Diffusion XL (SDXL) utilizzando il dataset Pick-a-Pic v2.

• Configurazione: Le immagini sono state valutate offline utilizzando il modello di reward PickScore. I modelli sono stati fine-tunati e valutati sulla generazione testo-immagine (HPD, Parti-prompts), sulla generazione composita (GenEval) e sulla modifica di immagini basata su istruzioni (InstructPix2Pix).
• Prestazioni: Diffusion LAIR ha superato basi solide (Diffusion-DPO, DSPO, Diffusion-KTO, MaPO, InPO) nella maggior parte delle metriche.
  • Su SD1.5, ha ottenuto i punteggi più alti su PickScore, HPS, Estetica, CLIP e ImageReward per entrambi i dataset Parti-prompt e HPD.
  • Su SDXL, ha superato similmente le basi, inclusi MaPO e Diffusion-DPO.
  • Generalizzazione: Notevolmente, sebbene addestrato solo con PickScore, il modello ha mostrato miglioramenti significativi in altre metriche (ad es. ImageReward), suggerendo che l'obiettivo listwise migliora la qualità generale dei campioni oltre il modello di reward specifico utilizzato.
  • Composizione e Modifica: Il metodo ha dimostrato guadagni coerenti nei tassi di vittoria su GenEval (ragionamento compositivo) e InstructPix2Pix (modifica di immagini) rispetto alle basi a coppie.
• Efficienza: Il processo di addestramento è stato significativamente più efficiente delle basi RL online e comparabile o più efficiente rispetto ad altri metodi offline, richiedendo circa 24 ore GPU per SD1.5 e 140 ore GPU A100 per SDXL.

Significato e Affermazioni
Il documento afferma che Diffusion LAIR offre una direzione semplice ed efficace per allineare i modelli di diffusione superando i limiti della riduzione a coppie. I suoi contributi chiave sono:

1. Apprendimento Listwise: È il primo metodo a sfruttare direttamente gruppi di immagini valutate con reward per l'ottimizzazione delle preferenze nella diffusione, utilizzando la struttura completa della qualità relativa dei campioni.
2. Chiarezza Teorica: Fornisce un ottimo in forma chiusa limitato, offrendo intuizioni su come la regolarizzazione controlli lo spostamento di distribuzione e l'entità dell'aggiornamento.
3. Superiorità Empirica: Dimostra prestazioni solide nella generazione, composizione e modifica, dimostrando che l'ottimizzazione listwise e consapevole del reward produce miglioramenti più ampi rispetto ai metodi a coppie.

Gli autori notano le limitazioni, in particolare la dipendenza dalla qualità del modello di reward utilizzato per l'addestramento e le ipotesi idealizzate nel limite KL teorico. Tuttavia, ipotizzano che il metodo sia robusto rispetto alle etichette di preferenze rumorose comuni nei dataset annotati dall'uomo.