Back

Action Matching: Learning Stochastic Dynamics from Samples

Kirill NeklyudovRob BrekelmansDaniel SeveroAlireza Makhzani

article2023ICML108 citations

Proposes Action Matching, a tractable framework for learning continuous population dynamics directly from uncorrelated snapshot samples across time without requiring optimal transport solvers or backpropagation through differential equations.

Listen

In many scientific and engineering domains—such as single-cell biology, quantum mechanics, and generative modeling—researchers need to reconstruct how continuous systems evolve over time. However, tracking individual particles or cells continuously is often physically impossible or destructive to the samples, leaving practitioners with only disconnected, cross-sectional snapshots taken at discrete time intervals. Existing approaches to infer continuous trajectories from such data often rely on restrictive analytical assumptions, struggle with scalability, or require computationally demanding simulations and backpropagation through complex solvers during model training.

The article demonstrates that continuous dynamic trajectories can be effectively learned from temporal snapshot data using a novel framework called Action Matching. The authors evaluate whether a single, tractable training objective can accurately recover continuous trajectories, simulate realistic physical and biological systems, and generate high-dimensional data without needing simulation during the training phase.

To achieve this, the authors mathematically formulate trajectory learning as identifying an optimal, curl-free gradient field that satisfies the physical continuity equation of density evolution. They prove that minimizing a practical training objective, which uses only sample batches across time, is mathematically equivalent to minimizing the kinetic energy gap between the model and the true underlying dynamics. The authors extend this foundational model into specialized variants: an entropic version that models stochastic diffusion processes (such as cellular Brownian motion), an unbalanced version that accounts for the creation and destruction of mass (such as cell division and death), and a generalized version supporting arbitrary convex cost functions. The framework was evaluated across synthetic benchmarks, real-world single-cell RNA sequencing data, quantum wave-function simulations, and image generation tasks using standard deep learning architectures.

The evaluation produced four key findings. First, in single-cell trajectory inference benchmarks, entropic Action Matching maintained robust tracking accuracy regardless of the number of intermediate time steps, outperforming existing flow-based models and matching the performance of leading specialized solvers. Second, in quantum simulations of an excited hydrogen atom, the method closely tracked the true physical evolution, achieving an average discrepancy roughly two orders of magnitude lower than standard score-based sampling baselines. Third, in generative modeling benchmarks on the CIFAR-10 image dataset, the approach successfully generated high-quality images and executed conditional super-resolution and colorization tasks without requiring knowledge of the ground-truth process. Finally, the framework achieved these results using substantially fewer function evaluations than traditional score-matching baselines—requiring 132 evaluations compared to 1,090—while eliminating the need for iterative differential equation solvers during training.

These findings indicate that Action Matching provides a computationally efficient, scalable alternative for trajectory inference and generative modeling. By removing the need to simulate entire trajectories during training, the framework substantially reduces training compute costs, accelerates experimentation cycles, and lowers the algorithmic complexity of modeling time-evolving systems. The extensions to unbalanced and stochastic dynamics make it especially relevant for biological and physical applications where mass variations and random motion are intrinsic.

Organizations analyzing snapshot population data or developing continuous-flow generative models should consider piloting Action Matching as a computationally lighter alternative to simulation-heavy methods. Implementation should leverage the entropic extension when underlying Brownian motion is present, or the unbalanced variant when modeling biological populations with cell proliferation and death. Practitioners should also adopt the authors' time-reweighting and importance-sampling strategies to ensure numerical stability when training on high-dimensional data distributions.

Confidence in the mathematical foundations and empirical performance of the framework is high across the tested domains. However, users should exercise caution when working with highly sparse or singular initial data distributions, where deterministic velocity fields can encounter numerical instabilities. In such cases, appropriate weighting schedules and time-sampling schemes remain essential prerequisites for stable training.

context-awarenessmeta-learningstochasticAction Matchingstochastic dynamicsstochastic interpolants
Neklyudov et al (2023).pdf

No sufficiently relevant recommendations were found.

Cover for Action Matching: Learning Stochastic Dynamics from Samples

Abstract

Machine learning offers several techniques where systems learn and improve automatically from experience without being clearly programmed. Some set ups which require very fast adaptation needs less training samples. In meta-learning, we focus on problems dealing with tasks where a learning algorithm has to quickly adapt itself with limited number of labelled samples to execute new tasks extracted from similar distribution. We introduce our approach to leverage large scale deep neural networks along with context-aware parameter generation mechanism using affine transformations of embeddings achieved through convolutional feature maps -> scaled Learned Init...

Topics

context-awarenessmeta-learningstochasticstochastic dynamicstraining samplematching algorithmmatching methodsonline learningconvolutional neural networksdeep learning

