Riepilogo Tecnico: Diff-Instruct con Ricompensa Diffusa (Didr)

Enunciato del Problema

I recenti progressi nei generatori testo-immagine (T2I) in un singolo passaggio hanno reso possibile la sintesi in tempo reale mappando il rumore latente direttamente sulle immagini in un'unica passata in avanti. Tuttavia, allineare questi generatori alle preferenze umane tramite Apprendimento per Rinforzo (RL) presenta una sfida strutturale fondamentale.

I metodi precedenti per RLHF in un singolo passaggio (ad es. Diff-Instruct*, Diff-Instruct++) combinano l'ottimizzazione della ricompensa nello spazio delle immagini con l'adattamento della distribuzione nello spazio latente. Questi approcci applicano tipicamente i segnali di preferenza (ricompense) esclusivamente al punto finale dell'immagine pulita ($x_0$), imponendo al contempo una regolarizzazione Kullback-Leibler (KL) contro la distribuzione di riferimento lungo l'intera traiettoria di diffusione rumorosa. Gli autori identificano una disallineamento strutturale critico in questo paradigma: la ricompensa agisce solo sullo stato terminale, mentre il regolarizzatore si estende su tutti i livelli di rumore.

Questa discrepanza porta a una modalità di fallimento definita dominazione della ricompensa terminale. Poiché il rumore in avanti maschera naturalmente i dettagli dell'immagine, la penalità di regolarizzazione KL contro la divergenza di modalità si indebolisce significativamente a livelli di rumore elevati. Di conseguenza, gli ottimizzatori sfruttano questa vulnerabilità della traiettoria, abbandonando la distribuzione di riferimento per collassare su regioni ad alta ricompensa senza incorrere in penalità KL sufficienti. Ciò si traduce in punteggi di ricompensa migliorati, ma spesso a scapito della fedeltà dell'immagine e della stabilità distribuzionale.

Metodologia: Diff-Instruct con Ricompensa Diffusa (Didr)

Per risolvere questo disallineamento, gli autori propongono Didr, un framework di allineamento a livello di traiettoria privo di dati, derivato dalla minimizzazione della KL Integrale (IKL). L'idea centrale è propagare sistematicamente la distribuzione delle immagini pulite ottimizzata per la ricompensa RLHF su tutti i livelli di rumore della traiettoria di diffusione.

1. L'Obiettivo di Traiettoria Principiato

Il metodo inizia stabilendo l'obiettivo di allineamento corretto per le immagini pulite sotto RLHF regolarizzato KL. La distribuzione ottimale delle immagini pulite $q^*(x_0|c)$ è la densità inclinata dalla ricompensa:
$$q^*(x_0|c) \propto q_0(x_0|c) \exp(r(x_0, c)/\tau)$$
dove $q_0$ è la distribuzione di riferimento, $r$ è la funzione di ricompensa e $\tau$ è un parametro di temperatura.

Invece di applicare le ricompense solo a $x_0$, Didr "solleva" questo obiettivo ottimale all'intera traiettoria diffondendo $q^*$ attraverso il processo in avanti di riferimento. Ciò genera una famiglia di marginali inclinati dalla ricompensa $\{q^*_t\}_{t \in [0, T]}$ definiti a ogni livello di rumore:
$$q^*_t(x_t|c) = \int q_t(x_t|x_0) q^*(x_0|c) dx_0$$
Gli autori dimostrano che la minimizzazione della divergenza KL Integrale tra le marginali del generatore e queste marginali inclinate dalla ricompensa condivide lo stesso minimizzatore dell'obiettivo standard RLHF regolarizzato KL. Ciò garantisce che ricompensa e regolarizzazione siano bilanciate a ogni livello di rumore, prevenendo la dominazione della ricompensa terminale.

2. Il Punteggio della Ricompensa Diffusa (DRS)

