Self-Consistent Velocity Matching of Probability Flows

TL;DR

This paper introduces a scalable, discretization-free framework for solving mass-conserving PDEs like the Fokker-Planck equation, leveraging a self-consistent velocity field approach that outperforms existing methods in accuracy and efficiency.

Contribution

The authors propose a novel iterative, discretization-free method for PDEs that ensures self-consistency of the velocity field, enabling high-dimensional solutions without traditional discretization.

Findings

Accurately recovers analytical solutions when available

Outperforms existing methods in high-dimensional settings

Requires less training time

Abstract

We present a discretization-free scalable framework for solving a large class of mass-conserving partial differential equations (PDEs), including the time-dependent Fokker-Planck equation and the Wasserstein gradient flow. The main observation is that the time-varying velocity field of the PDE solution needs to be self-consistent: it must satisfy a fixed-point equation involving the probability flow characterized by the same velocity field. Instead of directly minimizing the residual of the fixed-point equation with neural parameterization, we use an iterative formulation with a biased gradient estimator that bypasses significant computational obstacles with strong empirical performance. Compared to existing approaches, our method does not suffer from temporal or spatial discretization, covers a wider range of PDEs, and scales to high dimensions. Experimentally, our method recovers analytical solutions accurately when they are available and achieves superior performance in high dimensions with less training time compared to alternatives.