Rethinking Preference Alignment for Diffusion Models with Classifier-Free Guidance

Questo lavoro propone un metodo di allineamento delle preferenze per i modelli di diffusione che, ispirandosi alla guida senza classificatore, migliora la generalizzazione e il controllo durante l'inferenza decouplando l'apprendimento in due moduli addestrati su dati positivi e negativi per generare un vettore di guida contrastiva, ottenendo risultati superiori su Stable Diffusion senza necessità di riaddestrare il modello base.

L'essenziale

Immagina di avere un artista digitale geniale, ma un po' "selvaggio". Questo artista (il Modello di Diffusione) è stato addestrato guardando milioni di immagini su internet. Sa disegnare benissimo, ma a volte non capisce esattamente cosa vuoi: se gli chiedi un "gatto felice", potrebbe disegnare un gatto triste o un gatto che sembra un'automobile.

Per insegnargli a fare esattamente ciò che vogliamo, gli esseri umani hanno creato dei metodi per "allinearlo" ai nostri gusti. Il metodo più famoso finora si chiama DPO (Ottimizzazione Diretta delle Preferenze).

Ecco il problema: il DPO è come un insegnante che urla all'artista: "NO! Questo non è un gatto felice! Ridisegnalo!". Se l'insegnante insiste troppo, l'artista si confonde, smette di essere creativo e inizia a copiare solo le immagini che gli sono state mostrate, perdendo la sua magia originale. In termini tecnici, il modello "si sovraincarna" (overfitting) e smette di funzionare bene su nuove richieste.

La Soluzione: La Guida Senza Classificatore (PGD e cPGD)

Gli autori di questo paper hanno avuto un'idea geniale: invece di urlare all'artista mentre lavora, perché non dargli un foglio di istruzioni da tenere accanto mentre dipinge?

Hanno creato due nuovi metodi, chiamati PGD e cPGD, che funzionano come un sistema di "guida" intelligente durante la creazione dell'immagine.

1. PGD: Il "Filtro di Correzione"

Immagina che l'artista stia dipingendo un quadro.

• Il Modello Base: È l'artista che dipinge liberamente.
• Il Modello "Preferito": È un piccolo assistente che guarda il quadro e dice: "Ehi, qui c'è qualcosa che non piace agli umani, correggilo un po'".

Invece di riaddestrare tutto l'artista (che è costoso e rischioso), usiamo l'assistente solo alla fine, mentre l'immagine viene generata. L'assistente non cambia l'artista, ma aggiunge una "spinta" leggera verso ciò che piace agli umani. È come se l'artista dipingesse, e tu gli dicessi: "Mettici un po' più di rosso qui, e rendi il sorriso più ampio lì".

2. cPGD: Il "Duo di Giudici" (Contrastive)

Qui la cosa diventa ancora più intelligente. Invece di avere un solo assistente che dice cosa fare, ne abbiamo due:

1. Il Giudice Positivo: Guarda le immagini che piacciono e dice: "Fai così!".
2. Il Giudice Negativo: Guarda le immagini che non piacciono e dice: "Non fare così!".

Il sistema prende la differenza tra i due: "Fai come il primo, ma evita assolutamente quello che dice il secondo".
È come se avessi due esperti che ti guidano: uno ti indica la strada giusta, l'altro ti avvisa delle buche da evitare. Il risultato è un'immagine che è molto più precisa, più creativa e meno confusa rispetto ai metodi precedenti.

Perché è meglio? (Le Analogie)

