偏微分方程的圖神經網路解法

摘要
偏微分方程在科學和工程領域扮演著重要作用,而尋找偏微分方程的解一直是富有挑戰性的問題,需要複雜的數值求解技巧。由於在描述非歐氏系統時的高度適用性,圖神經網路給偏微分方程的求解提供了新的有趣思路。本文介紹了兩種具體方法,一種使用圖神經網路描述偏微分方程中關鍵的時間求導項,另一種則直接利用圖中的結點描述離散化後系統中的粒子。由於二者間的密切關聯,偏微分方程也能對圖神經網路的設計進行指導和幫助,相關內容將在最後的部分呈現。
背景介紹
偏微分方程(partial differential equations, or PDEs)在科學和工程領域扮演著重要作用,因為大部分自然和人造複雜系統中的控制方程(governing equation)都是偏微分方程。但絕大部分偏微分方程的求解一直是一個富有挑戰的問題,這往往涉及複雜的數值求解技巧,且相當耗時,尤其是當方程引數或是邊界、初始條件部分未知的時候。
人工智慧近年來在科研領域取得了巨大的成功,其中深度學習的發展更是對科研和人們的日常生活產生了深遠影響。除了歸功於計算資源(例如GPU)的快速發展、海量資料的採集等,深度學習的各種方法,包括典型的卷積神經網路(CNN)對於歐氏資料如影像、文字或影片等在提取其特徵時的有效性也是其成功的原因之一。
儘管深度學習在歐氏空間中取得了很大的成功,但越來越多資料呈現出非歐氏結構,這些資料需要使用圖(Graph)來描述,它包含頂點和邊兩個要素。在典型的在電子商務領域,圖可以描述使用者和產品、使用者和使用者等之間的關聯,並對使用者的行為做出預測或推薦。
在偏微分方程所在的科學領域中,我們也可以利用圖對非歐氏資料或形態的描述能力使用基於圖的深度學習方法來學習並解決一些問題。
用圖神經網路表徵時間導數項解法
一個典型的偏微分方程可以寫為

它表示一個系統關於空間座標x和時間t的狀態u隨時間的變化依賴其自身以及它關於空間座標x的一階或更高階的導數。這一類偏微分方程是一大類科學問題的基礎,它在具有傳播特徵的系統如聲波、流體、熱擴散甚至疾病發展中都有十分廣泛的應用。
在[1]中,作者希望利用圖神經網路來描述方程中u關於座標導數的函式F。首先利用線方法(method of lines, or MOL)將原方程進行離散化,選取方程描述的系統域Ω中的N個位置結點並將F在這些空間節點上進行離散化。將這些空間結點排列成一個觀測向量

F則可以近似表達為只在某個離散點上的

原方程(1)也相應改寫為如下的常微分方程組(ordinary differential equations, or ODEs):

值得注意的是,

表示座標

處的系統狀態,而

是座標

周圍的一系列領域點的下標集合。這表明為了描述座標

處的系統狀態變化函式

,我們不僅需要此處的系統狀態,還需要與之相鄰的一些離散座標處的系統狀態,這體現了系統本身對於相鄰座標間的狀態的交流與作用的依賴。這一依賴當然是符合自然規律本身的,而它也正好符合圖結點之間的鄰域依賴特徵。不僅如此,

中的鄰域相關結點數量無法預知,且對每個離散結點來說也未必相同,因而

需要對任意數量的相關結點都適用,並且與它們的排列順序無關。

圖神經網路(graph neural networks, or GNN)正好具備這些性質,因而可以採用基於資訊傳遞(message passing)機制的GNN來描述這一離散

函式:

這一網路與絕對座標不再有關,以此構建的模型也具有空間的平移或旋轉變換的不變性,因而對於有不同結點位置的的系統有很好的可擴充套件性。
這一資訊傳遞機制具體為:
其中

表示GNN中第k+1層的結點i的狀態,

表示它的鄰域傳遞來的資訊和,即每個節點基於自身資訊和鄰域中的結點資訊之和來計算自己在下一層中的狀態。我們可以用所定義的GNN中最後的第K層結點狀態來替代

:

而最終關於每個離散結點上的常微分方程的解則容易求得:

中的引數θ為資訊傳遞機制中的

函式中引數的集合,具體訓練時,將觀測狀態

網路對應輸出狀態的均值差作為代價函式即可。

