Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme

TL;DR

This paper introduces iJKOnet, a novel method combining inverse optimization with the JKO scheme to learn population dynamics from snapshots, offering theoretical guarantees and improved performance without restrictive neural network architectures.

Contribution

The paper presents iJKOnet, an end-to-end adversarial training approach that integrates inverse optimization with the JKO scheme, eliminating the need for input-convex neural networks and providing theoretical guarantees.

Findings

Demonstrates improved performance over prior JKO-based methods

Provides theoretical guarantees for the proposed approach

Does not require restrictive neural network architectures

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. The code of $\texttt{iJKOnet}$ is available at https://github.com/MuXauJl11110/iJKOnet.