• Evitare l'Overfitting (La memoria di elefante): I vecchi metodi (DPO) facevano memorizzare all'artista le risposte perfette, così quando gli chiedevi qualcosa di nuovo, si bloccava. I nuovi metodi (PGD/cPGD) lasciano all'artista la sua libertà creativa, aggiungendo solo una "bussola" per orientarlo. L'artista non impara a memoria, impara a capire la direzione.
• Plug-and-Play (Il gioco dei mattoncini): La cosa fantastica è che questi "assistenti" (i modelli di guida) sono come moduli separati. Una volta addestrati, puoi attaccarli a qualsiasi artista digitale, anche a uno diverso da quello su cui sono stati creati. Non devi ricominciare da capo ogni volta.
• Il Compromesso Perfetto: Spesso, quando si cerca di migliorare la qualità, si perde la varietà (tutti i gatti diventano uguali). Questi metodi riescono a ottenere il meglio di entrambi i mondi: immagini che piacciono molto agli umani (alta "ricompensa") ma che sono ancora diverse e creative tra loro (alta "diversità").

In sintesi

Gli autori hanno detto: "Smettiamola di addestrare l'artista a forza di urla. Diamogli invece due guide silenziose che lo correggono mentre lavora".

Il risultato? Immagini generate dall'intelligenza artificiale che sono più belle, più fedeli alle richieste e meno "robotiche" rispetto a prima. È come passare da un insegnante severo che ti corregge ogni errore a un mentore esperto che ti sussurra consigli mentre dipingi il tuo capolavoro.

Sommario tecnico

1. Il Problema

L'allineamento dei modelli di diffusione su larga scala (text-to-image) con le preferenze umane sfumate rimane una sfida significativa. Sebbene i modelli pre-addestrati su dataset di Internet siano capaci di generare immagini di alta qualità, spesso deviano dalle preferenze reali degli umani in termini di estetica e rispetto delle istruzioni.

L'approccio attuale dominante per l'allineamento è l'Ottimizzazione Diretta delle Preferenze (DPO), adattata ai modelli di diffusione (Diffusion-DPO). Tuttavia, DPO presenta limiti critici:

• Sovradattamento (Overfitting): I modelli tendono a memorizzare i dati di addestramento invece di generalizzare.
• Collasso delle modalità: La distribuzione generata può perdere diversità.
• Mancanza di robustezza: Le prestazioni peggiorano su prompt fuori distribuzione (OOD).
• Dipendenza dai dataset di pre-addestramento: Per regolarizzare DPO e prevenire l'oblio catastrofico, spesso servirebbero i dataset originali di pre-addestramento, che sono spesso inaccessibili per modelli su larga scala.

2. Metodologia Proposta

Gli autori propongono di reinterpretare il problema dell'allineamento delle preferenze non come un problema di ottimizzazione dei parametri del modello, ma come un problema di guida al momento dell'inferenza (test-time guidance), ispirandosi al Classifier-Free Guidance (CFG).

L'idea centrale è trattare un modello finetunato sulle preferenze come un "segnale di controllo esterno" che viene combinato con il modello base (prior) durante il campionamento, senza necessariamente ri-addestrare pesantemente il modello base.

La metodologia si articola in due varianti principali:

A. Preference-Guided Diffusion (PGD)

• Concetto: Il modello finetunato con DPO ($\pi_{DPO}$) viene trattato come un modello condizionato su un segnale di preferenza virtuale, mentre il modello base pre-addestrato ($\pi_{ref}$) funge da modello incondizionato (prior).
• Meccanismo: Durante l'inferenza, si applica una formula simile al CFG:
  $$\nabla \log \pi_{PGD}(x) = \nabla \log \pi_{ref}(x) + w \cdot (\nabla \log \pi_{DPO}(x) - \nabla \log \pi_{ref}(x))$$
  dove $w$ è un peso di guida.