Keywords

Action Matchingstochastic dynamicsstochastic interpolantsstochastic controlstochastic differential equationsFlow Matchingvariational actionNeural SDEentropic action

Table of Contents

  • 1. Introduction
  • 2. Action Matching
  • 2.1. Continuity Equation
  • 2.2. Action Matching Loss
  • 2.3. Learning, Sampling, and Likelihood Evaluation
  • 3. Extensions of Action Matching
  • 3.1. Entropic Action Matching
  • 3.2. Unbalanced Action Matching
  • 3.3. Action Matching with Convex Costs
  • 4. Applications of Action Matching
  • 4.1. Population Dynamics in Biology
  • 4.2. Quantum System Simulation
  • 4.3. Generative Modeling
  • 4.3.1. Empirical Study for Generative Modeling
  • 5. Related Works
  • 5.1. Connection with Flow Matching Methods
  • 6. Conclusion
  • References
  • A. Proofs
  • A.1. Action Matching Proofs
  • A.2. Entropic Action Matching Proofs
  • A.3. Unbalanced Action Matching Proofs
  • B. Action Matching and Optimal Transport
  • B.1. Action Matching as Infinitesimal Optimal Transport
  • B.2. Action Matching with Lagrangian Costs
  • B.3. Entropic Action Matching and Entropy-Regularized Optimal Transport
  • B.4. Unbalanced Action Matching and Unbalanced Optimal Transport
  • B.5. Action Matching and Wasserstein Gradient Flows
  • C. Generative Modeling in Practice
  • D. Sparse Data Regime
  • D.1. Delta Function Data Distribution
  • D.2. Mixture of Delta Functions Data Distribution
  • E. Experiments Details
  • E.1. Schrödinger Equation Simulation
  • E.2. Generative Modeling
  • E.3. Unbalanced Action Matching
  • F. Generated Images

