Sintesi Tecnica: Migliorare il Rectified Flow con le Condizioni al Contorno

Definizione del Problema

Il Rectified Flow (RF), noto anche come Flow Matching, è emerso come un potente framework di modellazione generativa che apprende un campo di velocità per trasformare una distribuzione di rumore in una distribuzione di dati target tramite un'equazione differenziale ordinaria (ODE). Nonostante il suo successo in compiti come la generazione di testo-immagine e video, gli autori identificano un limite critico nelle implementazioni standard di RF: il campo di velocità appreso spesso non soddisfa le condizioni al contorno derivate teoricamente.

Nello specifico, sotto l'interpolazione lineare tra rumore ($t=0$) e dati ($t=1$), il campo di velocità ottimale dovrebbe soddisfare:

1. Contorno Sinistro ($t=0$): $\mathbf{v}(\mathbf{x}, 0) = \mathbb{E}[\mathbf{X}_1] - \mathbf{x}$
2. Contorno Destro ($t=1$): $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$

Le parametrizzazioni di reti neurali standard del campo di velocità non impongono intrinsecamente questi vincoli. Come visualizzato nel paper, la velocità predetta a $t=1$ devia spesso dalla mappa identità attesa. Questa violazione è particolarmente dannosa durante il campionamento stocastico (ad esempio, utilizzando la dinamica di Langevin). Poiché la stima della funzione di score (tramite la formula di Tweedie) comporta una divisione per $(1-t)$, una velocità imprecisa vicino a $t=1$ causa la divergenza della stima dello score. Ciò porta a un campionamento instabile, producendo immagini eccessivamente levigate o artefatti "cartonesti", specialmente quando si utilizzano step size o scale di rumore più grandi.

Metodologia

Per affrontare questo problema, gli autori propongono il Boundary-enforced Rectified Flow Model (Boundary RF Model). Questo approccio incorpora esplicitamente i vincoli di contorno nel campo di velocità attraverso due varianti distinte:

1. Modello Boundary RF basato su Maschera (Mask-based)

Questa variante impone entrambe le condizioni al contorno ($t=0$ e $t=1$) strutturando il campo di velocità come una combinazione lineare di termini di contorno e l'output di una rete neurale apprendibile:
$$\mathbf{v}(\mathbf{x}, t) = g(t) \cdot (C - \mathbf{x}) + f(t) \cdot \mathbf{x} + h(t) \cdot m_\theta(\mathbf{x}, t)$$
Dove:

• $m_\theta(\mathbf{x}, t)$ è una rete neurale standard (ad esempio, U-Net o DiT).
• $C = \mathbb{E}[\mathbf{X}_1]$.
• $f(t), g(t), h(t)$ sono funzioni scalari progettate per annullarsi o attivarsi in specifici contorni (ad esempio, $g(1)=f(0)=h(0)=h(1)=0$ e $g(0)=f(1)=1$).
• Gli autori suggeriscono una scelta "Standard Cosine": $g(t) = \cos(\frac{\pi}{2}t)$, $f(t) = \sin(\frac{\pi}{2}t)$, e $h(t) = \sin(\pi t)$.

Questa costruzione garantisce che le condizioni al contorno siano soddisfatte per progettazione, indipendentemente dall'output della rete neurale.

2. Modello Boundary RF basato su Sottrazione (Subtraction-based)

Motivati dall'osservazione che l'imposizione del solo vincolo a $t=1$ produce significativi benefici empirici, gli autori propongono una variante più semplice a singolo contorno:
$$\mathbf{v}(\mathbf{x}, t) = \mathbf{x} + m_\theta(\mathbf{x}, t) - m_\theta(\mathbf{x}, 1)$$
Questa formulazione soddisfa intrinsecamente $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$ perché il termine $m_\theta(\mathbf{x}, 1)$ si cancella a $t=1$. Sebbene questo approccio offra semplicità di design e flessibilità, comporta un costo computazionale, richiedendo un passaggio forward aggiuntivo per calcolare $m_\theta(\mathbf{x}, 1)$ per ogni valutazione.

Impatto sul Campionamento Stocastico

