Research map / Probabilistic and scientific modeling / Molecules and proteins
Learning diffusion at lightspeed
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Learns the energy functional (potential, interaction, internal) governing diffusion from population snapshots by replacing the JKO bilevel problem with first-order optimality conditions and a simple quadratic loss, with a closed-form solution for linear parametrizations.
Problem
Learning the energy functional driving diffusion processes from population snapshot data, which existing JKO-based methods (JKOnet) solve via a computationally expensive infinite-dimensional bilevel optimization and only for potential energies.
Approach
Replace the JKO step with its first-order necessary optimality condition evaluated on optimal transport couplings computed beforehand, yielding a single-level quadratic loss over potential, interaction, and internal energy terms; linear parametrizations admit a closed-form least-squares solution.
Key results
- On the scRNA-seq embryoid body data, JKOnet*-time-varying reaches EMD 0.624 ± 0.007, better than JKOnet (1.363 ± 0.214) and the best listed prior method DMSB (0.67).
- JKOnet* trains in a few minutes on a single GPU while the listed methods take hours, per the authors' reported comparisons.
- Linear parametrization yields a closed-form optimum needing a single epoch, and the nonlinear JKOnet* has lower time per epoch and faster convergence than JKOnet.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use time-varying JKOnet_V* instead of JKOnet | Embryoid body scRNA-seq prediction · Embryoid body scRNA-seq | EMD: 1.363 ± 0.214 -> 0.624±0.007 | Helped |
| Use time-varying JKOnet_V* instead of JKOnet-vanilla | Embryoid body scRNA-seq prediction · Embryoid body scRNA-seq | EMD: 3.237 ± 1.135 -> 0.624±0.007 | Helped |
| Use time-varying JKOnet_V* instead of DMSB | Embryoid body scRNA-seq prediction · Embryoid body scRNA-seq | EMD: 0.67 ± -- -> 0.624±0.007 | Helped |
| Use an implicit time-varying loss instead of an explicit loss (ablation) | JKOnet* time-varying potential learning · Synthetic time-varying potential | potential recovery | Helped |
| Pair the implicit loss with implicit rather than explicit prediction (ablation) | JKOnet* time-varying potential learning · Synthetic time-varying potential | potential recovery | Helped |
| Compute optimal couplings with Sinkhorn-type algorithms instead of plain linear programming, when couplings are close to correct (ablation) | JKOnet* · Synthetic potential recovery | EMD | No clear effect |
Surprising findings
- Simply replacing the ICNN in JKOnet with a vanilla MLP causes training instability (NaN divergence), suggesting the optimal-transport theory connection matters for stability.
- An explicit prediction/loss scheme for time-varying potentials shifts predictions in time, requiring an implicit scheme instead.
Limitations
- Architecture remains vanilla MLP; image or high-capacity domains not investigated.
- Feature choice for linear parametrization not studied in detail.
- Requires optimal transport couplings computed beforehand, which scales with number of particles via linear programming.
Experimental setup
Datasets: synthetic particle populations generated by Euler-Maruyama diffusion embryoid body scRNA-seq (Moon et al.) PCA-processed
Benchmarks: synthetic potential energy recovery embryoid body scRNA-seq trajectory prediction
Metrics: one-step-ahead earth mover's distance time per epoch convergence ratio
Baselines: JKOnet JKOnet-vanilla TrajectoryNet Regularized CNF DSB I-CFM SB-CFM OT-CFM NLSB MIOFLOW DMSB
Builds on: JKOnet JKO scheme Monge gap Wasserstein gradient flow
Closest papers by meaning
- Towards Non-Asymptotic Convergence for Diffusion-Based Generative Models (ICLR 2024 · similarity 0.87)
- Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics (ICLR 2025 · similarity 0.87)
- Learning the Infinitesimal Generator of Stochastic Diffusion Processes (NeurIPS 2024 · similarity 0.87)
- Simple Guidance Mechanisms for Discrete Diffusion Models (ICLR 2025 · similarity 0.87)
- Neural Flow Diffusion Models: Learnable Forward Process for Improved Diffusion Modelling (NeurIPS 2024 · similarity 0.87)
- On scalable and efficient training of diffusion samplers (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.