Dynamical Measure Transport and Neural PDE Solvers for Sampling

TL;DR

This paper introduces a unified PDE-based framework for sampling via measure transport, leveraging physics-informed neural networks (PINNs) to improve efficiency and accuracy without requiring data samples or normalization constants.

Contribution

It proposes a novel PDE-based approach for sampling that integrates prior methods and employs PINNs for efficient, discretization-free PDE solution approximation.

Findings

PINNs enable efficient PDE solution approximation for sampling.

The method achieves better mode coverage than existing approaches.

High-accuracy sampling can be obtained with fine-tuning using Gauss-Newton methods.

Abstract

The task of sampling from a probability density can be approached as transporting a tractable density function to the target, known as dynamical measure transport. In this work, we tackle it through a principled unified framework using deterministic or stochastic evolutions described by partial differential equations (PDEs). This framework incorporates prior trajectory-based sampling methods, such as diffusion models or Schr\"odinger bridges, without relying on the concept of time-reversals. Moreover, it allows us to propose novel numerical methods for solving the transport task and thus sampling from complicated targets without the need for the normalization constant or data samples. We employ physics-informed neural networks (PINNs) to approximate the respective PDE solutions, implying both conceptional and computational advantages. In particular, PINNs allow for simulation- and discretization-free optimization and can be trained very efficiently, leading to significantly better mode coverage in the sampling task compared to alternative methods. Moreover, they can readily be fine-tuned with Gauss-Newton methods to achieve high accuracy in sampling.