Aligning Few-Step Diffusion Models with Dense Reward Difference Learning

Il paper propone SDPO, un nuovo framework di apprendimento per rinforzo che allinea i modelli di diffusione a pochi passi agli obiettivi specifici tramite un campionamento a doppio stato e un'apprendimento della differenza di ricompensa densa, superando così le limitazioni dei metodi RL esistenti in termini di qualità del campione e stabilità dell'ottimizzazione.

L'essenziale

Immagina di avere un artista digitale (il modello di diffusione) che sa disegnare quadri bellissimi, ma ci mette molto tempo. Per renderlo veloce, gli abbiamo insegnato a fare "scorciatoie": invece di dipingere il quadro passo dopo passo per 50 minuti, ora lo fa in 4 minuti (o addirittura in 1 minuto!). Questo è il mondo dei modelli a pochi passaggi (few-step diffusion models).

Il problema? Quando l'artista corre troppo, tende a fare errori o a non ascoltare bene le tue richieste specifiche (ad esempio, "voglio un gatto cyberpunk che sembri reale e non un mostro").

Gli scienziati hanno provato a usare un sistema di premi e punizioni (Reinforcement Learning) per insegnargli a fare meglio. Ma qui c'era un grosso ostacolo: i metodi vecchi funzionavano come se l'artista dovesse finire l'intero quadro prima di ricevere un voto. Se l'artista correva (pochi passaggi), il voto arrivava troppo tardi o era confuso, e l'artista si frustrava, peggiorando invece che migliorando.

Ecco come SDPO (il nuovo metodo proposto in questo paper) risolve il problema, usando tre idee geniali:

1. Il "Fotografo Fantasma" (Campionamento a Doppio Stato)

Immagina che mentre l'artista dipinge, ci sia un fotografo fantasma che scatta una foto di come il quadro potrebbe essere finito in quel preciso istante, anche se il quadro è ancora pieno di macchie di colore.

• Vecchio metodo: Aspettava che il quadro fosse finito per dire "Bravo" o "Brutto".
• Metodo SDPO: Guarda la "foto fantasma" ad ogni singolo colpo di pennello. Se il colpo di pennello attuale sembra promettente, il fotografo dà un feedback immediato. Questo permette di correggere l'artista mentre lavora, non dopo aver sbagliato tutto.

2. Il "Saggio Indovino" (Previsione dei Premi)

