NormGuard: Reward-Preserving Norm Constraints in Flow-Matching Reinforcement Learning

Sintesi Tecnica: NormGuard: Vincoli di Norma che Preservano la Ricompensa nel Reinforcement Learning per il Flow-Matching

Problematica
Il post-training tramite reinforcement learning (RL) è diventato un approccio standard per allineare i modelli generativi basati su flow alle preferenze umane. Tuttavia, questo processo porta frequentemente alla "sovra-ottimizzazione della ricompensa" (reward over-optimization), dove i modelli sfruttano le imperfezioni del proxy di ricompensa, causando un degrado della qualità percettiva (ad esempio, eccessiva nitidezza, spostamenti di colore, illuminazione innaturale) che il modello di ricompensa non riesce a penalizzare. Le strategie di mitigazione esistenti, come l'early stopping e la regolarizzazione KL, trattano la deriva del post-training come una singola quantità aggregata. Esse vincolano la deviazione totale della velocità del modello rispetto a un riferimento, ma non riescono a distinguere tra le componenti della deriva che trasportano segnali di ricompensa e quelle che causano artefatti visivi.

Osservazioni Chiave e Diagnostica
Gli autori identificano una specifica firma strutturale di questa deriva: il fine-tuning tramite RL gonfia costantemente la norma della velocità per passo ($\|v_\theta\|$) del 5% - 15% rispetto al riferimento pre-addestrato ($v_{ref}$) attraverso diversi metodi (NFT, AWM, DPO) e modelli (SD3.5, FLUX).

Il paper investiga due potenziali soluzioni prima di proporne una nuova:

1. Rinormalizzazione al Momento dell'Inferenza: Ispirandosi al Classifier-Free Guidance (CFG), dove la riscalatura della velocità durante l'inferenza mitiga gli artefatti, gli autori hanno testato la riscalatura della velocità del modello RL-fine-tuned per riportarla alla norma del riferimento durante l'inferenza. Hanno scoperto che questo approccio fallisce per i modelli addestrati con RL. A differenza del CFG, dove l'inflazione è una combinazione esplicita di teste di velocità all'inferenza, l'inflazione indotta da RL è co-adattata nei pesi del modello. Riscalare all'inferenza preserva la ricompensa ma non riesce a correggere gli artefatti visivi (eccessiva nitidezza, illuminazione innaturale).
2. Analisi della Sensibilità Adjoint: Gli autori hanno analizzato se la soppressione della magnitudo della velocità durante l'addestramento avrebbe rimosso un segnale di ricompensa coerente. Utilizzando un'analisi di sensibilità adjoint su 6.400 campioni, hanno scoperto che il cambiamento di ricompensa del primo ordine sotto la scalatura moltiplicativa della velocità presenta rapporti segnale-rumore molto elevati (da 3× a 100×) ed eterogeneità di segno tra i diversi prompt. La media a livello di batch rimane vicina allo zero, suggerendo che sopprimere l'inflazione della norma non è probabile che rimuova sistematicamente i guadagni di ricompensa.

Metodologia: NormGuard
Sulla base del riscontro che le correzioni al momento dell'inferenza sono inefficaci e che la soppressione della norma non comporta un costo di ricompensa consistente, gli autori propongono NormGuard, un regolarizzatore in fase di addestramento.

• Meccanismo: NormGuard è una penalità hinge unidirezionale che si attiva solo quando la norma della velocità corrente supera la norma del modello di riferimento. Viene aggiunta additivamente a qualsiasi perdita base locale sulla velocità (come quelle utilizzate in NFT, AWM e DPO).
• Formulazione: Per uno stato campionato $(x_t, t, c)$, il regolarizzatore è definito come:
  $$R_{norm}(\theta) = \lambda \cdot \max\left(0, \frac{\|v_\theta(x_t, t, c)\|^2 - \|v_{ref}(x_t, t, c)\|^2}{\|v_{ref}(x_t, t, c)\|^2}\right)$$
  dove $\lambda > 0$ controlla l'intensità. La penalità lascia gli aggiornamenti non vincolati finché $\|v_\theta\| \leq \|v_{ref}\|$, penalizzando solo l'eccesso di norma.
• Compatibilità: Poiché la penalità agisce nello stesso spazio di output della velocità della perdita base, essa si compone naturalmente con gli obiettivi di post-training esistenti senza richiedere vincoli separati nello spazio dei parametri.

Risultati Sperimentali
Gli autori hanno valutato NormGuard su due modelli base (SD3.5-Medium, FLUX.2-4B), tre metodi di post-training (NFT, AWM, DPO) e due modelli di ricompensa (PickScore, HPSv2).

• Qualità Percettiva: Utilizzando giudici MLLM (Qwen3.5 e GPT-4.1) per valutare la qualità dell'immagine su assi come la fedeltà della texture e il rilevamento degli artefatti, NormGuard ha superato costantemente il baseline non regolarizzato in tutte le sette configurazioni. I tassi di vittoria (win rates) sono stati significativamente più alti (ad esempio, 64% contro 35% per SD3.5-M + NFT).
• Realismo Forense: NormGuard ha migliorato i punteggi di realismo forense (RealScore) in sei casi su sette, riducendo artefatti come l'illuminazione innaturale e l'eccessiva nitidezza.
• Preservazione della Ricompensa: Il metodo ha ampiamente preservato i guadagni di ricompensa del baseline. Le differenze in PickScore e HPSv2 sono state minime (variando da -0.004 a +0.011), indicando che il controllo dell'inflazione radiale non ostacola i cambiamenti direzionali responsabili dell'allineamento della ricompensa.
• Inferenza a Pochi Step: I benefici di NormGuard si sono amplificati al diminuire del numero di step di inferenza. Il divario nei win rate degli MLLM è passato dal 9% a 28 step al 20% a 4 step, suggerendo che il regolarizzatore stabilizza il processo di integrazione dell'ODE contro gli errori di discretizzazione causati dalle velocità gonfie.
• Ablazioni: I miglioramenti non sono stati spiegati dall'early stopping (il modello regolarizzato ha superato i checkpoint precedenti del baseline) ed erano complementari alla regolarizzazione KL, affrontando modalità di fallimento distinte.

Significato e Rivendicazioni
Il paper sostiene che la deriva aggregata penalizzata dalla standard regolarizzazione KL può essere fruttuosamente decomposta: l'inflazione della norma guida gli artefatti percettivi pur portando poco segnale di ricompensa, mentre il riallineamento direzionale costituisce la maggior parte del guadagno di ricamente.

NormGuard offre una diagnosi e una soluzione mirata per la sovra-ottimizzazione della ricompensa nei modelli di flow-matching. Sopprimendo la crescita eccessiva della norma della velocità durante l'addestramento, migliora il fotorealismo e riduce gli artefatti senza sacrificare l'allineamento della ricompensa. Gli autori presentano questo come un intervento semplice ed efficace che affronta una specifica firma strutturale della deriva RL, distinta dai vincoli globali della regolarizzazione KL o dalle inefficaci correzioni all'inferenza del CFG.

Limitazioni
Gli autori osservano che l'origine dinamica dell'inflazione radiale rimane una questione aperta. La loro affermazione di neutralità della ricompensa si basa su segnali del primo ordine a livello di batch; le perturbazioni radiali per singolo campione influenzano la ricompensa, ma questi effetti si annullano mediamente tra i vari prompt. Inoltre, l'attuale analisi si applica a obiettivi locali sulla velocità (NFT, AWM, DPO), e l'estensione di questa prospettiva a obiettivi a livello di traiettoria come Flow-GRPO è lasciata ai lavori futuri.