Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme

Each language version is independently generated for its own context, not a direct translation.

De Kunst van het Voorspellen van Menigevorming: Hoe iJKOnet de Onzichtbare Krachten Ontdekt

Stel je voor dat je naar een grote menigte mensen kijkt op een plein. Je ziet foto's van deze menigte op verschillende tijdstippen: 's ochtends, 's middags en 's avonds. Maar er is een groot probleem: je kunt de individuele mensen niet volgen. Ze verdwijnen na elke foto en komen in de volgende foto als nieuwe, willekeurige gezichten terug. Je weet niet wie naar wie is verplaatst.

De vraag is nu: Wat drijft deze menigte?
Zijn ze aan het rennen omdat er een brand is (een externe kracht)? Zitten ze in groepjes te kletsen omdat ze elkaar aantrekken (interactie)? Of bewegen ze willekeurig rond als een drijvende massa (diffusie)?

In de wetenschap noemen we dit het leren van populatiedynamica. Het is een enorm lastig raadsel, vooral in gebieden zoals biologie (hoe cellen zich ontwikkelen), meteorologie (hoe stormen bewegen) of financiën (hoe prijzen schommelen).

Het Probleem: De Ontbrekende Schakel

Vroeger probeerden wetenschappers dit op te lossen door te raden welke "energie" of "kracht" de beweging veroorzaakt. Ze gebruikten een slimme wiskundige truc genaamd het JKO-schema (vernoemd naar de wetenschappers Jordan, Kinderlehrer en Otto).

Je kunt het JKO-schema vergelijken met een helling.

De mensen op het plein willen altijd naar beneden rollen, naar de laagste vallei (de laagste energie).
Maar ze kunnen niet direct naar beneden springen; ze moeten stap voor stap bewegen, en elke stap kost energie.
Het JKO-schema helpt om die stap-voor-stap beweging te berekenen als je weet hoe de helling eruitziet.

Maar hier zit de klem: De meeste bestaande methoden wisten de vorm van de helling (de energie) niet. Ze moesten de helling leren op basis van de foto's. En dat was erg lastig, omdat de oude methoden vaak vastliepen in complexe wiskunde of specifieke, starre bouwplannen voor hun computerprogramma's (neural networks) vereisten.

De Oplossing: iJKOnet (Het Omgekeerde Raadsel)

De auteurs van dit paper hebben een nieuwe methode bedacht, genaamd iJKOnet. Ze gebruiken een slimme combinatie van twee ideeën:

Omgekeerde Optimalisatie (Inverse Optimization):
Stel je voor dat je een bal ziet rollen over een onbekend terrein. In plaats van te proberen het terrein te tekenen en te kijken waar de bal heen gaat, doe je het andersom:
- Je kijkt naar de beweging van de bal.
- Je vraagt je af: "Welke vorm van het terrein zou ervoor zorgen dat de bal precies deze beweging maakt?"
- iJKOnet doet precies dit. Het kijkt naar de foto's van de menigte en vraagt zich af: "Welke onzichtbare krachten (energie) zouden nodig zijn om deze veranderingen te verklaren?"
Het JKO-Schema als Leraar:
Het systeem probeert een "foute" helling te tekenen. Als de helling verkeerd is, zal de voorspelde beweging van de menigte niet overeenkomen met de echte foto's.
- Het systeem corrigeert de helling steeds opnieuw.
- Het doet dit door een soort wedstrijd (een min-max spelletje) te spelen tussen twee delen van het programma:
  - De ene kant probeert de helling zo goed mogelijk te tekenen.
  - De andere kant probeert de beweging van de mensen zo goed mogelijk te simuleren op die helling.
- Als ze elkaar goed aanvullen, hebben we de juiste krachten gevonden!

Waarom is dit zo speciaal? (De Creatieve Analogie)

1. Geen Starre Bouwplannen
Vroeger moesten de computerprogramma's die dit deden (zoals JKOnet) gebouwd zijn volgens een heel strikt, star ontwerp (zoals een Lego-set waar je maar één manier op kunt bouwen). Als je iets anders wilde, kon het niet.
iJKOnet is als LEGO voor volwassenen: je kunt er alles mee bouwen. Je kunt de bouwplannen (de neurale netwerken) vrij kiezen, zolang ze maar slim genoeg zijn. Dit maakt het veel makkelijker om complexe, echte situaties (zoals 100 verschillende eigenschappen van cellen tegelijk) te modelleren.

2. Het Voorspellen van de Toekomst
Stel je voor dat je een film hebt van een storm die over een stad trekt, maar je hebt alleen de beelden van de straten op maandag en woensdag. Je mist dinsdag.

