Accelerating Fokker–Planck Regularization For Force Field Learning

Diffusion models can learn molecular equilibrium distributions without force labels, but they often predict inaccurate forces that cause subsequent Langevin simulations to diverge. While standard Fokker–Planck (FP) regularization mitigates this inconsistency by constraining the density evolution, FP regularization heavily relies on expensive Hessian computations, which makes training computationally prohibitive. To address this, we introduce a novel weak FP regularization that achieves both consistency and computational efficiency. Rather than relying on expensive higher-order derivatives, our method approximates standard FP using a small set of efficiently-computable test functions. To optimize this objective correctly using stochastic gradients, we then derive an unbiased U-statistic estimator for minibatched moments. Through extensive experiments on real-world and analytic potentials using augmented as well as SE(3)-equivariant graph transformers, we demonstrate that our objective achieves the accuracy of the state-of-the-art FP regularization at a fraction of the memory and runtime.