ParetoSlider: Diffusion Models Post-Training for Continuous Reward Control

1. Il Problema

L'allineamento dei modelli generativi (come i modelli di diffusione) con le preferenze umane tramite Reinforcement Learning (RL) è diventato lo standard. Tuttavia, la maggior parte delle attuali metodologie si basa su un singolo reward scalare.
Quando sono coinvolti criteri multipli e potenzialmente conflittuali (ad esempio, nell'editing di immagini: aderire al prompt di modifica vs. preservare l'immagine originale), la pratica corrente è la "early scalarization" (scalarizzazione anticipata). Questo approccio combina i reward in una somma ponderata fissa prima dell'addestramento.
Le limitazioni di questo metodo sono:

• Rigidità: Il modello è vincolato a un singolo punto di compromesso (trade-off) definito durante l'addestramento.
• Mancanza di controllo a inferenza: Non è possibile modificare il bilanciamento tra gli obiettivi senza riaddestrare il modello o mantenere checkpoint separati per ogni combinazione di pesi.
• Inefficienza: Mantenere modelli separati per ogni trade-off scala linearmente con il numero di obiettivi, rendendo il processo costoso in termini di memoria e calcolo.

2. Metodologia: ParetoSlider

Gli autori introducono ParetoSlider, un framework di Reinforcement Learning Multi-Obiettivo (MORL) progettato per addestrare un singolo modello di diffusione che approssima l'intera frontiera di Pareto. Questo permette agli utenti di navigare continuamente tra i compromessi ottimali durante l'inferenza.

Componenti Chiave:

1. Condizionamento sulle Preferenze ($\omega$):

   • Viene introdotto un vettore di preferenze $\omega = [\omega_1, ..., \omega_M]$ che risiede sul semplice di probabilità (i pesi sono non negativi e sommano a 1).
   • Questo vettore viene fornito al modello come segnale di condizionamento durante l'addestramento e l'inferenza.
   • Il modello impara a mappare qualsiasi vettore $\omega$ specifico al punto corrispondente sulla frontiera di Pareto, senza bisogno di riaddestramento.

2. Strategia di "Late Scalarization" (Scalarizzazione Tardiva):

   • Un problema comune nel MORL è che reward con scale numeriche diverse possono "dirottare" il segnale di gradiente, oscurando le preferenze dell'utente.
   • ParetoSlider evita la scalarizzazione immediata. Invece, calcola i vantaggi (advantages) per ogni obiettivo di reward indipendentemente, normalizzandoli all'interno di un gruppo di campioni.
   • Solo alla fase finale di aggregazione della loss, i singoli loss per obiettivo vengono combinati utilizzando i pesi $\omega$. Questo garantisce che l'ottimizzazione rispetti fedelmente il vettore di preferenza fornito, indipendentemente dalla scala grezza dei reward.

3. Architettura e Adattamento:

   • Il metodo si basa su DiffusionNFT, un framework di RL fine-tuning efficiente per modelli flow-matching che ottimizza il processo forward invece di quello inverso.
   • Il vettore $\omega$ viene iniettato nell'architettura del modello (es. SD3.5, FluxKontext, LTX-2) attraverso moduli di condizionamento leggeri (MLP e modulatori AdaLN) addestrati congiuntamente con LoRA.
   • L'iniezione avviene in modo graduale per non disturbare il modello pre-addestrato.

4. Processo di Addestramento:

   • Per ogni prompt, il modello genera un gruppo di campioni condizionati sia dal prompt che da un vettore $\omega$ campionato casualmente (inclusi i vertici e i bordi del semplice per coprire l'intero spazio).
   • Ogni campione viene valutato da $M$ modelli di reward.
   • Vengono calcolati vantaggi normalizzati per ogni reward, mappati in pesi di interpolazione, e le loss vengono aggregate tramite $\omega$ prima dell'aggiornamento del gradiente.

3. Contributi Chiave

• Controllo Continuo a Inferenza: Per la prima volta nei modelli di diffusione, viene dimostrato come un singolo modello possa gestire un'intera gamma continua di compromessi tra reward conflittuali, agendo come un "slider" interattivo.
• Efficienza dei Parametri: Elimina la necessità di mantenere checkpoint multipli, riducendo drasticamente l'overhead di memoria e calcolo rispetto ai metodi basati su interpolazione di pesi o modelli multipli.
• Late Scalarization: Propone una strategia di aggregazione della loss che previene il dominio di reward con scale elevate, garantendo un'apprendimento più stabile e fedele alle preferenze espresse.
• Validazione Multi-Dominio: Il metodo è stato applicato e validato su tre compiti diversi: sintesi testo-immagine (T2I), editing immagine basato su istruzioni (I2I) e generazione video testo-video (T2V).

4. Risultati Sperimentali

Il paper valuta ParetoSlider su tre backbone all'avanguardia: Stable Diffusion 3.5, FluxKontext e LTX-2.

• Confronto con Baseline:

  • Rispetto ai metodi a pesi fissi (che richiedono modelli separati per ogni trade-off), ParetoSlider raggiunge prestazioni uguali o superiori in ogni punto della frontiera, offrendo al contempo la flessibilità del controllo continuo.
  • Supera metodi di controllo a inferenza come la riscrittura dei prompt (Prompt Rewriting) o la guida tramite classifier-free (CFG), che offrono solo punti di controllo grossolani o isolati.
  • Supera approcci basati su selezione Pareto durante l'addestramento (es. Flow-Multi) che non offrono controllo a inferenza.

• Qualità e Transizioni:

  • Le transizioni qualitative mostrano che variando $\omega$ si ottengono cambiamenti fluidi e coerenti (es. da fotorealismo a schizzo, o da preservazione dell'immagine originale a forte adesione al prompt di editing).
  • Le analisi quantitative (Hypervolume indicator) confermano che ParetoSlider domina le altre metodologie, coprendo una frontiera di Pareto più ampia e ben distribuita.

• Ablazioni:

  • Le ablation study confermano che la strategia di late scalarization è superiore alla scalarizzazione anticipata o a metodi come Smooth Tchebycheff, poiché produce transizioni più lisce e riduce la tendenza a collassare verso un singolo obiettivo dominante.
  • L'architettura di condizionamento "Shared" (condivisione del segnale di preferenza attraverso i blocchi) si è dimostrata più efficace di approcci basati su token o ibridi.

5. Significato e Impatto

ParetoSlider rappresenta un cambio di paradigma nell'allineamento multi-obiettivo dei modelli generativi visivi.

• Usabilità: Abilita interfacce utente intuitive dove utenti non esperti possono regolare finemente il comportamento del modello (es. "quanto devo cambiare l'immagine rispetto all'originale?") senza competenze tecniche.
• Scalabilità: Risolve il problema della scalabilità nell'addestramento per obiettivi multipli, permettendo di gestire spazi di reward complessi con un'unica istanza di modello.
• Fondamentale per l'Editing: È particolarmente rilevante per l'editing di immagini e video, dove il bilanciamento tra fedeltà alla fonte e creatività è cruciale e soggettivo.

In sintesi, ParetoSlider trasforma il compromesso tra obiettivi conflittuali da una scelta statica fatta dagli sviluppatori durante l'addestramento a una funzione dinamica e controllabile dall'utente al momento della generazione.