无惧分辨率变化，顽强求解PDE家族：加州理工学院等提出傅里叶神经算子方法

这篇由加州理工学院 Zongyi Li、Anima Anandkumar，以及普渡大学（Purdue University）Kamyar Azizzadenesheli 等人提交的论文，目前正在接受 ICLR 2021 大会的审阅。

本文的作者之一 Anima Anandkumar 是加州理工学院教授，也是英伟达机器学习研究的负责人。

传统意义上，神经网络主要学习有限维欧式空间之间的映射。近期，这一做法被推广到神经算子，它主要学习函数空间之间的映射。对于偏微分方程（PDE）而言，神经算子直接学习任意函数参数依赖到解的映射。因而，与解决一个方程实例的经典方法不同，神经算子学习整个 PDE 家族。

近日， 来自加州理工学院和普渡大学的研究者通过直接在傅里叶空间中对积分核进行参数化，构造了新的神经算子——傅里叶神经算子（FNO），帮助实现表达力强且高效的架构。

论文地址：https://arxiv.org/abs/2010.08895v1

研究人员在伯格斯方程（Burgers’ equatio）、达西流动（Darcy flow）和纳维 - 斯托克斯方程（Navier-Stokes equation）上进行实验， FNO 超越现有的神经网络方法取得了 SOTA 性能 ，且与传统的 PDE 求解器相比，其速度快了三个数量级。

如何求解 PDE？

科学和工程学中的许多问题都涉及针对某些参数的不同值反复求解复杂的偏微分方程（PDE）系统，例如分子动力学、微观力学等。此类系统通常需要精细离散化，以捕获需要建模的现象。因此，传统的有限元法（finite element method，FEM）和有限差分法（finite difference method，FDM）速度较慢，有时效率低下。例如，当设计机翼之类的材料时，我们需要解决相关的逆问题，其中需要对前向模型进行数千次评估。对此，存在一种快速解决问题的方法。

机器学习方法通过提供近似于传统方法的快速求解器，成为革新许多科学学科的关键所在。 但是经典的神经网络是在有限维度的空间之间进行映射，因此它们只能学习与特定离散化相关的解 。这在实际应用中通常无法克服，因此我们更需要开发  mesh-invariant 神经网络 。

在该论文中，研究者首先概述了两种基于主流神经网络的求解 PDE 方法：有限维算子和神经 FEM（Neural-FEM）。

有限维算子

有限维算子方法将解算子参数化为有限维欧几里得空间之间的深度卷积网络。按照定义，这类方法并不是独立于 mesh 的，并且需要针对不同的分辨率和离散化进行修改和调整，以实现一致的误差（如果可能的话）。此外，这些方法受限于训练数据的离散化规模和几何形状，因此它们无法对新的点找到解。

而这篇论文提出的方法正相反， 其误差对于网格分辨率具备不变性，并能够在 mesh 之间传递解。

Neural-FEM

Neural-FEM 方法将解函数直接参数化为神经网络。该方法旨在为 PDE 的特定实例建模，而不是解算子。它具备 mesh 独立性并且准确，但是对于函数参数 / 系数的任何给定新实例，它都需要训练新的神经网络。

这种方法与经典方法（如有限元法）非常相似，用神经网络空间代替局部基函数有限集的线性范围。Neural-FEM 方法遇到的计算问题也与经典方法相同：每个新实例都需要解决优化问题。此外，该方法还受限于此设置：基础 PDE 已知。

研究贡献

该研究提出一种能够学习无限维函数空间之间映射的新型深度学习架构——傅里叶神经算子，通过傅里叶空间中的线性变换实例化积分算子（参见图 1a）。

该研究的主要贡献有：

• 傅里叶神经算子方法共享相同的学得网络参数，而与出于计算目的在输入和输出空间上使用的离散化无关。

• 对于 parametric PDE，傅里叶神经算子始终优于所有现有的深度学习方法。其误差率在伯格斯方程上降低了 30%，在达西流动问题上降低了 60%，在纳维 - 斯托克斯方程（雷诺数为 10000 的湍流状态）上降低了 30%（如图 1b 所示）。在学习整个时间序列的映射时，该方法在雷诺数为 1000 时，达到了 < 1% 的误差，在雷诺数为 10000 时，误差为 8%。

