FP-Diffusion: Improving Score-based Diffusion Models by Enforcing the Underlying Score Fokker-Planck Equation

TL;DR

This paper introduces FP-Diffusion, a method that enhances score-based diffusion models by enforcing the score Fokker-Planck equation, leading to improved likelihood and consistency in generative modeling.

Contribution

The paper derives the score FPE for noise-conditional scores and proposes a regularization technique to enforce this equation, improving model performance.

Findings

Scores learned via DSM often violate the score FPE.

Enforcing the score FPE improves likelihood and conservativity.

Regularized models outperform baseline diffusion models on various datasets.

Abstract

Score-based generative models (SGMs) learn a family of noise-conditional score functions corresponding to the data density perturbed with increasingly large amounts of noise. These perturbed data densities are linked together by the Fokker-Planck equation (FPE), a partial differential equation (PDE) governing the spatial-temporal evolution of a density undergoing a diffusion process. In this work, we derive a corresponding equation called the score FPE that characterizes the noise-conditional scores of the perturbed data densities (i.e., their gradients). Surprisingly, despite the impressive empirical performance, we observe that scores learned through denoising score matching (DSM) fail to fulfill the underlying score FPE, which is an inherent self-consistency property of the ground truth score. We prove that satisfying the score FPE is desirable as it improves the likelihood and the degree of conservativity. Hence, we propose to regularize the DSM objective to enforce satisfaction of the score FPE, and we show the effectiveness of this approach across various datasets.