Fine-Tuning Diffusion Models via Intermediate Distribution Shaping

Questo lavoro introduce un quadro teorico e pratico per l'adattamento fine dei modelli di diffusione e flusso attraverso la modellazione delle distribuzioni a livelli di rumore intermedi, proponendo l'algoritmo P-GRAFT per la massimizzazione della ricompensa e la correzione inversa del rumore per migliorare le prestazioni senza ricompense esplicite, con risultati empirici superiori su diverse attività generative.

L'essenziale

Immagina di avere un artista digitale molto talentuoso, addestrato su milioni di quadri, che sa disegnare quasi tutto. Tuttavia, quando gli chiedi di disegnare qualcosa di specifico (ad esempio, "un gatto che fa yoga"), a volte sbaglia: il gatto ha sei zampe, o fa yoga in modo strano.

Il problema è che correggere questo artista è difficile. Se gli dici semplicemente "fai meglio", potrebbe diventare confuso o perdere il suo stile originale.

Questo articolo presenta due nuovi metodi intelligenti per "aggiustare" questi artisti digitali (chiamati modelli di diffusione e modelli di flusso) senza doverli riaddestrare da zero o usare metodi complicati e costosi.

Ecco come funzionano, spiegati con metafore semplici:

1. Il Problema: L'Artista Confuso

I modelli attuali funzionano come un processo di "pulizia" di un quadro sporco. Partono da un foglio pieno di rumore (polvere) e, passo dopo passo, rimuovono la polvere fino a rivelare l'immagine finale.

• Il problema: Quando l'artista impara a fare cose nuove (ad esempio, allinearsi meglio alle richieste), spesso si "allontana troppo" dal suo stile originale o impara male perché il processo di apprendimento è troppo rumoroso e instabile. È come se un allenatore cercasse di correggere un calciatore guardando solo l'ultimo secondo della partita, ignorando tutto il movimento che ha portato al gol.

2. La Soluzione 1: "P-GRAFT" (Il Metodo del "Metà Viaggio")

Il primo metodo si chiama P-GRAFT. Immagina che l'artista stia dipingendo un quadro.

• Il vecchio modo: L'allenatore guarda il quadro completamente finito. Se non piace, dice: "Riprova tutto dall'inizio". Questo è costoso e l'artista si confonde perché non sa dove ha sbagliato esattamente.
• Il nuovo modo (P-GRAFT): L'allenatore guarda il quadro quando è per metà finito.
  • Immagina di fermare il processo di pulizia a metà strada. A quel punto, l'immagine è ancora un po' sfocata, ma si capisce già l'idea.
  • L'allenatore dice: "Ok, a metà strada la forma è buona, ma la fine sarà sbagliata. Correggiamo solo la parte iniziale del processo di pulizia".
  • Perché funziona? È come insegnare a qualcuno a nuotare. È più facile correggere la posizione del corpo quando è ancora in acqua calma (metà strada) piuttosto che quando è già sotto pressione e sta annegando (fine del processo). Questo riduce il "rumore" e rende l'apprendimento più stabile e preciso.

Risultato: Usando questo metodo, l'artista (ad esempio Stable Diffusion) disegna immagini che seguono le istruzioni molto meglio, con meno errori e meno calcoli necessari.

3. La Soluzione 2: "Correzione del Rumore Inverso" (Il Metodo del "Retrocedere")

Il secondo metodo si applica a un tipo di artista diverso (i modelli di flusso) e non richiede nemmeno di dire all'artista cosa è "giusto" o "sbagliato" (non servono premi o punizioni).

• L'idea: Immagina che l'artista abbia un difetto nascosto: quando inizia a dipingere, il "rumore" iniziale (il foglio sporco) non è perfetto. È come se il pennello fosse sporco di un colore sbagliato prima ancora di toccare la tela.
• Il trucco: Invece di insegnare all'artista a dipingere meglio, gli chiediamo di fare il contrario.
  1. Prendiamo un'immagine perfetta che l'artista ha già creato.
  2. Gli chiediamo di "riavvolgere il nastro" e trasformare quell'immagine perfetta di nuovo in rumore.
  3. Osserviamo: "Oh, guarda! Il rumore che esce non è un rumore normale, è un rumore 'strano' o 'distorto'".
  4. Invece di usare il rumore normale per iniziare, insegniamo a un piccolo assistente a generare proprio quel "rumore strano" corretto.
