偏微分方程（PDEs）被稱為物理學(xué)的語(yǔ)言，因?yàn)樗鼈兛梢栽趶V泛的時(shí)間 - 空間尺度上對(duì)各種各樣的物理現(xiàn)象進(jìn)行數(shù)學(xué)建模。常用的有限差分、有限元等數(shù)值方法通常用于近似或模擬偏微分方程。

然而，這些方法計(jì)算成本高昂，特別是對(duì)于多查詢問(wèn)題更是如此，因而人們?cè)O(shè)計(jì)了各種數(shù)據(jù)驅(qū)動(dòng)的機(jī)器學(xué)習(xí)（ML）方法來(lái)模擬偏微分方程。其中，算子學(xué)習(xí)（ operator learning）算法近年來(lái)受到越來(lái)越多的關(guān)注。

然而，現(xiàn)有的算子學(xué)習(xí)方法樣本效率并不高，因?yàn)樗鼈冃枰罅康挠?xùn)練樣例才能以期望的準(zhǔn)確率學(xué)習(xí)目標(biāo)解算子（如圖 1 所示）。這阻礙了算子學(xué)習(xí)的廣泛使用，因?yàn)橥ㄟ^(guò)數(shù)值模擬或底層物理系統(tǒng)的測(cè)量來(lái)生成特定任務(wù)的訓(xùn)練數(shù)據(jù)非常昂貴。

研究者不禁提出，如何才能顯著減少 PDE 學(xué)習(xí)所需的訓(xùn)練樣本數(shù)量？

來(lái)自蘇黎世聯(lián)邦理工學(xué)院等機(jī)構(gòu)的研究者提出了 Poseidon，這是一種用于學(xué)習(xí) PDE 解算子的基礎(chǔ)模型。該模型基于多尺度 operator transformer，可實(shí)現(xiàn)連續(xù)時(shí)間評(píng)估。

研究者將 Poseidon 在大規(guī)模數(shù)據(jù)集上進(jìn)行了預(yù)訓(xùn)練，然后對(duì)其進(jìn)行評(píng)估。具體而言，他們?cè)?15 項(xiàng)具有挑戰(zhàn)性的下游任務(wù)上對(duì) Poseidon 進(jìn)行了評(píng)估，這些任務(wù)涵蓋線性和非線性、時(shí)間相關(guān)以及橢圓、拋物線、雙曲線和混合型 PDE。

結(jié)果表明，Poseidon 在樣本效率和準(zhǔn)確率方面都遠(yuǎn)遠(yuǎn)超過(guò)基線，展現(xiàn)出優(yōu)異的性能。

Poseidon 還可以很好地泛化到預(yù)訓(xùn)練期間未見過(guò)的物理學(xué)問(wèn)題。此外，Poseidon 可以根據(jù)模型和數(shù)據(jù)大小進(jìn)行擴(kuò)展，無(wú)論是預(yù)訓(xùn)練還是下游任務(wù)?？偨Y(jié)來(lái)看，本文展示了 Poseidon 的驚人能力，它能夠在預(yù)訓(xùn)練期間從非常小的一組 PDE 中學(xué)習(xí)有效表示，從而很好地?cái)U(kuò)展到下游未見過(guò)和不相關(guān)的 PDE，證明了其作為有效通用 PDE 基礎(chǔ)模型的潛力。

這些結(jié)果首次肯定了 PDE 基礎(chǔ)模型的可行性這一基本問(wèn)題，并為進(jìn)一步開發(fā)和部署 Poseidon 作為高效的通用 PDE 基礎(chǔ)模型鋪平了道路。 

最后，Poseidon 模型以及底層預(yù)訓(xùn)練和下游數(shù)據(jù)集都是開源的。 

• 論文地址：https://arxiv.org/pdf/2405.19101
• 項(xiàng)目地址：https://github.com/camlab-ethz/poseidon
• 論文標(biāo)題：Poseidon: Efficient Foundation Models for PDEs

方法介紹

問(wèn)題描述：該研究將偏微分方程表示為：

然后假設(shè)，可以得到與時(shí)間無(wú)關(guān)的 PDE 的解：