這一方法的優點在於,雖然學習的是離散時刻的系統狀態,但是可以預測連續時間的狀態,與基於純離散時間的模型相比當資料的時間間隔較大時更有優勢。
上述方法是使用GNN表達PDE中的時間微分函式,那麼能否基於圖的本身結構來直接表達原先由PDE描述的系統並學習和預測系統隨時間的演化呢?DeepMind的一個工作很好地展示了這一思路[2]。
圖神經網路直接表達離散化粒子系統
我們知道流體的運動遵循著名的Navier-Stokes方程:
該方程的非線性特性使得它異常難解,而這也是複雜的湍流現象產生的原因。人們藉助各種數值計算方法來模擬流體的行為,常見的有基於網格離散化的物質點方法(material point method, or MPM)和光滑粒子流體動力學(smooth particle hydrodynamics, or SPH). 但這些方法對算力都有著巨大的需求,運算成本很高,如圖1中擷取的某機翼上的複雜湍流模擬需要消耗3500萬核心小時的計算資源與75TB的巨大儲存。
圖1. 某機翼上的複雜湍流模擬[7]
這些數值計算方法本質是將流體離散化,比如SPH方法將流體處理成了一組粒子離散元素,這些粒子具有空間位置,攜帶了其所在位置的流體資訊如質量、速度、密度、能量等。這些粒子之間的作用顯然具有很強的鄰域相關性,藉助這些特點,可以構造一種圖神經網路,使得網路中的每個結點代表該處的離散流體粒子,如圖2中最左側圖所示,再基於一定的鄰域半徑定義相鄰結點和連線相鄰結點的邊,而粒子(結點)與邊的資訊採用多層感知機進行編碼至隱空間。
圖2. 用圖神經網路表達離散粒子系統[2]
該GNN也遵循資訊傳遞機制進行演化。在之前所述的例子中,GNN包含了多層的結點和邊的資訊,每進行一次資訊傳遞這些資訊就傳遞至網路中的下一層,每層之間的資訊傳遞的引數是不同的,需要整體進行學習。而這裡GNN在每次資訊傳遞時採用的引數相同,每進行一次資訊傳遞的過程代表了一次特定時間間隔的系統狀態演化。可以這樣做的原因是這裡的流體系統的離散化粒子間的作用對於時間具有不變性,即任何一次資訊傳遞對應的相同時間間隔前後的系統演化都遵循統一的物理規律,因此我們可以用相同的資訊傳遞引數來表徵這一規律。這樣做的另一個好處是不需要將演化過程中每一次資訊傳遞的狀態都進行儲存。
對經過了資訊傳遞後的網路再進行解碼至結點對應粒子的物理狀態,就獲得了若干時間步後的離散化後的流體系統狀態,將這個狀態與採用經典數值方法(例如SPH等)計算的基準結果進行比較就可以定義該任務的代價函式。作者展示的例子中學習好一個典型的流體系統演化所需的時間步數量為幾百到一千左右。有意思的是,雖然利用SPH等經典數值方法求得的系統狀態演化是隨時間逐漸進行的,作者在使用上述GNN學習該過程的時候卻將已求得的離散時間間隔的一系列系統狀態預先打亂,這樣可以進一步增強模型對於任意時刻系統狀態的下一步狀態的預測能力,使得該預測能力更具有一般性。
圖3是該模型預測的流體和其他粒子系統(如沙粒)的運動結果,更精確的對比表明該模型比之前基於卷積神經網路等深度學習方法[3]在此任務上的預測結果更接近參考值,同時相對於產生參考值的傳統數值方法有10倍到100倍的提速。
圖3. 該模型預測的流體和沙粒的運動結果[2]
偏微分方程對圖神經網路的啟示
另一個有趣的問題是,PDE能否對圖神經網路做出啟示和指導呢?
圖卷積網路(GCN)中相比於卷積神經網路(CNN)往往是淺層的,而隨著網路的層數增多容易產生過平滑的問題,也就是圖上的結點和邊的值趨向於一致而不可分辨,使得網路失去了相應的功能。另外GCN在其指定的應用中調優後,若遷移至別的應用往往表現不佳。表1的例子中,效果出色的點雲分類網路DGCNN在另一個引用圖結點分類問題Cora中表現就很差。
表1.DCGNN在不同任務下的表現[6]
近期已經有很多工作指出了PDEs和CNNs的潛在關聯[4,5],即可以將CNNs等深度神經網路本身看成由PDE指導的動力系統,而CNNs中的每一層即是PDE離散化後的一步時間演進。這是因為CNNs中的卷積核可以看成多個線性求導運算元的組合,例如下方的3X3卷積核可以展開成後面9個子核的線性組合,而每個子核則對應一種離散的網格求導運算元,即
CNNs作用在影像上,而影像可以看成是一個特殊的圖,圖上的每個結點與其相鄰的結點相連線,由於CNNs中卷積核的權重是不變的,因而CNNs可以看成是所有邊都不變的特殊的圖神經網路。
類似地[6],在圖神經網路中我們可以定義求導運算元G
其中f是定義在圖上的函式,

