Branched Schrödinger Bridge Matching

Dit paper introduceert BranchSBM, een nieuw generatief model dat de beperkingen van bestaande methoden overbrugt door vertakte trajecten te leren die populaties van een gemeenschappelijke oorsprong naar meerdere eindverdelingen laten divergeren, wat essentieel is voor toepassingen zoals het modelleren van celfate-bifurcaties.

De kern

Stel je voor dat je een groep mensen hebt die allemaal in één punt beginnen, en je wilt voorspellen waar ze over een uur zullen zijn.

In de wereld van kunstmatige intelligentie (AI) is dit een heel bekend probleem. Meestal denken AI-modellen dat iedereen een rechte lijn volgt of dat ze allemaal in één grote, rommelige hoop eindigen. Maar in het echte leven is dat niet hoe het werkt.

Denk aan een stamcel in je lichaam. Die begint als één soort "leeg" cel. Maar als je die cel een medicijn geeft, kan hij zich op verschillende manieren ontwikkelen: hij kan een bloedcel worden, een spiercel, of misschien zelfs een kankercel. Hij splitst zich op in meerdere paden tegelijk. Of denk aan een drukke menigte op een plein: als er een gevaar is, splitst de menigte zich op in verschillende richtingen om te ontsnappen.

Deze paper introduceert een nieuwe AI-methode genaamd BranchSBM (Branched Schrödinger Bridge Matching). Hier is hoe het werkt, vertaald naar simpele taal:

1. Het oude probleem: De "Rechte Lijn" valkuil

Tot nu toe gebruikten AI-modellen voor dit soort voorspellingen vaak een methode die werkt als een trein op één spoor. Als je twee punten wilt verbinden (begin en eind), rijdt de trein rechtstreeks daarheen.

• Het probleem: Als je beginpunt één groep is, maar je eindpunt twee verschillende groepen zijn (bijvoorbeeld bloedcellen én spiercellen), probeert de trein dan vaak de "gemiddelde" weg te nemen. Het resultaat? Een rommelige mix van beide, of de trein kiest willekeurig één kant en negeert de andere. Dit noemen ze "mode collapse" (een instorting naar één optie). Het kan geen vertakkende paden zien.

2. De nieuwe oplossing: De "Vijver met Vertakkingen"

BranchSBM is als een slimme waterval of een stroompje dat zich splitst.
In plaats van één trein, stelt dit model zich voor dat er een stroom van water (de cellen of mensen) is die begint bij één bron.

• De vertakking: Het model leert niet alleen waar het water naartoe stroomt (de richting), maar ook hoeveel water er naar welk pad gaat.
• De groei: Het model kan ook beslissen of er meer "water" (massa) bij komt of dat het verdwijnt. In de biologie betekent dit dat sommige celgroepen groeien terwijl andere krimpen.

3. Hoe leert de computer dit? (De 4-fasen training)

De auteurs hebben een slimme manier bedacht om dit te leren, alsof je een kind leert een complex parcours te lopen:

• Fase 1: De Schets (De Interpolant)
  Eerst tekent de computer een ruwe schets van alle mogelijke routes. Het kijkt naar het begin en het eind, en bedenkt een "droompad" dat de kortste weg is, rekening houdend met obstakels (zoals bergen of muren in de data).
• Fase 2: De Stroom (De Drift)
  Nu leert het model hoe het water precies moet stromen om die schets te volgen. Het leert de "windrichting" voor elk punt op het pad.
• Fase 3: De Verdeling (De Groei)
  Dit is het magische deel. Het model leert nu: "Oké, bij dit punt in de tijd, moet 60% van het water linksaf gaan en 40% rechtsaf." Het leert ook hoeveel water er bij elke tak bijkomt of verdwijnt.
• Fase 4: De Perfectie (Samenwerken)
  Tot slot laat het model alles samenwerken. De windrichting en de verdeling van het water worden tegelijkertijd aangepast zodat het eindresultaat precies klopt met de werkelijkheid.

Waarom is dit belangrijk?

Stel je voor dat je een medicijn ontwikkelt tegen kanker. Je wilt weten: "Als ik dit medicijn geef, welke cellen genezen en welke worden kwaadaardig?"

• Oude modellen: Zeggen misschien: "De cellen worden een beetje van alles." (Onnauwkeurig).
• BranchSBM: Zegt: "60% van de cellen geneest, 30% wordt resistent, en 10% sterft." En het laat je precies zien hoe ze daar terechtkomen.

Samenvattend