• Il risultato: Quando l'artista principale usa questo "rumore corretto" per iniziare a dipingere, il risultato finale è molto più bello, anche se l'artista non è stato modificato. È come se avessimo pulito il pennello prima di iniziare a dipingere.

Perché è importante?

Questi due metodi sono come due nuovi strumenti nella cassetta degli attrezzi degli sviluppatori di intelligenza artificiale:

1. P-GRAFT ci permette di addestrare gli artisti in modo più intelligente, fermandoci a metà strada per correggere gli errori, ottenendo risultati migliori con meno sforzo.
2. La Correzione del Rumore ci permette di migliorare la qualità delle immagini semplicemente "pulendo" il punto di partenza, senza bisogno di spiegare all'IA cosa è bello o brutto.

In sintesi, invece di spingere l'artista a correre più veloce (che lo fa inciampare), questi metodi gli insegnano a camminare con più equilibrio, partendo da una base più solida.

Sommario tecnico

1. Il Problema

I modelli generativi pre-addestrati, in particolare i modelli di diffusione (Diffusion Models) e i modelli di flusso (Flow Models), richiedono spesso un adattamento specifico per compiti (fine-tuning) per correggere errori di apprendimento o allinearsi a preferenze umane tramite feedback di reward.
Le sfide principali identificate sono:

