PaperlayerView exact revision on arXiv

Source-linked AI summary

Action Matching: Learning Stochastic Dynamics from Samples

Kirill Neklyudov, Rob Brekelmans, Daniel Severo, Alireza Makhzani

arXiv:2210.06662v3cs.LG

TL;DR

Action Matching addresses learning continuous dynamics from independent samples of temporal marginals when full trajectories are unavailable. It learns a tractable model from those samples and extends the approach to stochastic and mass-changing dynamics. The paper reports competitive performance across applications in biology, physics, and generative modeling.

  • Problem

    Learning continuous dynamics from independent samples of temporal marginals is important when full individual trajectories are expensive or impossible to observe.

    Source: §1 · p. 1

  • Method

    Action Matching learns a gradient field from temporal-marginal samples with a tractable objective, with extensions for stochastic dynamics and probability-mass creation or destruction.

    Source: §1 · pp. 1–2, §3.1 · p. 4

  • Results

    Action Matching demonstrates competitive performance in applications across biology, physics, and generative modeling.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The method can model continuous dynamics from samples along distributional paths without requiring tractable intermediate densities or knowledge of the underlying stochastic dynamics.

    Source: §1 · p. 2, §6 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

Action Matching addresses learning dynamics from independent samples at different times, when complete trajectories may be unavailable. It provides a tractable objective and extensions for stochastic and mass-changing dynamics, with applications across biology, physics, and generative modeling.

1. Introduction

  • Independent samples at different times can support models that simulate trajectories even when full trajectories are unavailable.The introduction notes that observing entire trajectories can be expensive or physically impossible, including in quantum measurement and cell sequencing settings.Source: §1 · p. 1
  • Action Matching learns a gradient field from marginal samples using a tractable objective, without requiring knowledge of the underlying dynamics.The paper also states that the method does not require tractable intermediate densities or back-propagation through differential equations or optimal transport solvers.Source: §1 · pp. 1–2
  • The entropic, unbalanced, and convex-cost extensions address stochastic dynamics, creation or destruction of probability mass, and alternative convex costs.For a given distributional path, AM learns deterministic trajectories, eAM stochastic trajectories, and uAM weighted trajectories.Source: §1 · p. 2, Figure 1 · p. 2
  • The paper demonstrates competitive performance in trajectory inference for biological data, quantum-system evolution, and generative modeling.The introduction attributes this breadth to learning from samples without requiring tractable intermediate densities or knowledge of the underlying stochastic dynamics.Source: §1 · p. 2

2. Action Matching

Action Matching represents density evolution through a continuity equation and learns its action from samples along the path. The learned field supports forward or backward simulation, while the paper also describes a bound connecting action-gap error to Wasserstein-2 distance.

2.1. Continuity Equation

  • Under mild conditions, the continuity equation describes density evolution, and a gradient field can represent any continuous density path considered by the theorem.The paper notes that this includes diffusion and porous-medium evolutions, so restricting to gradient fields does not sacrifice expressivity for modeling these marginals.Source: §2.1 · pp. 2–3

2.2. Action Matching Loss

  • Action Matching trains a variational action by minimizing a tractable objective derived from the action gap, using samples from the marginals.The sample-based objective admits an unbiased, low-variance estimate; its intractable constant is the path's kinetic energy.Source: §2.2 · p. 3
  • The Wasserstein-2 distance between true and simulated marginals can be upper-bounded in terms of the action gap.The bound is stated for simulations using the learned gradient field.Source: §2.2 · p. 3

2.3. Learning, Sampling, and Likelihood Evaluation

  • The learned gradient field simulates dynamics by solving an ODE forward in time; when the initial density is known, the paper also describes evaluating sample likelihoods.Likelihood evaluation uses the continuous change-of-variables formula and backward ODE integration.Source: §2.3 · pp. 3–4

3. Extensions of Action Matching 3

Action Matching extensions broaden the dynamics it can model: entropic AM handles diffusion, unbalanced AM handles probability-mass creation and destruction, and convex-cost AM supports alternative kinetic costs.

3. Extensions of Action Matching 3

  • The extensions cover stochasticity, probability-mass creation or destruction, and alternative kinetic-energy costs.The unbalanced and entropic extensions can also be combined.Source: §3 · p. 4, §3.2 · p. 5

3.1. Entropic Action Matching

  • Entropic AM can recover diffusion dynamics with gradient-field drift and known diffusion, while at optimality tracing any absolutely continuous distributional path under mild conditions.The choice between entropic and deterministic AM should depend on prior knowledge of the underlying dynamics.Source: §3.1 · p. 4