Knowls

  1. Knowl 1 — Action Matching learns a sample-based potential for a prescribed marginal path

    model/method

    Let qtq_t be a time-indexed path of distributions on X⊆Rd\mathcal X\subseteq\mathbb R^d, for t∈[0,1]t\in[0,1], and let st(x)s_t(x) be a differentiable scalar potential. Action Matching (AM) trains sts_t by minimizing

    LAM(s)=Ex∼q0[s0(x)]−Ex∼q1[s1(x)]+∫01Ex∼qt[12∥∇xst(x)∥2+∂tst(x)]dt.L_{\mathrm{AM}}(s)=\mathbb E_{x\sim q_0}[s_0(x)]-\mathbb E_{x\sim q_1}[s_1(x)]+\int_0^1\mathbb E_{x\sim q_t}\left[\frac12\|\nabla_x s_t(x)\|^2+\partial_t s_t(x)\right]dt.

    For the ground-truth potential st∗s_t^*, the squared velocity error (the action gap) equals this objective plus a constant independent of the candidate ss. Thus minimizing LAML_{\mathrm{AM}} recovers the gradient field ∇xst∗\nabla_x s_t^* using samples from the endpoint and intermediate marginals, without requiring trajectories, tractable densities, simulation of the learned dynamics during training, or backpropagation through an ODE. In practice, each update samples endpoint batches from q0q_0 and q1q_1, samples tt uniformly on [0,1][0,1] and intermediate points from qtq_t, estimates the displayed objective, and updates the parameters of sts_t by gradient descent. The learned field generates samples by integrating x˙(t)=∇xst(x(t))\dot x(t)=\nabla_xs_t(x(t)) from x(0)∼q0x(0)\sim q_0.

  2. Knowl 2 — Every sufficiently regular marginal path has a unique gradient-field velocity

    theoretical result

    For a path qtq_t that is absolutely continuous in the 2-Wasserstein space P2(X)\mathcal P_2(\mathcal X) and satisfies the paper's mild regularity conditions, there is a unique velocity field of the form vt∗(x)=∇xst∗(x)v_t^*(x)=\nabla_xs_t^*(x) that transports the path through the continuity equation

    ∂tqt=−∇⋅(qt∇xst∗).\partial_t q_t=-\nabla\cdot(q_t\nabla_xs_t^*).

    The potential st∗s_t^* is unique up to an additive spatial constant. Consequently, representing dynamics with gradient fields does not restrict which such marginal paths can be modeled: the ODE x˙(t)=∇xst∗(x(t))\dot x(t)=\nabla_xs_t^*(x(t)) has marginals qtq_t. The claim concerns matching the distributional evolution; it does not assert that this deterministic ODE recovers the original individual trajectories when the observed path arose from stochastic dynamics.

  3. Knowl 3 — Action Matching objective controls error in simulated marginals

    theoretical result

    Let qtq_t be the true path generated by ∇xst∗\nabla_xs_t^*, and let q^t\hat q_t be generated from the same initial distribution by a learned field ∇xst\nabla_xs_t. Assume the learned field is continuously differentiable in time and state and is KK-Lipschitz in state uniformly over t∈[0,1]t\in[0,1]. For every τ∈[0,1]\tau\in[0,1], the squared 2-Wasserstein distance between the simulated and true marginals obeys

    W22(q^τ,qτ)≤e(1+2K)τ∫0τEx∼qt[∥∇xst(x)−∇xst∗(x)∥2]dt.W_2^2(\hat q_\tau,q_\tau)\leq e^{(1+2K)\tau}\int_0^\tau\mathbb E_{x\sim q_t}\left[\|\nabla_xs_t(x)-\nabla_xs_t^*(x)\|^2\right]dt.

    Here W2W_2 is the 2-Wasserstein distance, and the integral is the action gap accumulated up to time τ\tau. The result links accuracy of the learned vector field, measured under the true marginals, to accuracy of the distributions produced by integrating that field.

  4. Knowl 4 — Entropic Action Matching represents diffusion paths with a known noise schedule

    model/method

    For a prescribed diffusion coefficient σt\sigma_t and a marginal path qtq_t, entropic Action Matching (eAM) uses an entropic potential s~t∗\tilde s_t^* whose gradient is the drift in the Fokker–Planck equation

    ∂tqt=−∇⋅(qt∇xs~t∗)+σt22Δqt.\partial_tq_t=-\nabla\cdot(q_t\nabla_x\tilde s_t^*)+\frac{\sigma_t^2}{2}\Delta q_t.

    Thus samples can be propagated with the SDE dx(t)=∇xs~t∗(x(t))dt+σtdWtdx(t)=\nabla_x\tilde s_t^*(x(t))dt+\sigma_t dW_t, where WtW_t is a Wiener process. Under the stated regularity and boundary conditions, the entropic potential is unique up to an additive constant, even when the observed marginal path was not generated by this particular SDE. The tractable training objective is

    LeAM(s)=Eq0[s0]−Eq1[s1]+∫01Eqt[12∥∇xst∥2+∂tst+σt22Δst]dt.L_{\mathrm{eAM}}(s)=\mathbb E_{q_0}[s_0]-\mathbb E_{q_1}[s_1]+\int_0^1\mathbb E_{q_t}\left[\frac12\|\nabla_xs_t\|^2+\partial_ts_t+\frac{\sigma_t^2}{2}\Delta s_t\right]dt.

    It differs from standard AM by the diffusion correction involving the Laplacian and the known schedule σt\sigma_t. Minimizing it is equivalent, up to a candidate-independent constant, to minimizing the squared error between ∇xst\nabla_xs_t and the entropic target drift.

  5. Knowl 5 — Unbalanced Action Matching couples transport and mass growth through one potential

    model/method

    Unbalanced Action Matching (uAM) models changing total probability mass by transporting weighted particles and changing each particle's weight. If vtv_t is the velocity, gtg_t is the growth rate, and w(t)w(t) is a particle's importance weight, then x˙(t)=vt(x(t))\dot x(t)=v_t(x(t)) and ddtlog⁡w(t)=gt(x(t))\frac{d}{dt}\log w(t)=g_t(x(t)). The resulting density satisfies

    ∂tqt=−∇⋅(qtvt)+qtgt.\partial_tq_t=-\nabla\cdot(q_tv_t)+q_tg_t.

    For a given marginal path, the paper's unbalanced formulation uses a single potential s^t∗\hat s_t^* for both quantities: vt∗=∇xs^t∗v_t^*=\nabla_x\hat s_t^* and gt∗=s^t∗g_t^*=\hat s_t^*. Its tractable objective is

    LuAM(s)=Eq0[s0]−Eq1[s1]+∫01Eqt[12∥∇xst∥2+∂tst+12st2]dt.L_{\mathrm{uAM}}(s)=\mathbb E_{q_0}[s_0]-\mathbb E_{q_1}[s_1]+\int_0^1\mathbb E_{q_t}\left[\frac12\|\nabla_xs_t\|^2+\partial_ts_t+\frac12s_t^2\right]dt.

    Minimizing this objective fits both the transport field and the mass-growth rate without observing particle lineages. The paper also derives a scaled version in which the final term is λ2st2\frac{\lambda}{2}s_t^2 for a growth-cost weight λ\lambda.

  6. Knowl 6 — Convex-cost Action Matching generalizes the kinetic-energy objective

    model/method

    Let c(v)c(v) be a strictly convex cost on velocities v∈Rdv\in\mathbb R^d, and let its convex conjugate be c∗(a)=sup⁡v{⟨v,a⟩−c(v)}c^*(a)=\sup_v\{\langle v,a\rangle-c(v)\}. For a fixed marginal path qtq_t satisfying the continuity equation, the velocity field that minimizes ∫01Eqt[c(vt)]dt\int_0^1\mathbb E_{q_t}[c(v_t)]dt has the form vt∗=∇xc∗(∇xsˉt∗)v_t^*=\nabla_xc^*(\nabla_x\bar s_t^*). The corresponding tractable objective for a candidate potential sts_t is

    LcAM(s)=Eq0[s0]−Eq1[s1]+∫01Eqt[c∗(∇xst)+∂tst]dt.L_{c\mathrm{AM}}(s)=\mathbb E_{q_0}[s_0]-\mathbb E_{q_1}[s_1]+\int_0^1\mathbb E_{q_t}\left[c^*(\nabla_xs_t)+\partial_ts_t\right]dt.

    Sampling follows the velocity x˙(t)=∇xc∗(∇xst(x(t)))\dot x(t)=\nabla_xc^*(\nabla_xs_t(x(t))). For the squared Euclidean cost c(v)=12∥v∥2c(v)=\frac12\|v\|^2, its conjugate is the same quadratic, recovering standard Action Matching. For other strictly convex costs, the discrepancy from the optimal field is measured by a Bregman divergence generated by c∗c^*.

  7. Knowl 7 — Time weighting and adaptive time sampling address singular, high-variance training

    model/method

    For generative paths that expand a low-dimensional or discrete data distribution into a higher-dimensional distribution, the deterministic target velocity can become singular near an endpoint. The paper's practical AM procedure weights the action-gap integrand by a nonnegative time function ω(t)\omega(t) and samples training times from a proposal density p(t)p(t), correcting each sampled integrand by 1/p(t)1/p(t) so the estimate remains unbiased. With ζt(x)=12ω(t)∥∇xst(x)∥2+ω(t)∂tst(x)+st(x)∂tω(t)\zeta_t(x)=\frac12\omega(t)\|\nabla_xs_t(x)\|^2+\omega(t)\partial_ts_t(x)+s_t(x)\partial_t\omega(t), the weighted objective is

    Lω(s)=ω(0)Eq0[s0]−ω(1)Eq1[s1]+∫01Eqt[ζt(x)]dt.L_{\omega}(s)=\omega(0)\mathbb E_{q_0}[s_0]-\omega(1)\mathbb E_{q_1}[s_1]+\int_0^1\mathbb E_{q_t}[\zeta_t(x)]dt.

    The authors choose p(t)p(t) proportional to the standard deviation of ζt(x)\zeta_t(x) under x∼qtx\sim q_t, estimating and updating those variances with an exponential moving average and interpolating between estimates. For the example path xt=1−t x0+t ϵx_t=\sqrt{1-t}\,x_0+\sqrt t\,\epsilon, with x0x_0 drawn from a data distribution and ϵ\epsilon standard Gaussian, they use ω(t)=(1−t)t3/2\omega(t)=(1-t)t^{3/2} to cancel the endpoint singularities described in the paper.

  8. Knowl 8 — AM supports unconditional and conditional generation through chosen interpolant paths

    model/method

    For generative modeling, Action Matching can learn the dynamics of a path specified by interpolating a prior sample x0∼q0x_0\sim q_0 and a data sample x1∼q1x_1\sim q_1: xt=αt(x0)+βt(x1)x_t=\alpha_t(x_0)+\beta_t(x_1), where α0(x0)=x0\alpha_0(x_0)=x_0, β0(x1)=0\beta_0(x_1)=0, α1(x0)=0\alpha_1(x_0)=0, and β1(x1)=x1\beta_1(x_1)=x_1. This defines intermediate marginals implicitly through samples rather than through known densities. The paper uses standard-normal priors and illustrates three paths: unconditional generation with xt=(1−t)x0+tx1x_t=(1-t)x_0+tx_1; super-resolution with xt=m⊙x1+(1−m)⊙((1−t)x0+tx1)x_t=m\odot x_1+(1-m)\odot((1-t)x_0+tx_1), where the binary mask mm preserves observed pixels; and colorization with xt=(1−t)(10−1x0+gray⁡(x1))+tx1x_t=(1-t)(10^{-1}x_0+\operatorname{gray}(x_1))+tx_1, where gray⁡\operatorname{gray} converts an image to grayscale. The learned gradient field transports samples along the chosen path, so the same training principle applies to unconditional generation and tasks conditioned on observed image content.

  9. Knowl 9 — Entropic AM compares favorably on synthetic and embryoid-cell trajectories

    empirical result

    The biology experiments use entropic AM because the target cell dynamics include Brownian motion. On synthetic differentiation data with branching and merging, the authors train using 5, 10, or 15 observed time points. The page-6 plots compare generated and observed marginals using Wasserstein-2 distance and maximum mean discrepancy (MMD): eAM is below MIOFlow across the reported time-point settings, and its performance does not show the degradation with finer temporal sampling reported for MIOFlow. The page-6 trajectory visualization also shows eAM paths staying closer to the scattered data marginals.

    For real embryoid-body scRNA-seq data, the comparison is in a five-dimensional PCA representation, with data interpolated between observed time points to make training denser in time. The metric is W2(qti,q^ti)W_2(q_{t_i},\hat q_{t_i}), comparing test marginals with predictions propagated from the preceding test marginal; smaller is better. The concurrent-method values are reported as in the original comparison, while eAM is reported with mean and standard deviation.

    Distance OT-flow Trajectory-Net IPF Neural SDE NLSB eAM (ours)
    W2(qt1,q^t1)W_2(q_{t_1},\hat q_{t_1}) 0.75 0.64 0.65±0.0160.65\pm0.016 0.62±0.0160.62\pm0.016 0.63±0.0150.63\pm0.015 0.58±0.0150.58\pm0.015
    W2(qt2,q^t2)W_2(q_{t_2},\hat q_{t_2}) 0.93 0.87 0.78±0.0210.78\pm0.021 0.78±0.0210.78\pm0.021 0.75±0.0170.75\pm0.017 0.77±0.0160.77\pm0.016
    W2(qt3,q^t3)W_2(q_{t_3},\hat q_{t_3}) 0.93 0.78 0.76±0.0180.76\pm0.018 0.77±0.0170.77\pm0.017 0.75±0.0140.75\pm0.014 0.72±0.0070.72\pm0.007
    W2(qt4,q^t4)W_2(q_{t_4},\hat q_{t_4}) 0.88 0.89 0.76±0.0170.76\pm0.017 0.75±0.0170.75\pm0.017 0.72±0.0100.72\pm0.010 0.74±0.0170.74\pm0.017

    The real-data results put eAM on par with NLSB overall and show lower error than the other listed methods at three of the four test marginals.

  10. Knowl 10 — AM reproduces quantum-system marginals more accurately than score-based sampling

    empirical result

    The quantum experiment learns the evolving position distribution qt(x)=∣ψ(x,t)∣2q_t(x)=|\psi(x,t)|^2 for an excited hydrogen-atom state governed by the Schrödinger equation. Performance is average MMD over time between generated samples and the simulated data marginals; lower is better. Annealed Langevin Dynamics (ALD) is tested with learned score estimates from score matching (SM), sliced score matching, and the ground-truth scores. AM uses the same reported experimental setup and substantially lowers average MMD. The page-28 convergence plots show that SM and sliced SM can estimate marginal scores accurately, yet ALD sampling remains less faithful to the evolving distributions.

    Method Average MMD ↓\downarrow
    AM (ours) 5.7⋅10−4±3.1⋅10−45.7\cdot10^{-4}\pm3.1\cdot10^{-4}
    ALD + SM 4.8⋅10−2±4.8⋅10−34.8\cdot10^{-2}\pm4.8\cdot10^{-3}
    ALD + Sliced SM 4.7⋅10−2±4.0⋅10−34.7\cdot10^{-2}\pm4.0\cdot10^{-3}
    ALD + True Scores 3.6⋅10−2±4.1⋅10−43.6\cdot10^{-2}\pm4.1\cdot10^{-4}

    The Schrödinger simulation runs for 14⋅10314\cdot10^3 time units and collects 1,000 time steps. Models use a five-layer MLP with 256 hidden units per layer; MMD is averaged over ten evaluation times. Even ALD supplied with true scores performs far worse than AM, which the authors attribute to the sampling approximation used by ALD.

  11. Knowl 11 — CIFAR-10 results show competitive AM generation with fewer evaluations than the score baseline

    empirical result

    On CIFAR-10, the authors evaluate unconditional generation, super-resolution, and colorization. Models use the same architecture and are trained for 500,000 iterations. FID and IS are evaluated on 50,000 generated images; BPD is reported when available, and NFE is the number of function evaluations. VP-SDE and Flow Matching use additional knowledge of the generating dynamics, unlike AM. The results show AM substantially improves on the ALD plus sliced-score-matching baseline for unconditional generation with fewer evaluations; its FID and IS are less favorable than VP-SDE and Flow Matching, although the page-9 generated-image comparison is described as showing only minor qualitative differences between AM and VP-SDE. For conditional generation the score baseline does not produce meaningful results, while AM produces reported super-resolution and colorization results.

    Model BPD ↓\downarrow FID ↓\downarrow IS ↑\uparrow NFE ↓\downarrow
    VP-SDE (uses extra information) 3.25 3.71 9.12 199
    Flow Matching (uses extra information) 2.99 6.35 −- 142
    Baseline (ALD + SSM) −- 86.29 5.43 1090
    AM - generation (ours) 3.54 10.04 8.60 132
    Baseline (ALD + SSM), super-resolution −- n/a n/a n/a
    AM - superres (ours) −- 1.44 10.93 166
    AM - color (ours) −- 2.47 9.88 89

    The table's conditional baseline row records that the score-based method did not yield evaluable results for super-resolution; the reported dashes and n/a entries are retained as given.