De oude methoden probeerden de weg van dinsdag te raden, maar vaak liepen ze vast of gaven ze een rommelig beeld.
iJKOnet leert de "windkracht" (de energie) die de storm drijft. Zodra het die kracht begrijpt, kan het de storm op dinsdag, donderdag en vrijdag perfect voorspellen, zelfs als het weer heel chaotisch is.

Wat hebben ze ontdekt?

De onderzoekers hebben getoond dat hun methode beter werkt dan de vorige versies.

Ze hebben het getest op synthetische data (kunstmatige menigten).
Ze hebben het getest op echte biologie-data: hoe menselijke stamcellen zich ontwikkelen tot verschillende soorten cellen. Dit is cruciaal voor het begrijpen van ziektes en het ontwikkelen van medicijnen.
Hun methode was sneller, nauwkeuriger en kon complexere patronen vinden dan de concurrenten.

Samenvattend

Dit paper introduceert iJKOnet, een slimme manier om te achterhalen waarom dingen bewegen, zelfs als je niet kunt zien wie precies naar waar gaat. Het is alsof je de windkracht kunt meten door alleen naar de golven op zee te kijken, zonder de wind zelf te zien.

Door dit te doen, kunnen wetenschappers beter begrijpen hoe cellen groeien, hoe steden veranderen en hoe de wereld zich ontwikkelt, puur op basis van momentopnamen. Het is een grote stap voorwaarts in het vertalen van ruwe data naar diep inzicht.

Each language version is independently generated for its own context, not a direct translation.

Titel: Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme

Conferentie: ICLR 2026
Auteurs: Mikhail Persiianov, Jiawei Chen, Petr Mokrov, Alexander Tyurin, Evgeny Burnaev, Alexander Korotin.

1. Probleemstelling

Het paper adresseert de uitdaging om de onderliggende stochastische dynamica van een populatie te reconstrueren op basis van "snapshots" (momentopnamen) van de populatieverdeling op discrete tijdstippen.

Context: In veel wetenschappelijke domeinen (biologie, ecologie, epidemiologie, financiën) is het onmogelijk om individuele trajecten van deeltjes te volgen. In plaats daarvan heeft men alleen toegang tot marginaal gedistribueerde data op verschillende tijdstippen (bijv. destructieve metingen in single-cell genomics).
Huidige aanpak: Bestaande methoden modelleren deze evolutie vaak als een Wasserstein Gradient Flow (WGF) en gebruiken het JKO-scheme (Jordan-Kinderlehrer-Otto) voor tijdsdiscretisatie.
Beperkingen van bestaande methoden:
- JKOnet: Vereist ingewikkelde bi-level optimalisatie en is beperkt tot potentie-energie (kan geen stochastische dynamica modelleren).
- JKOnet*: Vervangt de optimalisatie door eerste-orde optimaliteitsvoorwaarden, maar vereist vooraf berekende optimale transportkoppelingen (OT), wat de schaalbaarheid beperkt en geen end-to-end training toestaat.

2. Methodologie: iJKOnet

De auteurs introduceren iJKOnet, een nieuwe aanpak die het JKO-framework combineert met inverse optimalisatie technieken.

Kernidee: Inverse Optimalisatie

In plaats van de JKO-stap te gebruiken om een verdeling te voorspellen (voorwaartse richting), wordt het probleem omgekeerd: gegeven de waargenomen verdelingen $\rho_k$ en $\rho_{k+1}$ , wordt de energiefunctional $J$ gezocht die deze evolutie het beste verklaart.

Het JKO-scheme stelt dat $\rho_{k+1}$ de oplossing is van een minimaliseringsprobleem voor een gegeven $J$ .
Als $J$ de ware functional is, dan moet de waarde van de JKO-objectief voor $\rho_{k+1}$ gelijk zijn aan de minimale waarde.
De auteurs formuleren een min-max doelstelling:
- Innerlijke minimalisatie: Vind de transportmap $T^k$ die de JKO-stap benadert voor een kandidaat-energie $J$ .
- Buitenste maximalisatie: Pas $J$ aan zodat de "gap" tussen de optimale waarde en de waarde bereikt met de ware data wordt gemaximaliseerd (of de gap tussen de theoretische en empirische JKO-stap wordt geminimaliseerd).

Formulering van de Loss

De doelstelling wordt herschreven als een adversarial training probleem:
$\max_{J} \min_{T^k} \sum_{k=0}^{K-1} \left[ J(T^k_\sharp \rho_k) - J(\rho_{k+1}) + \frac{1}{2\tau} \int \|x - T^k(x)\|^2 d\rho_k(x) \right]$
Hierbij:

