Multi-Marginal Flow Matching with Adversarially Learnt Interpolants

Technische Samenvatting: Multi-Marginal Flow Matching met Adversarieel Geleerde Interpolanten

Probleemstelling
Het artikel behandelt de uitdaging van het leren van de tijdsafhankelijke dynamica van een systeem, gegeven enkel gesamplede observaties (snapshots) op discrete tijdstippen, zonder toegang tot grondwaarheidstrajecten. Dit is een centrale taak in velden zoals de biologie (bijv. single-cell RNA-sequencing, celtracking) en de geneeskunde. Formeel is het doel om een tijdsafhankelijke vectorveld $v_t$ te extraheren die een gewone differentiaalvergelijking (ODE) $dx_t = v_t(x_t)dt$ aanstuurt, zodanig dat de integratiemapping de initiële marginale distributie $q_0$ naar een doelverdeling $q_1$ brengt, terwijl tussenliggende marginale restricties $p_{t_i} = q_{t_i}$ worden voldaan voor een set tijdstappen $0 = t_1 < \dots < t_K = 1$.

Bestaande benaderingen, zoals het toepassen van standaard Flow Matching (FM) op opeenvolgende paren van marginalen of het gebruik van stuksgewijze lineaire of cubische spline-interpolanten, resulteren vaak in niet-gladde interpolatiecurven. Deze problemen met gladheid leiden tot een hoge gradiëntvariantie in de leerdoelstelling en problemen bij het marginaliseren van de resulterende vectorvelden, met name in hoogdimensionale instellingen of wanneer data ruis bevat.

Methodologie: ALI-CFM
De auteurs stellen ALI-CFM voor (Adversarially Learnt Interpolants with Conditional Flow Matching), een tweestapsmethode die het leren van interpolanten ontkoppelt van het leren van het vectorveld.

1. Adversarieel Leren van Interpolanten (ALI):
   In plaats van interpolanten te dwingen om door specifieke datapunten te gaan (pointwise matching), leert ALI een neuraal interpolant $G_\phi(x_0, x_1, t)$ die de distributie van tussenliggende marginalen matcht. De interpolant wordt geparametriseerd als een lineaire interpolatie plus een neurale correctieterm:
   $$G_\phi(x_0, x_1, t) = (1-t)x_0 + tx_1 + t(1-t)f_\phi(x_0, x_1, t)$$
   Om af te dwingen dat de pushforward-distributie van de interpolant de geobserveerde tussenliggende marginalen $q_{t_i}$ matcht, gebruiken de auteurs een GAN-geïnspireerde adversariële doelstelling. Een discriminator $D_\gamma$ onderscheidt samples van de ware marginal $q_{t_i}$ en samples gegenereerd door de interpolant $G_\phi$. De generator (de interpolant) wordt getraind om de Jensen-Shannon divergentie tussen deze distributies te minimaliseren.

   Om de uniciteit van de interpolanten te waarborgen en willekeurig gekromde paden te voorkomen, introduceren de auteurs regularisaties:

   • Lineaire Referentie: Bestraft afwijking van een rechte lijn tussen gekoppelde eindpunten.
   • Stuksgewijs Lineaire Referentie: Regresseert de interpolant naar een stuksgewijs lineair pad gedefinieerd door Markov-keten Optimal Transport (OT) koppelingen.
   • Tweede Afgeleide Norm: Dwingt gladheid af door de kromming van de interpolant te bestraffen.

2. Marginalisatie via Conditional Flow Matching (CFM):
   Zodra de interpolanten $G_\phi$ zijn getraind, zijn deze vastgesteld. Een tweede neuraal netwerk $u_\theta$ wordt getraind om het vectorveld te benaderen dat deze interpolanten genereert. Dit gebeurt met de standaard Conditional Flow Matching doelstelling, waarbij het doelvectorveld de tijdsafgeleide is van de geleerde interpolant:
   $$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$$
   De koppeling $\pi$ tussen $q_0$ en $q_1$ kan onafhankelijk zijn of een Optimal Transport plan.

