Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme

Mikhail Persiianov, Jiawei Chen, Petr Mokrov, Alexander Tyurin, Evgeny Burnaev, Alexander Korotin

TL;DR

The paper tackles learning population dynamics from marginal distributions by embedding the learning problem in Wasserstein gradient flows via the JKO scheme. It introduces iJKOnet, an end-to-end adversarial framework that jointly learns a flexible energy functional ${\mathcal{J}}_\theta$ and transport maps without restrictive convexity priors, using an inverse-optimization objective that aligns JKO updates with observed marginals. A key theoretical result provides quality bounds for recovering the potential energy under mild smoothness and convexity assumptions, and experiments on synthetic and single-cell data demonstrate improved performance over prior JKO-based methods and competitive results against non-JKO baselines. The approach scales to higher dimensions and avoids precomputing OT couplings, offering a practical and principled tool for inferring governing dynamics from population-level observations. These contributions advance the applicability of energy-based dynamics modeling to biology, epidemiology, and beyond, where trajectory data are often unavailable.

Abstract

Learning population dynamics involves recovering the underlying process that governs particle evolution, given evolutionary snapshots of samples at discrete time points. Recent methods frame this as an energy minimization problem in probability space and leverage the celebrated JKO scheme for efficient time discretization. In this work, we introduce $\texttt{iJKOnet}$, an approach that combines the JKO framework with inverse optimization techniques to learn population dynamics. Our method relies on a conventional $\textit{end-to-end}$ adversarial training procedure and does not require restrictive architectural choices, e.g., input-convex neural networks. We establish theoretical guarantees for our methodology and demonstrate improved performance over prior JKO-based methods.

Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme

TL;DR

The paper tackles learning population dynamics from marginal distributions by embedding the learning problem in Wasserstein gradient flows via the JKO scheme. It introduces iJKOnet, an end-to-end adversarial framework that jointly learns a flexible energy functional and transport maps without restrictive convexity priors, using an inverse-optimization objective that aligns JKO updates with observed marginals. A key theoretical result provides quality bounds for recovering the potential energy under mild smoothness and convexity assumptions, and experiments on synthetic and single-cell data demonstrate improved performance over prior JKO-based methods and competitive results against non-JKO baselines. The approach scales to higher dimensions and avoids precomputing OT couplings, offering a practical and principled tool for inferring governing dynamics from population-level observations. These contributions advance the applicability of energy-based dynamics modeling to biology, epidemiology, and beyond, where trajectory data are often unavailable.

Abstract

Learning population dynamics involves recovering the underlying process that governs particle evolution, given evolutionary snapshots of samples at discrete time points. Recent methods frame this as an energy minimization problem in probability space and leverage the celebrated JKO scheme for efficient time discretization. In this work, we introduce , an approach that combines the JKO framework with inverse optimization techniques to learn population dynamics. Our method relies on a conventional adversarial training procedure and does not require restrictive architectural choices, e.g., input-convex neural networks. We establish theoretical guarantees for our methodology and demonstrate improved performance over prior JKO-based methods.

Paper Structure

This paper contains 35 sections, 2 theorems, 38 equations, 8 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. iJKOnet Method
  4. Problem statement
  5. Methodology and loss derivation
  6. Practical aspects: method parametrization and learning procedure
  7. Theoretical aspects
  8. Related Works
  9. Solving aggregation-diffusion equation with given energy
  10. Learning energy for aggregation-diffusion equation via JKO scheme
  11. Experiments
  12. Learning potential energy
  13. Learning population dynamics for single-cell data
  14. Discussion
  15. Extended Discussion
  16. ...and 20 more sections

Key Result

Theorem 3.1

Let $\varepsilon(V) \stackrel{\text{def}}{=} {\mathcal{L}}(V^\ast, T_{V^\ast}) - {\mathcal{L}}(V, T_V)$ be the gap between the optimal and optimized value of inverse JKO loss eq:loss with internal $\min_T$ problem solved exactly, i.e., $T_V \stackrel{\text{def}}{=} \min_{T} {\mathcal{L}}(V, T)$. Let

Figures (8)

  • Figure 1: iJKOnet working scheme: our method minimizes the gap between the optimal values of parametric ($\theta$) JKO functional and suboptimal values obtained at ground truth population measures.
  • Figure 2: Numerical results from \ref{['subsec:learning-potential-energy']} for the unpaired setup. The reported absolute values show that, while increasing the number of samples generally improves performance across metrics, certain potentials remain challenging, highlighting the difficulty of this setup.
  • Figure 3: Level curves of the true (green) and estimated (blue) potentials for the paired setup, following terpin2024learning. These results can be directly compared with those in terpin2024learning. Note that for the flat potential, the value range is near zero, as expected for the ground-truth potential.
  • Figure 5: Estimation error of the diffusion coefficient $\theta_3$ relative to the ground-truth values $\beta^\ast$. Blue bars correspond to $\texttt{iJKOnet}$ and orange bars to $\texttt{JKOnet}^\ast$. The Y-axis indicates the absolute deviation between the estimated $\theta_3$ and true $\beta^\ast$ values.
  • Figure : (a) Ground-truth ${\mathcal{W}}^\ast$
  • ...and 3 more figures

Theorems & Definitions (4)

  • Theorem 3.1: Quality bounds for recovered potential energy
  • Lemma D.1: Solution to the JKO problem with potential energy is unique
  • proof
  • proof