AI求解偏微分方程新基準登NeurIPS，發現JAX計算速度比PyTorch快6倍，LeCun轉發：這領域確實很火
包含更全面數據集
蕭簫 發自 凹非寺
量子位 | 公眾號 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，斯圖加特大學博士研究生，研究興趣是開發基于數據驅動和先驗物理知識的神經網絡模型。
論文地址：
https://arxiv.org/abs/2210.07182
PDEBench地址：
https://github.com/pdebench/PDEBench
參考鏈接：
[1]https://twitter.com/Mniepert/status/1581010273246523393
[2]https://mp.weixin.qq.com/s/Rbw2QFavSn8N7pPGS05o6w
- 首個GPT-4驅動的人形機器人！無需編程+零樣本學習，還可根據口頭反饋調整行為2023-12-13
- IDC霍錦潔：AI PC將顛覆性變革PC產業2023-12-08
- AI視覺字謎爆火！夢露轉180°秒變愛因斯坦，英偉達高級AI科學家：近期最酷的擴散模型2023-12-03
- 蘋果大模型最大動作：開源M芯專用ML框架，能跑70億大模型2023-12-07
