Learning diffusion at lightspeed
This addresses the need for efficient and accurate diffusion learning in generative models and natural processes, offering a novel method that bypasses complexity while enhancing capabilities, though it appears incremental as it builds on the JKO scheme.
The paper tackled the problem of learning diffusion terms from observational data, which existing models handle with complex bilevel optimization and only model drift, by proposing JKOnet*, a simple model that recovers potential, interaction, and internal energy components, achieving state-of-the-art accuracy at a fraction of the computational cost in real-world cellular process predictions.
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.