Riepilogo Tecnico: Reinforcement Learning Guidato da Discriminatore per il Flow Matching

Definizione del Problema

I modelli basati su flow e score-matching sono diventati il paradigma dominante nella modellazione generativa grazie ai loro obiettivi di addestramento basati su regressione e privi di simulazione. Tuttavia, questi modelli vengono tipicamente addestrati in una pipeline multi-stadio: prima, un modello base viene adattato a una distribuzione di dati $q$ tramite flow o score matching (FSM); in secondo luogo, viene applicato il reinforcement learning (RL) per orientare il modello verso regioni desiderabili definite da una ricompensa $r(x)$.

Il documento identifica un disallineamento strutturale in questo paradigma. Sebbene le perdite di FSM minimizzino l'errore di regressione $\ell_2$ sui campi di velocità o score sotto i marginali di addestramento, esse spesso non riescono a recuperare le proprietà intrinseche della distribuzione dei dati — come il realismo visivo, la coerenza della struttura degli oggetti o il movimento fisicamente corretto — che sono teoricamente presenti nei dati. Di conseguenza, i professionisti si affidano al RL basato sulle preferenze per recuperare queste proprietà. Questo approccio è problematico per due ragioni:

1. Fallimento Strutturale: È sorprendente che i modelli non riescano a imparare le proprietà del mondo reale direttamente dai dati tramite perdite di matching, ma richiedano il RL per recuperarle.
2. Costo e Confluenza: L'addestramento delle ricompense richiede tipicamente costosi dati di preferenza umana, il che conflaguna il realismo oggettivo dei dati con le inclinazioni soggettive degli annotatori.

Gli autori sostengono che il fallimento degli obiettivi di matching nel catturare queste proprietà derivi da un "mismatch tra addestramento e test" (dove gli errori si accumulano lungo la distribuzione di rollout del modello piuttosto che lungo i marginali di addestramento) e un "mismatch geometrico" (dove piccoli errori nello spazio di velocità possono portare a grandi errori nello spazio della qualità del campione/ricompensa).

Metodologia: Discriminator-Guided RL (DRL)

Il documento propone il Discriminator-Guided RL (DRL), un metodo per correggere i modelli di flow matching senza fare affidamento su dati di preferenza umana. L'intuizione centrale è che la ricompensa ottimale per allineare un modello base $p_{base}$ con una distribuzione target $q$ è il rapporto di densità logaritmica $\log(q(x)/p_{base}(x))$. Poiché questo rapporto è intrattabile, il DRL lo stima utilizzando un discriminatore.

La pipeline DRL consiste in due stadi:

Stadio 1: Stima della Ricompensa nello Spazio di Rappresentazione

Invece di addestrare un discriminatore direttamente nello spazio dei pixel (il che è statisticamente difficile e semanticamente inaffidabile), il DRL addestra un discriminatore in uno spazio di rappresentazione di apprendimento auto-supervisionato (SSL) pre-addestrato.

• Un encoder congelato $\phi$ (ad es. DINOv2) mappa le immagini in uno spazio di feature $Z$.
• Un discriminatore $D$ viene addestrato per distinguere tra campioni di dati reali $x \sim q$ e campioni del modello base $x_{fake} \sim p_{base}$ utilizzando le loro rappresentazioni codificate $\phi(x)$.
• La funzione di ricompensa è definita come il logit di questo discriminatore: $\hat{r}(x) = \log(D(\phi(x)) / (1 - D(\phi(x))))$.
• Giustificazione Teorica: Limitando la ricompensa allo spazio SSL, il DRL risolve un problema di ottimizzazione vincolata: trova la distribuzione $p$ più vicina a $p_{base}$ (in divergenza KL) tale che la pushforward di $p$ attraverso $\phi$ corrisponda alla pushforward di $q$ attraverso $\phi$. Ciò assicura che le correzioni siano effettuate lungo assi semanticamente significativi visibili all'encoder.

Stadio 2: RL Regolarizzato tramite KL via Adjoint Matching

La ricompensa stimata $\hat{r}$ viene utilizzata per perfezionare il modello base tramite RL regolarizzato tramite KL.

• L'obiettivo è massimizzare la ricompensa attesa minimizzando al contempo la divergenza KL tra le distribuzioni delle traiettorie del modello perfezionato e del modello base.
• Gli autori utilizzano l'Adjoint Matching, un algoritmo all'avanguardia per i modelli di flow che ottimizza l'obiettivo a livello di traiettoria tramite la propagazione del gradiente attraverso il processo di campionamento.
• Per garantire la stabilità dell'addestramento con l'SDE senza memoria richiesta per il tilting unbiased dell'endpoint, gli autori introducono un Integratore Lineare Locale. Questo integratore gestisce la rigidità numerica vicino a $t=0$ (dove i coefficienti di drift e diffusione divergono) in modo più efficace rispetto alle discretizzazioni standard Euler-Maruyama con offset.

