技術要約：学習済みシミュレータの確率論的レトロフィッティング（後付け）

問題定義

深層学習を用いた偏微分方程式（PDE）モデリングの支配的なパラダイムは、平均二乗誤差（MSE）または平均絶対誤差（MAE）を最小化する決定論的な学習フレームワークに依存している。この手法は、安定した、完全に観測されたシステムには効果的であるが、カオス的、ノイズを含む、あるいは部分的にしか観測されていない物理システムに対しては根本的に不適切である。粗い解像度で動作する決定論的モデルは、ダイナミクスに内在する不確実性を捉えることができず、「ぼやけた」予測に陥る傾向がある。これは、モデルが起こりうる将来の状態の平均値に収束してしまうことで、物理的に矛盾したロールアウト（展開）をもたらす。

拡散モデルやフローベースの生成モデルといった既存の代替案は、確率論的な能力を提供するものの、計算負荷が極めて高い。これらは専門的でリソース集約的な学習プロセスを必要とし、推論プロセスも反復的で遅いことが多い。さらに、パラメータ数と学習データの急速なスケーリングを考慮すると、高精度な確率論的基盤モデルを一から訓練することは、経済的および環境的に持続不可能になりつつある。既存の高精度な決定論的バックボーンに投じられた多大なリソースを放棄することなく、それらを効率的に確率論的モデルへと変換するための決定的な手法が欠落している。

手法

著者らは、学習済みの決定論的モデルを確率論的エミュレータへとレトロフィットするための、訓練効率の高い戦略を提案している。このアプローチはアーキテクチャに依存せず、最小限のコード修正で済む。コアとなる手法は、2つの主要なコンポーネントで構成される：

1. 条件付きレイヤー正規化（AdaLN）によるノイズ注入：
   決定論的アーキテクチャは、確率的なブランチを含むように修正される。グローバルなノイズベクトル $\epsilon \sim \mathcal{N}(0, I)$ がサンプリングされ、多層パーセプトロン（MLP）を通過してノイズ埋め込みが生成される。これらの埋め込みは、条件付きレイヤー正規化層を介して、変調パラメータ（$\gamma, \beta, \alpha, \delta$）へと投影される。これらのパラメータは、正規化された入力をスケールおよびシフトし、バックボーン（例：トランスフォーマーのアテンション・ブロックやU-Netのエンコーダ層）内の残差接続をゲート制御する。このメカニズムにより、モデルは学習済みの核となる特徴を大幅に変更することなく、単一の入力履歴から多様で鮮明な予測のアンサンブルを生成できる。

2. 連続ランク確率スコア（CRPS）による最適化：
   点ごとの誤差を最小化する代わりに、厳密に適切なスコアリングルールである連続ランク確率スコア（CRPS）を用いてモデルを微調整する。経験的CRPSは、$M$ 個の予測 $\{x_j\}$ と真値 $y$ に対して次のように計算される：
   $$\text{CRPS}(\{x_j\}_{j=1}^M, y) = \frac{1}{M}\sum_{j=1}^M |x_j - y| - \frac{1}{2M^2}\sum_{j=1}^M\sum_{k=1}^M |x_j - x_k|$$
   有限のアンサンブルにおけるバイアスに対処するため、著者らは不偏な「フェア（公正）」なCRPS推定器を採用している。第1項は個々のアンサンブルメンバーの正確性を促し（絶対誤差の最小化）、第2項は多様性を促し（アンサンブルメンバーを分散させる）、モード崩壊を防ぐ。

訓練戦略：
レトロフィッティングのプロセスは、CRPS目的関数と、事前学習で使用されたMSE/MAE損失との間の互換性を活用している。主な戦略は以下の通りである：

• バックボーンのパラメータを、事前学習済みの決定論的モデルから初期化する。
• 差分学習率を使用する：バックボーンには小さな学習率を適用して特徴を保持し、新たに導入されたノイズブランチには大きな学習率を適用する。
• 条件付き正規化層を小さな値に注意深く初期化し、モデルが決定論的な挙動の近くから開始されるようにすることで、確率性の学習をスムーズにする。

