Stepwise Credit Assignment for GRPO on Flow-Matching Models

1. Il Problema

L'apprendimento per rinforzo (RL), in particolare l'algoritmo GRPO (Group Relative Policy Optimization), ha dimostrato grande successo nel migliorare il ragionamento nei modelli linguistici (LLM). Tuttavia, la sua applicazione ai modelli di Flow-Matching per la generazione di immagini testo-su-immagine (come SD3.5) presenta sfide specifiche legate all'assegnazione uniforme del credito.

I metodi esistenti, come Flow-GRPO, assegnano lo stesso "vantaggio" (calcolato sulla qualità dell'immagine finale) a tutti i passaggi del processo di denoising. Questo approccio ignora due aspetti fondamentali della generazione diffusa:

1. Struttura Temporale Gerarchica: I passaggi iniziali determinano la composizione e la struttura a bassa frequenza (layout, oggetti), mentre i passaggi finali affinano i dettagli e le texture ad alta frequenza. Trattare tutti i passaggi allo stesso modo confonde queste fasi distinte.
2. Ricompensa di Errori Corretti: Un percorso di generazione potrebbe commettere errori gravi nei passaggi iniziali (es. colore sbagliato di un oggetto) che vengono successivamente corretti dai passaggi finali. Se l'immagine finale è di alta qualità, Flow-GRPO premia anche i passaggi iniziali errati, rafforzando comportamenti subottimali che potrebbero causare artefatti in altri contesti.

2. Metodologia: Stepwise-Flow-GRPO

Gli autori propongono Stepwise-Flow-GRPO, un approccio che assegna il credito a ogni singolo passaggio di denoising basandosi sul miglioramento marginale della ricompensa, piuttosto che sulla ricompensa finale totale.

Componenti Chiave:

• Stima delle Ricompense Intermedie: Poiché i modelli di ricompensa non possono essere valutati direttamente su stati rumorosi ($x_t$), il metodo utilizza la formula di Tweedie per stimare l'immagine pulita sottostante $\hat{x}_0(t) = E[x_0|x_t]$ a partire dallo stato rumoroso corrente. Per ottenere stime di alta qualità, viene utilizzata una breve traiettoria ODE (con $T'=5$ sottopassi) per denoizzare da $x_t$ a $\hat{x}_0(t)$.
• Vantaggi Basati sui Guadagni (Gain-based Advantages): Invece di usare la ricompensa finale, il metodo calcola il guadagno ($g_t$) per ogni passaggio:
  $$g_t = r_{t-1} - r_t$$
  dove $r_t$ è la ricompensa stimata dello stato intermedio. Questo misura quanto un passaggio specifico ha migliorato (o peggiorato) la qualità dell'immagine rispetto al passaggio precedente.
  • I passaggi che aumentano la ricompensa ricevono un rinforzo positivo.
  • I passaggi che riducono la ricompensa vengono penalizzati.
  • La somma telescopica dei guadagni equivale alla differenza tra la ricompensa finale e quella iniziale, garantendo l'allineamento con l'obiettivo globale.
• Normalizzazione Congiunta: Per preservare la struttura temporale (dove i guadagni iniziali sono tipicamente più grandi), i guadagni vengono normalizzati congiuntamente su tutti i passaggi e le traiettorie, invece di normalizzarli separatamente per ogni passo.
• SDE Ispirato a DDIM: Il paper introduce una nuova equazione differenziale stocastica (SDE) ispirata a DDIM. L'SDE originale di Flow-GRPO introduce rumore che degrada la qualità delle immagini intermedie, peggiorando il segnale di ricompensa. La nuova SDE bilancia l'esplorazione stocastica necessaria per il RL con la produzione di campioni più puliti, allineandosi meglio alla distribuzione di addestramento dei modelli di ricompensa.

3. Contributi Principali

1. Stepwise-Flow-GRPO: Introduzione di un meccanismo di assegnazione del credito basato sui guadagni relativi per ogni passaggio di diffusione, superando i limiti dell'assegnazione uniforme.
2. Ottimizzazione SDE: Sviluppo di una SDE di tipo DDIM che genera campioni meno rumorosi durante l'addestramento, migliorando la qualità del segnale di ricompensa senza sacrificare la stocasticità necessaria per l'esplorazione.
3. Efficienza Superiore: Dimostrazione empirica che l'approccio proposto raggiunge una convergenza più rapida e una maggiore efficienza dei campioni rispetto allo standard Flow-GRPO.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su dataset come GenEval (valutazione della comprensione composizionale) e PickScore, utilizzando modelli di ricompensa come PickScore, ImageReward e UnifiedReward.

• Efficienza dei Campioni: Stepwise-Flow-GRPO supera costantemente Flow-GRPO in termini di ricompensa ottenuta per iterazione di addestramento, convergendo più velocemente in tutte le configurazioni testate.
• Qualità Finale: Il metodo raggiunge punteggi GenEval superiori o equivalenti a Flow-GRPO, con miglioramenti significativi in compiti complessi come il conteggio degli oggetti, il posizionamento spaziale e il binding degli attributi (es. "un autobus sopra una barca" invece che "nel cielo").
• Efficienza Temporale: Nonostante il costo computazionale aggiuntivo per la stima delle immagini intermedie, la maggiore efficienza dei campioni porta a un tempo totale di addestramento (wall-clock time) inferiore per raggiungere obiettivi di ricompensa specifici.
• Stabilità: Il metodo dimostra una maggiore stabilità durante l'addestramento con modelli di ricompensa complessi (come UnifiedReward-7B), dove Flow-GRPO tende a divergere.
• Confronto con SOTA: Dopo 400 ore di addestramento, il modello proposto raggiunge un punteggio GenEval di 0.87, superando Flow-GRPO (0.72) e avvicinandosi alle prestazioni di modelli autoregressivi all'avanguardia come GPT-4o (0.84).

5. Significato e Impatto

Questo lavoro rappresenta un passo avanti significativo nell'applicazione del Reinforcement Learning alla generazione di immagini basata su Flow-Matching.

• Superamento del "Credit Assignment Problem": Risolve il problema fondamentale di come attribuire il merito o la colpa in processi generativi sequenziali, riconoscendo che le decisioni iniziali hanno un impatto sproporzionato sulla qualità finale.
• Efficienza Computazionale: Dimostra che un'allocazione più intelligente del credito può compensare il costo computazionale aggiuntivo, rendendo l'addestramento RL più pratico per modelli di grandi dimensioni.
• Generalizzabilità: L'approccio basato sui guadagni telescopici offre un nuovo paradigma che potrebbe essere applicato ad altri modelli generativi sequenziali, spostando l'attenzione dalla semplice ottimizzazione del risultato finale all'ottimizzazione del processo di generazione stesso.

In sintesi, Stepwise-Flow-GRPO trasforma l'addestramento RL da un processo "cieco" basato sul risultato finale a un processo guidato da feedback granulare e temporale, portando a immagini più coerenti, realistiche e allineate al prompt.