Riepilogo Tecnico: Reinforce Adjoint Matching (RAM)

Enunciato del Problema

I modelli di diffusione e di flow-matching raggiungono una generazione ad alta fedeltà perché il loro preaddestramento si basa sulla regressione supervisionata: un campione pulito viene analiticamente rumorizzato e il modello esegue una regressione rispetto a un target in forma chiusa. Tuttavia, allineare questi modelli a funzioni di ricompensa (ad esempio, per correttezza composizionale, rendering del testo o preferenze umane) tramite l'addestramento post-RL (Reinforcement Learning) interrompe questa struttura efficiente.

I metodi esistenti di addestramento post-RL affrontano ostacoli computazionali e strutturali significativi:

1. Metodi del Gradiente della Politica: Trattano il denoising come un Processo Decisionale di Markov, richiedendo rollout stocastici SDE e l'ottimizzazione di limiti inferiori variazionali. Ciò introduce un'alta varianza e richiede molti passaggi di campionamento.
2. Metodi di Controllo Ottimo Stocastico (Adioint): Mirano esattamente all'obiettivo di massimizzazione della ricompensa regolarizzata KL, ma richiedono scansioni adiointe all'indietro lungo le traiettorie SDE memorizzate e gradienti di ricompensa retropropagati attraverso il processo di campionamento. Questo è numericamente instabile e computazionalmente costoso, in particolare per i modelli di flow-matching dove i programmi di rumore divergono vicino alla fine rumorosa.
3. Funzioni di Perdita Surrogato: Gli approcci recenti sostituiscono le verosimiglianze intrattabili con surrogati (ad esempio, ELBO o punti finali contrastivi), sacrificando la struttura di regressione pulita che permette al preaddestramento di scalare.

La sfida fondamentale è derivare un obiettivo di addestramento post-RL che mantenga la semplicità e la scalabilità del framework di regressione del preaddestramento, ottimizzando efficacemente per le ricompense.

Metodologia: Reinforce Adjoint Matching (RAM)

Gli autori propongono Reinforce Adjoint Matching (RAM), un metodo che estende la struttura di regressione del preaddestramento all'addestramento post-RL senza richiedere rollout SDE, scansioni adiointe all'indietro o gradienti di ricompensa.

Fondamento Teorico

Il metodo si basa sulla massimizzazione della ricompensa regolarizzata KL formulata come un problema di controllo ottimo stocastico.