Imponendo $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$, il Boundary RF Model assicura che il numeratore nella formula di Tweedie ($t\mathbf{v}(\mathbf{x}, t) - \mathbf{x}$) tenda a zero quando $t \to 1$. Ciò impedisce alla funzione di score $\nabla \log \rho_t$ di divergere, stabilizzando così i campionatori stocastici (campionatori SDE) e permettendo una generazione robusta anche con scale di rumore o step size più grandi.

Contributi Chiave

1. Identificazione di un Limite: Il paper evidenzia come i campi di velocità non vincolati nel Rectified Flow violino le condizioni al contorno teoriche, portando ad traiettorie ODE imprecise e campionamento stocastico instabile.
2. Nuova Parametrizzazione: L'introduzione del Boundary RF Model, che offre due varianti pratiche (basata su maschera e basata su sottrazione) per imporre queste condizioni con modifiche minime al codice.
3. Stabilizzazione delle Funzioni di Score: Dimostrare che l'imposizione del contorno stabilizza direttamente la funzione di score vicino a $t=1$, mitigando i problemi di divergenza che affliggono il campionamento stocastico nel vanilla RF.
4. Validazione Empirica: Esperimenti completi che mostrano miglioramenti consistenti nella qualità della generazione e nella stabilità del campionamento attraverso diversi dataset e scale di modello.

Risultati Sperimentali

Gli autori hanno valutato i loro modelli su CIFAR-10 e ImageNet (256$\times$256 e 512$\times$512) utilizzando architetture Diffusion Transformer (DiT).

• Miglioramenti Quantitativi:

  • Su ImageNet 256$\times$256, il Modello Boundary RF basato su Sottrazione ha ottenuto un FID di 6.32 (campionamento ODE) e 6.38 (campionamento SDE), rispetto al baseline vanilla RF di 6.87 e 6.79, rispettivamente.
  • Il modello basato su Maschera ha ottenuto un FID di 6.63 (ODE) e 6.18 (SDE).
  • Miglioramenti sono stati osservati anche nell'Inception Score (IS), Precision e Recall su entrambi i dataset.
  • La variante basata su Sottrazione ha costantemente ottenuto i punteggi IS più elevati.

• Robustezza del Campionamento:

  • Negli esperimenti utilizzando l'Overshooting Sampler con diverse scale di rumore ($c$), i modelli Boundary RF hanno dimostrato una performance significativamente più stabile e una minore degradazione della qualità del campione rispetto al vanilla RF all'aumentare della stocasticità.
  • I risultati qualitativi mostrano che i modelli Boundary RF producono immagini più nitide con meno artefatti e migliore fedeltà degli oggetti, particolarmente nei setting di campionamento stocastico dove il vanilla RF tende a sovra-levigare.

• Scalabilità:

  • I benefici del Boundary RF Model sono stati coerenti scalando a modelli più grandi (DiT-L, DiT-XL) e risoluzioni più elevate (512$\times$512).

Significato e Rivendicazioni

Il paper sostiene che il Boundary RF Model offre una soluzione semplice e facilmente implementabile a un limite fondamentale del Rectified Flow. Imponendo esplicitamente le condizioni al contorno, il metodo:

1. Migliora la Qualità della Generazione: Supera costantemente il vanilla RF in termini di FID e altre metriche.
2. Abilita un Campionamento Stocastico Stabile: Risolve la divergenza della funzione di score vicino al tempo terminale, permettendo strategie di campionamento più robuste.
3. È Facile da Adattare: Può essere integrato nei modelli RF esistenti (come quelli che usano U-Net o DiT) con modifiche architettoniche minime.

Gli autori pongono questo lavoro come un necessario raffinamento per l'impiego pratico del Rectified Flow, in particolare in scenari che richiedono un campionamento stocastico di alta qualità. Suggeriscono che, sebbene il loro focus sia sul RF, il concetto centrale di imporre le condizioni al contorno potrebbe potenzialmente essere esteso ai modelli di diffusione con piccoli adattamenti. Il lavoro futuro proposto esplorerà l'imposizione del contorno adattiva e l'applicazione alla generazione di testo-video.