Differenziando l'obiettivo IKL, gli autori derivano la funzione di punteggio target. Il punteggio ottimale si decompone in due componenti:
$$s_{target} = s_{ref} + s_r$$
dove $s_{ref}$ è la funzione di punteggio di riferimento e $s_r$ è il Punteggio della Ricompensa Diffusa (DRS). Il DRS agisce come un termine di correzione guidato dalla ricompensa definito a ogni livello di rumore:
$$s_r(x_t, t, c) = \nabla_{x_t} \log \mathbb{E}_{x_0 \sim q(x_0|x_t, c)} \left[ \exp\left(\frac{r(x_0, c)}{\tau}\right) \right]$$
A differenza della guida basata su classificatori (che è un'approssimazione del primo ordine), il DRS mantiene l'intera inclinazione log-esponenziale, propagando le informazioni di preferenza solo dove l'immagine pulita è identificabile dall'osservazione rumorosa.

3. Il Proxy della Ricompensa Diffusa (DRP)

Poiché il DRS richiede un'attesa posteriore intrattabile $q(x_0|x_t, c)$, gli autori introducono il Proxy della Ricompensa Diffusa (DRP) per rendere il metodo pratico. Il DRP approssima l'attesa posteriore utilizzando catene di denoising a pochi passi differenziabili dal modello di riferimento congelato.

• Meccanismo: Da uno stato rumoroso $x_t$, il metodo esegue $K$ catene di denoising indipendenti a $S$ passi utilizzando il modello di riferimento congelato per generare campioni approssimati di immagini pulite $\{\hat{x}^{(k)}_0\}$.
• Stima del Gradiente: Il DRS è approssimato come una media ponderata con softmax dei gradienti di ricompensa lungo il percorso:
  $$\hat{s}_r(x_t, t, c) = \frac{1}{\tau} \sum_{k=1}^K \omega^{(k)} \nabla_{x_t} r(\hat{x}^{(k)}_0, c)$$
  dove i pesi $\omega^{(k)}$ sono proporzionali a $\exp(r(\hat{x}^{(k)}_0, c)/\tau)$.
• Efficienza: Questo approccio permette ai gradienti di ricompensa di propagarsi stabilmente all'indietro attraverso la traiettoria senza memorizzare le attivazioni del modello di riferimento o richiedere dati di addestramento sulle immagini.

4. Framework di Addestramento

Il processo di addestramento alterna due fasi:

1. Aggiornamento dell'Assistente Didattico (TA): Un modello di punteggio TA ($s_\psi$) viene addestrato tramite Denoising Score Matching (DSM) per tracciare le marginali in evoluzione del generatore, surrogando il punteggio implicito del generatore.
2. Aggiornamento del Generatore: Il generatore ($g_\theta$) viene aggiornato utilizzando il gradiente dell'obiettivo IKL, che confronta il punteggio TA con la somma del punteggio di riferimento corretto da CFG e la stima DRP.

Contributi Chiave

1. Identificazione della Dominazione della Ricompensa Terminale: Il documento identifica e analizza formalmente una modalità di fallimento strutturale nei precedenti metodi RLHF in un singolo passaggio, in cui gli ottimizzatori sfruttano i livelli di rumore per collassare su regioni ad alta ricompensa, violando l'equilibrio RLHF previsto.
2. Allineamento a Livello di Traiettoria Principiato: Gli autori propongono Didr, un framework che deriva un obiettivo a livello di traiettoria matematicamente solido propagando in avanti la densità delle immagini pulite ottimizzata per RLHF. Ciò garantisce che ricompensa e regolarizzazione siano bilanciate a ogni livello di rumore.
3. Derivazione Analitica del DRS: Il documento deriva analiticamente il Punteggio della Ricompensa Diffusa, mostrando che decompone esattamente l'obiettivo di traiettoria in un punteggio di riferimento più una correzione indotta dalla ricompensa.
4. Stimatore Pratico (DRP): Gli autori introducono il Proxy della Ricompensa Diffusa, uno stimatore privo di dati che calcola il DRS tramite denoising a pochi passi differenziabili, abilitando un addestramento stabile senza dati di immagini.

Risultati Empirici

Esperimenti estensivi dimostrano che Didr supera costantemente i baseline esistenti in un singolo passaggio:

• Dominanza di Pareto: Sul backbone SDXL (1024×1024), la variante standard Didr domina per Pareto tutti i baseline di allineamento in un singolo passaggio esistenti nel compromesso PickScore–FID, superando simultaneamente i modelli multi-step come FLUX-dev e SDXL-DPO su entrambi gli assi.
• Metriche di Preferenza: La variante con addestramento esteso (Didr_longer) ottiene il PickScore più alto (23.9) tra i modelli in un singolo passaggio valutati e migliora ImageReward del 8,9% rispetto al baseline precedente più forte (Diff-Instruct*). Guida anche in HPSv2.1 e Aesthetic Score.
• Generalizzazione: Didr trasferisce con successo al backbone Z-Image da 6 miliardi di parametri. In un singolo passaggio di generazione, supera il suo modello insegnante a 50 passi sulle metriche di preferenza, dimostrando scalabilità tra architetture.
• Miglioramenti Qualitativi: I confronti visivi mostrano che Didr produce dettagli fini più nitidi e colori più naturali rispetto ai precedenti metodi in un singolo passaggio, evitando la sovrassaturazione e gli artefatti pittorici spesso osservati nei baseline ottimizzati per la ricompensa.

Significato e Affermazioni

Il documento afferma che la propagazione della ricompensa a livello di traiettoria principiato è una strategia di allineamento efficace che risolve il disallineamento intrinseco dei metodi di ricompensa basati solo sul punto finale. Propagando rigorosamente la distribuzione ottimizzata per RLHF attraverso il processo di diffusione, Didr raggiunge il miglior compromesso tra allineamento alle preferenze e fedeltà dell'immagine tra i generatori in un singolo passaggio.

Gli autori sottolineano che il loro metodo non richiede dati di addestramento sulle immagini, basandosi esclusivamente su prompt testuali e un modello di riferimento congelato. Posizionano Didr come una soluzione che permette ai generatori in un singolo passaggio di eguagliare o superare l'allineamento alle preferenze dei modelli insegnanti multi-step mantenendo velocità di inferenza in tempo reale. Il lavoro suggerisce che la discrepanza strutturale identificata nei metodi precedenti è una limitazione fondamentale che può essere risolta solo allineando l'intera traiettoria, non solo lo stato terminale.