Learning population dynamics from snapshots using state graph methods

Discrete Action Matching: Learning Stochastic Dynamics from Samples via State Graphs

Machine LearningArtificial Intelligence

Summary

This paper tackles the hard problem of understanding how populations change over time when you only have separate snapshots at different moments, not continuous observations. The authors developed a method called Discrete Action Matching that works with finite sets of states using ideas from geometry on graphs. Their approach estimates how populations likely move from one state to another by learning patterns from these snapshots, then uses this to create models that can simulate the dynamics. They tested their method on both synthetic datasets and real biological data to see how well it reconstructs missing information and predicts transitions.

What this means in practice

  • For bioinformatics teams: Model cell population changes over time from discrete samples to improve understanding of developmental biology processes.
  • For computational biologists: Generate interpolated cellular state trajectories from partial time points to design more accurate dynamic models for complex biological systems.

Authors

Mikhail Persiianov, Alexander Korotin

Abstract

Learning population dynamics from unpaired temporal marginals is an ill-posed inverse problem that requires structural assumptions on the underlying dynamics. We introduce $\textit{Discrete Action Matching}$ (DAM), a finite-state counterpart of Action Matching based on discrete Wasserstein geometry. For a prescribed marginal path and transport geometry, we derive an action-minimization objective for its canonical minimum-kinetic-energy current. Our key observation is that the density dependence of the discrete action reduces to neighboring density ratios. Along an empirical interpolation of the snapshots, DAM first estimates these ratios and then learns an action potential. The learned fields also define a graph-supported Markov sampler. Experiments on controlled synthetic dynamics and real mouse gastrulation data evaluate marginal reconstruction and interpolation. Additional experiments approximate numerical surface-transport paths from paired samples.