Diffusion-PINN Sampler

TL;DR

The paper introduces the Diffusion-PINN Sampler, a new diffusion-based sampling method that uses physics-informed neural networks to accurately estimate the drift term in reverse SDEs, with proven convergence and effective performance on complex tasks.

Contribution

It proposes a novel diffusion sampling algorithm that leverages PINNs to solve the PDE of log-density, providing convergence guarantees and improved accuracy in complex sampling scenarios.

Findings

Effective in identifying mixing proportions with isolated components

Provides convergence guarantees for the sampling process

Demonstrates superior performance on various sampling tasks

Abstract

Recent success of diffusion models has inspired a surge of interest in developing sampling techniques using reverse diffusion processes. However, accurately estimating the drift term in the reverse stochastic differential equation (SDE) solely from the unnormalized target density poses significant challenges, hindering existing methods from achieving state-of-the-art performance. In this paper, we introduce the Diffusion-PINN Sampler (DPS), a novel diffusion-based sampling algorithm that estimates the drift term by solving the governing partial differential equation of the log-density of the underlying SDE marginals via physics-informed neural networks (PINN). We prove that the error of log-density approximation can be controlled by the PINN residual loss, enabling us to establish convergence guarantees of DPS. Experiments on a variety of sampling tasks demonstrate the effectiveness of our approach, particularly in accurately identifying mixing proportions when the target contains isolated components.