Score-based deterministic density sampling

Abstract

We propose a deterministic sampling framework using Score-Based Transport Modeling for sampling an unnormalized target density $\ddot{I}€$ given only its score $\nabla \log \ddot{I}€$. Our method approximates the Wasserstein gradient flow on $\mathrm{KL}(f_t\Vert \ddot{I}€)$ by learning the time-varying score $\nabla \log f_t$ on the fly using score matching. While having the same marginal distribution as Langevin dynamics, our method produces smooth deterministic trajectories, resulting in monotone noise-free convergence. We prove that our method dissipates relative entropy at the same rate as the exact gradient flow, provided sufficient training. Numerical experiments validate our theoretical findings: our method converges at the optimal rate, has smooth trajectories, and is often more sample efficient than its stochastic counterpart. Experiments on high-dimensional image data show that our method produces high-quality generations in as few as 15 steps and exhibits natural exploratory behavior. The memory and runtime scale linearly in the sample size.

13 pages, 2 tables, 11 figures. Key words: Deterministic sampling; score-based transport modeling; Wasserstein gradient flow; relative entropy; Fisher information; annealing; neural network; neural tangent kernel