Diffusion regulates a phenomenal number of 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
properly model only the drift of the system. We propose a new simple model,
JKOnet*, which bypasses altogether the complexity of existing architectures
while presenting significantly enhanced representational capacity: JKOnet*
recovers the potential, interaction, and internal energy components of the
underlying diffusion process. JKOnet* minimizes a simple quadratic loss, runs
at lightspeed, and drastically outperforms other baselines in practice.
Additionally, JKOnet* provides a closed-form optimal solution for linearly
parametrized functionals. 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, in light of few-weeks-old advancements in optimization in the
probability space.

JKOnet * 是一个简单的模型，能够学习扩散过程的漂流、潜力、相互作用和内部能量组成，通过简单的二次损失最小化，表现出更强的表示能力，优于其他基准模型。