Deze paper introduceert een nieuwe manier voor computers om verwarring en keuze te modelleren. In plaats van te denken dat iedereen één pad volgt, begrijpt BranchSBM dat het leven (en cellen, en menigten) vaak vertakt. Het is alsof je van een simpele lijntekening bent gegaan naar een dynamische, levende boom waar je precies kunt zien welke tak waar naartoe groeit en hoe groot die takken worden.

Dit helpt wetenschappers beter te begrijpen hoe ziektes ontstaan, hoe medicijnen werken, en hoe complexe systemen zich gedragen wanneer ze uit elkaar vallen in verschillende richtingen.

Technische samenvatting

Probleemstelling

Generatieve modellering heeft vaak tot doel om de tussenliggende trajecten te voorspellen tussen een initiële verdeling ($\pi_0$) en een doelverdeling ($\pi_1$). Bestaande methoden, zoals Flow Matching en Schrödinger Bridge Matching (SBM), modelleren doorgaans een enkele stochastische pad om een mapping tussen twee distributies te leren. Deze methoden hebben echter fundamentele beperkingen:

1. Unimodale beperking: Ze zijn inherent beperkt tot overgangen tussen twee verdelingen en kunnen geen vertakte (branched) of divergente evolutie modelleren vanuit één oorsprong naar meerdere distincte eindmodi.
2. Massa-behoud: Traditionele SBM veronderstelt vaak behoud van massa, wat onrealistisch is voor systemen waar groei of vernietiging van massa optreedt (bijv. celproliferatie of -dood).
3. Mode Collapse: Bij het modelleren van systemen die zich vertakken naar meerdere doelen (zoals celfate-bifurcaties), neigen bestaande methoden ertoe dat alle deeltjes dezelfde "laag-energie" route volgen, waardoor ze falen om de volledige diversiteit van de doelverdelingen te bereiken.

Dit probleem is cruciaal in domeinen zoals crowds navigatie over complexe oppervlakken en, vooral, in de biologie voor het modelleren van celdifferentiatie en de respons op perturbaties (zoals medicijnen), waar een homogene populatie zich splitst in meerdere fenotypisch verschillende uitkomsten.

Methodologie: BranchSBM

De auteurs introduceren Branched Schrödinger Bridge Matching (BranchSBM), een raamwerk dat de Generalized Schrödinger Bridge (GSB) uitbreidt naar een scenario met meerdere doelen en variabele massa.

1. Theoretische Formulering:

• Branched Generalized Schrödinger Bridge: Het probleem wordt gedefinieerd als het minimaliseren van de som van "Unbalanced Generalized Schrödinger Bridge" problemen over $K+1$ takken.
• Unbalanced Conditional Stochastic Optimal Control (CondSOC): Het probleem wordt herschreven als een som van CondSOC-objectieven. Dit omvat het leren van:
  • Tijd-afhankelijke snelheidsvelden ($u_{t,k}$): Deze sturen de deeltjes langs specifieke trajecten naar elke doelverdeling.
  • Groei-processen ($g_{t,k}$): Deze bepalen hoe massa wordt verdeeld over de takken. De primaire tak ($k=0$) start met een gewicht van 1, terwijl secundaire takken ($k>0$) starten met gewicht 0 en groeien naarmate massa wordt overgedragen.
• Massa-behoud: Voor gebalanceerde problemen wordt de som van de gewichten over alle takken constant gehouden ($\sum w_{t,k} = 1$), wat impliceert dat massa die de primaire tak verlaat, wordt herverdeeld naar de secundaire takken.

2. Leerarchitectuur (Multi-stadia Training):
BranchSBM maakt gebruik van een vier-stadia trainingsalgoritme om stabiliteit en schaalbaarheid te garanderen:

• Stadium 1: Neural Interpolant Training: Een neurale netwerk $\phi_{t,\eta}$ wordt getraind om de optimale interpolatie tussen eindpunten $(x_0, x_{1,k})$ te voorspellen, waarbij een state-cost $V_t$ (bijv. LAND of RBF metriek) wordt geminimaliseerd. Dit levert de "ground truth" snelheidsvelden op.
• Stadium 2: Flow Matching: Onafhankelijke drift-netwerken ($u_{\theta, t, k}$) worden getraind om de snelheidsvelden uit Stadium 1 te benaderen voor elke tak afzonderlijk.
• Stadium 3: Groei-netwerken Training: De drift-netwerken worden vastgezet (frozen). De groei-netwerken ($g_{\phi, t, k}$) worden getraind om de massa-distributie over de tijd te optimaliseren. Dit gebeurt door een gecombineerde loss te minimaliseren bestaande uit:
  • Energy Loss: Minimaliseert de kinetische energie en state-cost.
  • Weight Matching Loss: Zorg ervoor dat de voorspelde gewichten bij $t=1$ overeenkomen met de echte verdeling van de doelpopulaties.
  • Mass Conservation Loss: Garandeert dat de totale massa op elk tijdstip behouden blijft.
