Multi-Marginal Flow Matching with Adversarially Learnt Interpolants

Sintesi Tecnica: Multi-Marginal Flow Matching con Interpolanti Appresi Adversarialmente

Definizione del Problema
Il documento affronta la sfida di apprendere la dinamica temporale di un sistema partendo solo da osservazioni campionate (snapshot) in punti temporali discreti, senza avere accesso alle traiettorie reali (ground-truth). Questo è un compito centrale in campi come la biologia (ad esempio, sequenziamento dell'RNA a singola cellula, tracciamento cellulare) e la medicina. Formalmente, l'obiettivo è inferire un campo vettoriale tempo-dipendente $v_t$ che governa un'equazione differenziale ordinaria (ODE) $dx_t = v_t(x_t)dt$, tale che la mappa di integrazione spinga una distribuzione marginale iniziale $q_0$ verso una distribuzione target $q_1$, soddisfacendo al contempo i vincoli marginali intermedi $p_{t_i} = q_{t_i}$ per un insieme di passi temporali $0 = t_1 < \dots < t_K = 1$.

Gli approcci esistenti, come l'applicazione del Flow Matching (FM) standard a coppie consecutive di marginali o l'uso di interpolanti lineari o cubici a tratti (spline), spesso risultano in curve di interpolazione non lisce. Questi problemi di non-linearità portano a un'elevata varianza del gradiente nell'obiettivo di apprendimento e a difficoltà nella marginalizzazione dei campi vettoriali risultanti, in particolare in contesti ad alta dimensionalità o con dati rumorosi.

Metodologia: ALI-CFM
Gli autori propongono ALI-CFM (Adversarially Learnt Interpolants with Conditional Flow Matching), un metodo in due fasi che scinde l'apprendimento degli interpolanti dall'apprendimento del campo vettoriale.

1. Apprendimento Adversariale degli Interpolanti (ALI):
   Invece di forzare gli interpolanti a passare attraverso specifici punti dati (matching puntuale), ALI apprende un interpolante neurale $G_\phi(x_0, x_1, t)$ che si adatta alla distribuzione dei marginali intermedi. L'interpolante è parametrizzato come un'interpolazione lineare più un termine di correzione neurale:
   $$G_\phi(x_0, x_1, t) = (1-t)x_0 + tx_1 + t(1-t)f_\phi(x_0, x_1, t)$$
   Per imporre che la distribuzione di pushforward dell'interpolante corrisponda ai marginali intermedi osservati $q_{t_i}$, gli autori utilizzano un obiettivo adversariale ispirato alle GAN. Un discriminatore $D_\gamma$ distingue tra i campioni della vera marginale $q_{t_i}$ e i campioni generati dall'interpolante $G_\phi$. L'interpolante viene addestrato per minimizzare la divergenza di Jensen-Shannon tra queste distribuzioni.

   Per garantire l'unicità degli interpolanti e prevenire traiettorie arbitrariamente curve, gli autori introducono dei regolarizzatori:

   • Riferimento Lineare: penalizza la deviazione da una linea retta tra gli estremi accoppiati.
   • Riferimento Lineare a Tratti: effettua la regressione dell'interpolante verso un percorso lineare a tratti definito da accoppiamenti di Trasporto Ottimale (OT) a catena di Markov.
   • Norma della Seconda Derivata: impone la regolarità (smoothness) penalizzando la curvatura dell'interpolante.

2. Marginalizzazione tramite Conditional Flow Matching (CFM):
   Una volta addestrati gli interpolanti $G_\phi$, questi vengono fissati. Una seconda rete neurale $u_\theta$ viene addestrata per approssimare il campo vettoriale che genera tali interpolanti. Ciò avviene utilizzando l'obiettivo standard di Conditional Flow Matching (CFM), dove il campo vettoriale target è la derivata temporale dell'interpolante appreso:
   $$L_{CFM} = \mathbb{E}_{t, (x_0, x_1) \sim \pi} \| u_\theta(G_\phi(x_0, x_1, t)) - \frac{d}{dt}G_\phi(x_0, x_1, t) \|^2$$
   L'accoppiamento $\pi$ tra $q_0$ e $q_1$ può essere indipendente o un piano di Trasporto Ottimale (OT).

Contributi Chiave

1. Tecnica di Interpolazione Innovativa: Il documento introduce un metodo per apprendere gli interpolanti utilizzando un obiettivo adversariale che confronta le distribuzioni anziché i punti. Ciò consente di ottenere traiettorie lisce che sono uniche sotto deboli assunzioni (dimostrate negli allegati).
2. Robustezza al Rumore e alla Complessità: Il metodo è progettato per gestire dataset rumorosi e geometrie complesse e variabili nel tempo, dove i metodi di interpolazione puntuale (come spline o lineari a tratti) falliscono o diventano instabili.
3. Scalabilità e Versatilità: Gli autori dimostrano che ALI-CFM scala a dati ad alta dimensionalità e a un gran numero di distribuzioni marginali ($K$), superando i baseline esistenti in scenari specifici complessi pur rimanendo competitivo in altri.

Risultati Sperimentali
Gli autori valutano ALI-CFM su quattro compiti distinti:

• Distribuzione "Knot" Sintetica: Su un dataset sintetico con 1.200 distribuzioni marginali che formano un nodo, ALI-CFM è stato l'unico metodo in grado di catturare accuratamente la geometria tempo-dipendente. I metodi concorrenti (OT-MFM, spline cubiche, lineari a tratti) hanno prodotto interpolanti non lisce con un'elevata varianza del gradiente.
• Tracciamento Cellulare (Glioblastoma): Utilizzando dati di microscopia di una cella in movimento, ALI-CFM ha appreso con successo la dinamica spazio-temporale da dati rumorosi e sottocampionati. Al contrario, le traiettorie di OT-CFM hanno divergente, e OT-MFM tempo-indipendente non è riuscito a catturare le dipendenze temporali.
• Inferenza di Traiettoria Single-Cell (scRNA-seq): Sui dataset Embryoid Body (EB), Cite-seq e Multiome, ALI-CFM ha ottenuto prestazioni paragonabili ai migliori modelli attuali (OT-MFM, OT-CFM) in termini di Earth Mover's Distance (EMD). Gli autori osservano che, sebbene l'addestramento adversariale renda difficile superare strettamente i metodi che si adattano eccessivamente (overfitting) ai dati puntuali in questi specifici parametri, il metodo rimane altamente competitivo.
• Inferenza delle Coordinate Tumorali (Trascrittomica Spaziale): Su un dataset di cancro al seno con sezioni tissutali spazialmente ordinate, ALI-CFM ha superato significativamente tutti i baseline (inclusi OT-CFM, OT-MMFM e OT-MFM) nella previsione della distribuzione delle regioni tumorali in sezioni non utilizzate per l'addestramento (held-out). Gli autori attribuiscono questo successo alla capacità del metodo di gestire distribuzioni multimodali e le sfide dell'allineamento delle sezioni tissutali.

Significato e Rivendicazioni
Il documento sostiene che ALI-CFM offre un approccio concettualmente nuovo al multi-marginal flow matching, spostando il focus dal rigido assegnamento puntuale al matching delle distribuzioni. Questo spostamento consente l'apprendimento di interpolanti lisce e uniche, robuste al rumore e capaci di modellare dinamiche complesse dove i metodi precedenti falliscono. Gli autori evidenziano come il loro metodo sia particolarmente efficace in scenari con centinaia o migliaia di distribuzioni marginali e dati rumorosi, un regime in cui nessun altro algoritmo di flow matching ha dimostrato una prestazione comparabile. Sebbene il metodo sia in linea con lo stato dell'arte esistente nei benchmark standard della singola cella, la sua prestazione superiore sulla trascrittomica spaziale e sui compiti di geometria sintetica complessa ne sottolinea la versatilità e la scalabilità.