Sparse probes and murky physics: a case study of interpretability challenges in a foundation model for continuum dynamics

技術要約：疎なプローブと不透明な物理学

問題提起
生成AIエミュレータは、強力な物理理論やベンチマークが既に存在する科学領域において、ますます導入が進んでいる。これは重大な解釈性の課題を生じさせる。すなわち、基盤モデルが既知の連続体ダイナミクスを正常に再現している場合であっても、そのモデルが基礎となる物理的メカニズムを学習したのか、あるいは単に統計的な相関関係を記憶しただけなのかは不明確である。具体的には、研究者は、モデルの内部表現が既知の物理原則に対応しているのか、これらの表現がモデルの成功や失敗とどのように関連しているのか、そして「異なる」内部構造が真に情報量豊かなものなのか、それとも単に効果的なショートカット（近道）に過ぎないのかを判断する必要がある。本論文は、多様な物理シナリオで学習された13億パラメータのトランスフォーマーであるWalrus基盤モデルを用い、学習された内部表現と物理的直感との間のギャップに焦点を当てることで、これらの問いに取り組む。

手法
著者らは、物理原則に基づいたメカニズム的解釈性の手法を採用しており、具体的にはSparse Autoencoder（SAE）を用いてWalrusモデルの内部活性化をプローブする。

• モデルとデータ: 本研究は、古典流体を含む19の物理シナリオで学習されたWalrusに焦点を当てている。分析は、非圧縮ナビエ・ストークス方程式に従う2次元せん断流シミュレーション・スイートに限定される。データセットはレイノルズ数（$Re$）およびシュミット数（$Sc$）によって変化し、特に3つのシミュレーション（Sim3: $Re=10^4, Sc=0.1$、Sim50: $Re=10^4, Sc=0.2$、Sim56: $Re=5\times10^4, Sc=0.1$）に注意を払っている。
• 物理的接地: SAEによって生成される大規模な特徴セットを選別するため、著者らは物理的に根拠のある指標として、渦度の二乗の積分として定義されるエントロピー（enstrophy）（$E$）を利用する。また、注入と散逸のバランスを評価するために、エントロピーの時間微分（$\dot{E}$）も分析する。
• SAEの実装: Walrusトランスフォーマーの中間ブロック（$l=20$）の活性化（2,816個の活性化）に対して、Top-K Sparse Autoencoderを訓練した。SAEは8倍の拡張係数を用い、22,528個の特徴量を持つ。訓練には、訓練を安定させ、特徴量の崩壊を防ぐために、二部構成の損失関数（平均二乗誤差およびデッドニューロン用の補助損失）を用いた。
• 分析戦略: 著者らは、グローバルなエントロピー指標とのスペアマンの順位相関係数に基づいてSAE特徴量をランク付けする。その後、特徴量の活性化ヒートマップをトレーサー場およびエントロピーマップに対して可視化し、異なるシミュレーション設定やタイムステップ間での空間的なコヒーレンス（一貫性）と整合性を評価する。

主な結果
本調査は、モデルの内部表現の整合性と解釈可能性に関して、以下の知見を得た：

1. 特徴の特定と相関: 参照シミュレーション（Sim50）において、トップランクの特徴がエントロピーと高い相関（$\rho = 0.85$）を示した一方で、全22,528個の特徴のうち約10%が有意水準を超えていた。これは、エントロピーのような単一の物理指標との相関だけでは、メカニズム的に意味のある特徴を一意に分離するには不十分であることを示唆している。
2. 断続的な一貫性: 本研究は「区分的な一貫性（piecewise consistency）」の証拠を見出した。特定の機能（例：チャネル界面や渦領域で活性化する特定の機能）が、異なるせん断流の設定間で同様の役割を繰り返す。しかし、この構造は極めて断続的である。それは標準的な物理的分解に明確に適合せず、タイムステップが進むにつれて持続しないことが多い。
3. 拡散性と不一致: 数値シミュレーションとWalrusエミュレータの直接比較により、系統的な出力レベルの不一致が明らかになった。Walrusは中波数帯域のエネルギーを過小評価し、長波長帯域を過大評価する傾向があり、これは過剰な小スケール構造、すなわち「多すぎる」小スケール拡散の生成を示唆している。
4. 表現的 vs 物理的拡散: シミュレーションの進行に伴う特徴のコヒーレンスの喪失は、2つの要因に起因する可能性がある：
   • 物理的要因: Walrusが流れの場を過度に平滑化しており、参照シミュレーションが構造を保持しているスケールにおいて、エネルギースペクトルを散逸させている。
   • 表現的要因: 物理的構造が残っている場合でも、モデルがそれを少数の解釈可能な特徴に集中させるのではなく、多くの弱く活性化した特徴にわたって符号化している可能性がある。
5. 安定したマッピングの欠如: 単一のシミュレーション内のトップランクの特徴に限定して分析を行っても、馴染みのある物理量への安定したマッピングは回収されない。初期タイムステップで意味があるように見えるパターンは、流れが進むにつれてノイズへと解消されるか、予測不能に変化する。

意義と主張
本論文は、基盤モデルの解釈可能性を解決した、あるいはWalrusが流体力学の支配方程式を学習したことを証明したと主張するものではない。むしろ、科学的基盤モデルコミュニティに対する開かれた問いを提示するケーススタディとして機能している：

• 優先順位付け: 大部分の特徴が物理的指標と相関する場合、メカニズム的に意味のある特徴を堅牢に優先順位付けすることは困難である。
• アーティファクト vs 構造: 安定した内部構造と、単一レイヤーのプロービングやSAEの手法に固有の限界を含む分析上のアーティファクトを分離することには、大きな課題がある。
• ベンチマーク: 確立されたベンチマーク（出力レベルの精度など）は、内部表現が真に情報量豊かであるかどうかを判断するには不十分である。モデルは、内部表現が人間の物理的直感と一致していなくても、ダイナミクスを効果的に再現することができる。

著者らは、このようなモデルが何を「学習した」かという主張については慎重に扱うべきであると結論付けている。科学的基盤モデルへの信頼を築くには、事後的に数十億パラメータのトランスフォーマーを逆エンジニアリングしようとするのではなく、より単純で解釈可能なアーキテクチャから始め、体系的にスケールアップしていく必要があることを示唆している。本研究は、モデルが支配的な関係を捉えたのか、それとも単に統計的な規則性を記憶しただけなのかを区別することの重要性を強調している。