Coverage note — The optimal-transport, flow-matching, and Wasserstein-gradient-flow connections are omitted because they primarily interpret or compare the method rather than add a separate core procedure or result; detailed proofs and the standard likelihood-evaluation formula are also omitted as supporting material.

References

  1. 1.Albergo, M. S. and Vanden-Eijnden, E. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571, 2022.
  2. 2.Ambrosio, L., Gigli, N., and Savaré, G. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008.
  3. 3.Banerjee, A., Merugu, S., Dhillon, I. S., Ghosh, J., and Lafferty, J. Clustering with bregman divergences. Journal of machine learning research, 6(10), 2005.
  4. 4.Baradat, A. and Lavenant, H. Regularized unbalanced optimal transport as entropy minimization with respect to branching brownian motion. arXiv preprint arXiv:2111.01666, 2021.
  5. 5.Benamou, J.-D. and Brenier, Y. A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375–393, 2000.
  6. 6.Brenier, Y. Décomposition polaire et réarrangement monotone des champs de vecteurs. CR Acad. Sci. Paris Sér. I Math., 305:805–808, 1987.
  7. 7.Buenrostro, J. D., Wu, B., Litzenburger, U. M., Ruff, D., Gonzales, M. L., Snyder, M. P., Chang, H. Y., and Greenleaf, W. J. Single-cell chromatin accessibility reveals principles of regulatory variation. Nature, 523(7561):486–490, 2015.
  8. 8.Bunne, C., Stark, S. G., Gut, G., del Castillo, J. S., Lehmann, K.-V., Pelkmans, L., Krause, A., and Rätsch, G. Learning single-cell perturbation responses using neural optimal transport. bioRxiv, 2021.
  9. 9.Bunne, C., Papaxanthos, L., Krause, A., and Cuturi, M. Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 6511–6528. PMLR, 2022.
  10. 10.Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
  11. 11.Chen, Y., Georgiou, T. T., and Pavon, M. On the relation between optimal transport and Schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 169(2):671–691, 2016.
  12. 12.Chen, Y., Georgiou, T. T., and Pavon, M. Stochastic control liaisons: Richard sinkhorn meets gaspard monge on a schrodinger bridge. SIAM Review, 63(2):249–313, 2021.
  13. 13.Chizat, L., Peyré, G., Schmitzer, B., and Vialard, F.-X. An interpolating distance between optimal transport and fisher–rao metrics. Foundations of Computational Mathematics, 18(1):1–44, 2018a.
  14. 14.Chizat, L., Peyré, G., Schmitzer, B., and Vialard, F.-X. Scaling algorithms for unbalanced optimal transport problems. Mathematics of Computation, 87(314):2563–2609, 2018b.
  15. 15.Chizat, L., Peyré, G., Schmitzer, B., and Vialard, F.-X. Unbalanced optimal transport: Dynamic and Kantorovich formulations. Journal of Functional Analysis, 274(11):3090–3123, 2018c.
  16. 16.Chizat, L., Zhang, S., Heitz, M., and Schiebinger, G. Trajectory inference via mean-field langevin in path space. arXiv preprint arXiv:2205.07146, 2022.
  17. 17.Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
  18. 18.Cuturi, M., Meng-Papaxanthos, L., Tian, Y., Bunne, C., Davis, G., and Teboul, O. Optimal transport tools (ott): A jax toolbox for all things Wasserstein. arXiv preprint arXiv:2201.12324, 2022. URL https://github.com/ott-jax/ott.
  19. 19.De Bortoli, V., Thornton, J., Heng, J., and Doucet, A. Diffusion Schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695–17709, 2021.
  20. 20.Figalli, A. and Glaudo, F. An Invitation to Optimal Transport, Wasserstein Distances, and Gradient Flows. 2021.
  21. 21.Finlay, C., Jacobsen, J.-H., Nurbekyan, L., and Oberman, A. How to train your neural ode: the world of jacobian and kinetic regularization. In International conference on machine learning, pp. 3154–3164. PMLR, 2020.
  22. 22.Gangbo, W. and McCann, R. J. The geometry of optimal transportation. Acta Mathematica, 177(2):113–161, 1996.
  23. 23.Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723–773, 2012.
  24. 24.Griffiths, D. J. and Schroeter, D. F. Introduction to quantum mechanics. Cambridge university press, 2018.
  25. 25.Hashimoto, T., Gifford, D., and Jaakkola, T. Learning population-level diffusions with generative rnns. In International Conference on Machine Learning, pp. 2417–2426. PMLR, 2016.
  26. 26.Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
  27. 27.Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
  28. 28.Huguet, G., Magruder, D. S., Fasina, O., Tong, A., Kuchroo, M., Wolf, G., and Krishnaswamy, S. Manifold interpolating optimal-transport flows for trajectory inference. arXiv preprint arXiv:2206.14928, 2022.
  29. 29.Hyvärinen, A. and Dayan, P. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
  30. 30.Klein, A. M., Mazutis, L., Akartuna, I., Tallapragada, N., Veres, A., Li, V., Peshkin, L., Weitz, D. A., and Kirschner, M. W. Droplet barcoding for single-cell transcriptomics applied to embryonic stem cells. Cell, 161(5):1187–1201, 2015.
  31. 31.Kondratyev, S., Monsaingeon, L., and Vorotnikov, D. A new optimal transport distance on the space of finite radon measures. Advances in Differential Equations, 21(11/12):1117–1164, 2016.
  32. 32.Koshizuka, T. and Sato, I. Neural lagrangian Schrödinger bridge, 2022. URL https://arxiv.org/abs/2204.04853.
  33. 33.Lavenant, H., Zhang, S., Kim, Y.-H., and Schiebinger, G. Towards a mathematical theory of trajectory inference. arXiv preprint arXiv:2102.09204, 2021.
  34. 34.Léonard, C. A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete & Continuous Dynamical Systems, 34(4):1533, 2014.
  35. 35.Li, X., Wong, T.-K. L., Chen, R. T., and Duvenaud, D. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pp. 3870–3882. PMLR, 2020.
  36. 36.Liero, M., Mielke, A., and Savaré, G. Optimal transport in competition with reaction: The Hellinger–Kantorovich distance and geodesic curves. SIAM Journal on Mathematical Analysis, 48(4):2869–2911, 2016.
  37. 37.Liero, M., Mielke, A., and Savaré, G. Optimal entropy-transport problems and a new Hellinger–Kantorovich distance between positive measures. Inventiones mathematicae, 211(3):969–1117, 2018.
  38. 38.Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
  39. 39.Liu, Q. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577, 2022.
  40. 40.Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003, 2022.
  41. 41.Liu, X., Wu, L., Ye, M., et al. Learning diffusion bridges on constrained domains. In The Eleventh International Conference on Learning Representations, 2023.
  42. 42.Lübeck, F., Bunne, C., Gut, G., del Castillo, J. S., Pelkmans, L., and Alvarez-Melis, D. Neural unbalanced optimal transport via cycle-consistent semi-couplings. arXiv preprint arXiv:2209.15621, 2022.
  43. 43.Macosko, E. Z., Basu, A., Satija, R., Nemesh, J., Shekhar, K., Goldman, M., Tirosh, I., Bialas, A. R., Kamitaki, N., Martersteck, E. M., et al. Highly parallel genome-wide expression profiling of individual cells using nanoliter droplets. Cell, 161(5):1202–1214, 2015.
  44. 44.Mikami, T. and Thieullen, M. Duality theorem for the stochastic optimal control problem. Stochastic processes and their applications, 116(12):1815–1835, 2006.
  45. 45.Mikhailov, V. Partial differential equations. 1976.
  46. 46.Moon, K. R., van Dijk, D., Wang, Z., Gigante, S., Burkhardt, D. B., Chen, W. S., Yim, K., Elzen, A. v. d., Hirn, M. J., Coifman, R. R., et al. Visualizing structure and transitions in high-dimensional biological data. Nature biotechnology, 37(12):1482–1492, 2019.
  47. 47.Noé, F., Tkatchenko, A., Müller, K.-R., and Clementi, C. Machine learning for molecular simulation. Annual review of physical chemistry, 71:361–390, 2020.
  48. 48.Onken, D., Fung, S. W., Li, X., and Ruthotto, L. Ot-flow: Fast and accurate continuous normalizing flows via optimal transport. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9223–9232, 2021.
  49. 49.Otto, F. The geometry of dissipative evolution equations: the porous medium equation. 2001.
  50. 50.Peluchetti, S. Diffusion bridge mixture transports, schr" odinger bridge problems and generative modeling. arXiv preprint arXiv:2304.00917, 2023.
  51. 51.Pooladian, A.-A., Ben-Hamu, H., Domingo-Enrich, C., Amos, B., Lipman, Y., and Chen, R. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304.14772, 2023.
  52. 52.Ronneberger, O., Fischer, P., and Brox, T. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention, pp. 234–241. Springer, 2015.
  53. 53.Salimans, T. and Ho, J. Should ebms model the energy or the score? In Energy Based Models Workshop-ICLR 2021, 2021.
  54. 54.Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
  55. 55.Schiebinger, G. Reconstructing developmental landscapes and trajectories from single-cell data. Current Opinion in Systems Biology, 27:100351, 2021.
  56. 56.Schiebinger, G., Shu, J., Tabaka, M., Cleary, B., Subramanian, V., Solomon, A., Gould, J., Liu, S., Lin, S., Berube, P., et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928–943, 2019. URL https://www.cell.com/cms/10.1016/j.cell.2019.01.006/attachment/d251b88d-a356-436a-a9a7-b7e04b56b41f/mmc8.pdf.
  57. 57.Shi, Y., De Bortoli, V., Campbell, A., and Doucet, A. Diffusion schr" odinger bridge matching. arXiv preprint arXiv:2303.16852, 2023.
  58. 58.Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
  59. 59.Song, Y., Garg, S., Shi, J., and Ermon, S. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence, pp. 574–584. PMLR, 2020a.
  60. 60.Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020b.
  61. 61.Tong, A., Huang, J., Wolf, G., Van Dijk, D., and Krishnaswamy, S. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pp. 9526–9536. PMLR, 2020.
  62. 62.Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Fatras, K., Wolf, G., and Bengio, Y. Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint arXiv:2302.00482, 2023.
  63. 63.Vázquez, J. L. The porous medium equation: mathematical theory. Oxford University Press on Demand, 2007.
  64. 64.Villani, C. Optimal transport: old and new, volume 338. Springer, 2009.
  65. 65.Yang, L. and Karniadakis, G. E. Potential flow generator with l2 optimal transport regularity for generative models. IEEE Trans. Neural Networks Learn. Syst., 33(2):528–538, 2020. doi: 10.1109/TNNLS.2020.3028042. URL https://doi.org/10.1109/TNNLS.2020.3028042.

