[论文解读] Scalable Inference in SDEs by Direct Matching of the Fokker-Planck-Kolmogorov Equation

Arno Solin, Ella Tamir+1 more|arXiv (Cornell University)|Oct 29, 2021

Gaussian Processes and Bayesian Inference被引用 4

一句话总结

该论文提出了一种可扩展的、无需采样的随机微分方程（SDEs）推理方法，通过直接匹配福克-普朗克-柯尔莫哥洛夫（FPK）方程的一阶和二阶矩，结合高斯近似密度近似。通过将SDE解表述为关于矩的确定性常微分方程（ODE），该方法在时间序列和生成建模任务中保持了具有竞争力的性能，同时相比基于采样的方法（如Euler–Maruyama）实现了数量级的速度提升。

ABSTRACT

Simulation-based techniques such as variants of stochastic Runge-Kutta are the de facto approach for inference with stochastic differential equations (SDEs) in machine learning. These methods are general-purpose and used with parametric and non-parametric models, and neural SDEs. Stochastic Runge-Kutta relies on the use of sampling schemes that can be inefficient in high dimensions. We address this issue by revisiting the classical SDE literature and derive direct approximations to the (typically intractable) Fokker-Planck-Kolmogorov equation by matching moments. We show how this workflow is fast, scales to high-dimensional latent spaces, and is applicable to scarce-data applications, where a non-parametric SDE with a driving Gaussian process velocity field specifies the model.

研究动机与目标

• 为解决SDE中基于采样的推理在高维和数据稀缺场景下的计算低效问题。
• 开发一种直接的、非采样方法，以近似求解SDE中难以处理的福克-普朗克-柯尔莫哥洛夫（FPK）方程。
• 通过在时间边缘分布上施加高斯假设，实现在潜在空间SDE模型中的快速、可扩展推理。
• 通过将高斯过程（GP）先验与矩匹配推理相结合，拓展SDE在机器学习中的适用性。
• 提供一种确定性替代方案，以替代路径采样的随机Runge–Kutta格式，同时保持分布精度。

提出的方法

• 在漂移向量场上使用高斯过程先验，以实现对潜在动态的灵活、基于先验的建模。
• 通过匹配解分布的一阶矩（均值）和二阶矩（协方差），推导出福克-普朗克-柯尔莫哥洛夫（FPK）方程的闭式近似。
• 将所得的矩动力学作为确定性常微分方程（ODE）求解，避免了轨迹采样需求。
• 使用高斯近似密度近似来表示SDE的时间边缘分布，从而实现在高维潜在空间中的高效计算。
• 采用变分推理框架实现该方法，其中编码器输出初始矩，解码器使用矩匹配的SDE动力学。
• 通过似然最大化或变分目标优化模型，利用自动微分计算通过矩ODE的梯度。

实验结果

研究问题

• RQ1 对福克-普朗克-柯尔莫哥洛夫方程的直接矩匹配能否为SDE中的基于采样的推理提供一种可扩展的替代方案？
• RQ2 在高维和低数据场景下，矩匹配推理在准确性和速度上与Euler–Maruyama采样相比表现如何？
• RQ3 高斯近似密度近似在多大程度上能捕捉具有非线性漂移和扩散项的SDE的真实时间边缘分布？
• RQ4 结合GP先验的SDE与矩匹配能否在数据稀缺的应用中实现有效建模？
• RQ5 所提出的方法在降低计算成本的同时，是否在时间序列和生成建模基准上保持了具有竞争力的性能？

主要发现

• 在MOCAP数据集上，该矩匹配方法的测试均方误差（MSE）为7.55±0.05，优于使用相同架构的基线潜在SDE（4.03±0.20），表明在近似密度约束下仍具鲁棒性。
• 实际运行时间显示，该方法在GPU上单次推理仅需6.0±0.1秒，在CPU上仅需1.8±0.1秒，而Euler–Maruyama方法使用200条路径需37.1±0.1秒。
• 在Rotating MNIST任务中，该方法实现了0.051±0.007的MSE和52.7±9.5的NLPD，尽管未进行路径采样，仍表现出具有竞争力的性能。
• 该方法在高维潜在空间中表现出高效扩展性，每步计算成本随维度线性增长（O(d)），而基于采样的方法在维度增加时扩展性差。
• 与基于采样的SDE求解器相比，该方法将推理时间减少了数量级，使其适用于实时和大规模应用。
• 该方法在低数据场景下也表现出有效性，其中对漂移场施加的GP先验提供了强归纳偏置，提升了泛化能力。