数值仿真新工具 --基于傅里叶神经网络的偏微分方程解算器

发布时间:2020-12-26 11:32

偏微分方程(Partial Differential Equation,PDE)求解是数值仿真的核心课题,广泛应用于科学和工程领域。从分子动力学与湍流模拟到飞行模拟,乃至天气预测,许多实践问题因涉及变量关于自身某个维度的变化(如空间坐标或时间),求解过程均可抽象成偏微分方程组。这些方程组通常需要较为精细的离散化以进行数值模拟和逼近,传统的有限元、有限差分方法面对这些问题时计算量巨大。近年来,机器学习技术的发展提供了另一种近似逼近的思路和研究方向,许多基于神经网络(NN)模型的PDE求解方法被开发出来,比如基于多层感知机的PDE解算器[7]。但神经网络往往囿于每一层维度尺寸的限制,只能学习定义在特定分辨率的有限维空间上的映射,这无疑限制了模型对不同精度的适应能力。

 

背景介绍

首先考虑形如下式定义的PDE:

 

 

其中L为依赖于系数函数(coefficient function)a的一个线性微分算子,u为目标解函数(target solution function),在空间域D上满足给定的约束条件和初始条件。

格林函数是PDE求解常用的解析方法。假设有依赖系数a的格林函数G,使得

LG(x, y)=δ(x-y)

其中,δ为狄拉克函数,即对所有非零值,δ取值0,当且仅当δ(0)=1。则将格林函数与f求积分即可得到PDE的目标解函数u

 

然而,实际的PDE往往很难确定解析解,而常以数值迭代逼近的方式获取近似解,问题也就转化成了寻找一个从系数空间A到解空间U的映射G,使得对给定的a∈A,G(a)都可以得出一个关于真实解u∈U的近似。假设我们有N个观测/采样到的点(a,u),就可以通过比较计算G(a)与u误差的有监督算法来学习映射G。

 

图神经算子

图神经网络(GNN)是当前蓬勃发展的深度学习方法。由Kaiming He和Jure Leskovec等合作的Graph Structure of Neural Networks[1]一文曾对图网络的思路进行过提纲挈领的概况总结,我们之前也做过相关介绍(见文末“往期推荐”),GNN算法通常的思路如下:

a.设计图中点(node)和边(edge)的特征向量表达

b.计算信息传递函数(message function),通常是矩阵乘法(全连接层或卷积)

c.计算信息聚合函数(aggregation function),通常是求和再通过一层激活函数

d.将步骤b-c重复迭代T次来更新a中定义的特征向量

经典的NN算法大多都是处理有限维欧氏空间之间的映射,图神经算子(Graph Neural Operator,GNO)[2]将NN推广至无限维空间,以及用有限维空间去近似无限维的情况。pointwise的映射可适用于输入变量x空间中不同的点或格点形式,不仅降低甚至移除了对网格(mesh)采样点的依赖,还可以通过更少采样点的训练达到同等或更加优秀的泛化能力(generalization)。

格林函数的对称性类似于无向图中边的特性。由此,GNO提出了一个通过不断数值迭代来求解PDE的思路,表达如下:

 

 

原文并没有详细解释这个式子是怎么来的,但我们可以从GNN的视角来理解:

a).函数定义域D中的每个采样x即为图中的点(node)

b).点x的特征向量为v(x)

c).ν(希腊字母nu而非英文v)是点x的博雷尔测度(Borel measure),在图的概念中可以近似理解为点x的一个邻域(这个邻域决定了x的邻居y,也就间接定义出了一条从x到y的边)

d).有了边(edge),便可以定义边的特征向量(x, y, a(x), a(y)),a的定义同上节

e).权重矩阵W和κ扮演了信息传递函数的角色,其中κ被称为核积分算子

f).最外层的求和与σ函数(如ReLU)即是信息聚合函数

g).最后进行从t到t+1的迭代,去更新v的值以期一步步靠近最终解函数u

实际操作时,GNO对上述算法做了诸多细节上的调整,例如:

·将积分替换成求和,也就是网格化的离散处理

·通过人为调整邻域半径大小来控制邻居数量

·在进行关于v的主体迭代运算前后分别加入了一 层1x1的卷积P、Q,其中:

i.先将a进行高斯平滑(Gaussian smoothing)处理,再利用卷积P将平滑后的结果参与v0的初始化

ii.卷积Q用于将迭代完毕的vT映射到目标解函数u的空间

对更多细节感兴趣的读者可以参阅原文[2]。

 

傅里叶神经算子

有了GNO的尝试,傅里叶神经算子(Fourier Neural Operator,FNO)[3]将积分核κ的思想做了进一步的改进。狄拉克函数δ(x-y)的表达形式不禁令人联想到数学上的卷积运算。所以,如果假设GNO中需要学习的核积分算子κ

 

不再依赖于系数a(见实验部分讨论)而直接将其改写成κ(x-y)的形式,则转化为形如卷积的下式

由于卷积可以通过傅里叶变换后相乘,再进行逆变换回到原空间,上式演变为:

 

花体F即为傅里叶变换。事实上,既然κ是要学习的参数,与其在D空间学习后进行傅里叶变换,不如直接在傅里叶空间内学,所以FNO直接用一个定义在傅里叶空间中的卷积算子R代替上式中F(κ):

 

实际的迭代更新过程便从式1演化成了式2:

 

迭代公式(2)对应的网络结构请参阅图1:

 

图1,FNO迭代层的网络结构(图片来源:文献[3]图1)

 