模型架構(gòu)。Poseidon（圖 1 和圖 2）包括：i）可擴(kuò)展的 Operator Transformer 或 scOT，這是一種具有（移位）窗口或 Swin 注意力機(jī)制的多尺度視覺(jué) transformer，適用于算子學(xué)習(xí)；ii）一種新穎的 all2all 訓(xùn)練策略；iii）以及一個(gè)開源大型預(yù)訓(xùn)練數(shù)據(jù)集。

其中 scOT 是一種具有前置時(shí)間條件的分層多尺度視覺(jué) transformer，用來(lái)處理前置時(shí)間 t 和函數(shù)空間值初始數(shù)據(jù)輸入 a，以近似 PDE (2.1) 的解算子 S (t, a)。

接著如圖 2 (a) 所示，研究者通過(guò) SwinV2 transformer 塊對(duì) patch 嵌入的輸出進(jìn)行處理，每個(gè) transformer 塊的結(jié)構(gòu)表示為，得到：

通過(guò)在 (2.3) 中引入前置時(shí)間條件層范數(shù)，該研究提出了一種時(shí)間調(diào)節(jié)策略。

最后，如圖 2 (a) 所示，SwinV2  transformer 塊在 U-Net 類型的編碼器 - 解碼器架構(gòu)中以層級(jí)多尺度方式排列，通過(guò)使用 patch 合并（下采樣）和 patch 擴(kuò)展（上采樣）操作完成。

實(shí)驗(yàn)結(jié)果

預(yù)訓(xùn)練數(shù)據(jù)：研究者提供了包含 6 個(gè)算子的數(shù)據(jù)集，詳細(xì)信息如下所示。

下游任務(wù)：研究者在 15 個(gè)具有挑戰(zhàn)性的下游任務(wù)上進(jìn)行了實(shí)驗(yàn)，如表 4 所示。

模型：本文考慮了三種不同的 Poseidon 模型：i) Poseidon-T ≈ 21M 個(gè)參數(shù)，ii) Poseidon-B ≈ 158M 個(gè)參數(shù)，iii) Poseidon-L ≈ 629M 個(gè)參數(shù)。

實(shí)驗(yàn)結(jié)果顯示，Poseidon 在 15 個(gè)下游任務(wù)中都表現(xiàn)良好，明顯優(yōu)于 FNO（ Fourier Neural Operator ）（參考論文中的圖 7 - 圖 21，這里只展示圖 7 ）。

表 1 進(jìn)一步支持了這一點(diǎn)。

從表 9 可得，平均而言，Poseidon-L 僅需要 20 個(gè)樣本即可達(dá)到 FNO 的 1024 個(gè)樣本的誤差，并且在 13 個(gè)（15 個(gè)）任務(wù)中，Poseidon-L 所需的樣本比 FNO 少一個(gè)數(shù)量級(jí)。同樣，從表 1 和表 9 中可以看到，對(duì)于相同數(shù)量的樣本，Poseidon-L 的誤差明顯低于 FNO，增益范圍從 10% 到 25 倍不等 ，此外，Poseidon 可以很好地泛化到未見過(guò)的物理任務(wù)。

從表 1 和表 9 可以觀察到，在 15 項(xiàng)下游任務(wù)中，Poseidon 在 14 項(xiàng)上的表現(xiàn)明顯優(yōu)于 CNO-FM。平均而言，CNO-FM 需要大約 100 個(gè)特定于任務(wù)的示例才能達(dá)到 FNO 的 1024 個(gè)樣本的誤差水平，而 Poseidon 只需要大約 20 個(gè)。由于 CNO-FM 和 Poseidon 已在完全相同的數(shù)據(jù)集上進(jìn)行了預(yù)訓(xùn)練，因此這種性能差異很大程度上可以歸因于架構(gòu)差異，因?yàn)?CNO-FM 基于多尺度 CNN，而 Poseidon 的主干則是多尺度視覺(jué) transformer。

從圖 22 可以看出，隨著 Poseidon 模型大小的增加，預(yù)訓(xùn)練數(shù)據(jù)集上的訓(xùn)練和評(píng)估（驗(yàn)證）錯(cuò)誤都明顯減少。

了解更多結(jié)果，請(qǐng)參考原論文。