3.2. Unbalanced Action Matching

  • Unbalanced AM models probability-mass creation and destruction through a growth rate, and under mild conditions can represent any distributional path with an importance-weighted ODE.Its tractable objective is derived by decomposing the action gap into intractable constants and a tractable term.Source: §3.2 · pp. 4–5

3.3. Action Matching with Convex Costs

  • Convex-cost AM extends the objective to arbitrary strictly convex kinetic costs, with the density-tracing velocity given by the gradient of the cost’s convex conjugate applied to the action gradient.The quadratic cost recovers standard AM.Source: §3.3 · p. 5

4. Applications of Action Matching 4

The applications section studies Action Matching on population dynamics from biological data, quantum-system simulation, and generative modeling. Across these settings, it evaluates learned dynamics against observed marginals or established generative-model baselines.

4. Applications of Action Matching 4

  • The applications evaluate Action Matching on biological and physical datasets as well as generative modeling, where intermediate samples or densities may be unavailable.The section distinguishes dataset-based dynamics applications from generative modeling of a target distribution represented by samples.Source: §4 · p. 6

4.1. Population Dynamics in Biology

  • On synthetic cell-differentiation data, entropic Action Matching beats MIOFlow on Wasserstein-2 distance and MMD, and its performance does not degrade with finer time discretization.Its learned trajectories also stay closer to the data marginals; Figure 2 compares distribution distances and Figure 3 shows generated trajectories.Source: §4.1 · p. 6, Figure 2 · p. 6, Figure 3 · p. 6
  • On embryoid-body scRNA-seq data represented in five-dimensional PCA space, entropic Action Matching performs on par with NLSB and outperforms the other compared methods.Table 1 reports W2 distance between test marginals and predictions from previous test marginals.Source: §4.1 · p. 6, Table 1 · p. 7

4.2. Quantum System Simulation

  • For Schrödinger-equation simulation, Action Matching outperforms Score Matching and Sliced Score Matching and precisely simulates the true dynamics.Table 2 reports average MMD over time, including ALD with estimated or ground-truth scores; even ground-truth-score ALD does not match Action Matching.Source: §4.2 · p. 7, Table 2 · p. 7

4.3. Generative Modeling

  • For generative modeling, Action Matching learns a vector field for a chosen path from prior to data, supporting sampling and likelihood evaluation.The section considers unconditional generation, super-resolution, and conditional generation with coupled dimensions, including image colorization.Source: §4.3 · p. 7

4.3.1. EMPIRICAL STUDY FOR GENERATIVE MODELING

  • Action Matching performs much closer to VP-SDE and Flow Matching than to the ALD+SSM baseline, with fewer function evaluations; VP-SDE has better FID and IS, but generated images show minor qualitative differences.For conditional generation, ALD+SSM fails to learn meaningful scores.Source: §4.3.1 · p. 8, Table 3 · p. 8

5. Related Works

Action Matching differs from prior continuous-normalizing-flow and trajectory-inference methods by using a simulation-free objective evaluated at distributional endpoints. Compared with Flow Matching, it targets a gradient field and can be used when no conditional bridge vector field is available.

5. Related Works

  • Action Matching evaluates its objective at the endpoints without simulating the learned dynamics or backpropagating through differential equations, unlike prior continuous-normalizing-flow regularization methods.Its objective includes the action regularizer, and the authors describe the action and kinetic-energy optimizations as dual; adding both regularizers is redundant for this loss.Source: §5 · p. 8
  • Unlike trajectory-inference methods that optimize entropy-regularized unbalanced optimal-transport losses and backpropagate through OT solvers, Action Matching learns a time-dependent potential without simulating model samples during training.The cited approaches include continuous normalizing flows and JKONET, which uses a time-independent potential and learned optimal-transport maps.Source: §5 · p. 8

5.1. Connection with Flow Matching Methods

  • Although Flow Matching and Action Matching produce the same marginals, the AM target is a gradient field without a divergence-free component, whereas the FM target may include one.The FM target can reflect the endpoint distribution and bridge process; direct vector-field parameterization also saves a backward pass compared with differentiating a learned potential.Source: §5.1 · p. 9
  • Action Matching can use arbitrary marginal paths without a known conditional bridge vector field, making it applicable when linear interpolation is unsuitable or no bridge field is available.The paper applies it to scientific settings where dynamics are provided only as samples from temporal marginals.Source: §5.1 · p. 9

6. Conclusion

The conclusion presents Action Matching as a family of methods for learning continuous dynamics from samples along distributional paths. Its objective is reported as efficiently optimizable across applications, with competitive image-generative performance.

