[论文解读] Learning Neural PDE Solvers with Convergence Guarantees

Jun-Ting Hsieh, Shengjia Zhao+3 more|arXiv (Cornell University)|Jun 4, 2019

Model Reduction and Neural Networks被引用 56

一句话总结

简短结论：论文训练一个神经网络来修改现有线性偏微分方程求解器的更新，在保持收敛性保证的同时实现更快的收敛和跨几何形状及边界条件的泛化。

ABSTRACT

Partial differential equations (PDEs) are widely used across the physical and computational sciences. Decades of research and engineering went into designing fast iterative solution methods. Existing solvers are general purpose, but may be sub-optimal for specific classes of problems. In contrast to existing hand-crafted solutions, we propose an approach to learn a fast iterative solver tailored to a specific domain. We achieve this goal by learning to modify the updates of an existing solver using a deep neural network. Crucially, our approach is proven to preserve strong correctness and convergence guarantees. After training on a single geometry, our model generalizes to a wide variety of geometries and boundary conditions, and achieves 2-3 times speedup compared to state-of-the-art solvers.

研究动机与目标

• 通过学习对现有线性求解器的更新来提升求解速度，同时保持正确性以实现更高效的领域特定PDE求解器。
• 通过固定点保持来保证收敛到真实PDE解。
• 尽管在单一实例上训练，但展示对未见几何形状、边界条件和网格尺寸的泛化能力。

提出的方法

• 将学习得到的求解器表示为基线迭代器的参数化更新：u' = Ψ(u; G, f, b, n) + G H (Ψ(u; G, f, b, n) - u)，其中 H 为一个由卷积网络实现的学习线性算子。
• 使用线性（Jacobi 风格）基线 Ψ，并在设计上确保固定点仍为解（命题 1）。
• 用线性深度网络（Conv 或 U‑Net 架构）对 H 进行参数化，以近似 T(I−T)−1，从而加速收敛（定理 2 及第3.3节的解释）。
• 在单一几何/问题实例上进行训练，但在不同几何形状和边界条件上评估泛化能力（命题 2）。
• 提供两个模型族：Conv 模型（3×3 卷积）和基于 U‑Net 的多网格模型，以捕捉局部和多尺度修正。

实验结果

研究问题

• RQ1 一个学习得到的线性修正是否能够保持对正确不动点的收敛至基线 Ψ？
• RQ2 在训练中未见的不同几何形状、边界条件和网格尺寸上，学习的修正是否能加速收敛？
• RQ3 卷积和基于 U‑Net 的线性网络在多大程度上近似最优算子以加速收敛？
• RQ4 对于学习求解器的收敛性与泛化属性，存在哪些理论保证？

主要发现

模型	基线	平方层/运算	L形域层/运算	圆柱域层/运算	Square-Poisson 层/运算
Conv1	Jacobi	0.432 / 0.702	0.432 / 0.702	0.432 / 0.702	0.431 / 0.701
Conv2	Jacobi	0.286 / 0.524	0.286 / 0.524	0.286 / 0.524	0.285 / 0.522
Conv3	Jacobi	0.219 / 0.424	0.219 / 0.423	0.220 / 0.426	0.217 / 0.421
Conv4	Jacobi	0.224 / 0.449	0.224 / 0.449	0.224 / 0.448	0.222 / 0.444
U‑Net2	Multigrid2	0.091 / 0.205	0.090 / 0.203	0.091 / 0.204	0.079 / 0.178
U‑Net3	Multigrid3	0.220 / 0.494	0.213 / 0.479	0.201 / 0.453	0.185 / 0.417

• 学习迭代器 Φ_H 保留了基线 Ψ 的不动点，从而确保正确性（引理 1）。
• Φ_H 的谱范数是 H 的凸函数，使得产生 ρ(Φ_H) < 1 的 H 的集合成为一个凸开集（定理 2）。
• 在单一域上的训练也可以在未见几何和网格尺寸上获得收敛性和加速（命题 2）。
• 经验结果显示在正方形、L 形和圆柱域以及 square-Poisson（f ≠ 0）等场景下，显著优于 Jacobi 和 Multigrid 基线。
• 在 CPU 上，Conv3 模型相对于 Jacobi 在层数/运算次数上为 0.219–0.220 倍、在乘加次数上为 0.424–0.426 倍；Conv3 的层数约高效 5×、运算效率约高效 2.5×。
• U‑Net 模型在所有测试设置中都超越了相应的 Multigrid 基线，额外的 GPU 加速可带来高达 ~30× 相对于 CPU 基线的加速。