FreeBridge: Variational Schrödinger Bridges for Cellular Transition Dynamics

Technische Samenvatting: FreeBridge – Variabele Schrödinger-bruggen voor cellulaire transitiedynamiek

Probleemstelling
High-content beeldvormingsassays zijn essentieel voor het kwantificeren van cellulaire reacties op chemische en genetische perturbaties. Er bestaat echter een fundamentele beperking: cellen zijn chemisch gefixeerd op het moment van beeldacquisitie, waardoor individuele temporele trajecten onobserveerbaar blijven. Bijgevolg wordt perturbatiemodellering gereduceerd tot het afleiden van stochastisch transport tussen controle- en behandelde populaties, waarbij alleen de marginale eindverdelingen bekend zijn.

Hoewel recente generatieve modellen (bijv. diffusiemodellen, flow matching) een sterke uitlijning bereiken tussen controle- en geperturbueerde eindpunten, is endpoint-consistentie alleen onvoldoende om de intermediaire evolutie te bepalen. Oneindig veel stochastische processen kunnen identieke marginalen verbinden; zonder verdere beperkingen kunnen deze modellen regio's in de latente ruimte doorkruisen die niet worden ondersteund door geobserveerde single-cell morfologieën, wat leidt tot biologisch onwaarschijnlijke intermediaire toestanden. Bestaande methoden optimaliseren vaak distributieniveau-objectieven, waardoor intermediaire cellulaire structuren onderbepaald blijven, wat kan leiden tot het voldoen aan metrieken zoals FID of KID terwijl er artefacten ontstaan.

Methodologie: Het FreeBridge-framework
De auteurs stellen FreeBridge voor, een geconstrueerd stochastisch transportframework geformuleerd als een "Cell Engine" die de definitie van geldige cellulaire toestanden expliciet scheidt van het transitiemechanisme. De methode werkt onder supervisie van enkel de eindpunten en bestaat uit drie kerncomponenten:

1. State Unit (Expliciete Single-Cell Geometrie):

   • In plaats van beelden te behandelen als populatiebrede mengsels, definieert FreeBridge atomaire toestanden als geïnstantieerde segmentatie van single-cell representaties (met behulp van Cellpose).
   • Een vaste, vooraf getrainde encoder $\phi$ embedt deze crops in een latente ruimte $z = \phi(x)$. Deze encoder wordt bevroren tijdens het transportleerproces, wat een vaste latente geometrie vestigt die voorafgaand aan het leren van de transitierule is gedefinieerd.
   • Een empirische support bank $Z_{emp}$ wordt geconstrueerd uit de latente embeddings van geobserveerde cellen. Deze bank definieert de "manifold" van biologisch geldige morfologieën.

2. Transition Rule (Schrödinger Bridge Transport):

   • De transitie tussen de controleverdeling $\mu_z$ en de geperturbueerde verdeling $\nu_{c,z}$ wordt gemodelleerd als een gecontroleerde stochastische differentiaalvergelijking (SDE): $dZ_t = u_t(Z_t) dt + \sigma dW_t$.
   • Het doel is om een driftveld $u_t$ te vinden dat de initiële verdeling transporteert naar de doelverdeling bij $t=1$.
   • In tegenstelling tot standaard optimal transport, incorporeert deze formulering een taakspecifieke state cost om het pad te reguleren.

3. Geometric Regularization (Empirical Support Constraint):

   • Om te voorkomen dat het stochastische proces "leegtes" in de latente ruimte doorkruist, introduceert FreeBridge een support cost $V(z) = \lambda \cdot d_{emp}(z)^2$, waarbij $d_{emp}(z)$ de afstand is tot de dichtstbijzijnde buur in de empirische support bank $Z_{emp}$.
   • De totale objectief minimaliseert de som van de kinetische energie (drift magnitude) en de support cost:
     $$\min_u \mathbb{E} \left[ \int_0^1 \frac{1}{2\sigma^2}\|u_t(Z_t)\|^2 dt + \lambda \int_0^1 d_{emp}(Z_t)^2 dt \right]$$
   • Deze regularisatie verankert intermediaire toestanden aan de empirisch geobserveerde morfologische manifold zonder dat er gepaarde cel-niveau trajecten nodig zijn.

