[論文レビュー] UPS: Efficiently Building Foundation Models for PDE Solving via Cross-Modal Adaptation

Jun-Hong Shen, Tanya Marwah+1 more|arXiv (Cornell University)|Mar 11, 2024

BIM and Construction Integration被引用数 5

ひとこと要約

UPSは多様な時空間PDEデータを共通表現に統合し、事前学習済みLLMsをPDE解法へ適応させるトランスフォーマーベースのモデルを訓練することで、データ効率が高くPDEBenchタスク全体で最先端の性能を達成します。

ABSTRACT

We present Unified PDE Solvers (UPS), a data- and compute-efficient approach to developing unified neural operators for diverse families of spatiotemporal PDEs from various domains, dimensions, and resolutions. UPS embeds different PDEs into a shared representation space and processes them using a FNO-transformer architecture. Rather than training the network from scratch, which is data-demanding and computationally expensive, we warm-start the transformer from pretrained LLMs and perform explicit alignment to reduce the modality gap while improving data and compute efficiency. The cross-modal UPS achieves state-of-the-art results on a wide range of 1D and 2D PDE families from PDEBench, outperforming existing unified models using 4 times less data and 26 times less compute. Meanwhile, it is capable of few-shot transfer to unseen PDE families and coefficients.

研究の動機と目的

• 領域・次元・解像度を跨ぐ多様な時間依存PDEを解くためのデータ効率の高い基盤モデルアプローチを提案する。
• FNOベースのPDE埋め込み、LLMs、および予測ヘッドを組み合わせた統一データ表現と3部構成のアーキテクチャを提案する。
• PDE埋め込みとテキスト埋め込みを整合させる2段階のクロスモーダル適応を活用し、複数のPDEファミリに渡って微調整を行う。
• PDEBenchで性能とサンプル効率の向上を示し、少数ショット・ゼロショット転送能力を強力に発揮する。）

提案手法

• 多様なPDE軌跡を複数の次元・解像度・量を跨ぐ共有表現空間に統一する。
• PDE特化の埋め込みネットワークと Fourier Neural Operator (FNO) 層を用いてPDEデータから解像度不変特徴を抽出する。
• PDEに関するメタデータを、テキスト形式のPDE記述として埋め込み、PDE特徴と結合してLLM本体へ処理する前に組み込む。
• PDE埋め込みネットワーク、モデル本体としての事前学習済みLLM層、次ステップのPDE状態を予測する線形予測ヘッドの3部構成アーキテクチャを採用する。
• 2段階で訓練する: (i) モダリティ整列損失(MMDベース)とタスク損失(nRMSE)を用いた埋め込み事前学習; (ii) 複数のPDEファミリにまたがるマルチタスク微調整。
• RoBERTa、T5、CLIPなどのさまざまなLLMバックボーンとの互換性を示し、ニューラルオペレータの格子・解像度不変性を実証する。）

Figure 1 : UPS outperforms leading competitors on a wide range of tasks from PDEBench, including the single-task neural operators U-Net (Ronneberger et al., 2015 ) , FNO (Li et al., 2020a ) , and ORCA (Shen et al., 2023 ) , and the multi-task MPP-Base (McCabe et al., 2023 ) (MPP only applies to 2D t

実験結果

リサーチクエスチョン

• RQ1 統一表現とアーキテクチャは、異なるPDEファミリ・次元・解像度間のクロスドメイン転移を可能にするか。
• RQ2 事前学習済みLMMのクロスモーダル適応は、多様なPDE解法タスクでデータ効率と予測精度を向上させるか。
• RQ3 部分的なメタデータ（テキスト形式のPDE記述）が神経演算子フレームワークの学習と一般化にどう影響するか。
• RQ4 見落としたPDEファミリ、係数、解像度への少数ショット・ゼロショット転送におけるUPSの性能はどうなるか。

主な発見

# Params	Advection nRMSE	Burgers nRMSE	Diffusion-Sorption nRMSE	Navier-Stokes nRMSE	Shallow-Water nRMSE	Navier-Stokes Incomp nRMSE	Incomp Navier-Stokes nRMSE
149M	0.0033	0.0399	0.0009	0.0056	0.0019	0.0174	0.104
7.7M	0.67	0.34	0.15	0.72	0.083	5.1	0.1903
466K	0.011	0.042	0.0017	0.068	0.0044	0.36	0.0942
125M	0.0098	0.12	0.0016	0.062	0.006	0.3549	0.1529
116M	-	-	-	-	0.0024	0.0281	-
409M	-	-	-	-	0.0022	0.0208	-
8.5K	1.1	0.96	0.22	-	0.017	-	-

• UPSはPDEBenchのタスク10件中8件で最先端の性能を達成し、しばしばベースラインに対して大きなマージンを示す。
• 分布内タスクでは、UPS (RoBERTa-Base) が7データセット中6データセットで最小nRMSEを達成し、1Dおよび2Dで強い性能を示す。
• UPSは未知のPDEファミリ・係数・高解像度への強力な少数ショット・ゼロショット転送を示し、限られたデータ下でいくつかのベースラインを上回る。
• Stage-1の埋め込み事前学習（モダリティ整列とタスク学習を含む）は性能に不可欠で、メタデータの組み込みが結果を改善する。
• 本手法はデータ・計算効率が高く、PDEファミリごとに5K未満のトレーニング軌跡を使用し、シングルGPUでの訓練を実現する。

Figure 2 : To adapt pretrained LLMs for PDE solving, UPS first expresses PDE trajectories of different families, dimensions, channels, and resolutions using a unified data representation (left panel). It then processes the data with a unified transformer architecture that integrates domain-specific