$J_\theta$ : De parametrische energiefunctional (potentiaal, interactie, interne energie).
$T^k_\phi$ : De transportmap (neuraal netwerk) die $\rho_k$ naar de volgende stap duwt.
$\tau$ : De tijdstapgrootte.

Praktische Implementatie

Architectuur: In tegenstelling tot eerdere werken die Input-Convex Neural Networks (ICNNs) vereisten voor de transportmaps (vanwege convexiteitsbeperkingen), gebruikt iJKOnet standaard architecturen zoals MLP's of ResNets. Dit komt doordat de nieuwe loss-functie geen strikte convexiteitsbeperkingen oplegt, wat de schaalbaarheid en stabiliteit verbetert.
Energie Parametrisatie: De functional $J$ $J$ wordt gemodelleerd als een vrije energie met drie componenten:
1. Potentiaal-energie ( $V_\theta$ ).
2. Interactie-energie ( $W_\theta$ ).
3. Interne energie (gebaseerd op entropie, met een leerbare diffusiecoëfficiënt $\theta_3$ ).
Entropie Schatting: De entropieterm wordt geschat via de Kozachenko-Leonenko schatter en de verandering van variabelen formule, waarbij de Jacobiaan van de transportmap wordt gebruikt.

3. Belangrijkste Bijdragen

Methodologie: De eerste formulering van het herstel van energiefunctionals binnen het JKO-framework als een inverse optimalisatie taak, wat leidt tot een robuust min-max leerdoel.
Algoritme: Een end-to-end trainingsprocedure die geen restrictieve architecturale keuzes (zoals ICNNs) vereist. Dit maakt de methode schaalbaar naar hogere dimensies.
Theoretische Garanties: De auteurs bewijzen (Theorema 3.1) dat onder geschikte aannames (zoals strikte convexiteit en gladheid van de potentiaal), het minimaliseren van hun loss-functie leidt tot een nauwkeurige reconstructie van de onderliggende potentiaal-energie (tot op een additieve constante na).
Prestaties: Uitgebreide experimenten tonen aan dat iJKOnet superieur presteert ten opzichte van JKOnet en JKOnet*, zowel op synthetische data als op real-world single-cell RNA-seq datasets.

4. Resultaten

De methode werd geëvalueerd op synthetische datasets en echte biologische data (Embryoid Body dataset).

Synthetische Data (Potentiaal-energie):
- In de "unpaired" setup (waarbij samples tussen tijdstippen niet gekoppeld zijn, wat realistischer is maar moeilijker), overtreft iJKOnet JKOnet* significant op bijna alle geteste potentialen.
- De auteurs merkten op dat bestaande codebases vaak per ongeluk "paired" data gebruikten (waarbij deeltjes over tijd worden gevolgd), wat de prestaties kunstmatig opvoerde. iJKOnet is robuust in de echte "unpaired" setting.
Single-Cell Genomics (EB Dataset):
- In een 5D "leave-two-out" experiment (voorspellen van tijdstippen die niet in de training zaten) behaalde iJKOnet de beste resultaten (laagste Bures-Wasserstein en EMD afstanden) vergeleken met state-of-the-art methoden zoals DMSB, TrajectoryNet en MMSB.
- In een 100D "leave-one-out" experiment presteerde iJKOnet vergelijkbaar met DMSB, maar met een veel efficiëntere trainingsprocedure (geen caching van trajecten nodig).
Beperkingen: Het herstel van interactie-energie en interne energie (diffusie) blijft uitdagend en kan instabiliteit veroorzaken bij het gezamenlijk optimaliseren van alle parameters. De theoretische garanties zijn momenteel beperkt tot potentiaal-energie.

5. Betekenis en Impact

Dit werk is significant omdat het een brug slaat tussen inverse optimalisatie en stochastische dynamica.

Schaalbaarheid: Door af te zien van ICNNs en pre-berekende OT-plans, maakt iJKOnet het mogelijk om complexe populatiedynamica te leren in hogere dimensies, wat essentieel is voor moderne biologische en fysieke toepassingen.
End-to-End Learning: Het elimineert de noodzaak voor gescheiden optimalisatiestappen (zoals het vooraf berekenen van transportkoppelingen), wat de generalisatie verbetert en de implementatie vereenvoudigt.
Toepassingsbereik: De methode biedt een krachtig instrument voor het begrijpen van systemen waar individuele trajecten ontbreken, zoals celdifferentiatie, epidemieverspreiding en verkeersstromen.

Samenvattend biedt iJKOnet een nieuw, theoretisch onderbouwd en praktisch efficiënt framework voor het leren van populatiedynamica uit statische momentopnamen, en lost het belangrijke beperkingen op van eerdere JKO-gebaseerde methoden.