1. Struttura del Processo Ottimale: Sotto la regolarizzazione KL, il processo generativo ottimale inclina la distribuzione degli endpoint puliti ($X_0$) verso campioni ad alta ricompensa ($p^*_0 \propto p_{ref}^0 \exp(r)$) ma lascia invariata la legge condizionale degli stati rumorosi dato un endpoint pulito. Nello specifico, la traiettoria da un $X_0$ fisso a uno stato rumoroso $X_t$ rimane lo stesso kernel Gaussiano analitico utilizzato nel preaddestramento.
2. Condizione di Adjoint Matching: Il controllo ottimo $u_t$ soddisfa una condizione di punto fisso in cui il controllo è uguale al gradiente scalato negativamente della funzione valore (l'adiointe): $u_t = -\sigma_t \nabla_x V_t(x)$.
3. Decomposizione REINFORCE: Per evitare il calcolo di $\nabla r$ (che potrebbe non essere differenziabile), gli autori applicano l'identità log-derivata (REINFORCE) alla funzione valore. Ciò decompone l'adiointe in un termine di ricompensa pesato dal punteggio del ponte all'indietro e un termine di correzione del costo del percorso.

L'Obiettivo RAM

RAM approssima l'adiointe eliminando la correzione del costo del percorso (giustificata da bassa varianza e stabilità su larga scala) e sfruttando l'insight strutturale secondo cui il kernel di rumorizzazione rimane analitico.

1. Campionamento degli Endpoint: Invece di simulare un SDE controllato, RAM estrae un endpoint pulito $X_0$ dalla distribuzione del modello corrente utilizzando qualsiasi campionatore ODE pronto all'uso.
2. Rumorizzazione Analitica: L'$X_0$ campionato viene rumorizzato analiticamente per generare stati di addestramento $X_t = (1-t)X_0 + t\epsilon$, esattamente come nel preaddestramento.
3. Target in Forma Chiusa: Il target di regressione per il campo di velocità $v_\theta$ è derivato sostituendo il termine di ricompensa e il punteggio del ponte Bayesiano (che ha un'espressione in forma chiusa per la rumorizzazione lineare) nella condizione adiointe.

La funzione di perdita risultante è:
$$L_{RAM}(\theta) = \mathbb{E}_t \left[ \left\| v_\theta^t(X_t) - \text{sg}\left( v_{ref}^t(X_t) + r(X_0) \left( (\epsilon - X_0) - v_\theta^t(X_t) \right) \right) \right\|^2 \right]$$
dove $X_0 \sim p_\theta^0$, $\epsilon \sim \mathcal{N}(0, I)$, e $\text{sg}(\cdot)$ denota l'operatore stop-gradient.

Meccanismi Chiave di Efficienza

• Nessun Rollout SDE: Gli stati di addestramento sono generati rumorizzando un singolo endpoint $K$ volte con vettori di rumore indipendenti. Ciò produce $K$ campioni di addestramento condizionalmente indipendenti per ogni passaggio in avanti del modello, fornendo più segnale di gradiente per passo rispetto ai metodi che si basano su singole traiettorie SDE.
• Nessun Gradiente di Ricompensa: Il metodo funziona con ricompense black-box.
• Nessuna Scansione All'Indietro: La natura analitica del kernel di rumorizzazione elimina la necessità di integrazione adiointe all'indietro.

Contributi Chiave

1. Estensione Strutturale: Il documento dimostra che la struttura di regressione del preaddestramento si estende all'addestramento post-RL sotto regolarizzazione KL. Il processo ottimale cambia solo la distribuzione degli endpoint, non la legge condizionale di rumorizzazione.
2. Algoritmo RAM: Una nuova funzione di perdita di consistenza che corregge il target del preaddestramento con informazioni di ricompensa utilizzando un punteggio del ponte Bayesiano in forma chiusa, evitando i colli di bottiglia computazionali dei metodi esistenti.
3. Giustificazione Teorica: Una derivazione che mostra come il punto fisso di RAM corrisponda all'ottimo regolarizzato KL, con esattezza garantita per riferimenti Gaussiani sotto ricompense lineari e accordo del primo ordine in generale.
4. Analisi della Varianza: Un'analisi degli stimatori del costo del percorso che mostra come l'eliminazione della correzione del costo del percorso (come fa RAM) riduca significativamente la varianza rispetto agli stimatori adiointi a orizzonte completo o locali, rendendolo l'unico stimatore praticabile per l'addestramento post-scala immagine.

Risultati Sperimentali

Gli autori hanno valutato RAM su Stable Diffusion 3.5 Medium (SD3.5M) attraverso tre compiti:

1. Generazione di Immagini Composizionali (GenEval): Misurazione della correttezza di oggetti, attributi e relazioni spaziali.
2. Rendering Visivo del Testo (OCR): Misurazione della leggibilità del testo nelle immagini generate.
3. Allineamento alle Preferenze Umane (PickScore): Allineamento con le preferenze umane.

Prestazioni:

• Massimizzazione della Ricompensa: RAM ha ottenuto le ricompense di addestramento più elevate in tutti e tre i compiti, superando le linee di base tra cui Flow-GRPO, DiffusionNFT e AWM.
• Efficienza di Addestramento: RAM ha raggiunto la ricompensa massima di Flow-GRPO in 50 volte meno passaggi per GenEval, 48 volte meno per OCR e 34 volte meno per PickScore.
• Qualità dell'Immagine: A differenza di alcune linee di base che hanno sofferto di "reward hacking" (degrado della qualità estetica per le metriche del compito), RAM ha mantenuto o migliorato le metriche di qualità dell'immagine (Aesthetic, DeQA, ImageReward, HPSv2) massimizzando al contempo le ricompense del compito.
• Compatibilità CFG: RAM rimane compatibile con la Classifier-Free Guidance (CFG) dopo l'addestramento post-RL, mentre linee di base come DiffusionNFT e AWM hanno visto un calo significativo delle ricompense del compito quando è stata applicata la CFG.

Significato e Affermazioni

Il documento afferma che RAM rappresenta un cambiamento di paradigma nell'addestramento post-RL per modelli di diffusione e flow-matching. Sfruttando le proprietà strutturali specifiche del problema di controllo ottimo, RAM ripristina il paradigma di "semplice regressione" che ha reso scalabile il preaddestramento.

Gli autori affermano che:

• RAM elimina la necessità di costosi rollout SDE e scansioni adiointe all'indietro, che sono stati i principali colli di bottiglia per scalare l'RL alla generazione di immagini ad alta risoluzione.
• Il metodo è "semplice e scalabile", ottenendo risultati all'avanguardia con significativamente meno passaggi di addestramento e costi computazionali per passo comparabili ai metodi esistenti.
• L'approccio allinea con successo i modelli a ricompense complesse (componibilità, testo, preferenze) senza degrado visibile della qualità o reward hacking, dimostrando che la struttura del preaddestramento è abbastanza robusta da gestire l'allineamento post-RL quando vengono soddisfatte le condizioni teoriche corrette.