• Stadium 4: Joint Training: Zowel de drift- als de groei-netwerken worden samen finetuned, inclusief een reconstructie-loss om de eindverdeling nauwkeurig te matchen.

Belangrijkste Bijdragen

1. Definitie van het Branched GSB-probleem: De auteurs formuleren wiskundig hoe een initiële verdeling naar meerdere gewogen doelverdelingen kan worden getransporteerd via een som van onbalans CondSOC-objectieven.
2. BranchSBM Framework: Een nieuw matching-framework dat zowel divergerende snelheidsvelden als tak-specifieke groeisnelheden leert, waardoor het in staat is om populatie-divergentie te modelleren.
3. Theoretische garanties: Bewijzen voor de existentie van optimale groei-functies en de convergentie van het trainingsproces naar de optimale drift en groei-velden.
4. Unieke capaciteit: Het vermogen om dynamische vertakkings-trajecten te modelleren zonder tussenliggende tijdstippen te hoeven observeren, wat essentieel is voor biologische data waar tussenstappen vaak ontbreken.

Resultaten

BranchSBM werd geëvalueerd op drie verschillende domeinen en overtrof bestaande methoden (zoals Single-Branch SBM, DeepRUOT en CytoBridge) aanzienlijk:

1. 3D LiDAR Navigatie:

   • Opdracht: Navigeren over een bergachtig LiDAR-manifold van één startpunt naar twee verschillende doelgebieden.
   • Resultaat: BranchSBM leerde distincte, niet-lineaire paden die het manifold volgden. Het bereikte een veel lagere Wasserstein-afstand ($W_1 \approx 0.24$ vs $0.98$ voor Single-Branch SBM) en vermijdt mode collapse door massa correct te verdelen.

2. Differentiatie van Eencellige Populaties (Mouse Hematopoiesis & Pancreatic $\beta$-cells):

   • Opdracht: Modelleren van de differentiatie van homogene progenitorcellen naar meerdere celfates.
   • Resultaat: Single-Branch SBM faalde om de twee distincte eindfates te bereiken en produceerde een vage, gemengde verdeling. BranchSBM reconstrueerde echter zowel de tussenliggende als de eindtoestanden nauwkeurig, met significante verbeteringen in $W_1$ en $W_2$ afstanden. Het model slaagde erin om de biologisch betekenisvolle vertakkingsstructuur te leren, zelfs zonder tussenliggende data.

3. Modellering van Cel-Status Perturbaties (Clonidine & Trametinib):

   • Opdracht: Voorspellen van hoe een homogene celpopulatie reageert op medicijnen en zich splitst in heterogene subpopulaties.
   • Resultaat: In hoge dimensies (50-150 PCA-componenten) kon Single-Branch SBM slechts één cluster nabootsen (de dichtstbijzijnde bij de controle), terwijl BranchSBM succesvol trajecten naar alle divergente clusters (2 voor Clonidine, 3 voor Trametinib) genereerde. Dit toonde aan dat het model complexere, hogedimensionale perturbatie-effecten kan vangen.

Beteeknis en Impact

BranchSBM vertegenwoordigt een significante doorbraak in generatieve modellering voor systemen met divergente dynamiek.

• Biologische Relevantie: Het biedt een wiskundig onderbouwde methode om celfate-bifurcaties en perturbatie-responsen te modelleren, wat cruciaal is voor het begrijpen van ziekteprogressie en het ontwerpen van therapieën.
• Algemene Toepasbaarheid: Het raamwerk is niet beperkt tot biologie; het is ook toepasbaar op robotica (multi-path navigatie) en fysica.
• Efficiëntie: Het vereist alleen eindpunten (start en doel) voor training, wat het zeer geschikt maakt voor scenario's waar tussenliggende data schaars of onbeschikbaar is.
• Schaalbaarheid: Het model kan schalen naar een groot aantal takken (tot 11 in de experimenten) met een beperkte toename in rekentijd, wat het een praktische oplossing maakt voor complexe real-world systemen.

Kortom, BranchSBM lost het probleem op van het modelleren van multi-modale, onbalans stochastische processen door een uniek framework te bieden dat zowel de richting (drift) als de populatieverdeling (groei) van een systeem dynamisch leert.