• 在 256×256 网格上，用于求解纳维 - 斯托克斯方程的伪谱方法用时 2.2 秒，而傅里叶神经算子的推断时间仅为 0.005 秒。该方法不仅具有巨大的速度优势，在下游应用（如解决贝叶斯逆问题）中使用时，其准确率也不会下降，如图 3 所示。

神经算子

神经算子是一个迭代结构 v_0 |→ v_1 |→ . . . |→ v_T，其中，v_j (j = 0, 1, . . . , T − 1) 是一系列函数，每一个函数取值于 R^dv。首先通过局部（逐点）变换 P 将输入  a ∈ A 转换为更高维度的表示 v_0 = P(a)。这一局部变换 P 一般通过浅层全连接网络进行参数化，P : R^da → R^dv 单独对每个空间组件  a(x) ∈ R^da 执行。类似地，输出 u = Q(v_T ) 是 v_T 通过局部变换 Q : R^dv → R^du 后的投影。

在每次迭代中，更新 v_t |→ v_t+1 被定义为非局部积分算子 K 和局部非线性激活函数σ 的组合。

κ_φ 是从数据中学得的核函数。定义 1 和定义 2 构成了神经网络向无限维空间的泛化 [Li et al., 2020b]。如果我们移除对函数 a 的依赖，并使 κ_φ(x, y) = κ_φ(x−y)，则得到 (4) 是卷积算子。

研究者在傅里叶空间中直接参数化 κ_φ，并使用快速傅里叶变换（FFT）对 (4) 进行高效计算，从而得到在 PDE 问题上获得 SOTA 结果的新型快速架构。

傅里叶神经算子

研究人员提出用傅里叶空间中定义的卷积算子替换掉 (4) 中的核积分算子。令 F 表示对函数的傅里叶变换 f : D → R^dv，F^−1 表示逆变换，则得到：

j = 1, . . . , d_v，i = √−1 表示虚数单位。令公式 (4) 中的 κ_φ(x, y, a(x), a(y)) = κ_φ(x − y)，并应用卷积定理，得到：

从而在傅里叶空间中直接参数化 κ_φ：

实验

研究者在一维伯格斯方程、二维达西流动问题和二维纳维 - 斯托克斯方程上对比了 FNO 和多个有限维架构和基于算子的逼近方法。

伯格斯方程

该实验的结果参见图 2a 和表 1。

该研究提出的 FNO 方法取得了最低的相对误差，并且该误差值不随分辨率的变化而变化，而基于卷积神经网络的方法（FCN）的误差随着分辨率的增长而增大。与其他神经算子方法（如在物理空间中使用 Nyström 采样的 GNO 和 MGNO）相比，傅里叶神经算子方法在准确率和计算效率方面均更胜一筹。

达西流动

该实验结果参见图 2b 和表 2。FNO 方法的相对误差比其他方法几乎低了一个数量级，而且该误差值并不会随着分辨率的变化而变化。

纳维 - 斯托克斯方程

如表 3 所示，在具备足够数据时（ν = 1e−3, N = 1000 和 ν = 1e−4, N = 10000），FNO-3D 展现出了最优性能。对于数据不足的情况（ν = 1e−4, N = 1000 和 ν = 1e−5, N = 1000），其他方法误差均大于 15%，而 FNO-2D 的误差值最低。

此外，该研究在 64 × 64 × 20 数据上训练 FNO-3D，在 256 × 256 × 80 上进行评估，取得了不错的泛化效果。这表明该方法不仅 可泛化至不同的空间分辨率，对时间分辨率也具备泛化性 。

贝叶斯逆问题

如上图 3 所示，FNO 和传统的 PDE 求解器可以恢复几乎相同的后验均值。但是，FNO 只需 0.005s 即可评估一个实例，而经过优化的传统求解器仍需要 2.2s。使用 FNO 的 MCMC 一共用时两分半，而使用传统求解器的 MCMC 则用时超过 18 个小时。

第一作者简介

该研究的第一作者 Zongyi Li，目前是加州理工学院计算机和数学系的在读博士生。他的研究方向为机器学习、理论计算科学和应用数学。最近，他一直致力于为偏微分方程研究图神经网络。

在来到加州理工之前，Zongyi Li 毕业于圣路易斯华盛顿大学（Washington University in St. Louis），主修计算机科学和数学，导师为 Brendan Juba 和 Xiang Tang。