Learning diffusion at lightspeed

TL;DR

This paper introduces JKOnet*, a simplified yet powerful model for learning diffusion processes that outperforms existing methods in efficiency and accuracy, especially in modeling cellular dynamics from real data.

Contribution

JKOnet* offers a new, simpler approach to learning diffusion processes, capturing multiple components and providing closed-form solutions, with superior performance and computational efficiency.

Findings

JKOnet* outperforms baselines in sample efficiency and accuracy.

It provides a closed-form solution for linear functionals.

Achieves state-of-the-art results in cellular process prediction.

Abstract

Diffusion regulates numerous natural processes and the dynamics of many successful generative models. Existing models to learn the diffusion terms from observational data rely on complex bilevel optimization problems and model only the drift of the system. We propose a new simple model, JKOnet*, which bypasses the complexity of existing architectures while presenting significantly enhanced representational capabilities: JKOnet* recovers the potential, interaction, and internal energy components of the underlying diffusion process. JKOnet* minimizes a simple quadratic loss and outperforms other baselines in terms of sample efficiency, computational complexity, and accuracy. Additionally, JKOnet* provides a closed-form optimal solution for linearly parametrized functionals, and, when applied to predict the evolution of cellular processes from real-world data, it achieves state-of-the-art accuracy at a fraction of the computational cost of all existing methods. Our methodology is based on the interpretation of diffusion processes as energy-minimizing trajectories in the probability space via the so-called JKO scheme, which we study via its first-order optimality conditions.