您現在的位置是:首頁 > 攝影首頁攝影
AI求偏微分方程新基準登NeurIPS,JAX比PyTorch快6倍,LeCun轉發
手機怎麼給照片ps紋身
蕭簫 發自 凹非寺
量子位 | 公眾號 QbitAI
用AI求解偏微分方程,這段時間確實有點火。
但究竟什麼樣的AI求解效果最好,卻始終沒有一個統一的定論。
現在,終於有人為這個領域製作了一個名叫
PDEBench
的完整基準,論文登上了
NeurIPS 2022
。
PDEBench不僅能當成一個大型偏微分方程資料集,也能作為新AI求解偏微分方程的基準之一——
不少“老前輩”的預訓練模型程式碼都能在這裡找到,作為一個比對基礎。
例如去年大火了一陣的FNO,幾秒鐘求解出傳統方法需要計算18個小時的偏微分方程,程式碼就被放進了PDEBench中。
這個新基準一出,LeCun也激情轉發:這領域確實很火。
所以,AI求解偏微分方程的優勢是什麼,這一基準具體提出了哪些評估方法?
為啥用AI求解偏微分方程?
偏微分方程(PDE,Partial Differential Equation),是一個生活中常見的方程。
包括預報天氣、模擬飛機空氣動力、預測疾病傳播模型,都會用到這個方程。
目前
北大數學系“韋神”
韋東奕的研究方向之一,就是流體力學中的數學問題,其中就包括偏微分方程中的Navier-Stokes方程。
所以,為啥要
用AI來求解偏微分方程
?
訓練AI的本質,是找到一種儘可能逼近真實結果的模型。
用AI求解偏微分方程,其實也是找到一種
代理模型
,來模擬偏微分方程模型。
代理模型,指找到一種近似模型,在計算量更小的同時,確保計算結果與原來的偏微分方程儘可能相似。
這與
傳統的數值方法求解偏微分方程
有著異曲同工之妙。
傳統方法往往需要透過將連續問題離散化(類似在一個連續函式上切割出很多小點),來對方程進行近似求解。
然而,傳統的數值方法非常複雜,計算量也很大;採用AI方法訓練出來的模型,卻模擬得又快又好——
繼2017年華盛頓大學提出PDE-FIND後,2018年穀歌AI又提出了資料驅動求解偏微分方程的方法,都比傳統方法要快上不少,讓更多人開始關注到AI求解偏微分方程這一領域。
2019年,布朗大學應用數學團隊提出一種名叫
PINN
(物理激發的神經網路)的方法,徹底打開了AI在物理學領域的廣泛應用。
這篇論文在理論上雖然沒有PDE-FIND和谷歌AI的方法突破性強,卻給出了非常完整的程式碼體系,使得開發人員很容易上手,讓更多研究者開發出了不同的PINN,如今它也成為AI物理最常見的框架和詞彙之一。
△PINN
去年加州理工大學和普渡大學團隊發表的一項研究,更是將偏微分方程計算時間從傳統求解的18個小時降低為1秒鐘。
這篇論文提出了一種名為
FNO
(傅立葉神經運算元)的方法,基於傅立葉變換給神經網路加上“傅立葉層”,進一步節省了近似模擬運算元的計算量。
除此之外,也有不少研究人員透過訓練一些經典AI模型,來求解偏微分方程,如
U-Net
等。
不過,無論是FNO、U-Net還是PINN,都還是基於各自給出的基準來評估AI計算偏微分方程的效果。
有沒有一個更統一、更通用的框架來評估這個領域的新突破?
更全面的AI偏微分方程基準
在這樣的背景下,研究人員提出了一種名叫PDEBench的基準。
首先是基準中包含的
資料集
,目前這些資料集已經全部歸納到GitHub中:
這裡麵包括不少經典偏微分方程問題,如Navier-Stokes方程,達西流模型、淺水波模型等等。
隨後,PDEBench提出了幾個
指標
,來從不同角度更全面地對AI模型進行評估:
最後,PDEBench還包含了幾種經典模型的
預訓練模型
程式碼,並將它們作為評估其他模型的基準之一,包括上述提到的FNO、U-Net、PINN等。
例如研究團隊將這幾個模型分別基於各資料集進行了訓練,得出的均方根誤差(RMSE)如下,也說明它們在不同偏微分方程問題上的表現並不一樣:
除此之外,團隊還將資料格式進行了統一,同時針對PDEBench的可擴充套件性進行了最佳化,因此任何人都能參與進來,給這一基準加入更多的資料集、或是更多基準模型。
值得注意的是,團隊試了試分別在PyTorch和JAX兩種框架上執行幾種預訓練模型,發現JAX的速度大約是PyTorch的
6倍
。
看來以後搞相關研究可以試試JAX框架了。
作者介紹
作者們來自德國斯圖加特大學,歐洲NEC研發中心,還有澳大利亞聯邦科學與工業研究組織(CSIRO)旗下的Data61數字創新中心。
Makoto Takamoto,歐洲NEC研發中心高階研究員,畢業於京都大學,研究方向是影象處理、圖神經網路和科學機器學習。
Timothy Praditia,斯圖加特大學博士研究生,研究興趣是開發基於資料驅動和先驗物理知識的神經網路模型。