• Vantaggio: Permette di utilizzare un modello finetunato con poche iterazioni (evitando l'overfitting) e di controllare l'intensità dell'allineamento tramite $w$ al momento dell'uso.

B. Contrastive Preference-Guided Diffusion (cPGD)

Per migliorare ulteriormente la generalizzazione, gli autori decostruiscono l'apprendimento delle preferenze in due moduli separati, ispirandosi all'addestramento dei modelli di diffusione condizionati:

1. Addestramento Separato: Si addestrano due modelli indipendenti:
   • $\theta_+$: Addestrato solo sui campioni positivi (preferiti).
   • $\theta_-$: Addestrato solo sui campioni negativi (non preferiti).
2. Guida Contrastiva: All'inferenza, si costruisce un vettore di guida contrastivo sottraendo le previsioni del modello negativo da quelle del modello positivo:
   $$\nabla \log \pi_{cPGD}(x, t) = \nabla \log \pi_{ref}(x, t) + w \cdot (\nabla \log \pi(x, t; \theta_+) - \nabla \log \pi(x, t; \theta_-))$$

• Interpretazione Teorica: Gli autori mostrano che cPGD equivale a una ridistribuzione dinamica (reweighting) dei gradienti della funzione di perdita DPO. Invece di ottimizzare direttamente il rapporto di probabilità tra coppie positive/negative (che può essere instabile), si ottimizzano separatamente le distribuzioni e si combinano linearmente. Questo riduce il rischio di collasso delle modalità e migliora la stabilità.

3. Contributi Chiave

1. Nuova Prospettiva Teorica: Riformulano l'allineamento delle preferenze per i modelli di diffusione come un caso speciale di inferenza stile CFG, spostando l'onere dell'allineamento dall'addestramento all'inferenza.
2. Metodi PGD e cPGD: Introducono due metodi semplici che non richiedono il ri-addestramento completo del modello base, ma utilizzano moduli finetunati come segnali di guida.
3. Miglioramento della Generalizzazione: Dimostrano che separare l'apprendimento in moduli positivi e negativi (cPGD) mitiga i problemi di sovradattamento e collasso tipici di DPO.
4. Moduli Plug-and-Play: I moduli finetunati sono riutilizzabili e trasferibili tra diversi modelli base, offrendo una soluzione flessibile.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su Stable Diffusion 1.5 (SD1.5) e Stable Diffusion XL (SDXL) utilizzando i dataset Pick-a-Pic v2 e HPDv3.

• Performance Quantitative: PGD e cPGD superano costantemente le baseline (DPO, NPO, MaPO, KTO, ecc.) in termini di:
  • Win Rate: Percentuale di volte in cui l'immagine generata è preferita rispetto al modello base (miglioramenti significativi su tutti i benchmark).
  • Reward: Punteggi più alti su metriche come PickScore, HPSv2, HPSv3 e ImageReward.
  • Pareto Improvement: I metodi ottengono contemporaneamente reward più alti, FID più bassi (migliore preservazione del prior) e maggiore diversità dei campioni.
• Robustezza: I metodi mostrano una maggiore robustezza rispetto a DPO quando si utilizzano dataset di preferenza di qualità variabile o sottogruppi ad alta qualità.
• Efficienza: È possibile distillare la guida PGD/cPGD in un singolo checkpoint, riducendo il costo computazionale dell'inferenza multi-modello senza perdere significativamente in prestazioni.
• Trasferibilità: Il metodo è stato testato con successo su architetture diverse (es. KOALA) purché condividano lo stesso spazio latente (VAE), confermando la natura "plug-and-play".

5. Significato e Impatto

Questo lavoro offre una soluzione elegante al problema dell'allineamento delle preferenze nei modelli di diffusione, aggirando le limitazioni computazionali e di stabilità di DPO.

• Paradigma Shift: Sposta il focus dal "come addestrare" il modello per allinearlo al "come guidarlo" durante la generazione.
• Accessibilità: Poiché non richiede l'accesso ai dataset di pre-addestramento originali per la regolarizzazione, è applicabile a modelli proprietari o su larga scala.
• Flessibilità: La capacità di controllare l'intensità dell'allineamento tramite un parametro di guida ($w$) offre agli utenti un controllo fine sul compromesso tra fedeltà alle istruzioni, estetica e diversità.

In sintesi, il paper dimostra che trattare le preferenze umane come un segnale di guida condizionale (simile a come si usano i prompt di testo) è una strategia più robusta, scalabile ed efficace rispetto alle attuali tecniche di ottimizzazione diretta.