主な貢献

本論文の主な貢献は以下の4点である：

1. 統一された、アーキテクチャに依存しないレシピ： 著者らは、ピクセル空間および潜在空間モデリングのパラダイムを代表する3つの異なるアーキテクチャ（Walrus、Lola、Poseidon）および5つの物理システムに対して、同一のレトロフィッティング戦略が有効であることを示した。
2. 効率性と性能： 本手法は、決定論的な微調整と比較して大幅な向上を実現している。単一システムモデルにおいて、ロールアウトCRPSで20〜54%、分散正規化平方根平均二乗誤差（VRMSE）で最大30%の改善を示した。基盤モデルにおいては、決定論的な微調整に対して最大40%のCRPS改善と15%のVRMSE改善を示している。
3. 推論時の柔軟性： レトロフィットされたモデルは、テスト時のスケーリング特性を示す。予測スキルは、生成されるアンサンブルメンバーの数に応じて単調に向上するため、ユーザーは再学習することなく、計算量と精度のトレードオフを動的に調整できる。
4. アブレーション研究と洞察： 広範なアブレーション研究により、ハイパーパラメータ調整（学習率、アンサンブルサイズ、ノイズ次元数）に関する実用的な推奨事項を提供し、手法の堅牢性を実証した。

結果

本アプローチは、単一システムモデルおよびマルチシステム基盤モデル（HalfWalrus）を用いて検証された。データセットは The Well（レイリー・ベナール対流、オイラー多重クアドラント、剪断流、乱流放射層、粘弾性不安定性）に基づいている。

• 単一システムモデル： Walrus、Lola、Poseidonの各アーキテクチャにおいて、CRPSレトロフィッティングは決定論的な微調整を一貫して上回った。例えば、レイリー・ベナール・データセットにおいて、WalrusモデルはCRPSで43.1%の改善、VRMSEで21.5%の減少を達成した。定性的分析によれば、決定論的なベースラインが長いホライゾンで過度に平滑化された予測に陥る一方で、CRPSアンサンブルは高周波の乱流構造を維持していた。
• 基盤モデル： HalfWalrus基盤モデルに適用した場合、レトロフィッティングは、カオス的な領域（レイリー・ベナールおよび乱流放射層）において34〜40%のCRPS改善と、>13%のVRMSE改善をもたらした。カオス的な粘弾性領域において、確率論的モデルは、過度の平滑化に苦しむ決定論的ベースラインを大幅に上回った。
• ベースラインとの比較：
  • ディープアンサンブル： CRPSレトロフィッティングは、計算コストが4倍であるディープアンサンブルに対し、精度と較正（広がり対スキル比）の両面で上回った。
  • 拡散モデル： Lolaフレームワーク内での計算量を一致させた比較において、CRPSレトロフィッティングは、精度において拡散ベースのレトロフィッティングと同等またはそれを上回り、かつ単一ステップの推論（反復的なデノイジングに対して）を提供した。決定的なことに、CRPSレトロフィッティングは、拡散モデルよりもデータセットの不均一性（例：剪断流）に対してより堅牢であり、事前学習済みの決定論的表現を活用する能力においてより効果的であることが証明された。

意義と主張

本論文は、確率論的なPDEモデリングには、ゼロからの再学習を必要としないことを主張している。CRPSとノイズ注入に基づくレトロフィッティング戦略を採用することで、コミュニティは既存の決定論的バックボーンから、追加の訓練コストを抑えつつ確率論的な能力を引き出すことができる。

著者らは、このアプローチが、生成基盤モデルを一から訓練する際の膨大なコストに対する実行可能な代替案を提供すると主張している。これは、巨大なモデルを再訓練することの経済的および環境的な持続不可能性に対処しつつ、決定論的な効率性と確率論的な忠実度の間の溝を埋める統一されたフレームワークを提供するものである。この手法は、ダイナミカルシステムのための基盤モデルの台頭において、特に重要である。なぜなら、堅牢な決定論的チェックポイントを再利用して、カオス的な物理システムに内在する不確実性を扱うことを可能にするからである。