AI求解偏微分方程新基準(zhǔn)登NeurIPS，發(fā)現(xiàn)JAX計(jì)算速度比PyTorch快6倍，LeCun轉(zhuǎn)發(fā)：這領(lǐng)域確實(shí)很火

蕭簫 2022-10-27 12:36:45 來源：量子位

蕭簫 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI

用AI求解偏微分方程，這段時(shí)間確實(shí)有點(diǎn)火。

但究竟什么樣的AI求解效果最好，卻始終沒有一個(gè)統(tǒng)一的定論。

現(xiàn)在，終于有人為這個(gè)領(lǐng)域制作了一個(gè)名叫PDEBench的完整基準(zhǔn)，論文登上了NeurIPS 2022。

PDEBench不僅能當(dāng)成一個(gè)大型偏微分方程數(shù)據(jù)集，也能作為新AI求解偏微分方程的基準(zhǔn)之一——

不少“老前輩”的預(yù)訓(xùn)練模型代碼都能在這里找到，作為一個(gè)比對(duì)基礎(chǔ)。

例如去年大火了一陣的FNO，幾秒鐘求解出傳統(tǒng)方法需要計(jì)算18個(gè)小時(shí)的偏微分方程，代碼就被放進(jìn)了PDEBench中。

這個(gè)新基準(zhǔn)一出，LeCun也激情轉(zhuǎn)發(fā)：這領(lǐng)域確實(shí)很火。

所以，AI求解偏微分方程的優(yōu)勢是什么，這一基準(zhǔn)具體提出了哪些評(píng)估方法？

為啥用AI求解偏微分方程？

偏微分方程（PDE，Partial Differential Equation），是一個(gè)生活中常見的方程。

包括預(yù)報(bào)天氣、模擬飛機(jī)空氣動(dòng)力、預(yù)測疾病傳播模型，都會(huì)用到這個(gè)方程。

目前北大數(shù)學(xué)系“韋神”韋東奕的研究方向之一，就是流體力學(xué)中的數(shù)學(xué)問題，其中就包括偏微分方程中的Navier-Stokes方程。

所以，為啥要用AI來求解偏微分方程？

訓(xùn)練AI的本質(zhì)，是找到一種盡可能逼近真實(shí)結(jié)果的模型。

用AI求解偏微分方程，其實(shí)也是找到一種代理模型，來模擬偏微分方程模型。

代理模型，指找到一種近似模型，在計(jì)算量更小的同時(shí)，確保計(jì)算結(jié)果與原來的偏微分方程盡可能相似。

這與傳統(tǒng)的數(shù)值方法求解偏微分方程有著異曲同工之妙。

傳統(tǒng)方法往往需要通過將連續(xù)問題離散化（類似在一個(gè)連續(xù)函數(shù)上切割出很多小點(diǎn)），來對(duì)方程進(jìn)行近似求解。

然而，傳統(tǒng)的數(shù)值方法非常復(fù)雜，計(jì)算量也很大；采用AI方法訓(xùn)練出來的模型，卻模擬得又快又好——

繼2017年華盛頓大學(xué)提出PDE-FIND后，2018年谷歌AI又提出了數(shù)據(jù)驅(qū)動(dòng)求解偏微分方程的方法，都比傳統(tǒng)方法要快上不少，讓更多人開始關(guān)注到AI求解偏微分方程這一領(lǐng)域。

2019年，布朗大學(xué)應(yīng)用數(shù)學(xué)團(tuán)隊(duì)提出一種名叫PINN?（物理激發(fā)的神經(jīng)網(wǎng)絡(luò)）的方法，徹底打開了AI在物理學(xué)領(lǐng)域的廣泛應(yīng)用。

這篇論文在理論上雖然沒有PDE-FIND和谷歌AI的方法突破性強(qiáng)，卻給出了非常完整的代碼體系，使得開發(fā)人員很容易上手，讓更多研究者開發(fā)出了不同的PINN，如今它也成為AI物理最常見的框架和詞匯之一。

△PINN

去年加州理工大學(xué)和普渡大學(xué)團(tuán)隊(duì)發(fā)表的一項(xiàng)研究，更是將偏微分方程計(jì)算時(shí)間從傳統(tǒng)求解的18個(gè)小時(shí)降低為1秒鐘。

這篇論文提出了一種名為FNO?（傅里葉神經(jīng)算子）的方法，基于傅里葉變換給神經(jīng)網(wǎng)絡(luò)加上“傅里葉層”，進(jìn)一步節(jié)省了近似模擬算子的計(jì)算量。

除此之外，也有不少研究人員通過訓(xùn)練一些經(jīng)典AI模型，來求解偏微分方程，如U-Net等。

不過，無論是FNO、U-Net還是PINN，都還是基于各自給出的基準(zhǔn)來評(píng)估AI計(jì)算偏微分方程的效果。

有沒有一個(gè)更統(tǒng)一、更通用的框架來評(píng)估這個(gè)領(lǐng)域的新突破？

更全面的AI偏微分方程基準(zhǔn)

在這樣的背景下，研究人員提出了一種名叫PDEBench的基準(zhǔn)。

首先是基準(zhǔn)中包含的數(shù)據(jù)集，目前這些數(shù)據(jù)集已經(jīng)全部歸納到GitHub中：

這里面包括不少經(jīng)典偏微分方程問題，如Navier-Stokes方程，達(dá)西流模型、淺水波模型等等。

隨后，PDEBench提出了幾個(gè)指標(biāo)，來從不同角度更全面地對(duì)AI模型進(jìn)行評(píng)估：

最后，PDEBench還包含了幾種經(jīng)典模型的預(yù)訓(xùn)練模型代碼，并將它們作為評(píng)估其他模型的基準(zhǔn)之一，包括上述提到的FNO、U-Net、PINN等。

例如研究團(tuán)隊(duì)將這幾個(gè)模型分別基于各數(shù)據(jù)集進(jìn)行了訓(xùn)練，得出的均方根誤差（RMSE）如下，也說明它們在不同偏微分方程問題上的表現(xiàn)并不一樣：

除此之外，團(tuán)隊(duì)還將數(shù)據(jù)格式進(jìn)行了統(tǒng)一，同時(shí)針對(duì)PDEBench的可擴(kuò)展性進(jìn)行了優(yōu)化，因此任何人都能參與進(jìn)來，給這一基準(zhǔn)加入更多的數(shù)據(jù)集、或是更多基準(zhǔn)模型。

值得注意的是，團(tuán)隊(duì)試了試分別在PyTorch和JAX兩種框架上運(yùn)行幾種預(yù)訓(xùn)練模型，發(fā)現(xiàn)JAX的速度大約是PyTorch的6倍。

看來以后搞相關(guān)研究可以試試JAX框架了。

作者介紹

作者們來自德國斯圖加特大學(xué)，歐洲NEC研發(fā)中心，還有澳大利亞聯(lián)邦科學(xué)與工業(yè)研究組織（CSIRO）旗下的Data61數(shù)字創(chuàng)新中心。

Makoto Takamoto，歐洲NEC研發(fā)中心高級(jí)研究員，畢業(yè)于京都大學(xué)，研究方向是圖像處理、圖神經(jīng)網(wǎng)絡(luò)和科學(xué)機(jī)器學(xué)習(xí)。

Timothy Praditia，斯圖加特大學(xué)博士研究生，研究興趣是開發(fā)基于數(shù)據(jù)驅(qū)動(dòng)和先驗(yàn)物理知識(shí)的神經(jīng)網(wǎng)絡(luò)模型。

論文地址：
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