Training-Free Refinement of Flow Matching with Divergence-based Sampling

1. Il Problema: Conflitti di Velocità nei Modelli Flow Matching

I modelli basati su Flow Matching (FM) apprendono una distribuzione target modellando un campo di velocità marginale. Questo campo è definito come la media delle velocità specifiche per ogni campione (sample-wise velocities) che collegano una distribuzione prior semplice (es. rumore gaussiano) ai dati target.

Tuttavia, il paper identifica un limite fondamentale:

• Conflitti Direzionali: In stati intermedi, le traiettorie di diversi campioni possono incrociarsi. In questi punti di intersezione, le velocità specifiche per ogni campione possono puntare in direzioni opposte o verso modalità diverse (distribuzione multimodale).
• Velocità Media Fuorviante: Poiché il modello FM apprende una singola velocità marginale (la media), in presenza di conflitti, questa velocità media può puntare verso regioni a bassa densità di probabilità (spazi vuoti tra le modalità), invece di indirizzare il campione verso una modalità valida.
• Conseguenze: Questo fenomeno causa la deriva delle traiettorie generative verso regioni indesiderate, producendo immagini sfocate, degradate o con dettagli poco definiti.
• Limitazione delle Soluzioni Esistenti: I lavori precedenti hanno tentato di risolvere questo problema ridisegnando la formulazione di addestramento (es. modellando esplicitamente la distribuzione multimodale delle velocità), ma ciò richiede un addestramento costoso e impedisce l'uso diretto di modelli pre-addestrati potenti e già disponibili.

2. Metodologia: Flow Divergence Sampler (FDS)

Gli autori propongono FDS, un framework training-free (senza ri-addestramento) e plug-and-play che opera esclusivamente durante la fase di inferenza. L'idea centrale è correggere lo stato intermedio della traiettoria prima di ogni passo del solver, invece di modificare il campo di velocità appreso.

A. Rilevamento Teorico: La Divergenza come Surrogato

Il cuore teorico del metodo è la scoperta che l'ambiguità locale (la discrepanza tra la velocità marginale e le velocità specifiche dei campioni) può essere quantificata tramite la divergenza del campo di velocità marginale.

• Teorema 1: Gli autori dimostrano che il residuo ottimale del CFM (Conditional Flow Matching), che misura l'errore locale, è direttamente proporzionale alla divergenza del campo di velocità $\nabla \cdot u_t(x_t)$, più alcuni coefficienti dipendenti dallo schedule temporale.
• Surrogato Training-Free: Poiché la divergenza può essere calcolata direttamente dal modello pre-addestrato $u_\theta$ senza accesso ai dati di training o alle velocità vere, essa funge da indicatore affidabile per identificare regioni "ambigue" o ad alto rischio di degradazione.

B. Algoritmo di Campionamento

FDS modifica il processo di generazione standard introducendo un passo di raffinamento spaziale:

1. Perturbazione: Ad ogni passo temporale $t$, partendo dallo stato corrente $x_t$, vengono generati $M$ candidati vicini aggiungendo piccole perturbazioni casuali (rumore gaussiano scalato).
2. Valutazione: Per ogni candidato, viene calcolata la divergenza del campo di velocità del modello pre-addestrato.
3. Selezione: Viene selezionato il candidato con la divergenza minima (indicante una regione più affidabile e meno ambigua).
4. Raffinamento: Lo stato viene aggiornato a questo nuovo punto $\tilde{x}_t$ prima di eseguire il passo successivo del solver ODE (es. Euler o Heun).

C. Efficienza Computazionale

• Basso Overhead: Il metodo non richiede derivate di secondo ordine (che sarebbero costose). Utilizza un approccio di ordine zero (ricerca locale tramite perturbazione).
• Ottimizzazione Temporale: L'applicazione di FDS è limitata alle fasi iniziali della generazione ($t < 0.5$), dove i conflitti di traiettoria sono più critici, massimizzando l'efficienza.
• Configurazione Pratica: Gli esperimenti mostrano che un singolo passo di raffinamento ($N=1$) con un solo candidato ($M=1$) è spesso sufficiente per ottenere miglioramenti significativi, rendendo il metodo molto leggero.

3. Contributi Chiave

1. Caratterizzazione Teorica: Dimostrazione che la discrepanza tra velocità marginale e velocità specifiche può essere espressa in forma chiusa tramite la divergenza del campo di velocità, fornendo un surrogato pratico e privo di dati.
2. FDS (Flow Divergence Sampler): Introduzione di un metodo di campionamento plug-and-play che raffina gli stati intermedi verso regioni a bassa divergenza senza ri-addestrare il modello sottostante.
3. Versatilità e Performance: Validazione del metodo su diversi backbones (EDM, JiT, SD3), solver (Euler, Heun) e compiti (generazione condizionata, text-to-image, problemi inversi), mostrando miglioramenti costanti rispetto ai baselines.

4. Risultati Sperimentali

Gli esperimenti confermano l'efficacia di FDS in diverse configurazioni:

• Dati Sintetici (2D): FDS riduce la distanza di Wasserstein rispetto al target, guidando i campioni dalle regioni ad alta divergenza verso quelle corrette, evitando la dispersione fuori dalle modalità target.
• Generazione di Immagini (CIFAR-10 e ImageNet 256x256):
  • FDS supera i baselines "compute-matched" (baselines con lo stesso tempo di calcolo ma più passi di integrazione). Ad esempio, su ImageNet con solver Heun, aumentare i passi di integrazione porta a rendimenti decrescenti, mentre FDS migliora significativamente l'FID (Fréchet Inception Distance).
  • I risultati qualitativi mostrano dettagli più nitidi (es. semi su un hamburger, texture del pelo) rispetto ai metodi standard.
• Text-to-Image (SD3-Medium): Miglioramento della fedeltà visiva e dell'allineamento testo-immagine su dataset come DrawBench, con metriche come ImageReward e CLIP Score superiori.
• Problemi Inversi (Deblurring e Super-Risoluzione): FDS integrato con framework di guida (TFG) produce immagini più dettagliate e realistiche, preservando strutture fini (es. pupille, baffi) che i metodi basali tendono a sfocare.
• Confronto con Metodi Training-Based: FDS, pur essendo training-free, supera metodi che richiedono un addestramento specifico per risolvere gli incroci (come HRF), grazie alla capacità dei modelli pre-addestrati di mantenere una forte capacità rappresentativa che FDS sfrutta senza degradarla.

5. Significato e Impatto

Questo lavoro offre una nuova prospettiva sulla generazione basata su flussi:

• Shift Spaziale vs Temporale: Mentre i solutori avanzati tradizionali si concentrano sulla riduzione dell'errore di discretizzazione temporale (passi più piccoli), FDS introduce un aggiustamento spaziale attivo, correggendo la traiettoria stessa.
• Accessibilità: Essendo un metodo training-free e plug-and-play, rende immediatamente utilizzabili i miglioramenti di qualità su modelli pre-addestrati all'avanguardia senza i costi proibitivi del ri-addestramento.
• Robustezza: La capacità di funzionare su diversi domini (immagini, testo, problemi inversi) e con diversi solver lo rende una soluzione generale e pratica per migliorare la fedeltà della generazione in Flow Matching.

In sintesi, FDS risolve il problema fondamentale degli incroci di traiettoria nei modelli Flow Matching utilizzando la divergenza come segnale di guida, offrendo un metodo efficiente, economico e altamente efficace per migliorare la qualità della generazione.