參數偏微分方程的傅里葉神經算子
論文鏈接:https://arxiv.org/abs/2010.08895
項目鏈接:https://github.com/neuraloperator/neuraloperator
作者博客:https://zongyi-li.github.io/blog/2020/fourier-pde/
- 參數偏微分方程的傅里葉神經算子
- ABSTRACT
- 1 INTRODUCTION
- 2 LEARNING OPERATORS
- 3 NEURAL OPERATOR
- 4 FOURIER NEURAL OPERATOR
- 5 NUMERICAL EXPERIMENTS
- 5.1 伯格斯方程BURGERS’ EQUATION
- 5.2 達西滲流DARCY FLOW
- 5.3 納維-斯托克斯方程Navier-Stokes equations
- 5.4 Zero-shot超分辨率
- 5.5 貝葉斯逆問題
- 6 DISCUSSION AND CONCLUSION
- A APPENDIX
- A.1 符號表
- A.2 光譜分析
- A.3 數據生成
- A.4 伯格斯方程和達西滲流結果
- A.5 貝葉斯逆問題
ABSTRACT
神經網絡的經典發展主要集中在有限維歐幾里得空間之間的學習映射。最近,這被推廣到學習函數空間之間映射的神經算子。對于偏微分方程,神經算子直接學習從任意函數參數依賴到解的映射。因此,他們學習了整個偏微分方程家族,而不像經典方法只解一個方程實例。在這項工作中,我們通過直接在傅里葉空間中參數化積分核來制定一個新的神經算子,允許一個表達和高效的架構。我們對Burgers方程、Darcy流和Navier-Stokes方程進行了實驗。傅里葉神經算子是第一個成功模擬zero-shot超分辨率湍流的基于ML的方法。與傳統的PDE求解器相比,它的速度快了三個數量級。此外,在固定分辨率下,與以往基于學習的求解器相比,它具有更高的精度。
1 INTRODUCTION
科學和工程中的許多問題涉及到對某些參數的不同值反復求解復雜的偏微分方程系統。例子出現在分子動力學、微觀力學和湍流中。通常這樣的系統需要精細的離散化,以便捕捉被建模的現象。因此,傳統的數值求解器速度很慢,有時效率低下。例如,當設計諸如翼型之類的材料時,需要解決相關的逆問題,其中需要對正演模型進行數千次評估。一種快速的方法可以使這類問題變得可行。
傳統解決方案vs.數據驅動的方法。傳統的求解方法如有限元法(FEM)和有限差分法(FDM)是通過對空間進行離散來求解方程的。因此,它們對分辨率進行了權衡:粗網格速度快,但精度較低;精細網格是精確的,但速度很慢。如上所述,復雜的PDE系統通常需要非常精細的離散化,因此對于傳統的求解器來說非常具有挑戰性和耗時。另一方面,數據驅動的方法可以直接從數據中學習方程族的軌跡。因此,基于學習的方法可以比傳統的求解器快幾個數量級。
機器學習方法可能是科學學科革命的關鍵,因為它提供了近似或增強傳統解決方案的快速解決方案(Raissi等人,2019;Jiang等,2020;Greenfeld等人,2019;Kochkov et al, 2021)。然而,經典的神經網絡在有限維空間之間映射,因此只能學習與特定離散化相關的解決方案。這通常是實際應用的限制,因此需要發展網格不變神經網絡。我們首先概述了兩種主流的基于神經網絡的偏微分方程求解方法——有限維算子和神經有限元法。
有限維空間算子。這些方法將解算子參數化為有限維歐幾里得空間之間的深度卷積神經網絡(Guo et al ., 2016);朱& Zabaras (2018);Adler & Oktem (2017);Bhatnagar等人(2019);Khoo等人(2017)。根據定義,這些方法依賴于網格,并且需要針對不同的分辨率和離散化進行修改和調整,以實現一致的誤差(如果可能的話)。此外,這些方法僅限于訓練數據的離散化大小和幾何形狀,因此不可能在域中的新點查詢解。相反,對于我們的方法,我們展示了誤差對網格分辨率的不變性,以及在網格之間傳遞解決方案的能力。
Neural-FEM。第二種方法直接將解函數參數化為神經網絡(E & Yu, 2018;Raissi等人,2019;Bar & Sochen, 2019;Smith等人,2020;Pan & Duraisamy, 2020)。此方法旨在為PDE的一個特定實例建模,而不是為解決方案算子建模。它是網格無關的和準確的,但對于任何給定的新的函數參數/系數實例,它需要訓練一個新的神經網絡。該方法與有限元等經典方法非常相似,用神經網絡空間代替局部基函數的有限集的線性跨度。神經有限元法與經典方法存在著相同的計算問題:需要對每一個新實例求解優化問題。此外,該方法僅限于底層PDE已知的設置。
神經算子。最近,一項新的研究提出了用神經網絡學習無網格、無限維算子(Lu et al, 2019;Bhattacharya et al ., 2020;尼爾森&斯圖爾特,2020;李等,2020b;a;Patel et al, 2021)。神經算子通過產生一組可用于不同離散化的網絡參數,彌補了上述有限維算子方法的網格依賴性質。它具有在網格之間傳遞解的能力。此外,神經算子只需要訓練一次。獲得一個新的參數實例的解只需要網絡的前向傳遞,從而減輕了神經有限元方法中產生的主要計算問題。最后,神經算子不需要底層PDE的知識,只需要數據。到目前為止,由于計算積分算子的成本,神經算子還沒有產生有效的數值算法,可以在有限維環境中與卷積或循環神經網絡的成功相媲美。通過快速傅里葉變換,我們的工作減輕了這個問題。
傅里葉變換。傅里葉變換經常用于求解微分方程的頻譜方法,因為微分相當于傅里葉域中的乘法。傅里葉變換在深度學習的發展中也發揮了重要作用。理論上,它們出現在普遍近似定理的證明中(Hornik等人,1989),經驗上,它們已被用于加速卷積神經網絡(Mathieu等人,2013)。涉及傅里葉變換或使用正弦激活函數的神經網絡架構也被提出和研究(Bengio等人,2007;Mingo et al, 2004;Sitzmann et al, 2020)。最近,一些譜方法已經擴展到神經網絡(Fan et al .,20119a;b;Kashinath et al, 2020)。在這些工作的基礎上,我們提出了一個直接在傅里葉空間中定義的神經算子體系結構,具有準線性時間復雜度和最先進的近似能力。
我們的貢獻。我們引入傅里葉神經算子,這是一種新的深度學習架構,能夠學習無窮維函數空間之間的映射;積分算子被限制為卷積,并通過傅里葉域中的線性變換實例化。
- 傅里葉神經算子是第一個學習湍流狀態下Navier-Stokes方程家族的分辨率不變解算子的工作,其中以前基于圖的神經算子不收斂。
- 通過構造,該方法共享相同的學習網絡參數,而不管在輸入和輸出空間上使用的離散化。它可以實現零拍攝超分辨率:在較低分辨率上訓練,直接在較高分辨率上評估,如圖1所示。
- 所提出的方法始終優于所有現有的深度學習方法,即使將分辨率固定為64×64。它在Burgers方程上的錯誤率降低了30%,在Darcy Flow上降低了60%,在Navier Stokes(黏度ν = 1e?4的湍流狀態)上降低了30%。在學習整個時間序列的映射時,當黏度ν = 1e?3時,該方法的誤差< 1%,當黏度ν = 1e?4時,該方法的誤差為8%。
- 在256 × 256網格上,傅里葉神經算子的推理時間僅為0.05 s,而用于求解Navier-Stokes的偽譜方法的推理時間為2.20 s。盡管該方法具有巨大的速度優勢,但在下游應用(如解決Bayesian逆問題)中使用時,其精度不會下降,如圖6所示。
我們觀察到,所提出的框架可以近似于高度非線性、高頻模式和緩慢能量衰減的偏微分方程中的復雜算子上升。神經算子的力量來自于線性、全局積分算子(通過傅里葉變換)和非線性、局部激活函數的結合。類似于標準神經網絡通過結合線性乘法和非線性激活來近似高度非線性函數的方式,所提出的神經算子可以近似高度非線性算子。
2 LEARNING OPERATORS
我們的方法從觀察到的輸入輸出對的有限集合中學習兩個無限維空間之間的映射。設 D ? R d D\subset\mathbb{R}^d D?Rd 為有界開集,且 A = A ( D ; R d a ) \mathcal{A}=\mathcal{A}(D;\mathbb{R}^{d_a}) A=A(D;Rda?), U = U ( D ; R d u ) U=\mathcal{U}(D;\mathbb{R}^{d_u}) U=U(D;Rdu?)是分別取 R d a \mathbb{R}^{d_a} Rda?和 R d u \mathbb{R}^{d_u} Rdu?值的函數的可分離巴拿赫空間。此外,讓 G ? : A → U G^\dagger:\mathcal{A}\to\mathcal{U} G?:A→U是一個(典型的)非線性映射。我們研究作為參數偏微分方程的解算子出現的映射 G ? G^\dagger G?-參見第5節的例子。假設我們有觀測值 { a j , u j } j = 1 N \left\{a_j,u_j\right\}_{j=1}^N {aj?,uj?}j=1N?,其中 a j ~ μ a_j\sim\mu aj?~μ是A上支持的概率測度 μ \mu μ的 A \mathcal{A} A序列, u j = G ? ( a j ) u_j=G^\dagger(a_j) uj?=G?(aj?)可能被噪聲破壞。我們的目標是通過構造一個參數映射來建立 G ? G^\dagger G?的近似值:
G : A × Θ → U or?equivalently, G θ : A → U , θ ∈ Θ (1) G:\mathcal{A}\times\Theta\to\mathcal{U}\quad\text{or equivalently,}\quad G_\theta:\mathcal{A}\to\mathcal{U},\quad\theta\in\Theta \tag{1} G:A×Θ→Uor?equivalently,Gθ?:A→U,θ∈Θ(1)
對于某個有限維參數空間 Θ \Theta Θ,通過選擇 θ ? ∈ Θ \theta^\dagger\in\Theta θ?∈Θ,使得 G ( ? , θ ? ) = G θ ? ≈ G ? G(\cdot,\theta^\dagger)=G_{\theta^\dagger}\approx G^\dagger G(?,θ?)=Gθ??≈G?。這是在無限維中學習的自然框架,就像我們可以定義一個代價函數 C : U × U → R C:\mathcal{U}\times\mathcal{U}\to\mathbb{R} C:U×U→R然后求問題的最小值。
min ? θ ∈ Θ E a ~ μ [ C ( G ( a , θ ) , G ? ( a ) ) ] (2) \min_{\theta\in\Theta}\mathbb{E}_{a\sim\mu}[C(G(a,\theta),G^\dagger(a))] \tag{2} θ∈Θmin?Ea~μ?[C(G(a,θ),G?(a))](2)
這與經典的有限維環境直接相似(Vapnik, 1998)。在無限維的情況下,如何證明最小值的存在,仍然是一個具有挑戰性的開放性問題。我們將在測試列車設置中通過使用數據驅動的經驗近似來解決這個問題,該近似用于確定θ并測試近似的準確性。因為我們在無限維的情況下概念化了我們的方法,所以所有有限維的近似都有一個共同的參數集,這些參數集在無限維中是一致的。表示法表見附錄A.3。
可學習算子。近似運算符 G ? G^{\dagger} G?是一個不同的任務,通常比為參數 a ∈ A a\in\mathcal{A} a∈A的單個實例找到PDE的解 u ∈ U u\in\mathcal{U} u∈U更具挑戰性。大多數現有的方法,從經典的有限元素,有限差分,以及有限體積到現代機器學習方法,如物理信息神經網絡(PINN)(Raissi et al.,2019)針對后者,因此計算成本可能很高。這對于需要為許多不同的參數實例提供PDE解決方案的應用程序來說是不切實際的。另一方面,我們的方法直接接近運算符,因此更便宜、更快,與傳統的求解器相比,可以節省大量的計算量。有關貝葉斯逆問題的示例應用程序,請參見第5.5節。
離散化。由于我們的數據 a j a_j aj?和 u j u_j uj?通常都是函數,因此我們假設只能訪問逐點求值。令 D j = { x 1 , … , x n } ? D {D}_{j}= \{ x_{1}, \ldots , x_{n}\} \subset D Dj?={x1?,…,xn?}?D是域 D D D的 n n n點離散化,并假設我們有觀測值 a j ∣ D j ∈ R n × d a , u j ∣ D j ∈ R n × d v a_{j}| _{D_{j}}\in \mathbb{R} ^{n\times d_{a}}, u_{j}| _{D_{j}}\in \mathbb{R} ^{n\times d_{v}} aj?∣Dj??∈Rn×da?,uj?∣Dj??∈Rn×dv?,對于以 j j j為索引的輸入輸出對的有限集合。為了離散化不變量,神經算子可以對任意 x ∈ D x\in D x∈D產生答案 u ( x ) u(x) u(x),可能是 x ? D j x\notin D_j x∈/Dj?。這樣的性質是非常可取的,因為它允許在不同的網格幾何和離散之間傳遞解。
3 NEURAL OPERATOR
在(Li et al ., 2020b)中提出的神經算子被表述為迭代架構 v 0 ? v_0\mapsto v0?? v 1 ? … ? v T v_1\mapsto\ldots\mapsto v_T v1??…?vT? 其中 v j v_j vj? , j = 0 , 1 , … , T ? 1 j=0,1,\ldots,T-1 j=0,1,…,T?1是一個取 R v d \mathbb{R}^d_v Rvd?值的函數序列。如圖2 (a)所示,輸入 a ∈ A a\in\mathcal{A} a∈A首先通過局部變換 P P P提升到更高的維度表示 v 0 ( x ) = P ( a ( x ) ) v_0(x)=P(a(x)) v0?(x)=P(a(x)),該變換通常由淺全連接神經網絡參數化。然后我們應用幾個迭代的更新 v t ? v t + 1 v_{t}\mapsto v_{t+ 1} vt??vt+1?(定義見下)。輸出 u ( x ) = Q ( v T ( x ) ) u( x) = Q( v_{T}( x) ) u(x)=Q(vT?(x))是 v T v_{T} vT?通過局部變換 Q : R d v → R d u Q: \mathbb{R} ^{d_{v}}\to \mathbb{R} ^{d_{u}} Q:Rdv?→Rdu?的投影。在每次迭代中,更新 v t ? v t + 1 v_t\mapsto v_{t+1} vt??vt+1?被定義為非局部積分算子 K \mathcal{K} K和局部非線性激活函數 σ \sigma σ的組合。
定義1(迭代更新) 定義對表示 v t ? v t + 1 v_t\mapsto v_{t+1} vt??vt+1?的更新,
v t + 1 ( x ) : = σ ( W v t ( x ) + ( K ( a ; ? ) v t ) ( x ) ) , ? x ∈ D (2) v_{t+1}(x):=\sigma\Big(Wv_t(x)+\Big(K(a;\phi)v_t\Big)(x)\Big),\quad\forall x\in D \tag{2} vt+1?(x):=σ(Wvt?(x)+(K(a;?)vt?)(x)),?x∈D(2)
其中 K : A × Θ K → L ( U ( D ; R d v ) , U ( D ; R d v ) ) K: \mathcal{A} \times \Theta _{\mathcal{K} }\to \mathcal{L} ( U( D; \mathbb{R} ^{d_v}) , \mathcal{U} ( D; \mathbb{R} ^{d_v}) ) K:A×ΘK?→L(U(D;Rdv?),U(D;Rdv?)) 映射到 U ( D ; R d v ) \mathcal{U}(D;\mathbb{R}^{d_v}) U(D;Rdv?),參數化為 ? ∈ Θ K , W : R d v → R d v \phi\in\Theta_{\mathcal{K}},W:\mathbb{R}^{d_v}\to\mathbb{R}^{d_v} ?∈ΘK?,W:Rdv?→Rdv?是一個線性變換, σ : R → R \sigma:\mathbb{R}\to\mathbb{R} σ:R→R是一個非線性激活函數,其作用是按組件定義的。
我們選擇 K ( a ; ? ) \mathcal{K}(a;\phi) K(a;?)是由神經網絡參數化的核積分變換。
定義2(核積分算子 K \mathcal{K} K) 定義映射到公式(2)中的核積分算子
( K ( a ; ? ) v t ) ( x ) : = ∫ D κ ( x , y , a ( x ) , a ( y ) ; ? ) v t ( y ) d y , ? x ∈ D (3) \big(\mathcal{K}(a;\phi)v_t\big)(x):=\int_D\kappa\big(x,y,a(x),a(y);\phi\big)v_t(y)\mathrm{d}y,\quad\forall x\in D \tag{3} (K(a;?)vt?)(x):=∫D?κ(x,y,a(x),a(y);?)vt?(y)dy,?x∈D(3)
其中 κ ? : R 2 ( d + d a ) → R d v × d v \kappa_{\phi}:\mathbb{R}^{2(d+d_{a})}\to\mathbb{R}^{d_{v}\times d_{v}} κ??:R2(d+da?)→Rdv?×dv?是 ? ∈ Θ K \phi\in\Theta_{\mathcal{K}} ?∈ΘK?參數化的神經網絡。
這里 κ ? \kappa_{\phi} κ??扮演核函數的角色,我們從數據中學習。定義1和定義2共同構成了Li等人(2020b)首次提出的將神經網絡推廣到無限維空間的方法。注意,即使積分算子是線性的,神經算子也可以通過將線性積分算子與非線性激活函數組合來學習高度非線性的算子,類似于標準神經網絡。
如果我們去除對函數 a a a的依賴,并施加 κ ? ( x , y ) = κ ? ( x ? y ) \kappa_\phi(x,y)=\kappa_\phi(x-y) κ??(x,y)=κ??(x?y),我們得到公式(3)是一個卷積算子,從基本解的角度來看,這是一個自然的選擇。在下一節中,我們通過直接在傅里葉空間中參數化 κ ? \kappa_\phi κ??并使用快速傅里葉變換(FFT)來有效地計算公式(3)來利用這一事實。這導致了一個快速架構,可以獲得PDE問題的最新結果。
4 FOURIER NEURAL OPERATOR
我們建議用傅里葉空間中定義的卷積算子代替公式(3)中的核積分算子。設 F \mathcal{F} F表示函數 F F F的傅里葉變換 f : D → R d v {f}:D\to\mathbb{R}^{d_{v}} f:D→Rdv?和 F ? 1 {\mathcal{F}^{-1}} F?1是它的逆。
( F f ) j ( k ) = ∫ D f j ( x ) e ? 2 i π ? x , k ? d x , ( F ? 1 f ) j ( x ) = ∫ D f j ( k ) e 2 i π ? x , k ? d k (\mathcal{F}f)_j(k)=\int_Df_j(x)e^{-2i\pi\langle x,k\rangle}\mathrm{d}x,\quad(\mathcal{F}^{-1}f)_j(x)=\int_Df_j(k)e^{2i\pi\langle x,k\rangle}\mathrm{d}k (Ff)j?(k)=∫D?fj?(x)e?2iπ?x,k?dx,(F?1f)j?(x)=∫D?fj?(k)e2iπ?x,k?dk
對于 j = 1 , … , d v j=1,\ldots,d_v j=1,…,dv? ,其中 i = ? 1 i=\sqrt-1 i=??1是虛單位。令在公式(3)中的 κ ? ( x , y , a ( x ) , a ( y ) ) = κ ? ( x ? y ) \kappa_\phi(x,y,a(x),a(y))=\kappa_\phi(x-y) κ??(x,y,a(x),a(y))=κ??(x?y),并應用卷積定理得到
( K ( a ; ? ) v t ) ( x ) = F ? 1 ( F ( κ ? ) ? F ( v t ) ) ( x ) , ? x ∈ D . \big(\mathcal{K}(a;\phi)v_t\big)(x)=\mathcal{F}^{-1}\big(\mathcal{F}(\kappa_\phi)\cdot\mathcal{F}(v_t)\big)(x),\quad\forall x\in D. (K(a;?)vt?)(x)=F?1(F(κ??)?F(vt?))(x),?x∈D.
因此,我們建議在傅里葉空間中直接參數化 κ ? \kappa_\phi κ??。
**定義3(傅里葉積分算子 ( K (\mathcal{K} (K) ** 定義傅里葉積分算子:
( K ( ? ) v t ) ( x ) = F ? 1 ( R ? ? ( F v t ) ) ( x ) ? x ∈ D (4) \big(\mathcal{K}(\phi)v_t\big)(x)=\mathcal{F}^{-1}\big(R_\phi\cdot(\mathcal{F}v_t)\big)(x)\quad\forall x\in D \tag{4} (K(?)vt?)(x)=F?1(R???(Fvt?))(x)?x∈D(4)
其中, R φ R_φ Rφ?是周期函數 κ : D ˉ → R d v × d v \kappa:\bar{D}\to\mathbb{R}^{d_{v}\times d_{v}} κ:Dˉ→Rdv?×dv?由 ? ∈ Θ K \phi\in\Theta_{\mathcal{K}} ?∈ΘK?參數化。圖2 (b)給出了一個說明。
對于頻率模式 k ∈ D k\in D k∈D,我們有 ( F v t ) ( k ) ∈ C v d (\mathcal{F}v_t)(k)\in\mathbb{C}^d_v (Fvt?)(k)∈Cvd?和 R ? ( k ) ∈ C d v × d v R_\phi(k)\in\mathbb{C}^{d_v\times d_v} R??(k)∈Cdv?×dv?。注意,因為我們假設 κ \kappa κ是周期性的,它允許傅里葉級數展開,所以我們可以處理離散模式 k ∈ Z d k\in\mathbb{Z}^d k∈Zd。我們選擇一個有限維的參數化通過截斷在最大模態數的傅里葉級數 k max ? = ∣ Z k max ? ∣ = ∣ { k ∈ Z d : ∣ k j ∣ ≤ k max ? , j k_{\max }= | Z_{k_{\max }}| = | \{ k\in \mathbb{Z} ^d: | k_j| \leq k_{\max , j} kmax?=∣Zkmax??∣=∣{k∈Zd:∣kj?∣≤kmax,j?, for $j= 1, \ldots , d} | $ 。因此,我們直接參數化 R ? R_{\phi } R??為復值 ( k m a x × d v × d v ) ( k_{\mathrm{max}}\times d_{v}\times d_{v}) (kmax?×dv?×dv?)張量,包括截斷的傅立葉模的集合,因此從我們的符號中去掉 ? \phi ?。由于 κ \kappa κ是實值,我們施加共軛對稱性。我們注意到集合Zkmax并不是 v t v_{t} vt?的低頻模態的標準選擇。事實上,低頻模態通常是通過在 k ∈ Z d k\in\mathbb{Z}^{d} k∈Zd的1范數上放置一個上界來定義的。我們如上所述選擇 Z k max ? Z_{k_{\max}} Zkmax??,因為它允許有效的實現。
離散情況和FFT。假設域 D D D被 n ∈ N n\in\mathbb{N} n∈N個點離散,我們得到 v t ∈ R n × d v v_t\in\mathbb{R}^{n\times d_v} vt?∈Rn×dv?和 F ( v t ) ∈ C n × d v \mathcal{F}(v_t)\in\mathbb{C}^{n\times d_v} F(vt?)∈Cn×dv?。由于我們將 v t v_t vt?與一個只有 k m a x k_{\mathrm{max}} kmax?傅里葉模式的函數進行卷積,我們可以簡單地截斷較高的模式以獲得 F ( v t ) ∈ C k m a x × d v \mathcal{F} ( v_{t}) \in \mathbb{C} ^{k_{\mathrm{max}}\times d_{v}} F(vt?)∈Ckmax?×dv?。乘以權張量 R ∈ C k m a x × d v × d v R\in \mathbb{C} ^{k_{\mathrm{max}}\times d_{v}\times d_{v}} R∈Ckmax?×dv?×dv?就是
( R ? ( F v t ) ) k , l = ∑ j = 1 d v R k , l , j ( F v t ) k , j , k = 1 , … , k max ? , j = 1 , … , d v . (5) \left(R\cdot(\mathcal{F}v_t)\right)_{k,l}=\sum_{j=1}^{d_v}R_{k,l,j}(\mathcal{F}v_t)_{k,j},\quad k=1,\ldots,k_{\max},\quad j=1,\ldots,d_v. \tag{5} (R?(Fvt?))k,l?=j=1∑dv??Rk,l,j?(Fvt?)k,j?,k=1,…,kmax?,j=1,…,dv?.(5)
當離散化均勻且分辨率為 s 1 × ? × s d = n s_1\times\cdots\times s_d=n s1?×?×sd?=n時,可以用快速傅里葉變換代替 F \mathcal{F} F。對于 f ∈ R n × d v , k = ( k 1 , … , k d ) ∈ Z s 1 × ? × Z s d f\in\mathbb{R}^{n\times d_v},k=(k_1,\ldots,k_d)\in\mathbb{Z}_{s_1}\times\cdots\times\mathbb{Z}_{s_d} f∈Rn×dv?,k=(k1?,…,kd?)∈Zs1??×?×Zsd??, x = ( x 1 , … , x d ) ∈ D x=(x_1,\ldots,x_d)\in D x=(x1?,…,xd?)∈D, FFT F ^ \hat{\mathcal{F}} F^和它的逆 F ^ ? 1 \hat{\mathcal{F}}^{-1} F^?1定義為
( F ^ f ) l ( k ) = ∑ x 1 = 0 s 1 ? 1 ? ∑ x d = 0 s d ? 1 f l ( x 1 , … , x d ) e ? 2 i π ∑ j = 1 d x j k j a j , ( F ^ ? 1 f ) l ( x ) = ∑ k 1 = 0 s 1 ? 1 ? ∑ k d = 0 s d ? 1 f l ( k 1 , … , k d ) e 2 i π ∑ j = 1 d x j k j s j \begin{align} (\hat{\mathcal{F}}f)_{l}(k)&=\sum_{x_{1}=0}^{s_{1}-1}\cdots\sum_{x_{d}=0}^{s_{d}-1}f_{l}(x_{1},\ldots,x_{d})e^{-2i\pi\sum_{j=1}^{d}\frac{x_{j}k_{j}}{a_{j}}},\\ (\hat{\mathcal{F}}^{-1}f)_l(x)&=\sum_{k_1=0}^{s_1-1}\cdots\sum_{k_d=0}^{s_d-1}f_l(k_1,\ldots,k_d)e^{2i\pi\sum_{j=1}^d\frac{x_jk_j}{s_j}} \end{align} (F^f)l?(k)(F^?1f)l?(x)?=x1?=0∑s1??1??xd?=0∑sd??1?fl?(x1?,…,xd?)e?2iπ∑j=1d?aj?xj?kj??,=k1?=0∑s1??1??kd?=0∑sd??1?fl?(k1?,…,kd?)e2iπ∑j=1d?sj?xj?kj????
對于 l = 1 , . . . , d v l = 1,...,d_v l=1,...,dv?$。在這種情況下,截斷模式集變為
Z k m a x = { ( k 1 , … , k d ) ∈ Z s 1 × ? × Z s d ∣ k j ≤ k m a x , j o r s j ? k j ≤ k m a x , j , f o r j = 1 , … , d } . Z_{k_{\mathrm{max}}}=\{(k_1,\ldots,k_d)\in\mathbb{Z}_{s_1}\times\cdots\times\mathbb{Z}_{s_d}\mid k_j\leq k_{\mathrm{max},j}\mathrm{~or~}s_j-k_j\leq k_{\mathrm{max},j},\mathrm{~for~}j=1,\ldots,d\}. Zkmax??={(k1?,…,kd?)∈Zs1??×?×Zsd??∣kj?≤kmax,j??or?sj??kj?≤kmax,j?,?for?j=1,…,d}.
當實現時, R R R被視為 ( s 1 × ? × s d × d v × d v ) (s_1\times\cdots\times s_d\times d_v\times d_v) (s1?×?×sd?×dv?×dv?)張量, Z k m a x Z_{k_\mathrm{max}} Zkmax??的上述定義對應于 R R R的“corners”,這允許通過矩陣向量乘法直接并行實現公式(5)。在實踐中,我們發現選擇 k m a x , j = 12 k_{\mathrm{max},j}=12 kmax,j?=12,每個通道產生 k m a x = 1 2 d k_\mathrm{max}=12^d kmax?=12d個參數,足以滿足我們考慮的所有任務。
R R R的參數化。通常, R R R可以定義為依賴于 ( F a ) (\mathcal{F}a) (Fa)來并行公式(3)。實際上,我們可以定義 R ? : Z d × R d v → R d v × d v R_\phi:\mathbb{Z}^d\times\mathbb{R}^{d_v}\to\mathbb{R}^{d_v\times d_v} R??:Zd×Rdv?→Rdv?×dv? 作為一個參數函數,它映射 ( k , ( F a ) ( k ) ) (k,(\mathcal{F}a)(k)) (k,(Fa)(k))到適當的傅里葉模態的值。我們對 R ? R_\phi R??的線性和神經網絡參數化進行了實驗。我們發現線性參數化與之前描述的直接參數化具有相似的性能,而神經網絡的性能較差。這可能是由于空間 Z d \mathbb{Z}^d Zd的離散結構。我們在這項工作中的實驗集中在上述直接參數化上。
離散化的不變性。傅里葉層是離散不變的因為它們可以從任意離散的函數中學習并求值。由于參數是直接在傅里葉空間中學習的,因此在物理空間中解析函數就等于在 R d {R}^d Rd上處處定義良好的 e 2 π i ? x , k ? e^{2\pi i\langle x,k\rangle} e2πi?x,k?的基礎上進行投影。這使我們能夠實現如5.4節所示的Zero-shot超分辨率。此外,我們的體系結構在輸入和輸出的任何分辨率上都有一致的錯誤。另一方面,請注意,在圖3中,我們比較的標準CNN方法的誤差隨著分辨率的增加而增加。
準線性復雜度。權重張量 R R R包含 k m a x < n k_\mathrm{max}<n kmax?<n個模式,因此內部乘法的復雜度為 O ( k max ? ) O(k_{\max}) O(kmax?)。因此,大部分的計算成本在于計算傅里葉變換 F ( v t ) \mathcal{F}(v_t) F(vt?)及其逆。一般傅里葉變換的復雜度為 O ( n 2 ) O(n^2) O(n2),然而,由于我們截斷了序列,復雜度實際上是 O ( n k m a x ) O(nk_\mathrm{max}) O(nkmax?),而FFT的復雜度為 O ( n log ? n ) O(n\log n) O(nlogn)。通常,我們發現使用FFT非常有效。然而,需要統一的離散化。
5 NUMERICAL EXPERIMENTS
在本節中,我們將提出的傅里葉神經算子與多個有限維架構以及基于算子的近似方法在一維Burgers方程、二維Darcy流問題和二維Navier-Stokes方程上進行比較。數據生成過程分別在附錄A.3.1、A.3.2和A.3.3中討論。我們不與傳統的求解器(FEM/FDM)或神經-FEM類型的方法進行比較,因為我們的目標是產生可用于下游應用的有效算子近似。我們將在第5.5節中演示貝葉斯逆問題的一個這樣的應用。
我們通過將(2)和(4)中指定的四個傅里葉積分算子層與ReLU激活以及批處理歸一化疊加來構建傅里葉神經算子。除非另有說明,否則我們使用N = 1000個訓練實例和200個測試實例。我們使用Adam優化器訓練500個epoch,初始學習率為0.001,每100個epoch減半。我們設置一維問題 k max ? , j = 16 , d v = 64 k_{\max,j}=16,d_{v}=64 kmax,j?=16,dv?=64;對于二維問題, k max ? , j = 12 , d v = 32 k_{\max,j}=12,d_{v}=32 kmax,j?=12,dv?=32。低分辨率數據從高分辨率下采樣。所有的計算都在一個具有16GB內存的Nvidia V100 GPU上進行。
分辨率評估。傳統的PDE求解方法如FEM和FDM近似于單一函數,因此它們對連續體的誤差隨著分辨率的增加而減小。另一方面,只要解出所有相關信息,算子近似與數據離散的方式無關。分辨率不變運算符在不同分辨率下具有一致的錯誤率,如圖3所示。此外,分辨率不變運算符可以實現Zero-shot超分辨率,如第5.4節所示。
時間無關問題的基準(Burgers和Darcy):
NN:一個簡單的逐點前饋神經網絡。
RBM:經典的約簡基法(使用POD基) (DeVore, 2014)。
FCN:基于全卷積網絡的最先進的神經網絡架構(Zhu & Zabaras, 2018)。
PCANN:一種算子方法,使用PCA作為輸入和輸出數據的自編碼器,并用神經網絡對潛在空間進行插值(Bhattacharya等人,2020)。
GNO:原始圖神經算子(Li et al ., 2020b)。
MGNO:多極圖神經算子(Li et al ., 2020a)。
LNO:一種基于核低秩分解的神經算子方法 κ ( x , y ) : = ∑ i = 1 r ? j ( x ) ψ j ( y ) \kappa(x,y):=\sum_{i=1}^r\phi_j(x)\psi_j(y) κ(x,y):=∑i=1r??j?(x)ψj?(y),類似于(Lu et al, 2019)中提出的未堆疊DeepONet。
FNO:新提出的傅里葉神經算子。
時間相關問題的基準(Navier-Stokes):
ResNet:帶有殘差連接的18層二維卷積(He et al, 2016)。
U-Net:圖像到圖像回歸任務的流行選擇,由四個塊組成,具有二維卷積和反卷積(Ronneberger等人,2015)。
TF-Net:基于空間和時間卷積的組合,設計用于學習湍流的網絡(Wang et al ., 2020)。
FNO-2d:具有時域RNN結構的二維傅里葉神經算子。
FNO-3d:在時空中直接卷積的三維傅立葉神經算子。
5.1 伯格斯方程BURGERS’ EQUATION
一維Burgers方程是一種非線性偏微分方程,具有多種應用,包括模擬粘性流體的一維流動。它的形式是:
? t u ( x , t ) + ? x ( u 2 ( x , t ) / 2 ) = ν ? x x u ( x , t ) , x ∈ ( 0 , 1 ) , t ∈ ( 0 , 1 ] u ( x , 0 ) = u 0 ( x ) , x ∈ ( 0 , 1 ) (6) \partial_{t}u(x,t)+\partial_{x}(u^{2}(x,t)/2)=\nu\partial_{xx}u(x,t),\quad x\in(0,1),t\in(0,1] \\ u(x,0)=u_{0}(x),\quad x\in(0,1)\tag{6} ?t?u(x,t)+?x?(u2(x,t)/2)=ν?xx?u(x,t),x∈(0,1),t∈(0,1]u(x,0)=u0?(x),x∈(0,1)(6)
具有周期邊界條件,其中 u 0 ∈ L p e r 2 ( ( 0 , 1 ) ; R ) u_0\in L_{\mathrm{per}}^2( ( 0, 1) ; \mathbb{R} ) u0?∈Lper2?((0,1);R)為初始條件, ν ∈ R + \nu\in\mathbb{R}_+ ν∈R+? 為粘度系數。我們的目標是學習將初始條件映射到時刻1的解的算子, G ? : L p e r 2 ( ( 0 , 1 ) ; R ) → H p e r r ( ( 0 , 1 ) ; R ) G^\dagger : L_{\mathrm{per}}^2( ( 0, 1) ; \mathbb{R} ) \to H_{\mathrm{per}}^r( ( 0, 1) ; \mathbb{R} ) G?:Lper2?((0,1);R)→Hperr?((0,1);R)由 u 0 ? u ( ? , 1 ) u_0\mapsto u(\cdot,1) u0??u(?,1)定義,對于任意 r > 0 r>0 r>0。
我們的實驗結果見圖3 (a)和表3(附錄A.3.1)。與任何基準測試相比,我們提出的方法獲得了最低的相對誤差。此外,誤差隨分辨率不變,而基于卷積神經網絡的方法(FCN)的誤差隨分辨率增大。與在物理空間中使用Nystrom采樣的其他神經算子方法(如GNO和MGNO)相比,傅里葉神經算子更精確,計算效率更高。
5.2 達西滲流DARCY FLOW
考慮單元上二維達西流方程的穩態,即二階線性橢圓偏微分方程:
? ? ? ( a ( x ) ? u ( x ) ) = f ( x ) x ∈ ( 0 , 1 ) 2 u ( x ) = 0 x ∈ ? ( 0 , 1 ) 2 (7) -\nabla\cdot(a(x)\nabla u(x))=f(x)\quad x\in(0,1)^{2}\\u(x)=0\quad x\in\partial(0,1)^{2}\tag{7} ???(a(x)?u(x))=f(x)x∈(0,1)2u(x)=0x∈?(0,1)2(7)
具有Dirichlet邊界,其中 a ∈ L ∞ ( ( 0 , 1 ) 2 ; R + ) a\in L^\infty((0,1)^2;\mathbb{R}_+) a∈L∞((0,1)2;R+?)為擴散系數, f ∈ f\in f∈ L 2 ( ( 0 , 1 ) 2 ; R ) L^2((0,1)^2;\mathbb{R}) L2((0,1)2;R)為強制函數。這種PDE有許多應用,包括模擬地下流動的壓力,線性彈性材料的變形,以及導電材料的電勢。我們感興趣的是學習將擴散系數映射到解 G ? : L ∞ ( ( 0 , 1 ) 2 ; R + ) → H 0 1 ( ( 0 , 1 ) 2 ; R + ) G^\dagger:L^\infty((0,1)^2;\mathbb{R}_+)\to H_0^1((0,1)^2;\mathbb{R}_+) G?:L∞((0,1)2;R+?)→H01?((0,1)2;R+?)由 a ? u a\mapsto u a?u定義。注意,雖然PDE是線性的,但算子 G ? G^\dagger G?不是。
我們的實驗結果見圖3 (b)和表4(附錄A.3.2)。所提出的傅里葉神經算子的相對誤差比任何基準都低近一個數量級。我們再次觀察到誤差相對于分辨率的不變性。
5.3 納維-斯托克斯方程Navier-Stokes equations
我們考慮在單位環面上以渦量形式存在的粘性不可壓縮流體的二維Navier-Stokes方程:
? t w ( x , t ) + u ( x , t ) ? ? w ( x , t ) = ν Δ w ( x , t ) + f ( x ) , x ∈ ( 0 , 1 ) 2 , t ∈ ( 0 , T ] ? ? u ( x , t ) = 0 , x ∈ ( 0 , 1 ) 2 , t ∈ [ 0 , T ] w ( x , 0 ) = w 0 ( x ) , x ∈ ( 0 , 1 ) 2 (8) \begin{aligned} \partial_{t}w(x,t)+u(x,t)\cdot\nabla w(x,t)& =\nu\Delta w(x,t)+f(x), && x\in(0,1)^{2},t\in(0,T] \\ \nabla\cdot u(x,t)& =0, && x\in(0,1)^{2},t\in[0,T] \\ w(x,0)& =w_0(x), && x\in(0,1)^{2} \end{aligned}\tag{8} ?t?w(x,t)+u(x,t)??w(x,t)??u(x,t)w(x,0)?=νΔw(x,t)+f(x),=0,=w0?(x),??x∈(0,1)2,t∈(0,T]x∈(0,1)2,t∈[0,T]x∈(0,1)2?(8)
其中 u ∈ C ( [ 0 , T ] ; H p e r r ( ( 0 , 1 ) 2 ; R 2 ) ) u\in C([0,T];H_{\mathrm{per}}^{r}((0,1)^{2};\mathbb{R}^{2})) u∈C([0,T];Hperr?((0,1)2;R2))對于任意 r > 0 r>0 r>0 為速度場, w = ? × u w=\nabla\times u w=?×u為渦度, w 0 ∈ L p e r 2 ( ( 0 , 1 ) 2 ; R ) w_0\in L_{\mathrm{per}}^2( ( 0, 1) ^2; \mathbb{R} ) w0?∈Lper2?((0,1)2;R)為初始渦度, ν ∈ R + \nu\in\mathbb{R}_+ ν∈R+?為黏度系數, f ∈ f\in f∈ L p e r 2 ( ( 0 , 1 ) 2 ; R ) L_{\mathrm{per}}^2( ( 0, 1) ^2; \mathbb{R} ) Lper2?((0,1)2;R)為強制函數。我們感興趣的是學習將時間為10的渦度映射到之后某個時間的渦度的算子 T > 10 , G ? : C ( [ 0 , 10 ] ; H p e r r ( ( 0 , 1 ) 2 ; R ) ) → T> 10, G^\dagger : C( [ 0, 10] ; H_{\mathrm{per}}^r( ( 0, 1) ^2; \mathbb{R} ) ) \to T>10,G?:C([0,10];Hperr?((0,1)2;R))→ C ( ( 10 , T ] ; H p e r r ( ( 0 , 1 ) 2 ; R ) ) C( ( 10, T] ; H_{\mathrm{per}}^r( ( 0, 1) ^2; \mathbb{R} ) ) C((10,T];Hperr?((0,1)2;R)) 由 w ∣ ( 0 , 1 ) 2 × [ 0 , 10 ] ? w ∣ ( 0 , 1 ) 2 × ( 10 , T ] w|_(0,1)^2\times[0,10]\mapsto w|_{(0,1)^2\times(10,T]} w∣(?0,1)2×[0,10]?w∣(0,1)2×(10,T]?定義。已知渦度,很容易求出速度。雖然與速度相比,渦度更難建模,但它提供了更多的信息。通過對渦度問題的表述,神經網絡模型模擬了偽譜方法。我們實驗粘度 ν = 1 e ? 3 , 1 e ? 4 , 1 e ? 5 ν = 1e?3,1e?4,1e?5 ν=1e?3,1e?4,1e?5,隨著動態變混沌,最終時間T減小。由于基線方法不是分辨率不變的,我們將訓練和測試的分辨率都固定為64 × 64。
由表1可知,當數據充足時,FNO-3D的性能最佳( ν = 1 e ? 3 , N = 1000 ν = 1e?3,N = 1000 ν=1e?3,N=1000和$ ν = 1e?4,N = 10000 ) 。對于數據量不足的配置 ( )。對于數據量不足的配置( )。對于數據量不足的配置(ν = 1e?4,N = 1000 和 和 和 ν = 1e?5,N = 1000$),所有方法的誤差都>15%,其中FNO-2D達到最低。請注意,我們只給出空間分辨率為64 × 64的結果,因為我們比較的所有基準測試都是針對該分辨率設計的。增加它會降低它們的性能,而FNO會產生相同的錯誤。
2D和3D卷積。FNO-2D、U-Net、TF-Net和ResNet都在空間域中進行2D卷積,在時間域中循環傳播(2D+RNN)。算子將前10個時間步的解映射到下一個時間步(2D函數到2D函數)。另一方面,FNO-3D在時空中進行卷積。它將初始時間步長直接映射到完整的軌跡(3D函數到3D函數)。2D+RNN結構可以將解以固定區間長度?T的增量傳播到任意時間T,而Conv3D結構則固定為區間[0,T],但可以將解轉換為任意時間離散化。我們發現,與RNN結構的方法相比,3-d方法更具表現力,更容易訓練。
5.4 Zero-shot超分辨率
神經算子是網格不變性的,因此它可以在較低分辨率下進行訓練,并在較高分辨率下進行評估,而無需看到任何更高分辨率的數據(Zero-shot超分辨率)。圖1顯示了一個例子,我們在上面設置的64 × 64 × 20分辨率數據上訓練FNO-3D模型,(ν = 1e?4;N = 10000),轉換為256 × 256 × 80分辨率,展示時空超分辨率。傅里葉神經算子是基準(FNO-2D, U-Net, TF-Net和ResNet)中唯一可以實現Zero-shot超分辨率的模型。令人驚訝的是,它不僅可以在空間域而且可以在時間域實現超分辨率。
5.5 貝葉斯逆問題
在本實驗中,我們使用函數空間馬爾可夫鏈蒙特卡羅(MCMC)方法(Cotter et al, 2013)從給定時間T = 50的稀疏噪聲觀測值的Navier-Stokes初始渦度的后檢分布中提取樣本。我們將傅里葉神經算子作為替代模型與用于生成訓練測試數據的傳統求解器進行比較(兩者都在GPU上運行)。我們從后驗生成25,000個樣本(具有5,000個樣本老化期),需要對前向算子進行30,000次評估。
如圖6(附錄A.5)所示,FNO和傳統求解器恢復幾乎相同的后置均值,當向前推進時,可以很好地恢復Navier Stokes的后期動態。與之形成鮮明對比的是,FNO計算單個實例的時間為0.05 s,而傳統的求解器在優化使用最大可能的內部時間步長(不會導致爆炸)后,需要2.2s。對于使用FNO的MCMC來說,這相當于2.5分鐘,而對于傳統求解器來說,這超過了18個小時。即使我們考慮到數據生成和訓練時間(離線步驟)需要12個小時,使用FNO仍然更快!經過訓練后,FNO可以針對不同的初始條件和觀測值快速執行多個MCMC運行,而傳統的求解器每個實例需要18個小時。此外,由于FNO是可微的,它可以很容易地應用于PDE約束優化問題,而不需要伴隨方法。
光譜分析。由于我們參數化 R ? R_\phi R??的方式,由公式(4)輸出的函數每個通道最多有 k ^ max ? , j \hat{k}_{\max,j} k^max,j?個傅里葉模式。然而,這并不意味著傅里葉神經算子只能近似到 k max ? , j k_{\max,j} kmax,j?模式的函數。實際上,發生在積分算子和最終解碼器網絡 Q Q Q之間的激活函數恢復了高頻模式。作為一個例子,考慮粘度為 ν = 1 \nu=1 ν=1e-3的Navier-Stokes方程的解。將該函數截斷為20個傅立葉模式會產生約 2 % 2\% 2%的誤差,而我們的傅立葉神經算子學習參數依賴性并產生誤差 ≤ 1 % \leq1\% ≤1%的近似值,只有 k m a x , j = 12 k_{\mathrm{max},j}=12 kmax,j?=12 個參數化模式。
非周期邊界條件。傳統的傅里葉方法只適用于周期邊界條件。然而,傅里葉神經算子沒有這種限制。這是由于線性變換 W W W(偏置項)保持了非周期邊界的軌跡。以Darcy Flow和Navier-Stokes的時域具有非周期邊界條件為例,傅里葉神經算子仍能以極好的精度學習解算子。
6 DISCUSSION AND CONCLUSION
數據要求。數據驅動的方法依賴于數據的質量和數量。為了學習黏度為ν = 1e?4的Navier-Stokes方程,我們需要生成N = 10000對訓練對faj;用數值解算器。然而,對于更具挑戰性的PDE,生成一些訓練樣本可能已經非常昂貴了。未來的發展方向是將神經算子與數值求解相結合,使對數據的要求懸浮起來。
復發性結構。神經算子具有迭代結構,可以自然地表述為循環網絡,其中所有層共享相同的參數而不犧牲性能。(我們在實驗中沒有施加這個限制。)
計算機視覺。算子學習并不局限于偏微分方程。圖像自然可以被視為二維域上的實值函數,視頻只是添加了一個時間結構。因此,我們的方法是計算機視覺問題的自然選擇,其中離散化的不變性至關重要(Chi等人,2020)。
A APPENDIX
A.1 符號表
A.2 光譜分析
Navier Stokes方程數據的譜衰減如圖4所示。光譜衰減斜率為 k ? 5 / 3 k^{-{5}/{3}} k?5/3,與湍流區的能譜相匹配。我們注意到能譜不隨時間衰減。
我們還以截斷kmax的形式給出了光譜衰減,如圖5所示。我們注意到所有方程(Burgers, Darcy和Navier-Stokes ν≤1e?4)都表現出高頻模式。即使我們在傅里葉層截斷kmax = 12,傅里葉神經算子也可以恢復高頻模式。
A.3 數據生成
A.4 伯格斯方程和達西滲流結果
Burgers方程和Darcy Flow的詳細錯誤率如表3和表4所示。
A.5 貝葉斯逆問題
Navier-Stokes方程的Bayesian逆問題結果如圖6所示。可以看出,采用傅里葉神經算子作為替代算子的求解結果與傳統求解器的求解結果一樣好。