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 flow characterized by the same velocity field. By parameterizing the flow as a time-dependent neural network, we propose an end-to-end iterative optimization framework called self-consistent velocity matching to solve this class of PDEs. Compared to existing approaches, our method does not suffer from temporal or spatial discretization, covers a wide range of PDEs, and scales to high dimensions. Experimentally, our method recovers analytical solutions accurately when they are available and achieves comparable or better performance in high dimensions with less training time compared to recent large-scale JKO-based methods that are designed for solving a more restrictive family of PDEs.
翻译:我们提出了一种免离散化的可扩展框架,用于求解一大类质量守恒偏微分方程(PDE),包括含时福克-普朗克方程和Wasserstein梯度流。核心发现是:PDE解的时变速度场必须满足自洽性,即该速度场需满足一个涉及由同一速度场所表征流动的不动点方程。通过将流动参数化为时变神经网络,我们提出了一种名为"自洽速度匹配"的端到端迭代优化框架来求解此类PDE。与现有方法相比,本方法无需时间和空间离散化,可覆盖广泛的PDE类型,并适用于高维场景。实验表明,在解析解已知时,本方法能精确复现解析解;与近期为求解更受限PDE族设计的大规模JKO方法相比,本方法在高维场景中训练时间更少,性能相当或更优。