Self-Consistent Velocity Matching of Probability Flows
This addresses the computational challenges of solving PDEs in high dimensions for researchers in machine learning and physics, though it appears incremental as it builds on existing neural parameterization methods.
The paper tackles solving mass-conserving PDEs like the Fokker-Planck equation by proposing a discretization-free framework based on self-consistent velocity fields, achieving accurate recovery of analytical solutions and superior performance in high dimensions with less training time compared to alternatives.
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.