[Paper Review] Self-Consistent Velocity Matching of Probability Flows

Lingxiao Li, Samuel Hurault+1 more|arXiv (Cornell University)|Jan 31, 2023

Generative Adversarial Networks and Image Synthesis5 citations

TL;DR

This paper introduces Self-Consistent Velocity Matching (SCVM), a discretization-free, scalable framework for solving mass-conserving PDEs such as the Fokker-Planck equation and Wasserstein gradient flows. By enforcing self-consistency between the velocity field and probability flow via an iterative optimization with a biased gradient estimator, SCVM achieves superior accuracy and speed in high dimensions without temporal or spatial discretization.

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.

Motivation & Objective

• Address the limitations of existing PDE solvers that rely on spatial or temporal discretization, which restrict scalability and accuracy in high dimensions.
• Overcome the computational inefficiency of adjoint-based methods for self-consistency in Fokker-Planck and related PDEs.
• Develop a unified, differentiable framework that applies to a broad class of mass-conserving PDEs beyond just Fokker-Planck or gradient flows.
• Enable accurate and efficient learning of probability flows in high-dimensional spaces using neural parameterization without grid-based representations.
• Demonstrate the method's flexibility on complex dynamics such as flow splashing and smooth measure interpolation, where prior methods fail or produce artifacts.

Proposed method

• Formulate the solution of mass-conserving PDEs as a self-consistent fixed-point problem where the velocity field must match the flow it generates.
• Use an iterative optimization scheme that refines the velocity field at each step using a neural network parameterization of the probability flow.
• Employ a biased gradient estimator to bypass the need for adjoint differentiation, significantly improving training speed and scalability.
• Reformulate velocity-matching objectives to eliminate expensive terms like ∇log p_t in Fokker-Planck equations, improving computational efficiency.
• Support two parameterization modes: time-incremental pushforwards (TIPF) and continuous velocity fields (NODE), enabling flexibility across different dynamics.
• Maintain full differentiability and scalability by avoiding mesh-based discretization and relying only on neural network inference and optimization.

Figure 1 : Probability flow produced by SCVM-NODE for a 60-dimensional mixture of Gaussians at irregular time steps. Samples are projected onto the first two PCA components and kernel density estimation is used to generate the contours.

Experimental results

Research questions

• RQ1 Can a self-consistent velocity matching framework be designed to solve a broad class of mass-conserving PDEs without temporal or spatial discretization?
• RQ2 How can self-consistency be enforced efficiently in high-dimensional probability flows using neural networks and iterative optimization?
• RQ3 Can a biased gradient estimator outperform adjoint-based methods in training speed and performance for PDE-constrained flow learning?
• RQ4 To what extent can SCVM generalize to complex dynamics such as flow splashing and smooth measure interpolation beyond standard Fokker-Planck or gradient flow settings?
• RQ5 How does SCVM compare quantitatively and qualitatively to state-of-the-art methods like JKO, adjoint, and particle-based solvers in high-dimensional problems?

Key findings

• SCVM recovers analytical solutions of mass-conserving PDEs with high accuracy, matching or exceeding the performance of prior methods in low dimensions.
• In high-dimensional settings (e.g., d=10), SCVM achieves 30× faster training than the adjoint method (ADJ), with comparable or better Wasserstein-2 and KL divergence metrics.
• For the time-dependent Fokker-Planck equation with harmonic drift, SCVM-TIPF and SCVM-NODE achieve results on par with ADJ but with significantly reduced training time.
• SCVM produces continuous, artifact-free particle paths in flow splashing simulations, unlike DFE (jagged paths) and SDE-EM (crossing obstacles), and avoids numerical instability seen in ADJ.
• SCVM successfully models non-Fokker-Planck dynamics such as flocking behavior with interaction terms, demonstrating broad applicability beyond standard PDE classes.
• In smooth measure interpolation tasks, SCVM enables novel formulations (e.g., optimal transport-based) that are infeasible for competing methods, producing visually coherent animations of shape deformation.