Data-driven machine learning approaches are being increasingly used to solve partial differential equations (PDEs). They have shown particularly striking successes when training an operator, which takes as input a PDE in some family, and outputs its solution. However, the architectural design space, especially given structural knowledge of the PDE family of interest, is still poorly understood. We seek to remedy this gap by studying the benefits of weight-tied neural network architectures for steady-state PDEs. To achieve this, we first demonstrate that the solution of most steady-state PDEs can be expressed as a fixed point of a non-linear operator. Motivated by this observation, we propose FNO-DEQ, a deep equilibrium variant of the FNO architecture that directly solves for the solution of a steady-state PDE as the infinite-depth fixed point of an implicit operator layer using a black-box root solver and differentiates analytically through this fixed point resulting in $\mathcal{O}(1)$ training memory. Our experiments indicate that FNO-DEQ-based architectures outperform FNO-based baselines with $4\times$ the number of parameters in predicting the solution to steady-state PDEs such as Darcy Flow and steady-state incompressible Navier-Stokes. Finally, we show FNO-DEQ is more robust when trained with datasets with more noisy observations than the FNO-based baselines, demonstrating the benefits of using appropriate inductive biases in architectural design for different neural network based PDE solvers. Further, we show a universal approximation result that demonstrates that FNO-DEQ can approximate the solution to any steady-state PDE that can be written as a fixed point equation.

翻译：数据驱动的机器学习方法正越来越多地被用于求解偏微分方程（PDEs）。在训练算子——即输入某个族中的PDE并输出其解——时，这些方法已展现出尤为显著的成效。然而，针对所关注的PDE族的结构性知识，架构设计空间仍未被充分理解。为弥补这一不足，我们研究了权重共享神经网络架构在稳态PDE中的优势。为此，我们首先证明大多数稳态PDE的解可表示为非线性算子的不动点。基于这一观察，我们提出FNO-DEQ，一种深度平衡变体的傅里叶神经算子（FNO）架构，该架构将稳态PDE的解直接求解为隐式算子层的无限深度不动点，通过黑盒根求解器实现，并对此不动点进行解析求导，从而将训练内存复杂度控制在$\mathcal{O}(1)$。实验表明，基于FNO-DEQ的架构在预测达西流及稳态不可压缩纳维-斯托克斯方程等稳态PDE的解时，以$4\times$更少的参数量超越了基于FNO的基线模型。此外，我们发现与FNO基线相比，FNO-DEQ在噪声观测值更多的数据集上训练时具有更强的鲁棒性，这证明了在基于神经网络的PDE求解器架构设计中采用适当归纳偏置的优势。最后，我们给出了一个通用逼近性结果，证明FNO-DEQ能逼近任何可写为不动点方程的稳态PDE的解。