是結點離散化後第i個結點上的f值。W是圖中邊的連線權重矩陣。透過G我們可以進一步定義諸如圖上的Laplacian運算元

常用的還有邊上求平均運算元
等。有了這些基本運算元就可以定義後續的PDE到圖神經網路的對映,通常這足以構建很多問題所需的PDEs在圖上的表達,如熱擴散PDE
或是雙曲型波動PDE
亦或是兩者結合的具有擴散性的波動PDE
而這些方程右邊的項來自於圖的能量項

對於U的梯度

基於Leapfrog方法可以很容易地將上述方程離散化,上述(18)式基於圖的離散化表達則為
自然地,由於熱擴散方程的解隨著時間推移本身就會變得在空間上越來越平滑,與之結構對應的圖卷積網路隨著層數增多容易產生的過平滑問題也就很好解釋了。影像分類任務中由此造成的過平滑是無關緊要的,因為我們不關注單個結點的值,更多關注整體的一種平均。當我們需要關注結點之間的差異資訊時,過平滑就會使得網路效能很差。這時候如果仍然堅持既有的與擴散型PDE相似的GCN就不再合適了,而與雙曲型波動PDE相似的GCN則可以更好地完成任務。如果應用如(19)式的同時包含兩者特性的GCN網路,則可以大大提升網路在不同應用中的可擴充套件性。
作者提供的模型應用了(19)式的混合擴散與波動型PDEs的GCN,並將α也作為一個需要學習的引數,則該網路在兩種完全不同型別的任務中都能取得SOTA的表現。在圖4中左側的結點分類任務中,α隨著訓練自動退縮回接近0的狀態,表明了此問題中擴散現象佔據的絕對主導,而由此帶來圖上特徵的過平滑也自然不會影像網路的效能。右側的例子中隨著訓練的進行α則增長到了接近1,這是由於此在FAUST資料集上的形態對應問題中波動現象又佔據了絕對主導。
圖4. 不同任務中不同型別圖神經網路的自動區分和匹配[6]
小結
我們看到,使用圖神經網路可以幫助求解偏微分方程,而不同形式的偏微分方程又能指導圖神經網路的設計。這些方法之所以可行,在於這些場景下它們分別描述的是同一類系統,系統遵循相似的內在演化機制。偏微分方程對演化機制有精確的數學描述,圖神經網路則將這些機制抽象成了結點(或邊)之間的資訊傳遞。由於圖對自然界和人類社會中廣泛存在的非歐氏資料或系統的描述有著天然的結構適用性,我們期待圖神經網路對科學研究和社會經濟等的進一步推動。
參考文獻
[1] Valerii Iakovlev, et. al., “Learning Continuous-time PDEs From Sparse Data with Graph Neural Networks”, arXiv:2006.08956.
[2] Alvaro Sanchez-Gonzalez, et. al., “Learning to Simulate Complex Physics with Graph Networks”, 2020 International Conference on Machine Learning Volume: 1, pp 8459-8468.
[3] Ummenhofer, et. al., “Lagrangian Fluid Simulation with Continuous Convolutions”, In International Conference on Learning Representations, 2020.
[4] Eldad Haber and Lars Ruthotto, “Stable architectures for deep neural networks”, Inverse Problems, 34 (1), 2017.
[5] Bo Chang, et. al., “Reversible architectures for arbitrarily deep residual neural networks”, Thirty-Second AAAI Conference on Artificial Intelligence, 2018.
[6] Moshe Eliasof and Eldad Haber, “PDE-GCN: Novel Architectures for Graph Neural Networks Motivated by Partial Differential Equations”, arXiv: 2108.01938.
[7] Hosseini S M, Vinuesa R, Schlatter P, et al. "Direct numerical simulation of the flow around a wing section at moderate Reynolds number[J]", International Journal of Heat and Fluid Flow, 2016, 61: 117-128.
關於壁仞科技研究院
壁仞科技研究院作為壁仞科技的前沿研究部門,旨在研究新型智慧計算系統的關鍵技術,重點關注新型架構,先進編譯技術和設計方法學,並將逐漸拓展研究方向,探索未來智慧系統的各種可能。壁仞科技研究院秉持開放的原則,將積極投入各類產學研合作並參與開源社群的建設,為相關領域的技術進步做出自己的貢獻。
掃碼關注我們

相關文章