Belangrijkste Bijdragen

• Formulering: Introduceert een Schrödinger Bridge-formulering voor single-cell transitiemodellering die de specificatie van de staat expliciet ontkoppelt van het stochastische transport.
• Geometric Grounding: Stelt een nieuwe regularisatieterm voor op basis van non-parametrische dichtheidsschatting (nearest-neighbor afstand tot een empirische support bank) om intermediaire trajecten binnen biologisch plausibele regio's te beperken.
• Unified Protocol: Vestigt een uniform evaluatieprotocol over meerdere datasets (BBBC021, RxRx1, JUMP) met behulp van geïnstantieerde single-cell crops, wat een eerlijke vergelijking tussen methoden garandeert.
• Prestaties: Toont aan dat het beperken van de geometrie niet alleen de endpoint-getrouwheid verbetert, maar ook de semantische retentie van perturbatiemechanismen (Mechanism-of-Action).

Resultaten
Geëvalueerd op de BBBC021, RxRx1 en JUMP datasets, demonstreert FreeBridge het volgende:

• Endpoint Fidelity: Bereikt competitieve of verbeterde Fréchet Inception Distance (FID) en Kernel Inception Distance (KID) scores vergeleken met state-of-the-art baselines (PhenDiff, IMPA, CellFlux). Op BBBC021 bereikt het de laagste FID (15.35) en KID (1.10).
• Intermediate Consistency: Vermindert aanzienlijk de Support Violation Rate ($R_{viol}$), een metriek die het aandeel intermediaire toestanden meet die buiten de empirische morfologische support vallen. Op BBBC021 reduceert FreeBridge $R_{viol}$ naar 0.11, vergeleken met 0.38 voor PhenDiff en 0.31 voor CellFlux.
• Semantic Retention: Ablatie-studies laten zien dat het verwijderen van de geometrische regularisatie de Mechanism-of-Action (MoA) classificatie nauwkeurigheid verlaagt (van 72.08% naar 65.08% op BBBC021), wat aangeeft dat het beperken van de intermediaire evolutie helpt om mechanistische signalen te behouden.
• Kwalitatieve Analyse: Visualisaties bevestigen dat FreeBridge verbinding-specifieke morfologische kenmerken (bijv. nucleaire compactie, cytoplasmatische krimp) beter behoudt dan baselines, die vaak last hebben van textuursmoothing of incomplete patronen.

Betekenis en Claims
Het artikel claimt dat het modelleren van biologische transities meer vereist dan alleen boundary alignment; het vereist de expliciete geometrische gronding van de intermediaire evolutie. Door stochastisch transport te verankeren aan een empirisch ondersteunde cellulaire manifold, biedt FreeBridge een principieel aanpak voor het genereren van biologisch interpreteerbare perturbatiedynamiek. De auteurs stellen dat deze formulering een robuuste basis biedt voor "virtuele cel" simulatie, waarmee het kritieke gat wordt gedicht waar endpoint-aligned modellen falen in het vastleggen van realistische intermediaire cellulaire toestanden. Het werk benadrukt dat zonder geometrische beperkingen, generatieve modellen statistisch geldige eindpunten kunnen produceren die echter biologisch niet-ondersteunde regio's in de latente ruimte doorkruisen.

Beperkingen
De auteurs erkennen dat FreeBridge afhankelijk is van een bevroren encoder; morfologische factoren die afwezig zijn in de pretraining van de encoder kunnen niet enkel door transport learning worden hersteld. Daarnaast schaalt de berekening van de support-afstand kwadratisch ($O(B^2)$) met de minibatch grootte, wat benaderde nearest-neighbor zoekopdrachten kan vereisen voor grotere support banken.