Title: Consistent Sampling and Simulation: Molecular Dynamics with Energy-Based Diffusion Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: Energy-based Diffusion Models

Molecular Dynamics

Coarse-Graining

Boltzmann Emulator

Fokker-Planck Equation

Mixture of Experts
Abstract: In recent years, diffusion models trained on equilibrium molecular distributions have proven effective for sampling biomolecules. Beyond direct sampling, the score of such a model can also be used to derive the forces that act on molecular systems. However, while classical diffusion sampling usually recovers the training distribution, the corresponding energy-based interpretation of the learned score is often inconsistent with this distribution, even for low-dimensional toy systems. We trace this inconsistency to inaccuracies of the learned score at very small diffusion timesteps, where the model must capture the correct evolution of the data distribution. In this regime, diffusion models fail to satisfy the Fokker-Planck equation, which governs the evolution of the score. We interpret this deviation as one source of the observed inconsistencies and propose an energy-based diffusion model with a Fokker-Planck-derived regularization term to enforce consistency. We demonstrate our approach by sampling and simulating multiple biomolecular systems, including fast-folding proteins, and by introducing a state-of-the-art transferable Boltzmann emulator for dipeptides that supports simulation and achieves improved consistency and efficient sampling. Our code, model weights, and self-contained JAX and PyTorch notebooks are available at https://github.com/noegroup/ScoreMD.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a regularization loss for diffusion model that is based on Fokker–Planck-equation (FPE) to enforce consistency between sampling and simulation modes in molecular dynamics (MD) simulations. The authors identify that diffusion models, when trained on equilibrium molecular distributions, exhibit inconsistencies between diffusion reverse process (backward inference) and use the first score function for Langevin simulation. To address this, they introduce a computationally efficient regularization method that ensures consistency. They validate their approach on a simple toy system, alanine dipeptide, and a transferable Boltzmann emulator for dipeptides demonstrating improved consistency and efficient sampling.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\n- Introducing Fokker-Planck equation (FPE) based regularization loss for diffusion models to learn to sample consistently with the energy-based model formulation.\n- The experiments demonstrate that with the proposed loss, the diffusion works well by generating consistent sample using langevin dynamics (LD) simulator, on MB potential, ALDP and dipeptide tasks.\n\nWeakness:\n- While the paper empirically demonstrates improvements on toy systems, the theoretical underpinning connecting regularization to broader implications of diffusion models is insufficiently developed, which may not support the main claim of diffusion models in the paper.\n- The writing is not clear enough and seems to lack of necessary explanation (see equations).\n- The significance of the methodological advance seems unclear. I am not convinced how incorporating the FP loss improves the diffusion sampling from the experimental results.\n\nQUESTIONS:\n1. For the definition of $L_{FP}$ (line 132), which does the D^{-2} denote or indicate? \n2. How do you obtain the $\\log p^{\\theta}(x, t)$ in equation 13 as your practical estimation? This term is not trivial for most of the diffusion models. Which is the concrete form of residual error (eq (8))?\n3. In equation (6), i do not quite understand the derivation in the following sense: the EBM that correlates p(x) and U(x) has its support over the whole configuration space of x (exp is non-negative). However, the learned score at t=0 $s^{\\theta}(x)$ does not necessarily capture the gradient field with arbitrary x, especially the region far from the data manifold. I am hesitant about the correctness of such derivation. Following such derivation, one can even simply train the network via denoising score matching to estimate the \"force\" around the data manifold. Training diffusion models on multi-tasking across multiple time steps yet throwing away the score network for time t>0 seems weird for sampling. Maybe you can elaborate or justify it furthermore.\n4. What does it mean by $NNET(x,t)$ in line 155 and line 156?\n5. In practical concern, why do we actually need to do “sim” instead of just “iid” for diffusion sampling? (\"sim\", \"iid\" are following the notation in the experimental section)\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses the inconsistency between molecular dynamics (MD) simulations and traditional sampling in diffusion models by introducing a Fokker–Planck-based regularization to the diffusion loss. To reduce computational costs, the authors use efficient approximations and a time-based mixture of experts architecture. Experiments on toy systems, alanine dipeptide, and transferable dipeptide settings demonstrate that their approach improves consistency between sampling and simulation. In short, the contribution of this paper lies in solving the inconsistency, i.e., mismatch between how we generate samples from a model and how we use it to simulate time-dependent behavior.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. (Clarity) clear and organized presentation\n\nThe paper is clearly written and well organized, making it easy to follow. Most experiments are well designed and support the authors’ claims across various settings.\n\n**Weakness**\n\n1. (Quality) Quantitative transferable dipeptides results\n\nFor the transferable dipeptide, quantitative and qualitative results only for dipeptide AC has been presented in the main paper. Though there are qualitative results on other test dipeptides in the appendix, I do not see any quantitative results on other peptides.\n\n(Minor) Moreover, additional qualitative results for dipeptides could be helpful. Specially, most dipeptide show similar landscape to the Ramachandran plot in Figure 5 and 12, while dipeptides containing proline and glycine shows some different Ramachandran plot trends. Though there is one in Figure 14, one or two additional test peptides results with different Ramachandran plot trends could strengthen the claim for the generalization of the proposed method in the transferable dipeptide settings.\n\n2. (Quality) Efficiency of training and inference\n\nWhile the authors have emphasized that the proposed method achieves efficient training and inference, I feel that the experiment results fall short on this. The author have claimed efficiency in many aspects, such as \n\n- line 58 (contribution 2) achieve efficient training and inference\n- line 80 efficiently evaluate the loss\n- line 150 efficient approximation of the loss\n- line 183 MOE approach to allocate model capacity more efficiently\n\nThere are not results supporting this claim in the main paper, while there is table 7 in the appendix reports the training and inference time for the proposed method. However, methods including Fokker-Planck regularization terms shows several longer training time than the original model. \n\n3. (Significance) Dynamic results\n\nThe authors have experimented only to dipeptide system scales, while a major baseline in the paper “Two For One” has done experiments to fast folding protein system scales (chignolin, BBA, etc). Experiments on systems with this scales allows dynamics analysis as in Figure 5 of the Two For One paper. Testing two or three fast folding protein could strengthen the paper’s claim a lot, since this dynamics analysis is closely related to why we run simulations.\n\nQUESTIONS:\n1. Mixture intervals\n\nA simple question in Section 5. The Mixture of Experts (MoE) model is trained on diffusion time intervals (0, 0.1), [0.1, 0.6), and [0.6, 1.0). Could the authors clarify the intuition or practical rationale behind choosing these specific cutoffs? Were they based on empirical findings, model behavior (e.g., Fokker–Planck error), or heuristic choices?\n\n2. High variance of Mixture of experts in simulations results\n\nIn Tables 1 and 2, the Mixture of Experts model shows relatively high variance in the simulation metrics compared to other methods. This seems like an interesting and potentially important observation. Could the authors offer an explanation or intuition for this behavior?\n\n3. Experiment metric (JS)\n\nFrom section 5, it seems that the JS divergence is computed for the energy distribution. Could the authors state why this is done, instead of JS divergence on dihedral angles / TIC(time-lagged independent component) or pairwise distance as in Two For One? A simple results on JS divergence as in Two For One or why they have chosen energy would be good be enough to resolve my question.\n\nLIMITATIONS:\nyes, the authors adequately addressed the limitations and potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors look at the problem of molecular dynamics, in particular sampling from boltzmann distributions, through the use of diffusion models and langevin sampling. They highlight that current score-based diffusion models, while immensely successful at modeling images, are not as efficient at modeling MD simulations which can be consequently improved with a combination of (a) energy parameterization as opposed to score parameterization, (b) mixture-of-experts formulation for the architecture, and (c) Fokker-planck regularization to enable consistency. Their experiments indicate that the combination of these approaches leads to better simulation results, while still maintaining good iid sampling.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n- The authors provide tractable means of leveraging fokker-planck regularization through finite-difference based estimation of divergence and derivative terms.\n- They showcase that such regularization leads to improved performance for both simulation and iid sampling, when combined with mixture-of-experts architecture.\n\n**Weaknesses**\n\n- The authors conjecture that learning a score function through denoising score matching leads to unreliable estimates closer to the data distribution. However, this has been a well known property in literature, and is in particular showed in detail in [1].\n- The authors aim to solve the main problem of simulation which requires score at $t=0$ for which they leverage the learned score. However, typically for these applications energy, and thus forces, are available at $t=0$ so what is the aim for using the learned score from denoising score matching? I believe an insight into this was lacking in the current draft.\n- While it is evident that enforcing the fokker-planck equation at intermediate time points would lead to consistent sampling, what is the reason for why it would lead to better estimation of score at $t=0$ since it is only the boundary point?\n\n[1] De Bortoli, Valentin, et al. \"Target score matching.\" arXiv preprint arXiv:2402.08667 (2024).\n\nQUESTIONS:\n- It was unclear how the authors successfully train energy-based parameterization since it would require taking the gradient of the model w.r.t the state as well as consequent gradient w.r.t the parameters of the model. Is this operation efficient or are there any approximations leveraged?\n- Are equations 10-12 somehow related to the hutchinson trace estimator or the estimator used in [1]?\n- What is the purpose of using a mixture-of-experts system instead of just increasing the model size/capacity with number of parameters?\n\n[1] Albergo, Michael S., and Eric Vanden-Eijnden. \"Nets: A non-equilibrium transport sampler.\" arXiv preprint arXiv:2410.02711 (2024).\n\nLIMITATIONS:\nNo additional limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 1, "ORIGINALITY": 1, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores the connection between the score of diffusion models and the potential energy surface of a molecular system. The authors show that there is a mismatch between samples generated via reverse diffusion and Langevin dynamics using the score as a force field. To fix this, the authors augment the standard denoising score matching loss with a Fokker–Planck residual penalty that biases the learned energy/time evolution to satisfy the Fokker–Planck equation. They further reduce computational cost by partitioning diffusion time into segments and training a Mixture-of-Experts. They show empirically that they can train consistent Boltzmann Generators.\n\nSTRENGTHS AND WEAKNESSES:\n-\tThe clarity of the paper is very strong as is the overall quality of the text.\n-\tFigure 2 nicely illustrates the need for the added FP regularisation. \n-\tThe results are convincing and the free energy plots informative. \n-\tThe transferability of the approach could be probed more. Is it transferrable across diverse molecular datasets? Rather than just similar dipeptides? \n-\tThe authors mention that the MoE approach is more expensive, but a detailed cost comparison of multiple inferences is not provided. \n-\tThe use of only a single MoE and the associated non smoothness of the potential energy surface is not thoroughly investigated and appears to me as a significant limitation of the approach. The artifacts may appear in kinetic observables and may produce artifacts in other relevant observables not investigated in this approach. Could the authors quantitatively investigate the non-smoothness of the potential energy surface, by probing the forces as the model changes between mixture of experts? How significant can these differences be? \n-\tHow does the proposed method compare to recently proposed long-stride MD simulations such as FlashMD, BoostMD, TrajCast, Timewarp?\n\nQUESTIONS:\n1.\tCould you provide a more detailed analysis of the non-smoothness of the resulting PES? (see above for more detail on the question) How smooth are geometry optimisations for example? \n2.\tCould you test a more diverse dataset, potentially of small molecules? With more diversity in the functional groups present. Does the FP regularisation truly transferrable learn more consistent scores? This is of interest as there is no guarantee of the score satisfying the FP equation from an architectural perspective.\n\nLIMITATIONS:\n-\tOverall adequately addressed, however the problems associated with a non-smooth PES do not seem adequately investigated.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}