• Intrattabilità del KL: A differenza dei modelli autoregressivi, la verosimiglianza marginale necessaria per implementare la regolarizzazione KL (cruciale per la stabilità dell'addestramento tramite RL come PPO) è intrattabile nei modelli di diffusione.
• Limiti delle attuali strategie: Ignorare il termine KL porta a instabilità su larga scala, mentre l'uso di vincoli KL su traiettorie intere (trajectory KL) porta a risultati subottimali e a problemi di bias nella funzione valore iniziale.
• Difficoltà di apprendimento: Le fasi finali della generazione (dove il rumore è basso e il segnale è alto) sono spesso più difficili da modellare a causa della complessità della distribuzione dei dati target, mentre le fasi iniziali (rumore puro) sono più semplici.

2. Metodologia Proposta

Gli autori propongono un approccio unificato basato sul Campionamento per Rifiuto (Rejection Sampling) e sulla Modellazione di Distribuzioni Intermedie.

A. GRAFT (Generalized Rejection sAmpling Fine-Tuning)

Il lavoro unifica varianti esistenti come RAFT e RSO sotto un framework chiamato GRAFT.

• Concetto: Dimostrano teoricamente che il campionamento per rifiuto (come Top-K o Best-of-N) esegue implicitamente una massimizzazione del reward regolarizzata da KL, anche senza calcolare esplicitamente la verosimiglianza marginale.
• Meccanismo: Si generano campioni dal modello di riferimento, si assegnano reward, e si selezionano solo i campioni ad alto reward (o si pesano) per creare un dataset di addestramento. Questo processo sposta implicitamente la distribuzione verso quella desiderata con un reward "rimodellato" (reshaped reward).

B. P-GRAFT (Partial-GRAFT)

Questa è la contribuzione principale per i modelli di diffusione. Invece di addestrare il modello su tutto il percorso di denoising (da $T$ a $0$), P-GRAFT addestra il modello solo fino a un passo intermedio di denoising ($t_{NI}$).

• Strategia:
  1. Si generano traiettorie complete dal modello di riferimento.
  2. Si calcolano i reward sui campioni finali ($X_0$).
  3. Si applica il campionamento per rifiuto (GRS) sugli stati intermedi ($X_t$) basandosi sui reward finali.
  4. Si addestra il modello fine-tuned solo per la fase da $T$ a $t_{NI}$.
  5. Per la fase da $t_{NI}$ a $0$, si utilizza il modello di riferimento originale.
• Giustificazione Teorica (Bias-Variance Tradeoff):
  • Variance: I reward calcolati sugli stati intermedi sono più rumorosi (alta varianza) rispetto a quelli finali.
  • Bias: Tuttavia, la funzione di score (score function) da apprendere negli stati intermedi (vicini al rumore gaussiano) è molto più semplice e vicina a quella gaussiana rispetto alla distribuzione complessa dei dati finali.
  • Risultato: Addestrando solo sulla parte "facile" del percorso (dove il bias di apprendimento è basso), si ottiene un miglioramento netto nonostante la varianza dei reward.

C. Inverse Noise Correction (per Flow Models)

Per i modelli di flusso (che usano ODE deterministiche), gli autori introducono un metodo per correggere gli errori del modello pre-addestrato senza reward espliciti.

• Idea: Sfruttando la reversibilità dei modelli di flusso, si può mappare un campione di dati ($X_1$) indietro al rumore iniziale ($X_0$). Se il modello pre-addestrato è imperfetto, la distribuzione del "rumore inverso" ($prev_1$) ottenuta da dati reali non sarà una Gaussiana standard.
• Algoritmo:
  1. Si usa il modello pre-addestrato per invertire i dati reali e ottenere campioni dalla distribuzione del "rumore inverso".
  2. Si addestra un piccolo modello "Noise Corrector" per generare questo rumore inverso partendo da una Gaussiana standard.
  3. Durante l'inferenza, si genera prima il rumore corretto con il Noise Corrector e poi si usa il modello principale per generare l'immagine.
• Vantaggio: Corregge la distribuzione iniziale, migliorando la qualità finale con meno calcoli (FLOPs).

3. Risultati Sperimentali

I metodi sono stati valutati su quattro domini: generazione testo-immagine (T2I), layout, molecole e generazione di immagini unconditional.

• Text-to-Image (Stable Diffusion v2):
  • P-GRAFT supera significativamente i metodi basati su Policy Gradient (come DDPO) e il modello base SDv2.
  • Su benchmark come GenAI-Bench e T2ICompBench++, P-GRAFT (con $t_{NI} = 0.25N$) ottiene il miglior punteggio VQAScore (es. 71.94 vs 66.87 del baseline), con un miglioramento relativo dell'8.81% rispetto al modello base.
  • Dimostra una migliore generalizzazione su prompt non visti.
• Layout e Molecole:
  • Applicato a modelli di diffusione discreto-continuo (IGD), P-GRAFT migliora l'allineamento nei layout e la stabilità delle molecole generate.
  • Per le molecole, l'uso di una strategia di rifiuto con deduplicazione previene il collasso della modalità (mode collapse), mantenendo alta la diversità.
• Inverse Noise Correction:
  • Su CelebA-HQ e LSUN-Church, il metodo migliora il punteggio FID riducendo al contempo i FLOPs per immagine.
  • Un modello corretto con 100 step di inferenza supera un modello pre-addestrato che ne usa 1000, dimostrando un'efficienza computazionale superiore.

4. Contributi Chiave

1. Unificazione Teorica (GRAFT): Dimostrazione che il campionamento per rifiuto equivale a una massimizzazione del reward regolarizzata da KL, risolvendo il problema dell'intrattabilità del KL marginale nei modelli di diffusione.
2. P-GRAFT: Un framework innovativo che sfrutta il trade-off bias-varianza addestrando solo sulle fasi iniziali (più facili da apprendere) della diffusione, ottenendo risultati superiori rispetto all'addestramento completo.
3. Inverse Noise Correction: Un metodo efficiente per correggere i modelli di flusso senza bisogno di reward, sfruttando la reversibilità per apprendere una distribuzione di rumore corretta.
4. Performance SOTA: Risultati empirici che superano gli attuali metodi di Policy Gradient (PPO/DDPO) su benchmark standard, con un'efficienza computazionale superiore.

5. Significato

Questo lavoro offre una soluzione elegante al problema della fine-tuning dei modelli di diffusione, aggirando le limitazioni computazionali e teoriche dei metodi RL tradizionali (PPO).

• Efficienza: P-GRAFT riduce la complessità dell'apprendimento focalizzandosi sulle fasi dove la funzione di score è più semplice.
• Versatilità: Il framework è applicabile sia a modelli di diffusione continui che a modelli di flusso e persino a modelli ibridi discreto-continuo.
• Impatto Pratico: Fornisce un metodo pratico per migliorare la qualità della generazione (allineamento prompt-immagine, stabilità molecolare) con costi computazionali ridotti, rendendo il fine-tuning di modelli su larga scala più accessibile e stabile.