Citation

MLA
Neklyudov, K., et al. “Action Matching: Learning Stochastic Dynamics from Samples”. International Conference on Machine Learning, vol. 202, 2023, pp. 25858–89, https://proceedings.mlr.press/v202/neklyudov23a.html.
APA
Neklyudov, K., Brekelmans, R., Severo, D., & Makhzani, A. (2023). Action Matching: Learning Stochastic Dynamics from Samples. International Conference on Machine Learning, 202, 25858–25889. https://proceedings.mlr.press/v202/neklyudov23a.html
Chicago
Neklyudov, K., R. Brekelmans, D. Severo, and A. Makhzani. 2023. “Action Matching: Learning Stochastic Dynamics from Samples”. International Conference on Machine Learning 202: 25858–89. https://proceedings.mlr.press/v202/neklyudov23a.html.
Harvard
Neklyudov, K. et al. (2023) “Action Matching: Learning Stochastic Dynamics from Samples”, International Conference on Machine Learning. PMLR, pp. 25858–25889. Available at: https://proceedings.mlr.press/v202/neklyudov23a.html.
Vancouver
1. Neklyudov K, Brekelmans R, Severo D, Makhzani A (2023) Action Matching: Learning Stochastic Dynamics from Samples. In: International Conference on Machine Learning. PMLR, pp 25858–25889

BibTeX

@InProceedings{pmlr-v202-neklyudov23a,
  title = 	 {Action Matching: Learning Stochastic Dynamics from Samples},
  author =       {Neklyudov, Kirill and Brekelmans, Rob and Severo, Daniel and Makhzani, Alireza},
  booktitle = 	 {Proceedings of the 40th International Conference on Machine Learning},
  pages = 	 {25858--25889},
  year = 	 {2023},
  editor = 	 {Krause, Andreas and Brunskill, Emma and Cho, Kyunghyun and Engelhardt, Barbara and Sabato, Sivan and Scarlett, Jonathan},
  volume = 	 {202},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {23--29 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://proceedings.mlr.press/v202/neklyudov23a/neklyudov23a.pdf},
  url = 	 {https://proceedings.mlr.press/v202/neklyudov23a.html},
  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.}
}
Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF
License: https://creativecommons.org/licenses/by/4.0/