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

2026-10-04

AI Summary

This study addresses the ill-posed inverse problem of learning population dynamics from unpaired temporal marginal distributions. To this end, it proposes a discrete action matching framework that extends continuous action matching to finite state spaces. By leveraging discrete Wasserstein geometry, the authors derive a minimum kinetic energy flow objective and simplify discrete action computation through neighbor density ratios. Furthermore, a graph-supported Markov sampler is constructed to learn the potential function. Integrating variational inference with graph theory, this work effectively achieves marginal reconstruction, temporal interpolation, and surface transport path approximation on both synthetic datasets and real-world mouse gastrulation data. Ultimately, it establishes a novel paradigm for dynamic modeling in discrete state spaces.

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.

Problem

• population dynamics
• unpaired temporal marginals
• ill-posed inverse problem
• discrete Wasserstein geometry
• stochastic dynamics

Innovation

• Discrete Action Matching
• Discrete Wasserstein geometry
• Action minimization
• Markov sampler
• Population dynamics