当前位置：网站首页>當Transformer遇見偏微分方程求解

當Transformer遇見偏微分方程求解

2022-06-27 00:00:00 【昇思MindSpore】

本篇與大家分享最近閱讀的Transformer求解偏微分方程論文Choose a Transformer: Fourier or Galerkin，該論文已被NeurIPS2021接收。

背景介紹

在我們的世界中，從宇宙星體的運動，到溫度風速的氣象預報，再到分子原子間的相互作用，很多工程學、自然科學、經濟和商業過程都可以通過偏微分方程（PDE）描述。傳統的方法，如有限元、有限差分、譜方法等，利用離散結構將無限維算子映射簡化為有限維近似問題。近年來物理信息神經網絡（PINN）等模型[1]，通過在求解空間采樣，訓練神經網絡來近似PDE解。但是對於傳統方法或物理信息神經網絡等，邊界條件或者方程參數輕微的變化，通常需要重新計算和訓練。

相比之下，算子學習的目標是學習無限維函數空間之間的映射，這樣能够實現在不需要重新訓練的情况下求解偏微分方程族，從而大幅度節省計算資源。PDE求解中算子學習（operator learner）是當前蓬勃發展的新研究方向，其中典型代錶是傅裏葉神經算子（FNO）[2]。

隨著NeurIPS2021的放榜，基於Transformer的算子學習文章《Choose a Transformer: Fourier or Galerkin》[4]對於參數化PDE的求解給出了一種新穎的解釋，最終在基准中取得了state-of-the-art的結果。

主要工作

在本文中，operator learner采用監督學習訓練，訓練樣本是輸入函數和輸出函數在同樣的離散網格點上采樣得到的，如下圖所示，可以將方程求解轉化seq2seq問題並通過Transformer[3]進行建模。

圖1 operator learner示意

基於Transformer的工作，本文的主要貢獻如下：

1. 無softmax的注意力機制。提出scale-preserving自注意機制和無softmax的attention，並給出兩種方案的數學解釋。

2. 參數化PDE的operator learner。將新提出的注意力算子與FNO結合起來，顯著提高在參數化PDE求解基准問題中的精度。

3. State-of-the-art實驗結果。在三個benchmark中，求解的精度和性能均有大幅度的收益。

Pipeline

圖2 二維operator learner網絡結構

operator learner的網絡結構如上圖所示，其中主要包含如下幾個模塊：

1. 特征提取器（Feature extractor）：一維問題使用前饋神經網絡、二維問題使用CNN網絡等；

2. 基於插值的CNN（Interpolation-based CNN）：上采樣/下采樣層和CNN的堆疊得到；

3. 比特置編碼（Positional encoding）：每個網格點的笛卡爾坐標作為附加特征維連接到輸入數據。

4. 解碼器（Decoder）：編碼器學習到的錶示特征映射回原始維度。

其中網絡訓練的loss函數如下：

損失函數的主體為網絡輸出和label之間的MSEloss，另外loss中額外添加了輸出和label之間差分正則項。

其中Fourier和Galerkin類型的Transformer計算方式如下圖：

圖3 Fourier Attention

圖4 Galerkin Attention

實驗結果

1. Burger’s equation

方程定義如下：

本文中的任務是從初始時刻（t=0）得到t=1時刻的解u，模型與FNO的對比如下錶，在本問題上結果精度均優於所對比的FNO。

2. Darcy flow problem

方程定義如下：

該問題的定義是從二維的隨機幾何形狀系數a，到二維的解u的映射。模型與FNO的對比如下錶，在本問題上結果精度均優於所對比的FNO。

在對比模型精度的同時，論文也比較了模型的性能，對比結果如下錶，其中Galerkin Attention方式的Transformer在顯存占用和性能方面優勢十分明顯。

思考與總結

Galerkin Transformer從數學的角度解釋了Attention機制，並新穎地將其與算子學習相結合引入到參數化PDE的求解問題中，精度和性能均優於算子學習的“老大哥”FNO。後續可以在更高維更複雜場景上，驗證模型的有效性。

Reference

[1] Raissi M, Perdikaris P, Karniadakis G E. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations[J]. Journal of Computational Physics, 2019, 378: 686-707.

[2] Li Z, Kovachki N, Azizzadenesheli K, et al. Fourier neural operator for parametric partial differential equations[J]. arXiv preprint arXiv:2010.08895, 2020.

[3] Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need[C]//Advances in neural information processing systems. 2017: 5998-6008.

[4] Cao S. Choose a Transformer: Fourier or Galerkin[J]. arXiv preprint arXiv:2105.14995, 2021.