Chiedere un voto per ogni singolo colpo di pennello sarebbe costosissimo e lento (come chiedere a un critico d'arte di valutare ogni singolo tratto di matita).

• La soluzione SDPO: Invece di chiedere un voto a ogni passo, ne chiede solo tre: all'inizio, alla fine e in un punto "strategico" (dove il quadro cambia più drasticamente).
• Poi, usa un Saggio Indovino basato sulla somiglianza: "Se questo tratto di pennello assomiglia molto a quello che ho già valutato, allora probabilmente vale lo stesso". In questo modo, ricostruisce un feedback continuo e dettagliato senza dover pagare (o calcolare) un voto per ogni singolo istante. È come indovinare il sapore di un piatto intero assaggiando solo tre ingredienti chiave.

3. La "Mappa del Tesoro" (Apprendimento delle Differenze)

Invece di dire all'artista: "Questo quadro è un 7 su 10", SDPO gli dice: "Guarda questo quadro (A) e questo quadro (B). Il quadro A è migliore di B perché in questo passaggio specifico hai fatto così".

• Si concentra sulle differenze tra un tentativo e l'altro, passo dopo passo.
• Inoltre, dà più importanza ai primi passi (come i primi tratti di un disegno che definiscono la forma) rispetto agli ultimi ritocchi, usando una sorta di "peso temporale".
• Per evitare che l'artista impari a memoria solo l'ordine in cui gli hai mostrato i quadri, SDPO mescola l'ordine dei passaggi durante l'allenamento, costringendo l'artista a capire la logica vera e non a memorizzare la sequenza.

Il Risultato?

Grazie a SDPO, l'artista digitale veloce (che prima produceva immagini sfocate o strane quando gli si chiedeva di correre) ora:

1. Impara molto più velocemente (serve meno materiale di allenamento).
2. Ascolta meglio le richieste (il gatto cyberpunk sembra davvero cyberpunk).
3. Funziona anche quando deve correre alla massima velocità (1 o 2 passaggi), cosa che i metodi precedenti non riuscivano a fare senza impazzire.

In sintesi: SDPO è come un allenatore sportivo che non aspetta la fine della partita per dare consigli, ma guarda ogni singolo movimento in tempo reale, usa la sua esperienza per prevedere il risultato dei movimenti intermedi e corregge l'atleta passo dopo passo, rendendolo un campione anche quando deve correre a velocità supersonica.

Sommario tecnico

1. Il Problema

I modelli diffusivi a pochi passi (few-step diffusion models) permettono una sintesi di immagini ad alta risoluzione estremamente efficiente, riducendo drasticamente il numero di passaggi di denoising necessari (es. 1-4 passi invece di 20-50). Tuttavia, questi modelli spesso non sono allineati con obiettivi specifici a valle, come la qualità estetica o le preferenze utente, che vengono solitamente valutati tramite funzioni di ricompensa.

L'apprendimento per rinforzo (RL) è un approccio promettente per l'allineamento, ma le metodologie esistenti presentano limiti critici quando applicate a regimi a pochi passi:

• Spazi degli stati limitati: Con pochi passi, le traiettorie di campionamento sono brevi, offrendo una diversità di segnali insufficiente per un'ottimizzazione stabile.
• Qualità subottimale: I campioni intermedi a pochi passi sono spesso di bassa qualità, rendendo difficile l'ottimizzazione basata su ricompense sparse (valutate solo all'ultimo passo).
• Instabilità nel Mixed-Step: Tentare di usare traiettorie di lunghezze diverse (mixed-step) per aumentare la copertura degli stati introduce un'alta varianza che destabilizza l'ottimizzazione della politica.
• Sovra-adattamento: L'uso di traiettorie lunghe per addestrare modelli a pochi passi porta spesso a un sovra-adattamento ai passaggi finali, degradando le prestazioni nell'inferenza a pochi passi.

2. Metodologia: SDPO (Stepwise Diffusion Policy Optimization)

Gli autori propongono SDPO, un nuovo framework di RL fine-tuning progettato specificamente per i modelli diffusivi a pochi passi. SDPO integra tre componenti principali:

A. Campionamento a Doppio Stato (Dual-State Trajectory Sampling)

Per superare la mancanza di feedback denso, SDPO introduce un meccanismo che traccia simultaneamente due stati a ogni passo $t$:

1. Lo stato rumoroso ($x_t$).
2. Lo stato pulito previsto ($\hat{x}_0^t$), ovvero una stima intermedia dell'immagine finale priva di rumore.
   Grazie alla forte capacità di denoising in un singolo passo dei modelli distillati, $\hat{x}_0^t$ funge da surrogato affidabile per l'output finale di un processo di denoising di $t$ passi. Questo permette di assegnare ricompense dense a ogni passo intermedio, mappando output di traiettorie di lunghezze diverse su una sequenza condivisa di stati puliti intermedi, riducendo la varianza e mantenendo dinamiche di denoising coerenti.

B. Predizione della Ricompensa Densa basata sulla Similarità Latente

Poiché interrogare la funzione di ricompensa a ogni passo è computazionalmente proibitivo, SDPO utilizza una strategia di predizione:

• Si interrogano le ricompense solo in tre punti chiave: primo passo, ultimo passo e un passo ancla (anchor) adattivamente selezionato.
• Il passo ancla è scelto per massimizzare l'informazione (minimizzando la similarità coseno con gli estremi nello spazio latente).
• Le ricompense per i passi non interrogati vengono stimate tramite interpolazione pesata dalla similarità latente, basandosi sull'assunzione di continuità di Lipschitz della funzione di ricompensa rispetto allo spazio latente.

