技術要約: Mosaic: 微分可能な物理ソルバーのためのベンチマーク・スイート

問題提起

微分可能な偏微分方程式（PDE）ソルバーは、ソルバー・イン・ザ・ループの機械学習、勾配ベースの最適制御、および逆問題を実現する基盤技術として期待されているが、正しく使用可能な勾配を得るための実質的なコストは不透明なままである。微分可能ソルバーの景観は、言語、自動微分（AD）フレームワーク、および数値手法によって断片化されている。実務者は以下の事項について大きな不確実性に直面している：

• 統合の労力: 即席のアプリケーションから低レベルのカーネル・ツールキットに至るまで、ソルバーを統合するために必要な専門知識の差異。
• 勾配の品質: 勾配の正確性、数値的コンディショニング、および安定性に関する系統的な測定の欠如。
• 構造的互換性: ソルバー固有の制約（例：周期境界条件のみ、互換性のないペナルティ手法）が、特定のスタックへの多大なコミットメントの後に初めて表面化すること。

既存のベンチマークは、ニューラル・サロゲートのフォワード精度（例：APEBench, PDEBench）に焦点を当てるか、あるいは単一のフレームワーク内でのADバックエンドの比較（例：PhiFlow）に留まっており、独立したソルバー実装間での勾配特性を系統的に測定できていない。

手法

著者らは、ソルバーの勾配へのアクセスを標準化し、ソルバー間のクロス比較を可能にするために設計された拡張可能なベンチマーク・フレームワークである Mosaic を導入する。この手法は、3つの主要なアーティファクトに基づいている：

1. コンテナ化されたソルバー・インターフェース (Tesseract):

   • ソルバーは、言語やAD戦略に関わらず統一されたAPIを公開するDockerコンテナ（Tessera）としてパッケージ化される。
   • このインターフェースは、2つのエントリポイントを定義する：apply(inputs) → outputs（フォワード・パス）および vector_jacobian_product(inputs, cotangents) → vjp（リバースモード勾配）。
   • この抽象化により、評価コードを書き換えることなく、設定文字列を通じてソルバーを入れ替えることが可能になる。フォワード専用のソルバー（例：OpenFOAM, deal.II）は、リファレンス・ベースラインとしての有限差分勾配を提供するために自動的にラップされる。

2. ベンチマーク・タスク:

   • 難易度順に並べられた4つの物理ドメインが評価される：
     • 熱伝導 (H): 熱伝導率の逆問題（ネガティブ・コントロール；滑らかな物理現象によりADのバグを分離）。
     • 構造力学 (S): SIMPトポロジー最適化によるコンプライアンス最小化。
     • 2D流体 (F2): 円柱周囲の流入最適化による抗力最小化（非周期、障害物あり）。
     • 3D流体 (F3): 周期的なテイラー・グリーン渦ボックスにおける初期条件の回復。
   • タスクは学習成分を持たないスタンドアロンの最適化問題として設計されており、失敗の原因がモデルの近似誤差ではなく、ソルバーの勾配メカニズムにあることを確実にできる。

3. 評価プロトコル:

   • 勾配の正確性: 最適な摂動サイズにおける中央差分に対する余弦類似度および相対 $\ell_2$ 誤差を用いて測定。
   • コンディショニング: ヤコビアンの特異値スペクトルを介して分析。
   • パフォーマンス: フォワードおよびVJPパスのウォールクロック時間、メモリ消費量、およびVJP/フォワード・コスト比。
   • 最適化の収束: AdamおよびL-BFGSオプティマイザを用いた固定イテレーション予算による、グローバル最適解への収束測定。

本スイートには、多様なAD戦略（ソース変換、テープベースのアジョイント、手導出のアジョイント）および数値手法（有限差分、有限体積、有限要素、格子ボルツマン、スペクトル法）にわたる14のソルバー・バックエンドが含まれている。

主要な結果

14のソルバーを4つのドメインで評価した結果、いくつかの重要な知見が得られた：

• 勾配の正確性とコンディショニング: 勾配を生成するすべての微分可能ソルバーは、有限差分に対して高い正確性（余弦類似度 > 0.999）を達成した。しかし、ヤコビアンのコンディショニングは数桁の差がある。投影ベースのソルバーは、格子ボルツマン法やスペクトル法と比較して急峻なスペクトル減衰（悪いコンディショニング）を示す傾向があり、これが二次最適化の安定性に影響を与える。
• 計算コスト: 計算コストには桁違いの差異が存在する。ソース変換AD（例：JAXベースのソルバー）は、ほぼ一定のVJP/フォワード比を維持するが、テープベースのアジョイント（例：FEniCS, Firedrake）は問題サイズとともに急激にスケールアップする。
• 構造的不適合性: 最も複雑なタスク（障害物のある2D抗力最小化）において、境界条件またはペナルティ手法の不適合により、7つの流体バックエンドのうち3つしか問題を表現できなかった。この「脱落」は、勾配の品質が問題になる前に発生する。
• 最適化の結果: コンディショニングとコストに大きな違いがあるにもかかわらず、正常に実行されたすべてのソルバーは、良条件の問題（熱、構造）において同様の最適解に収束した。3D初期条件回復タスクでは、L-BFGS（ソレノイダル投影あり）がAdamおよび非投影L-BFGSを大幅に上回り、初期条件誤差を40%超から0.5%未満へと減少させた。
• 実装の労力: 統合の複雑さは大きく異なる。現代的なADフレームワーク（JAX, PyTorch）上に構築されたソルラーは、勾配のために追加のコードをほとんど必要としないことが多いが、低レベルのカーネル・ツールキット（例：NVIDIA Warp）は、アジョイント実装のために数百行の専用コードを必要とする。

意義と主張

本論文は、Mosaicを単なる単一ソルバーの新しい経験的結果の源泉としてではなく、強化学習におけるOpenAI Gymのような、微分可能物理研究への参入障壁を下げるためのインフラストラクチャとして位置づけている。

• 標準化: Mosaicは、独立したソルバー実装間で勾配にアクセスし、比較し、信頼するための最初の標準化されたインターフェースを提供する。
• 実務家へのガイダンス: 本ベンチマークは、勾配の正確性そのものよりも、実用的な障壁（メモリ制限、数値的安定性、セットアップの互換性）を浮き彫りにする。これは、良条件の問題においては、最も高速で統合の労力が少ないソルバーが最適な選択であることを示唆している。
• コミュニティの拡張性: フレームワークはオープンソースであり、進化するように設計されている。これにより、研究者は特定のソルバー・スタックにコミットする前に失敗モードをテストすることができ、統一されたインターフェース上でソルバー・イン・ザ・ループのパイプラインを構築することが可能になる。
• ADの妥当性: 結果は、手導出のアジョイントがより高速になり得る一方で、現代的なフレームワークを介したリバースモードADが、実装コストを抑えつつ十分な精度を提供する、ベストな「労力対能力」のトレードオフを提供することを検証している。

著者らは、勾配の品質は一般に十分であるが、ソルバーのバックエンドの選択が、勾配自体の正しさよりも、勾配ベースのワークフローが成功するか失敗するかを決定付ける決定的な要因であると結論付けている。