Riepilogo Tecnico: Adjoint Matching Efficiente per il Fine-tuning di Modelli Diffusivi

Enunciato del Problema

Il fine-tuning basato su ricompensa è diventato l'approccio standard per allineare modelli diffusivi e di flusso preaddestrati alle preferenze umane nella generazione da testo a immagine (T2I). Tra questi, Adjoint Matching (AM) offre una formulazione rigorosa, inquadrando il fine-tuning basato su ricompensa come un problema di Controllo Ottimo Stocastico (SOC). AM mira a trasformare un modello preaddestrato in una distribuzione target $p^\star_1(x) \propto e^{\beta r(x)}p_{data}(x)$, dove $r(x)$ è una funzione di ricompensa.

Tuttavia, AM soffre di significative inefficienze computazionali che ne ostacolano l'applicazione pratica:

1. Collo di bottiglia della simulazione in avanti: AM richiede la simulazione di traiettorie generative complete utilizzando un solver di equazioni differenziali stocastiche (SDE) sotto dinamiche "senza memoria". Ciò necessita di un elevato numero di valutazioni della funzione (NFE) e richiede l'archiviazione dell'intera traiettoria in memoria.
2. Collo di bottiglia della simulazione all'indietro: Per calcolare il segnale di controllo, AM richiede una simulazione all'indietro di un'Equazione Differenziale Ordinaria (ODE) dello stato aggiunto lungo la traiettoria campionata. Ciò comporta il calcolo di prodotti Jacobiano-Vettore (JVP) attraverso la rete preaddestrata ad ogni passo di integrazione, operazione computazionalmente costosa.

Gli autori osservano che questi colli di bottiglia sono intrinsecamente legati alla deriva di base non banale ereditata dal modello preaddestrato nelle formulazioni standard di AM.

Metodologia: Adjoint Matching Efficiente (EAM)

Il documento propone Efficient Adjoint Matching (EAM), che riformula il problema SOC per eliminare entrambi i colli di bottiglia computazionali preservando al contempo la garanzia teorica di raggiungere la distribuzione target inclinata dalla ricompensa.

1. Riprogettazione delle Dinamiche di Base

EAM introduce due condizioni strutturali per le dinamiche di base:

• Linearità: La deriva di base $b(x,t)$ è riprogettata per essere lineare in $x$ (cioè $b(x,t) = D(t)x$).
• Assenza di memoria: La dinamica di base rimane senza memoria, permettendo di campionare stati intermedi direttamente dall'estremo finale.

Fissando il coefficiente di diffusione $\sigma(t)$ per corrispondere alla configurazione standard, gli autori derivano una specifica famiglia di derive lineari ammissibili (Eq. 15) che soddisfano queste condizioni. Questa deriva lineare consente una soluzione in forma chiusa per lo stato aggiunto, eliminando la necessità di simulazione ODE all'indietro e di valutazioni JVP.

2. Correzione del Costo Terminale

Cambiare la deriva di base altera la distribuzione terminale del sistema non controllato. Per garantire che la distribuzione controllata ottimale corrisponda ancora alla distribuzione target desiderata inclinata dalla ricompensa $p^\star_1$, EAM deriva un costo terminale corretto $g(x)$. Il gradiente di questo costo, $\nabla g(x)$, include termini per il punteggio preaddestrato, il nuovo punteggio terminale di base e il gradiente della ricompensa (Eq. 16).

3. Algoritmo di Addestramento

Il ciclo di addestramento EAM (Algoritmo 1) procede come segue:

1. Generazione dell'estremo finale: Invece di un rollout completo SDE, l'immagine estremo finale $X_1$ è generata utilizzando un solver deterministico ODE a pochi passi (ad esempio, 10 NFE).
2. Costruzione degli stati intermedi: Gli stati intermedi $X_t$ sono campionati direttamente dal kernel di rumore originale $q_t(\cdot | X_1)$ condizionato all'estremo finale generato, bypassando la necessità di simulazione SDE in avanti.
3. Stato aggiunto in forma chiusa: Lo stato aggiunto è calcolato tramite una scalatura in forma chiusa del gradiente dell'estremo finale $\nabla g(X_1)$, rimuovendo completamente la simulazione all'indietro.
4. Parametrizzazione del controllo: Il controllo $u(x,t)$ è parametrizzato tramite un modello di velocità addestrabile (ad esempio, LoRA) inizializzato dal modello preaddestrato, permettendo al fine-tuning di iniziare dalle dinamiche preaddestrate.

Contributi Chiave

• Proposta dell'algoritmo: L'introduzione di EAM, un algoritmo di fine-tuning basato sul gradiente di ricompensa che raggiunge gli stessi obiettivi teorici di AM ma con un costo computazionale significativamente ridotto.
• Caratterizzazione teorica: Una caratterizzazione delle derive di base lineari che soddisfano i requisiti per la costruzione efficiente di traiettorie (assenza di memoria) e il calcolo in forma chiusa dello stato aggiunto.
• Validazione empirica: Dimostrazione che EAM converge fino a 4 volte più velocemente di AM su benchmark T2I standard, mantenendo o superando le sue prestazioni su più metriche.

Risultati

Gli autori hanno valutato EAM su Stable Diffusion 3.5-Medium (SD3.5-M) utilizzando il dataset Pick-a-Pic, confrontandolo con AM e con la baseline SD3.5-M.

• Efficienza: EAM riduce il costo di addestramento per iterazione fino a 4 volte. Ciò è ottenuto sostituendo le simulazioni SDE a 40 passi con il campionamento dell'estremo finale ODE a 10 passi ed eliminando le simulazioni aggiunte all'indietro.
• Prestazioni: Su benchmark inclusi PickScore, ImageReward, HPSv2.1, CLIPScore e Aesthetics, EAM corrisponde costantemente o supera AM.
  • In impostazioni a singola ricompensa (PickScore), EAM ha ottenuto un PickScore di 23.16 rispetto ai 23.00 di AM.
  • In impostazioni a ricompensa multipla (PickScore + HPSv2.1 + Aesthetics), EAM ha ottenuto un PickScore di 22.93 e un punteggio Aesthetics di 6.56, superando AM (6.21).
• Miglioramenti qualitativi: I confronti visivi indicano che EAM preserva meglio i dettagli fini e le composizioni complesse rispetto ad AM, che talvolta appiana i dettagli. Ciò è attribuito alla coerenza di EAM tra le traiettorie di addestramento (rumore condizionato all'estremo finale) e le traiettorie di inferenza.

Significato e Affermazioni

Il documento afferma che i limiti primari di Adjoint Matching non sono intrinseci alla formulazione SOC stessa, ma sono artefatti della specifica scelta della deriva di base. Riprogettando la deriva di base per essere lineare e senza memoria, e aggiustando di conseguenza il costo terminale, EAM rimuove i due maggiori colli di bottiglia computazionali (simulazione SDE in avanti e simulazione aggiunta all'indietro) senza sacrificare le garanzie teoriche di allineamento.

Gli autori posizionano EAM come una soluzione pratica e scalabile per il fine-tuning basato su ricompensa che mantiene l'efficienza del campione dei metodi basati su gradiente riducendo drasticamente l'overhead computazionale, rendendo più accessibile l'allineamento di alta qualità dei modelli diffusivi. Il lavoro non afferma di risolvere tutte le sfide di ottimizzazione (ad esempio, il replay del buffer o gli schedule adattivi sono indicati come lavoro futuro), ma stabilisce un nuovo punto di riferimento per l'efficienza in questo dominio.