Sintesi Tecnica: CoCA – Reward a livello di step gratuitamente nei modelli di diffusione T2I sottoposti a RL-Finetuning

1. Problema

I modelli di diffusione text-to-image (T2I) allo stato dell'arte si affidano sempre più all'apprendimento per rinforzo (RL) per il fine-tuning al fine di allinearsi alle preferenze umane (ad esempio, appeal estetico, correttezza semantica). Tuttavia, questi metodi affrontano un limite critico: la sparsità del reward.

Nelle formulazioni standard di RL per la diffusione, il processo è modellato come un Processo Decisionale di Markov (MDP) in cui l'agente riceve un singolo reward terminale solo dopo l'ultimo step di denoising ($t=0$). I metodi di policy-gradient esistenti (ad esempio, REINFORCE, DDPO) distribuiscono questo reward scalare uniformemente attraverso tutti i $T$ step di denoising. Ciò crea un disallineamento:

• Impatto disuguale: Gli step iniziali sono decisivi per stabilire la struttura globale dell'immagine, mentre gli step successivi servono principalmente a raffinare i dettagli fini.
• Assegnazione del credito inefficiente: Il weighting uniforme non riesce a identificare quali specifici step abbiano contribuito significativamente alla qualità finale, portando a una efficienza campionaria subottimale e a una precisione limitata nell'attribuire i reward alle azioni generative.
• Limitazioni delle soluzioni esistenti: I tentativi precedenti per affrontare questo problema includono il discounting temporale fisso (che si basa su assunzioni a priori che non riescono a generalizzare su prompt diversi) o l'addestramento di reti ausiliarie (ad esempio, critic step-wise) per predire i reward intermedi, il che introduce un significativo overhead computazionale e richiede ulteriore addestramento.

2. Metodologia: CoCA (Contribution-based Credit Assignment)

Gli autori propongono CoCA, un framework che ridistribuisce il reward terminale lungo la traiettoria di denoising per generare segnali densi a livello di step senza addestrare reti ausiliarie.

Meccanismo Core

CoCA stima il contributo di ogni step di denoising misurando la progressione latente verso la immagine pulita realizzata.

1. Proxy della Progressione Latente: Invece di assumere un'importanza fissa, CoCA traccia la similarità tra lo stato latente intermedio ($x_{T-t}$) e il latente terminale realizzato ($x_0$) nello spazio di embedding.
   • Calcola la similarità del coseno: $Sim_t = \frac{\langle x_{T-t}, x_0 \rangle}{\|x_{T-t}\| \cdot \|x_0\|}$.
   • Il contributo di uno step è derivato dal cambiamento in questa similarità ($\Delta Sim_t$).
2. Smoothing a Finestra (Window Smoothing): Per ridurre le fluttuazioni nella similarità per step, la traiettoria viene segmentata in finestre non sovrapposte di dimensione fissa $W$. Viene calcolata la similarità media all'interno di una finestra, e il punteggio di progressione è la differenza tra finestre consecutive ($\Delta \bar{Sim}_i$).
3. Ridistribuzione del Reward: Il reward terminale $r(x_0, c)$ viene ridistribuito a ciascuno step $t$ sulla base dei punteggi di progressione normalizzati:
   $$\hat{R}(s_t, a_t) = w_t \cdot r(x_0, c)$$
   dove $w_t$ è proporzionale al contributo stimato di quello step.

Stabilizzazione e Garanzie Teoriche

• Normalizzazione del Reward: Per stabilizzare l'addestramento, CoCA applica la normalizzazione due volte:
  1. Normalizzazione per-prompt prima della ridistribuzione per mantenere la coerenza del ranking tra i prompt.
  2. Normalizzazione per-prompt per-timestep dopo la ridistribuzione per garantire scale di reward coerenti tra i timestep e ridurre la varianza.
