PairSAE: Mechanistic Interpretability from Pair Representations in Protein Co-Folding

技術要約：PairSAE – タンパク質の共折り畳みにおけるペア表現からのメカニスティックな解釈可能性

問題提起
AlphaFoldやBoltzモデルのような構造生物学のための基盤モデルは、生体分子の構造や複合体の予測において驚異的な精度を達成している。しかし、それらの出力の背後にある内部特徴を理解することは依然として大きな課題である。Sparse Autoencoder（SAE）は、トランスフォーマーベースの言語モデル（LLM）やタンパク質言語モデル（pLM）において、「重畳（superposed）」された特徴をモノセマンティック（単一の意味を持つ）な概念へと写像することで、これらを解きほぐすことに成功してきたが、構造予測モデルへの転用は容易ではない。

標準的なLLMとは異なり、構造予測モデルは、配列埋め込みとペアワイズ相互作用テンソルの両方を維持する**ペアフォーマー（pairformer）**アーキテクチャを利用している。単純にSAEをこれらのモデルに適用すると、2つの決定的な問題が生じる：

1. 二次的な爆発（Quadratic Blow-up）： ペア表現の辞書を個別に学習することは、配列長に対して二次的にスケールするため、計算量的に極めて困難になる。
2. 概念の隠蔽： 特徴量は、配列空間とペア空間の両方に共同で重畳されている可能性がある。これらを別々に扱うことは、両方の表現に分散している概念を捉え損なうことになり、また、独立して扱うことは学習された特徴の結合的な性質を不明瞭にする。

手法：PairSAE
著者らは、配列とペアの両方の表現を共同で再構成する、共有されたトークンレベルの特徴を学習するために設計されたフレームワークであるPairSAEを提案している。この手法は、主に3つの段階で構成される：

1. NモードSVDによるペアワイズ相互作用の要約：
   ペアワイズテンソル（$Z \in \mathbb{R}^{N_{tok} \times N_{tok} \times n_z}$）の二次的なスケーリングを回避するため、著者らはNモード特異値分解（SVD）を実行する。彼らは、各トークンの行および列の相互作用における役割をエンコードする左特異ベクトル（モード1およびモード2の行列 $U^{(1)}$ および $U^{(2)}$）を抽出する。これらは上位 $r$ 個の特異ベクトルに切り詰められ、トークンごとの要約埋め込み $m_i$ を形成するために結合される。このステップにより、二次的な複雑さを増すことなく、相互作用の役割を保持したまま、ペアワイズ情報をトークンレベルの表現へと圧縮する。

2. 共有オートエンコーダ・アーキテクチャ：
   PairSAEエンコーダは、元の配列埋め込み（$s_i$）とSVD由来の相互作用要約（$m_i$）の結合を入力として受け取る。そして、スパース性を誘導する非線形関数（BatchTopKとReLUの合成）を用いて、スパースな潜在特徴ベクトル $h_i$ を生成する。
   決定的なのは、デコーダがこの共有された特徴セット $h$ を用いて、以下の両方を再構成することである：

• 配列埋め込み：$\hat{s}_i = D_s h_i + b_{dec}^s$
• ペアワイズ埋め込み：$\hat{z}_{i,j} = D_z^{row} h_i + D_z^{col} h_j + b_{dec}^z$
  これにより、モデルは局所的な配列コンテキストとペアワイズ相互作用の両方を説明する、統一された特徴基底を学習することを強制される。

3. 損失関数と訓練：
   モデルは、入れ子状の幅（$c_1 < c_2 < c_3 = D$）にわたってデコーダと特徴ベクトルをスライスし、堅牢性を確保するMatryoshka SAE損失を用いて訓練される。総損失には、両方の表現に対する再構成項（等しく重み付けされる）と、「死んだ（dead）」特徴を防ぐための補助損失が含まれる。訓練は、ペアワイズ損失項のモンテカルロ近似を計算し、各 $i$ に対して単一の $j$ をサンプリングすることで加速される。

主な貢献

• PairSAEアーキテクチャ： NモードSVDを通じてペアワイズテンソルを要約し、共有された特徴セットから配列空間とペア空間の両方へとデコードすることにより、ペアフォーマーベースのモデルにスパース辞書学習を適用する新しい手法。
• 解釈可能な特徴の発見： 本手法は、構造生物学モデルに従来のSAEを適用した際の落とし穴を回避しつつ、生物物理学的概念に整合する特徴を正常に回収することに成功した。
• 結合親和性への適用： 本フレームワークは、タンパク質-リガンド複合体（PLINDERデータセット）におけるBoltz-2の活性化に適用され、潜在表現と結合親和性の予測を結びつけている。

結果
著者らは、タンパク質-リガンド複合体のBoltz-2表現を用いてPairSAEを評価した：

• 線形プロービング： PairSAEの特徴は、UniProt残基アノテーションに対して高い予測能力を示した。レイヤー33およびレイヤー64（第3リサイクリングステップ）の特徴は、トークンレベルで29.1%および24.0%のF1スコアを、複合体レベルで53.2%および49.6%のF1スコアを達成し、ESM2-650Mのニューロンに基づく分類器（それぞれ0.8%および19.7%を達成）を大幅に上回った。
• 概念の整合性： 可視化により、膜貫通領域、ジスルフィド結合、プロテアーゼ活性部位などの特定の構造モチーフに反応する特定の特徴が示され、UniProtのグラウンドトゥルースに対して高いF1スコアを示した。
• 親和性予測： 最大プーリングされたPairSAE特徴を用いたLASSO回帰により、テストセット（PoseBusters）においてBoltz-2の親和性値を $R^2 = 0.528$ で予測した。高親和性複合体に強く相関し、リガンドに特異的に反応して不活性な複合体との顕著なグループ差を示す特徴（例：特徴2299）が特定された。

意義と主張
本論文は、PairSAEが構造生物学のための基盤モデルの潜在空間を、解釈可能な構造的概念へと正常に結びつけたことを主張している。ペアフォーマーの特定のアーキテクチャ上の制約に対処することで、PairSAEは、従来のSAEを制限する二次的な複雑さや特徴の重畳の問題に陥ることなく、折り畳みや結合に関して「モデルが何を知っているか」を明確にする。

著者らは、本研究を構造生物学における原理的な評価と発見への一歩として位置づけており、研究者が予測の信頼性と生物物理学的な妥当性を評価することを可能にするものである。彼らは、手法が配列特徴とペア特徴を既知のアノテーションに効果的にマッピングしている一方で、詳細なリガンドのアノテーションが不足しているため、リガンド活性化特徴を特定の概念にマッピングすることには依然として課題が残ると指摘している。今後の展望として、より多くのレイヤーへのスケールアップ、これらの特徴を自動化パイプラインへの統合、および解釈可能なタンパク質設計のためのステアリング手法の開発を挙げている。