Belangrijkste Bijdragen

1. Nieuwe Interpolatietechniek: Het artikel introduceert een methode om interpolanten te leren met behulp van een adversariële doelstelling die distributies matcht in plaats van punten. Dit maakt het mogelijk om gladde trajecten te leren die uniek zijn onder milde aannames (bewezen in de appendixen).
2. Robuustheid tegen Ruis en Complexiteit: De methode is ontworpen om met ruizige datasets en complexe, tijdsafhankelijke geometrieën om te gaan waar puntgewijze interpolatiemethoden (zoals splines of stuksgewijs lineaire methoden) falen of instabiel worden.
3. Schaalbaarheid en Veelzijdigheid: De auteurs demonstreren dat ALI-CFM schaalt naar hoogdimensionale data en een groot aantal marginale distributies ($K$), waarbij het bestaande baselines overtreft in specifieke complexe scenario's terwijl het in andere scenario's concurrerend blijft.

Experimentele Resultaten
De auteurs evalueren ALI-CFM op vier verschillende taken:

• Synthetische "Knot" Distributie: Op een synthetische dataset met 1.200 marginale distributies die een knoop vormen, was ALI-CFM de enige methode die in staat was de tijdsafhankelijke geometrie accuraat te vatten. Concurrerende methoden (OT-MFM, cubische splines, stuksgewijs lineair) produceerden niet-gladde interpolanten met een hoge gradiëntvariantie.
• Celtracking (Glioblastoom): Met behulp van microscopiegegevens van een bewegende cel, slaagde ALI-CFM erin de spatio-temporele dynamica te leren van ruizige, ondergesamplede data. In tegen testelling hiermee divergeerden OT-CFM trajecten, en faalde de tijd-onafhankelijke OT-MFM in het vatten van temporele afhankelijkheden.
• Single-Cell Trajectory Inference (scRNA-seq): Op de Embryoid Body (EB), Cite-seq en Multiome datasets presteerde ALI-CFM op gelijke voet met state-of-the-art baselines (OT-MFM, OT-CFM) in termen van Earth Mover's Distance (EMD). De auteurs merken op dat hoewel adversariële training het moeilijk maakt om in deze specifieke metrieken strikt beter te presteren dan methoden die overfitten op puntgewijze data, de methode zeer competitief blijft.
• Tumor Coördinaat Inferentie (Spatial Transcriptomics): Op een borstkankercatasets met ruimtelijk geordende weefselsecties presteerde ALI-CFM aanzienlijk beter dan alle baselines (inclusclusief OT-CFM, OT-MMFM en OT-MFM) in het voorspellen van de distributie van tumorregio's in niet-geteste secties. De auteurs schrijven dit succes toe aan het vermogen van de methode om multimodale distributies en de uitdagingen van het uitlijnen van weefselsecties aan te kunnen.

Betekenis en Claims
Het artikel claimt dat ALI-CFF een conceptueel nieuwe aanpak biedt voor multi-marginal flow matching door de focus te verschuiven van strikte puntgewijze toewijzing naar distributiematching. Deze verschuiving maakt het mogelijk om gladde, unieke interpolanten te leren die robuust zijn tegen ruis en in staat zijn complexe dynamica te modelleren waar eerdere methoden falen. De auteurs benadrukken dat hun methode bijzonder effectief is in scenario's met honderden of duizenden marginale distributies en ruizige data, een regime waarin geen andere flow matching algoritme vergelijkbare prestaties heeft laten zien. Hoewel de methode op standaard single-cell benchmarks op gelijke voet presteert met bestaande state-of-the-art, onderstreept de superieure prestatie op spatial transcriptomics en synthetische complexe geometrie taken de veelzijdigheid en schaalbaarheid van de methode.