Action Matching: Learning Stochastic Dynamics from Samples (2210.06662v3)

Abstract: Learning the continuous dynamics of a system from snapshots of its temporal marginals is a problem which appears throughout natural sciences and machine learning, including in quantum systems, single-cell biological data, and generative modeling. In these settings, we assume access to cross-sectional samples that are uncorrelated over time, rather than full trajectories of samples. In order to better understand the systems under observation, we would like to learn a model of the underlying process that allows us to propagate samples in time and thereby simulate entire individual trajectories. In this work, we propose Action Matching, a method for learning a rich family of dynamics using only independent samples from its time evolution. We derive a tractable training objective, which does not rely on explicit assumptions about the underlying dynamics and does not require back-propagation through differential equations or optimal transport solvers. Inspired by connections with optimal transport, we derive extensions of Action Matching to learn stochastic differential equations and dynamics involving creation and destruction of probability mass. Finally, we showcase applications of Action Matching by achieving competitive performance in a diverse set of experiments from biology, physics, and generative modeling.

• The paper introduces the Action Matching method to learn dynamics by minimizing the action gap from temporal marginals rather than complete trajectories.
• It leverages connections with optimal transport and extends to stochastic differential equations, including adaptations for entropic and unbalanced dynamics.
• Empirical results show AM’s superior performance in biological, quantum, and generative tasks, achieving lower Wasserstein distances and competitive scores.

Action Matching: Learning Stochastic Dynamics from Samples

This paper presents a method called Action Matching (AM), developed for learning continuous stochastic dynamics from the snapshots of a system's temporal marginals rather than from full trajectory data. Such scenarios arise in fields like quantum mechanics and single-cell biology, where observing individual trajectories is inherently challenging or unfeasible. The aim is to construct a model that can propagate samples over time to simulate entire trajectories.

Methodology

The core of the paper revolves around the notion of learning the underlying dynamics of a system through the minimization of the "action-gap"—the difference between a parameterized action and the true optimal action associated with the system dynamics. The authors propose a tractable training objective for AM that leverages samples from temporal marginals without requiring any explicit assumptions about an underlying dynamic model. Crucially, this objective does not necessitate backpropagation through differential equations or optimal transport solvers.

The method builds on connections to optimal transport theory, extending to scenarios involving stochastic differential equations (SDEs) and dynamics with probabilistic mass creation and destruction. The authors introduce extensions such as entropic AM (eAM) to tackle stochastic trajectories and unbalanced AM (uAM) to handle mass conservation issues, augmenting AM's versatility in modeling diverse dynamical systems.

Key Results

Empirical validations across various domains highlight AM's substantial efficacy and adaptability:

1. Biological Systems: AM outperforms traditional methods for trajectory inference in synthetic and real single-cell datasets. It offers lower Wasserstein distances between predicted and actual marginals than existing techniques, indicating better trajectory approximation.
2. Quantum Systems: In simulating the Schrödinger dynamics of quantum systems, AM demonstrates superior performance against score-based methods and even succeeds where exact score dynamics approximations (e.g., Annealed Langevin Dynamics) fail to converge.
3. Generative Modeling: When applied to generative tasks like image synthesis, AM exhibits comparable results to state-of-the-art models like Variance Preserving SDEs and achieves notable image quality as evidenced by competitive FID and IS scores. It effectively models dynamics from data to target distributions, showcasing applications in conditional generative tasks such as image super-resolution and colorization.

Implications and Future Directions

The authors position AM as a pivotal tool for scientific applications that require learning dynamics from temporal marginal distributions. By removing the need for complete trajectory data, AM opens new possibilities where data acquisition constraints exist. Moreover, given its extension capabilities, AM provides a robust framework for addressing dynamics with randomness and mass-variability features.

Looking ahead, AM can inspire further research in both theoretical and practical directions. Theoretically, exploring its connections with other dynamical systems' theory, such as Floquet dynamics, could offer deeper insights. Practically, implementations could be enhanced with potential action augmentation techniques. Moreover, since AM's applicability spans beyond classical dynamics, potential intersections with areas like reinforcement learning or climate modeling—where temporal dynamics play a critical role—invite exploration and could yield substantial advancements.

Overall, Action Matching introduces a significant paradigm for learning system dynamics under constraints, promising broad implications across various domains requiring dynamic modeling from limited data views.