为简化计算(图3中黄色框部分),作者对傅里叶变换F进行了截断(truncate,即用有限阶变换去近似真实的频域信号)和共轭对称(conjugate symmetry,因为是实值函数)的改进,最终用快速傅里叶变换(FFT)完成了核卷积算子的运算。

如果以L代表式1与式2的迭代层逻辑,GNO与FNO完整的NN模型可概括如下:

 

图2,GNO/FNO模型框图概况

 

其中,P,Q分别为在不同维度间变换的局部映射,可分别视作编码层与解码层。FNO方法同GNO类似,也是先处理从a到v的初始化,在完成了关于v的迭代更新之后,由解码层将v映射到u的空间内。值得一提的是,正是这一层解码层的存在,令FNO可以有效学习正确的边界条件,从而摆脱了“传统傅里叶方法仅对周期性边界条件有效”的限制。不同于GNO,FNO在L中移除了对a的依赖,进一步简化了计算。求解Navier-Stokes方程时,FNO在inference阶段的速度远胜常用的pseudo-spectral数值方法(0.005秒对比2.2秒)。

 

实验解读

作者测试了三类PDE下FNO的表现:Burgers’ Equation(BE),Darcy Flow方程(DF)和Navier-Stokes方程(NS)。实验显示核积分算子κ对参数a的依赖性取决于问题本身。如果a仅仅是物理模型中的初始条件,例如BE和NS方程,依赖性是不必要的;但对类似于DF的情况,例如a是个几何参数,加入a也许会有帮助。

2D的NS方程常用于描述不可压缩的黏性流动:

 

 

作者对NS问题尝试了两种思路:

1).将x和t分开处理,例如用ResNet[4]、U-Net[5]、TF-Net[6],以及2D-FNO处理二维变量x,再配上一个RNN结构处理时间t

2).3D-FNO直接在时空三维中模拟全部x、t

两者对于时间的适应性不同,2D+RNN的模式可以将解函数传播至任意长度的T,但时间增量必须是固定的;3D模式则必须固定在[0,T]的区间上建模,但可应对任意采样精度的时间离散化。3D-FNO在64x64x20分辨率的数据上训练后完美适应了256x256x80分辨率的测试数据,体现出优秀的超分辨率(super-resolution)效果(图3)。此外作者声称3D-FNO相较于2D-FNO更加容易训练且更有表现力。

 

图3,NS方程模型,两行分别为数值解的groud truth与prediction(图片来源:文献[3]图1)

 

作者最后讨论了激活函数ReLU是施加于空间域D还是傅里叶域的问题。如果ReLU定义在傅里叶域里,pointwise操作相当于空间中的卷积,这失去了它作为“激活函数”的意义;定义于D的ReLU则能帮助模型恢复部分因截断操作而损失的高频段信号。更多细节请参阅[3]及其附录。

 

总结与思考

从解算器网络[7]的计算结构来看,FNO方法的一个创新点是基于FFT的NN算子:先将积分函数转换为卷积,再转至频域操作利用FFT加速。与传统CNN因数据与卷积核尺寸不匹配而计算效率不高的问题不同,FNO可以很好地利用FFT特性,提升整体计算性能。同时,FFT使得模型具备离散不变性(discretization-invariant),即使为了降低运算开销而截取k阶近似。传统的PDE数值解法,如有限元/有限差分算法,为了降低误差往往需要更加精细的采样;FNO则无论采样的分辨率粗细均保持极低的误差值,这使得FNO可以更好应对数据量少的情况。事实上,傅里叶空间由自己的一组基明确定义而不受原始空间域D的影响,所以无论在D上的采样分辨率如何,通过FFT投射到傅里叶空间后,参数矩阵R均可从容应对,达到超分辨率的效果。

神经网络和偏微分方程的相关性远不止NN用于PDE求解,而是有着更深刻的内在联系。事实上,神经网络模型可以看作一个动力学系统,从而可以用微分方程来描述。比如ResNet就可看作一个动力学系统的离散时间近似。鄂维南院士等学者对此提出了许多有趣的观点,相关工作可参考[8,9]。

 

参考文献

1、You, Jiaxuan, et al. "Graph structure of neural networks." International Conference on Machine Learning. PMLR, 2020.

2、Li, Zongyi, et al. "Neural operator: Graph kernel network for partial differential equations." arXiv preprint arXiv:2003.03485 (2020).

3、Li, Zongyi, et al. "Fourier neural operator for parametric partial differential equations." arXiv preprint arXiv:2010.08895 (2020).

4、He, Kaiming, et al. "Deep residual learning for image recognition." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.

5、Ronneberger, Olaf, Philipp Fischer, and Thomas Brox. "U-net: Convolutional networks for biomedical image segmentation." International Conference on Medical image computing and computer-assisted intervention. Springer, Cham, 2015.

6、Wang, Rui, et al. "Towards physics-informed deep learning for turbulent flow prediction." Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2020.

7、Hennigh, Oliver, et al. "NVIDIA SimNet^{TM}: an AI-accelerated multi-physics simulation framework." arXiv preprint arXiv:2012.07938 (2020).

8、Lu, Yiping, et al. "Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations." International Conference on Machine Learning. PMLR, 2018.

9、Weinan, E. "A proposal on machine learning via dynamical systems." Communications in Mathematics and Statistics 5.1 (2017): 1-11.

 

上一个: 贝叶斯方法与深度学习的结合及应用(1)

下一个: 在壁仞科技研究院,我们思考的问题

近期文章

AI 智能体:应用前景与对算力需求的影响

人工智能技术的迅猛发展,尤其是 AI 智能体这一核心部分。本文讨论了AI智能体的定义、应用前景和其对算力需求的影响。

2023-11-13

查看更多