6. Conclusion

  • Action Matching learns continuous dynamics from samples along arbitrary absolutely continuous paths on the 2-Wasserstein space using a tractable objective based on intermediate-distribution samples.The method is extended to model stochastic state dynamics, create or destroy probability mass, and incorporate convex state-space costs.Source: §6 · p. 9
  • Its extensions address stochastic state dynamics, creation and destruction of probability mass, and convex costs on the state space.These are the entropic, unbalanced, and c-Action Matching variants, respectively.Source: §6 · p. 9
  • Action Matching is reported to be efficiently optimized across applications in the physical and natural sciences and to achieve competitive image-generative performance.The conclusion notes that the prior and density path can be chosen flexibly for the generative-modeling task.Source: §6 · p. 9

A.1. Action Matching Proofs

The appendix establishes the action representation and relates the tractable Action Matching objective to the action gap. It also gives a bound connecting that gap to the distance between learned and target marginals.

A.1. Action Matching Proofs

  • Under mild conditions, a continuous density evolution has a unique action up to a constant, whose gradient supplies a vector field for the dynamics.The existence argument uses an elliptic PDE with a Neumann boundary condition.Source: §A.1 · p. 13
  • The Action Matching objective is the tractable term in a decomposition of the action gap, which the method minimizes.The appendix also interprets minimizing this objective as maximizing a variational lower bound on kinetic energy.Source: §A.1 · pp. 13–14
  • The appendix bounds the squared Wasserstein-2 distance between learned and target marginals using the action gap.The stated bound applies to the dynamics induced by the learned gradient field under the proposition's regularity conditions.Source: §A.1 · p. 14

A.2. Entropic Action Matching Proofs

The entropic appendix defines an entropic action for dynamics with a specified diffusion term and gives a tractable objective for minimizing the entropic action gap. It interprets that objective as a variational lower bound on entropic kinetic energy.

A.2. Entropic Action Matching Proofs

  • For continuous dynamics with a specified σ_t, the appendix defines a unique entropic action up to a constant, with the density evolution satisfying the Fokker–Planck equation.The appendix includes a proof of existence and uniqueness for this action.Source: §A.2 · p. 15
  • The entropic Action Matching objective is the tractable term in a decomposition of the entropic action gap and is minimized.Minimizing it is interpreted as maximizing a variational lower bound on entropic kinetic energy.Source: §A.2 · pp. 15–16
  • The appendix notes that a boundary term can be dropped if the density q_t vanishes on the boundary.Source: §A.2 · p. 16

A.3. Unbalanced Action Matching Proofs

The unbalanced formulation learns dynamics that transport probability while also allowing mass to grow or shrink. Its objective is tractable up to constants and has an optimal-transport interpretation.

A.3. Unbalanced Action Matching Proofs

  • An unbalanced action and growth term satisfy an importance-weighted continuity equation for the evolving density.The stated result assumes a continuous dynamic whose density evolution satisfies mild conditions.Source: §A.3 · p. 17
  • The unbalanced Action Matching objective decomposes the action gap into intractable constants and a tractable term to minimize.Minimizing that term maximizes a variational lower bound on the sum of unbalanced kinetic and growth energies.Source: §A.3 · pp. 17–18

B. Action Matching and Optimal Transport

  • The unbalanced action links mass transport by a vector field with mass creation or destruction through a growth rate.The resulting unbalanced objective is minimized over a variational action.Source: B. Action Matching and Optimal Transport · p. 18
  • Action Matching, entropic Action Matching, and unbalanced Action Matching connect to local optimal-transport problems between infinitesimally close distributions.The paper also considers general convex costs and interprets Action Matching through a Wasserstein-2 gradient flow.Source: B. Action Matching and Optimal Transport · p. 18

B.1. Action Matching as Infinitesimal Optimal Transport

Action Matching can be understood as a local optimal-transport procedure along a supplied density path. The appendix extends this view to general costs, stochastic dynamics, and changing mass.

B.1. Action Matching as Infinitesimal Optimal Transport

  • The Action Matching objective selects a continuity-equation vector field that minimizes infinitesimal squared displacement, yielding a local transport map between neighboring densities.The supplied path is fixed and need not be the global optimal-transport path between its endpoints.Source: §B.1 · p. 18

B.2. Action Matching with Lagrangian Costs

  • For a strictly velocity-convex Lagrangian, the paper formulates a corresponding dynamical optimal-transport problem and derives an Action Matching objective.The standard squared Euclidean cost recovers the usual objective; the optimal value reflects infinitesimal squared Wasserstein-2 distance along the path.Source: §B.2 · p. 19, §B.1 · p. 19