• Consistenza del Policy-Gradient: Gli autori dimostrano che se i reward per step ridistribuiti sommano al reward terminale originale per ogni traiettoria ($\sum \hat{R}_t = r(x_0, c)$), il gradiente della policy atteso rimane non distorto (unbiased). Pertanto, CoCA preserva i punti stazionari dell'obiettivo originale fornendo al contempo segnali di apprendimento più densi.

3. Contributi Chiave

1. Framework Innovativo: Introduzione di CoCA, un meccanismo di assegnazione del credito basato sul contributo che affronta il disallineamento tra aggiornamenti di policy uguali e importanza variabile degli step nei modelli di diffusione sottoposti a RL-finetuning.
2. Nessun Addestramento Ausiliario: A differenza di altri metodi step-level, CoCA non richiede reti aggiuntive (critic) o annotazioni di reward; opera puramente come un meccanismo di ridistribuzione del reward condizionato sulle traiettorie realizzate.
3. Osservazione Empirica: Il paper dimostra che i pattern di progressione latente a livello di timestep variano sostanzialmente tra diversi prompt e traiettorie, spiegando perché gli schedule di weighting temporale fissi falliscono nel generalizzare e motivando la necessità di un'assegnazione adattiva.
4. Generalità Applicativa: Il framework è agnostico rispetto all'algoritmo, integrandosi con successo sia con pipeline di ottimizzazione basate su REINFORCE (DDPO) che su GRPO (DanceGRPO).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su quattro funzioni di reward per le preferenze umane (Aesthetic, ImageReward, PickScore, HPSv2) e due setting di ottimizzazione (DDPO e DanceGRPO).

• Efficienza Campionaria: CoCA migliora l'efficienza di ottimizzazione del 25% - 100% rispetto ai baseline a livello di traiettoria (DDPO) e ai baseline step-level esistenti (TDPO, UCA). Raggiunge curve di apprendimento più ripide e una convergenza più rapida.
• Performance: Nel setting REINFORCE, CoCA ha ottenuto i punteggi più alti in tutte le metriche, superando sia i baseline a livello di traiettoria che quelli a livello di step. Nel setting GRPO, ha mostrato risultati superiori, particolarmente in scenari di ottimizzazione multi-reward.
• Generalizzazione: CoCA ha dimostrato una migliore generalizzazione su prompt non visti e reward function non viste (ad esempio, migliorando PickScore quando addestrato su HPSv2), suggerendo che i miglioramenti derivino da un reale potenziamento della qualità dell'immagine piuttosto che da reward hacking.
• Studi di Ablazione:
  • Dimensione della Finestra: Una dimensione della finestra moderata ($W=5$) ha prodotto le prestazioni migliori, bilanciando stabilità e granularità.
  • Metrica di Distanza: La similarità del coseno ha superato la distanza $\ell_2$ per misurare la progressione latente.
  • Normalizzazione: Combinare la normalizzazione pre- e post-ridistribuzione ha fornito l'addestramento più stabile.

5. Significato e Rivendicazioni

Il paper sostiene che CoCA offra un miglioramento "gratuito" nel RL-finetuning risolvendo il problema dell'assegnazione del credito inerente ai setting di reward sparso. Il suo significato risiede in:

• Efficienza: Riduzione drastica del numero di query al reward (e quindi di calcolo/dati) richiesti per allineare i modelli generativi alle preferenze umane.
• Semplicità: Fornire una soluzione plug-and-play generale che non aumenta la complessità del modello né richiede addestramento ausiliario.
• Solidità Teorica: Garantire che la ridistribuzione dei reward non distorca l'obiettivo di ottimizzazione, mantenendo l'integrità della policy gradient originale.

Gli autori concludono che, fornendo segnali di apprendimento più densi e allineati all'output, CoCA abbassa la barriera per applicazioni benefiche nella creazione di contenuti e nel design assistito, pur riconoscendo che una maggiore efficienza potrebbe anche facilitare l'uso improprio o amplificare i bias se non gestita con attenzione.