Mechanistic Interpretability and Causal Feature Steering of Neural Quantum States via Sparse Autoencoders

技術要約：疎な自己符号化器によるニューラル量子状態のメカニズム的解釈可能性と因果的特徴制御

問題提起
ニューラル量子状態（NQS）は、スピン系やフェルミオン系における基底状態の表現および実時間ダイナミクスの両方において成功を収めており、強力な変分アンザッツとして台頭している。しかし、その経験的な有効性にもかかわらず、NQSは依然として「ブラックボックス」のままである。特定の物理量に対して明示的な最適化を行わず、変分目的関数（エネルギー最小化など）のみに基づいて学習されるため、NQSがどのようにして秩序パラメータや相関関数といった量を正確に捉えているのかは不明である。この不透明さは、数値的なベンチマークの取得が困難な領域において、モデルの信頼性と堅牢性を阻害する要因となる。既存のアプローチは、主に外部的なアーキテクチャの制約やパラメータ解析に焦点を当てており、ネットワークが内部でどのように物理情報を組織化しているかという事後的な（post hoc）分析が欠落している。

手法
著者らは、メカニズム的解釈可能性（mechanistic interpretability）のツールである**疎な自己符号化器（Sparse Autoencoders: SAEs）**を用いた体系的なフレームワークを提案している。これは、NQSの内部活性化を、疎で解釈可能な特徴へと分解するものである。

1. NQSアーキテクチャ: 本研究では、自己回帰型トランスフォーマーベースのNQSを利用する。波動関数の振幅 $\psi(\sigma)$ は、ボルン確率 $p(\sigma)$ と位相 $\phi(\sigma)$ に分解される。トランスフォーマーは因果的シーケンスモデルとして機能し、スピン配置をサイトごとに予測する。
2. SAEの適用: SAEは、最終層のトランスフォーマーの**残差ストリーム（residual stream）**に適用される。これは、特徴方向の集合 $\{f_k\}$ と疎な係数 $z_k$ からなる過完備な辞書を学習し、活性化ベクトル $h$ を $h \approx \sum z_k f_k$ という疎な線形結合として近似する。
3. 教師なし学習: 決定的なことに、SAEは残差ストリームの活性化に対して完全に教師なしで学習される。物理的なラベル（磁化の値など）は一切受け取らず、活性化の再構成（スパース性ペナルティとしての $L_1$ ノルムを伴う）のみを最適化する。
4. 特徴量と物理量の相関: 学習後、著者らはスピン鎖にわたって疎な係数を「平均プーリング」することで、スカラー値の特徴強度を得る。これらを、同一のサンプリングされた配置から計算された物理量（磁化、スタッガード磁化、半鎖相関関数など）と相関させる。
5. 因果的介入（ステアリング）: 因果関係を確立するために、著者らは標的を絞った介入を行う。特定の物理量（例：磁化）を制御するために、他の要素は固定したまま、特定の項の活性化強度をスケーリング（$z_{k^*} \to \alpha z_{k^*}$）し、修正された残差ストリームを再構成して、結果として生じる物理量の変化を観察する。この際、介入が局所的に留まるよう、変分エネルギーを監視する。

主な貢献と結果

• 解釈可能な特徴の発見: 1次元横磁場イジング模型（TFIM）において、SAEは、磁化（$M_z$）、絶対磁化（$|M_z|$）、および半鎖相関関数（$C_{zz}$）を含む物理量とほぼ完璧に相関（$|r| > 0.97$）する疎な特徴を自律的に発見した。この発見は、SAEがこれらの量に関する事前知識を持たないにもかかわらず、常磁性相、強磁性相、および臨界点にわたって実現された。
• 因果的制御: 著者らは、これらの特徴がNQSの予測において因果的な役割を果たしていることを示した。単一の特徴強度（$\alpha$）をスケーリングすることで、対応する物理量（例：磁化）を滑らかかつ単調に制御（ステアリング）できる一方で、変分エネルギーの変化は極めて小さかった（相対偏差 $< 0.02\%$）。これにより、これらの特徴が単なる相関関係によるアーティファクトではなく、物理的予測を制御する内部座標であることが確認された。
• 実時間ダイナミクス: このフレームワークは、TFIMにおける実時間ダイナミクスへと拡張された。単一の静的な特徴方向が、全進化軌道を通じて時間依存的な磁化 $M_z(t)$ を追跡できることが判明した。さらに、時間の経過に伴うこの特徴の活性化強度をスケーリングすることで、磁化のダイナミクスを系統的に変化させることができ、NQSが静的な内部座標系の中にダイナミカルな情報をエンコードしていることが証明された。
• 2次元系への汎用性: この手法は、2次元ハイゼンベルグ反強磁性模型（非可積分な2次元系）においても検証された。SAEは、スタッガード磁化（サブ格子を区別する空間的に構造化された秩序パラメータ）と強く相関する特徴を特定した。この特徴のステアリングによってスタッガード磁化を制御することに成功し、本手法が1次元系や可解な系を超えて適用可能であることを示した。

意義と主張
論文は、NQSは単なる波動関数の振幅の関数近似器ではなく、豊かで解釈可能な物理情報の内部表現をエンコードしていると主張している。

• 診断ツール: このアプローチは、訓練されたNQSが物理的に妥当な表現を学習しているかどうかを評価する方法を提供する。秩序量を因果的に制御する特徴の存在は、学習の成功を示す証拠となり、逆にその欠如はモデルの失敗を診断する手がかりとなる。
• 介入ツール: 本研究は、NQSの制御された特徴レベルの操作への経路を確立する。残差ストリーム内の特定のスカラ係数を調整することで、ネットワーク全体を再学習または微調整することなく、特定の物理量をチューニングできる。
• 双方向の交換: 著者らは、量子多体系物理学とメカニズム的解釈可能性との間の相乗的な関係を提唱している。NQSは、解釈可能性の主張をテストするための「正解（グラウンドトゥルース）」となる物理量を提供し、一方で解釈可能性のツールは、変分波動関数という「ブラックボックス」を開くための手法を提供することで、より信頼性が高く透明な量子シミュレーションへと導く。

著者らは、本研究の範囲について謙虚な姿勢を維持しており、本研究はトランスフォーマー・アーキテクチャおよび最終層の残差ストリームに限定されていることを指摘している。これには、他のアーキテクチャ（RBMやRNNなど）や、層ごとのセルフアテンション機構に関する未解決の問いが残されている。