This paper presents a learnable solver tailored to iteratively solve sparse linear systems from discretized partial differential equations (PDEs). Unlike traditional approaches relying on specialized expertise, our solver streamlines the algorithm design process for a class of PDEs through training, which requires only training data of coefficient distributions. The proposed method is anchored by three core principles: (1) a multilevel hierarchy to promote rapid convergence, (2) adherence to linearity concerning the right-hand-side of equations, and (3) weights sharing across different levels to facilitate adaptability to various problem sizes. Built on these foundational principles and considering the similar computation pattern of the convolutional neural network (CNN) as multigrid components, we introduce a network adept at solving linear systems from PDEs with heterogeneous coefficients, discretized on structured grids. Notably, our proposed solver possesses the ability to generalize over right-hand-side terms, PDE coefficients, and grid sizes, thereby ensuring its training is purely offline. To evaluate its effectiveness, we train the solver on convection-diffusion equations featuring heterogeneous diffusion coefficients. The solver exhibits swift convergence to high accuracy over a range of grid sizes, extending from $31 \times 31$ to $4095 \times 4095$. Remarkably, our method outperforms the classical Geometric Multigrid (GMG) solver, demonstrating a speedup of approximately 3 to 8 times. Furthermore, our numerical investigation into the solver's capacity to generalize to untrained coefficient distributions reveals promising outcomes.
翻译:本文提出了一种可学习的求解器,专门用于迭代求解离散化偏微分方程产生的稀疏线性系统。与传统依赖专业经验的方法不同,我们的求解器通过训练简化了一类偏微分方程的算法设计过程,仅需系数分布的训练数据即可完成。该方法基于三个核心原则:(1)多层次结构促进快速收敛;(2)保持与方程右端项的线性关系;(3)跨层级共享权重以适应不同问题规模。基于这些基本原则,并考虑到卷积神经网络与多重网格组件相似的计算模式,我们引入了一种适用于结构化网格上具有异质系数偏微分方程线性系统求解的网络。值得注意的是,所提求解器具备对右端项、偏微分方程系数和网格规模的泛化能力,从而确保训练可完全离线进行。为评估效果,我们在具有异质扩散系数的对流扩散方程上训练该求解器。该求解器在$31 \times 31$至$4095 \times 4095$的网格规模范围内均能快速收敛到高精度。特别地,我们的方法比经典几何多重网格求解器快约3至8倍。此外,对求解器在未训练系数分布上泛化能力的数值研究表明了令人鼓舞的结果。