Contributi Chiave

1. Analisi Teorica dei Limiti del Matching: Il documento fornisce un'analisi formale mostrando che le perdite di flow/score matching non garantiscono il controllo sui gap di ricompensa (proprietà di interesse) a causa degli shift di distribuzione tra addestramento e test e del disallineamento tra la geometria dello spazio di velocità e la geometria dello spazio di ricompensa dei campioni.
2. Discriminator-Guided RL (DRL): Un nuovo framework che utilizza un discriminatore addestrato in uno spazio di rappresentazione SSL congelato per generare un segnale di ricompensa. Ciò consente al modello di recuperare le proprietà dei dati senza etichette di preferenza umana.
3. Integratore Lineare Locale: Un nuovo integratore numerico per l'SDE senza memoria utilizzato nell'addestramento RL, che previene l'esplosione della norma nei tempi iniziali e migliora sia la qualità del campionamento che la stabilità dell'addestramento.
4. Validazione Empirica: Dimostrazione che il DRL migliora significativamente l'allineamento distributivo e la qualità delle immagini attraverso diverse architetture (SiT, JiT, REPA, RAE) e spazi di feature (DINOv2, DINOv3, SigLIP, Inception).

Risultati

Gli autori hanno valutato il DRL su quattro modelli di flow pre-addestrati su ImageNet. Le scoperte principali includono:

• Allineamento Distributivo: Il DRL riduce significativamente la Distanza di Fréchet (FD). Ad esempio, su SiT, la FD senza guidance nello spazio DINOv3 è scesa da 88.2 a 19.3. In tutti i modelli e gli spazi di feature, il DRL ha costantemente migliorato l'allineamento senza utilizzare il CFG (Classifier-Free Guidance).
• Qualità dell'Immagine senza Dati di Preferenza: Il DRL ha migliorato i punteggi sui modelli di ricompensa basati sulle preferenze umane tenuti in hold-out (ImageReward, HPSv2, PickScore, Aesthetics), nonostante non sia mai stato addestrato su di essi. Ciò suggerisce che il DRL recupera con successo le proprietà di "realismo" che gli standard di matching-loss trascurano.
• Base Migliorata per l'Allineamento delle Preferenze: Quando il DRL viene utilizzato come punto di partenza per il successivo RL basato sulle preferenze (PRL), produce una frontiera di Pareto superiore tra ricompensa e fedeltà dell'immagine. Il DRL assorbe le correzioni strutturali recuperabili dai dati, permettendo al PRL di concentrarsi sulle genuine preferenze soggettive piuttosto che sulla correzione di artefatti di basso livello (ad es. sovrasaturazione, eccessiva luminosità).
• Gap di Distillazione: Gli esperimenti hanno mostrato che un modello studente addestrato tramite standard flow matching su campioni di un insegnante DRL non poteva replicare le prestazioni dell'insegnante, anche dopo aver visto 50 milioni di campioni. Ciò conferma che l'obiettivo di RL stesso, e non solo i campioni, è necessario per recuperare queste proprietà.

Significato e Rivendicazioni

Il documento sostiene che l'addestramento post-RL per i modelli di flow debba essere visto non solo come uno strumento per ottimizzare le preferenze esterne, ma come un meccanismo complementare per recuperare la struttura nei dati di addestramento che è catturata imperfettamente dagli standard obiettivi di matching.

Utilizzando un discriminatore in uno spazio di rappresentazione pre-addestrato, il DRL fornisce un segnale di ricompensa che è:

1. Allineato con le Proprietà dei Dati: Mira direttamente alla distribuzione dei dati piuttosto che alle preferenze soggettive umane.
2. Geometricamente Solido: Aggira gli ostacoli della geometria dello spazio di velocità ottimizzando direttamente sul panorama della ricompensa.
3. Efficiente: Richiede significativamente meno calcolo rispetto al pre-addestramento ed evita il costo di raccogliere dati di preferenza.

Gli autori concludono che gli obiettivi di matching e il RL non sono metodi con obiettivi conflittuali, ma piuttosto due strumenti complementari per massimizzare l'utilità dei dati di addestramento. Il lavoro suggerisce che la "ricompensa" per una generazione di alta qualità è spesso "sempre stata nei dati", ma richiede un meccanismo on-policy guidato dalla ricompensa per essere pienamente realizzata.