B.3. Entropic Action Matching and Entropy-Regularized Optimal Transport

  • Entropic Action Matching uses a Fokker–Planck constraint to model stochastic dynamics, with the minimizing kinetic-energy vector field satisfying that constraint.This vector field is identified with the unique gradient field satisfying the Fokker–Planck equation for the fixed path.Source: §B.3 · pp. 21–22

B.4. Unbalanced Action Matching and Unbalanced Optimal Transport

  • Unbalanced Action Matching links transport by a vector field with mass creation or destruction through a growth rate, producing an objective to minimize over the variational action.At each time, the optimal objective reflects infinitesimal squared Wasserstein Fisher–Rao distance along an absolutely continuous curve.Source: §B.4 · pp. 22–23

B.5. Action Matching and Wasserstein Gradient Flows

The paper interprets learned dynamics as Wasserstein gradient flows and discusses practical adaptations for generative modeling. It also analyzes the sparse-data setting, including a single-point distribution.

B.5. Action Matching and Wasserstein Gradient Flows

  • Action Matching parameterizes dynamics as the gradient flow of a time-dependent linear functional on the Wasserstein-2 manifold.This contrasts with JKONET, which is limited to time-homogeneous linear functionals; the paper describes Action Matching as learning the functional traced by the given density path.Source: §B.5 · pp. 23–24

C. Generative Modeling in Practice

  • For complicated dynamics such as image generation, naive Action Matching may converge poorly because the objective estimate has high variance.The practical procedure uses a time-sampling distribution as an importance-sampling proposal; the authors report that choosing it proportional to integrand standard deviation reduces variance.Source: C. Generative Modeling in Practice · pp. 24–25
  • When the initial data distribution is concentrated near a low-dimensional manifold and the final distribution has higher intrinsic dimensionality, the deterministic ground-truth velocity field can become singular.The required velocity becomes very large, infinite in the limit, to spread the manifold into higher dimensions.Source: C. Generative Modeling in Practice · p. 24

D. Sparse Data Regime; D.1. Delta Function Data Distribution

  • In the single-point sparse-data case, the derived vector field is a gradient field and therefore is the unique field Action Matching would recover.The section studies continuity-equation fields for delta-function data and Gaussian conditional distributions.Source: §D.1 · p. 26, D. Sparse Data Regime · p. 26

D.2. Mixture of Delta Functions Data Distribution

The section gives a solution for dynamics with a mixture of delta-function data. Figure 5 describes evaluation by sample-distribution discrepancy and score-estimation error.

D.2. Mixture of Delta Functions Data Distribution

  • The passage concludes that expression (108) is indeed a solution.Source: §D.2 · p. 27
  • Figure 5 reports average MMD between generated samples and training data over time, alongside squared score-estimation error for score-based methods.The caption specifies the plotted quantities but not their values or which method performs best.Source: Figure 5 · p. 28

E. Experiments Details

The experiments describe a Schrödinger-equation simulation, image-generation setups, unbalanced mass changes, and generated-image examples.

E.1. Schrödinger Equation Simulation

  • The Schrödinger-equation dataset evolves the initial state for T = 14 · 10^3 time units and samples every 14 time units across 10^3 steps.Methods use a five-layer MLP with 256 hidden units per layer; Action Matching defines its vector field as the gradient of its scalar model.Source: §E.1 · p. 28
  • Although SM and SSM accurately recover ground-truth marginal-distribution scores, they cannot efficiently sample ground-truth dynamics; Figure 5 compares Action Matching and baselines using Average MMD over time.MMD compares generated samples from qt with training samples at 10 timesteps, t = k/10T for k = 1,...,10, and is averaged over those timesteps.Source: §E.1 · p. 28

E.2. Generative Modeling

  • The image-generation setup parameterizes Action Matching with a U-net-based scalar model and uses VP-SDE as the diffusion process.The baseline uses 10 sampling steps per distribution and 100 additional steps for the final distribution.Source: §E.2 · p. 29

E.3. Unbalanced Action Matching

  • Action Matching models probability-mass changes with a growth term instead of transporting particles between mixture modes.The toy demonstration starts from ground-truth samples and reweights particles over time.Source: §E.3 · p. 30

F. Generated Images

  • The generated-image examples show baseline and Action Matching outputs from noise, alongside Action Matching examples for upscaling and grayscale colorization.Figure 6 illustrates different noising processes interpolating from a prior at t = 0 to a target at t = 1.Source: Figure 6 · p. 30, Figure 8 · p. 31, Figure 9 · p. 31

Related papers

Loading 2210.06662v3…