Sintesi Tecnica: Gradual Fine-Tuning per Modelli di Flow Matching

Problematica

Il fine-tuning di modelli di flow matching pre-addestrati presenta una sfida significativa in scenari caratterizzati da dati limitati, distribuzioni in evoluzione o vincoli computazionali. Sebbene i recenti progressi nel fine-tuning basato su ricompense (ad esempio, Adjoint Matching, approcci di apprendimento per rinforzo) abbiano mostrato potenziale, essi soffrono spesso di limitazioni teoriche e instabilità pratiche. Nello specifico, i metodi esistenti falliscono frequentemente nel garantire simultaneamente la correttezza teorica, la stabilità dell'addestramento, l'efficienza dell'inferenza e la preservazione della diversità.

Un contesto critico, seppur poco esplorato, è il fine-tuning quando sono disponibili campioni diretti della distribuzione target (ad esempio, nell'adattamento di dominio o nel drift del dataset), eliminando la necessità di funzioni di ricompensa esterne. Il fine-tuning standard in questo contesto tratta spesso il modello pre-addestrato solo come un'inizializzazione, ignorando le sue dinamiche apprese (come i percorsi di probabilità brevi), il che porta a un'alta instabilità e al potenziale collasso della moda (mode collapse). È necessario un framework principato che bilanci l'adattamento a una nuova distribuzione target con la preservazione delle proprietà benefiche del modello pre-addestrato.

Metodologia: Gradual Fine-Tuning (GFT)

Gli autori propongono il Gradual Fine-Tuning (GFT), un framework che utilizza una regolarizzazione per annealing per creare una sequenza di obiettivi intermedi controllata dalla temperatura. Questo approccio interpola fluidamente tra le dinamiche pre-addestrate e le dinamiche target.

Formulazione Teorica

GFT è derivato da un problema di ottimizzazione sullo spazio dei percorsi indotto da flussi stocastici. L'obiettivo minimizza la divergenza di Kullback-Leibler (KL) tra il processo fine-tuned $P_\theta$ e il processo target $P_q$, regolarizzata dalla divergenza KL tra $P_\theta$ e il processo pre-addestrato $P_{\theta_0}$:

$$\min_{\theta} \text{KL}(P_\theta \| P_q) + \beta \text{KL}(P_\theta \| P_{\theta_0})$$

dove $\beta \in \mathbb{R}^+$ è un parametro di temperatura.

Teorema 4.1 stabilisce che la minimizzazione di questo obiettivo produce un campo di deriva ottimale a forma chiusa $v^*_\theta$ che è una combinazione convessa del campo vettoriale target $v_q$ e del campo vettoriale pre-addestrato $v_{\theta_0}$:

$$v^*_\theta(X_t, t) = \left( \frac{1}{1 + \beta} \right) v_q(X_t, t) + \left( \frac{\beta}{1 + \beta} \right) v_{\theta_0}(X_t, t)$$

Regolarizzazione per Annealing: Per raggiungere la convergenza, GFT impiega uno schema di raffreddamento dipendente dal tempo in cui $\beta$ viene gradualmente ridotto (annealed) da un valore iniziale elevato verso $0^+$.

• $\beta$ elevato: Penalizza fortemente la deviazione dal modello pre-addestrato, prevenendo il collasso prematuro e garantendo la stabilità.
• $\beta$ basso ($\to 0^+$): Rilassa il vincolo, permettendo al modello di recuperare asintoticamente il campo vettoriale target $v_q$.

Riformulazione Trattabile

Poiché il campo vettoriale target $v_q$ è sconosciuto e l'obiettivo è definito sugli spazi di misura dei percorsi, gli autori applicano il Teorema di Girsanov per riformulare l'obiettivo in termini di termini di deriva. Ciò risulta in una funzione di perdita trattabile che coinvolge la distanza $L_2$ al quadrato tra le derive.

Inoltre, per consentire un addestramento e un'inferenza efficienti, GFT utilizza i principi del Conditional Flow Matching (CFM). Condizionando su accoppiamenti arbitrari sorgente-target (ad esempio, piani di Trasporto Ottimale), la perdita diventa:

$$\mathcal{L}_\pi(\theta) = \mathbb{E}_{Z \sim \pi, X_t \sim P(\cdot|Z), t \sim U[0,1]} \left[ \frac{1}{2} \left( \|v_\theta - v_q(\cdot|Z)\|^2 + \beta \|v_\theta - v_{\theta_0}\|^2 \right) \right]$$

Questa formulazione preserva le garanzie teoriche dell'obiettivo incondizionato consentendo al contempo l'uso di accoppiamenti efficienti (come l'OT) per accelerare la generazione.

Estensione all'Inferenza a Pochi Step

Il framework si estende naturalmente ai modelli Mean Flow (MF), che consentono l'inferenza a pochi step. Gli autori dimostrano che il campo di velocità media ottimale per i modelli MF segue anch'esso una forma di combinazione convessa, permettendo a GFT di essere applicato alla generazione a uno o pochi step senza sacrificare la correttezza teorica.

Contributi Chiave

Il documento delinea quattro contributi primari:

1. Convergenza Teorica: Dimostrazione che GFT converge alla distribuzione target sotto ipotesi standard, con un tasso di convergenza esplicito dipendente dalla temperatura $\beta$.
2. Equivalenza di Accoppiamento: Stabilire che GFT produce risultati equivalenti quando condizionato su accoppiamenti arbitrari sorgente-target, consentendo l'uso di tecniche di addestramento (come l'OT) che migliorano l'efficienza dell'inferenza.
3. Estensione a Pochi Step: Dimostrazione della naturale estensione di GFT al fine-tuning per metodi di inferenza a pochi step (ad esempio, Mean Flow) senza compromettere le garanzie teoriche.
4. Superiorità Empirica: Validazione sperimentale completa che mostra come GFT raggiunga una qualità di generazione comparabile o superiore ad altri metodi, fornendo al contempo una migliore stabilità dell'addestramento, diversità di generazione e percorsi di probabilità costantemente più brevi.

Risultati Sperimentali

Gli autori hanno valutato GFT in molteplici setting, inclusi l'adattamento cross-domain e in-domain sui dataset benchmark WILDS (Camelyon17, RxRx1, FMoW) e dati simulati 2D.

• Stabilità e Convergenza: GFT ha dimostrato una stabilità significativamente maggiore rispetto al fine-tuning CFM standard e ai metodi basati su ricompensa (ad esempio, Adjoint Matching). Metriche come la varianza istantanea e la correlazione di Spearman tra epoche e FID hanno indicato una convergenza monotona e stabile.
• Qualità della Generazione: Nell'adattamento cross-domain, GFT ha raggiunto punteggi FID comparabili al fine-tuning standard ma con lunghezze medie dei percorsi inferiori, indicando un'inferenza più efficiente. Nell'adattamento in-domain, GFT ha eguagliato l'accuratezza del full fine-tuning mantenendo al contempo percorsi più brevi.
• Diversità: Su target Gaussiani bimodali 2D, GFT e il CFM standard sono riusciti a coprire entrambi i modi con alta diversità. Al contrario, i metodi basati su ricompensa come Adjoint Matching hanno sofferto di collasso della moda, generando campioni strettamente raggruppati all'interno di un singolo modo.
• Efficienza: GFT ha mantenuto l'efficienza di addestramento simulation-free del CFM. Mentre Adjoint Matching richiedeva costose simulazioni ODE per ogni step (risultando in tempi di addestramento un ordine di grandezza superiori), GFT è rimasto efficiente.
• Inferenza a Pochi Step: Quando applicato ai modelli Mean Flow, GFT ha ottenuto un FID comparabile ai modelli iterativi ma con una riduzione di 50-100 volte delle valutazioni del modello (NFE) durante l'inferenza.
• Controllo della Regolarizzazione: Gli esperimenti variando la temperatura minima $\beta$ hanno rivelato un compromesso controllabile tra il FID finale e la lunghezza media del percorso, permettendo agli utenti di sintonizzare esplicitamente l'equilibrio tra accuratezza della generazione e velocità di inferenza.

Significato e Rivendicazioni

Il paper posiziona GFT come un'alternativa semplice, ma teoricamente fondata e praticamente efficace per l'adattamento scalabile dei modelli di flow matching sotto distribuzione shift.

Gli autori affermano che GFT colma lacune critiche nella letteratura esistente soddisfacendo tutti i desiderata chiave per un metodo di fine-tuning ideale:

• Teoricamente Corretto: Converge provabilmente alla corretta distribuzione target con un tasso controllabile.
• Stabile: Mostra una convergenza altamente monotona, evitando le oscillazioni e l'instabilità comuni nel fine-tuning standard.
• Efficiente: Mantiene l'addestramento simulation-free del CFM e supporta l'inferenza efficiente tramite accoppiamenti OT e metodi a pochi step.
• Diverso: Preserva la diversità del modello pre-addestrato, evitando il collasso della moda osservato negli approcci basati su ricompensa.

Colmando il divario tra il pre-addestramento di alta qualità e il fine-tuning, GFT offre una soluzione robusta per adattare i modelli generativi a domini complessi guidati dai dati dove è richiesto un controllo sfumato sul risultato del fine-tuning. Il lavoro stabilisce una base per l'adattamento scalabile dei modelli generativi basati su flow.