Excited to share our latest preprint: ๐๐ฒ๐ฎ๐ฟ๐ป๐ถ๐ป๐ด ๐๐ฎ๐บ๐ถ๐น๐๐ผ๐ป๐ถ๐ฎ๐ป ๐๐น๐ผ๐ ๐ ๐ฎ๐ฝ๐: ๐ ๐ฒ๐ฎ๐ป ๐๐น๐ผ๐ ๐๐ผ๐ป๐๐ถ๐๐๐ฒ๐ป๐ฐ๐ ๐ณ๐ผ๐ฟ ๐๐ฎ๐ฟ๐ด๐ฒ-๐ง๐ถ๐บ๐ฒ๐๐๐ฒ๐ฝ ๐ ๐ผ๐น๐ฒ๐ฐ๐๐น๐ฎ๐ฟ ๐๐๐ป๐ฎ๐บ๐ถ๐ฐ๐ ๐
(7/n) Check out our paper and code:
Paper: arxiv.org/abs/2506.17139
Code + models: github.com/noegroup/Sco...
And also, our self-contained notebooks!
Colab (JAX):
colab.research.google.com/drive/1r3DGO...
Colab (PyTorch):
colab.research.google.com/drive/1rbcND...
#NeurIPS2025 #Diffusion #MD
(6/n)
Done with a brilliant team: Hao Wu, Leon Klein, Stephan Gรผnnemann, and @franknoe.bsky.social .
Comparison of equilibrium distributions obtained by iid sampling and Langevin simulation (sim) across different systems and methods. While classical iid sampling recovers the reference equilibrium distribution, performing simulation with the learned score reveals inconsistencies when models are not trained with Fokker-Planck regularization, i.e., p(x) != p_0(x). Regularized models achieve consistent behavior across systems.
(5/n) With this, we can run coarse-grained Langevin dynamics directly, without the need for any priors or force labels.
This works across biomolecular systems including fast-folding proteins like Chignolin and BBA.
Here is a comparison with and without our regularization:
A model trained with FokkerโPlanck regularization is self-consistent, and aligns the learned score at t = 0 with the distribution recovered by diffusion sampling.
(4/n) Our solution:
We train an energy-based diffusion model and regularize it to satisfy the FokkerโPlanck equation.
This enforces consistency between:
- The density recovered via denoising
- The potential energy learned at t = 0
Result: the same model can be used for sampling AND simulation.
(3/n) The root issue is that at very small diffusion times, diffusion models are inaccurate.
The loss is large, and the models violate the Fokker-Planck equation, meaning the evolution of the modelโs density and its score disagree.
When that happens, the recovered energy ๐ผ(x) is not meaningful.
Training diffusion models on a 2D toy example reveals inconsistencies. While classical iid diffusion sampling (i.e., denoising) correctly reproduces both modes, evaluating the score at t = 0 to estimate the unnormalized density yields a third mode and an incorrect mass distribution. Such a diffusion model would produce incorrect dynamics while producing correct samples.
(2/n) The problem: classical diffusion models learn scores that reproduce equilibrium samples, but the corresponding energy-based parameterization is not consistent.
So if you try to use the learned energy to derive forces, the dynamics are wrong, even if the samples themselves look fine.
(1/n) Can diffusion models simulate molecular dynamics instead of just generating independent samples?
In our NeurIPS 2025 paper, we train energy-based diffusion models that can do both:
- Generate independent samples
- Learn the underlying potential ๐ผ
๐งต๐
Paper: arxiv.org/abs/2506.17139
(6/n)
Done with a brilliant team: Hao Wu, Leon Klein, Stephan Gรผnnemann, and @franknoe.bsky.social .