C. Apprendimento della Differenza di Ricompensa Densa (Dense Reward Difference Learning)

Invece di ottimizzare il ritorno aggregato dell'intera traiettoria, SDPO formula un obiettivo che allinea la differenza di ricompensa densa con la differenza dei log-rapporti di verosimiglianza a livello di singolo passo.

• Stima dell'Advantage Passo-Passo: Incorpora dipendenze temporali a lungo termine calcolando un ritorno scontato ($\hat{G}_t$) e normalizzando per passo e prompt per ottenere stime di vantaggio ($\hat{A}_t$).
• Pesatura Temporale: Introduce una pesatura esponenziale che dà priorità ai passi iniziali (bassi), cruciali per i modelli a pochi passi.
• Aggiornamenti del Gradiente Mescolati (Step-Shuffled): Per evitare l'overfitting all'ordine fisso dei passi, gli aggiornamenti del gradiente vengono eseguiti su indici di passo mescolati all'interno del batch.

3. Contributi Chiave

1. Meccanismo di Campionamento a Doppio Stato: Abilita il feedback di ricompensa denso e l'ottimizzazione mista a bassa varianza per modelli a pochi passi.
2. Strategia di Predizione Efficiente: Riduce drasticamente il costo computazionale delle query di ricompensa mantenendo una guida granulare e robusta.
3. Obiettivo di Apprendimento della Differenza Densa: Permette aggiornamenti della politica più frequenti e granulari rispetto ai metodi basati su traiettoria completa.
4. Framework Unificato SDPO: Combina vantaggio passo-passo, pesatura temporale e aggiornamenti mescolati per un'ottimizzazione stabile ed efficiente, specialmente in regimi estremamente bassi (1-4 passi).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli come SD-Turbo e Latent Consistency Models (LCM), utilizzando diverse funzioni di ricompensa (Aesthetic Score, PickScore, ImageReward, HPSv2).

• Efficienza del Campione: SDPO supera significativamente metodi esistenti come DDPO, PRDP/REBEL e D3PO, raggiungendo punteggi di ricompensa più alti con meno campioni di addestramento.
• Prestazioni a Pochi Passi: In scenari di 1, 2 e 4 passi, SDPO produce immagini di qualità superiore e meglio allineate rispetto ai baselines, che spesso falliscono o producono immagini sfocate.
• Generalizzazione: Il modello mantiene alte prestazioni su prompt non visti durante l'addestramento, dimostrando una forte capacità di generalizzazione.
• Stabilità: A differenza di altri metodi che mostrano instabilità o collasso della ricompensa quando addestrati con traiettorie a 1 o 2 passi, SDPO mantiene curve di ottimizzazione stabili.
• Estensibilità: Il metodo è stato validato con successo anche su modelli per la generazione di immagini multivista (Text-to-Multiview).

5. Significato e Impatto

Questo lavoro risolve una delle principali sfide nell'adozione pratica dei modelli diffusivi a pochi passi: la loro difficoltà ad essere allineati con preferenze umane complesse tramite RL.

• Efficienza Computazionale: Riducendo il numero di query di ricompensa e permettendo l'uso di modelli a inferenza ultra-rapida (1-4 passi) senza sacrificare la qualità, SDPO rende l'allineamento RL economicamente e temporalmente fattibile.
• Stabilità Teorica: La proposta di un meccanismo di campionamento che unifica dinamiche coerenti su traiettorie miste offre una soluzione teorica solida al problema della varianza nell'RL per la generazione di immagini.
• Applicabilità: Il framework è generico e può essere applicato a una vasta gamma di modelli distillati e compiti di generazione, aprendo la strada a sistemi di generazione di immagini in tempo reale e di alta qualità.

In sintesi, SDPO rappresenta un avanzamento significativo nel campo dell'allineamento dei modelli generativi, rendendo possibile l'uso pratico di modelli diffusivi estremamente veloci mantenendo un controllo fine sulla qualità e sulle preferenze dell'utente.