技術概要：ENSEMBITS

問題定義

タンパク質言語モデル（PLM）と構造予測ツールは、主に静的な構造表現または単純なアミノ酸配列に依存してきた。既存のタンパク質構造トークナイザー（PST）は静的スナップショットの局所幾何学を効果的に捉えるが、タンパク質アンサンブルに内在する相関運動と代替的コンフォメーション状態を符号化することはできない。生物学的機能は、単一の構造ではなく、タンパク質がサンプリングするコンフォメーションの分布（例：基質アクセスを制御する触媒ループ）にコード化されることが多い。分野はアンサンブル生成へと移行しつつあるが、トランスフォーマーベースのアーキテクチャで使用するために、これらの動的アンサンブルをトークン化できる離散語彙が存在しない。さらに、動的データはしばしば希薄であり、既存の手法は残基間での相関運動を共同モデル化することに失敗するか、推論時に利用可能とは限らない完全な軌跡を必要とする。

手法：ENSEMBITS

ENSEMBITSは、タンパク質コンフォメーションアンサンブル用に設計された最初のトークナイザーとして導入される。これは、構造フレーム（コンフォマー）の順序なし多重集合を、残差ベクトル量子化変分オートエンコーダー（RVQ-VAE）アーキテクチャを介して離散トークンにマッピングする。

1. 問題定式化

このモデルは、アンサンブルのトークン化を多重集合からトークンへの圧縮問題として扱う。

• 入力: 単一のタンパク質に対する $P$ フレームの順序なし多重集合 $P = \{x_1, \dots, x_P\}$。ここで、各フレームは原子ごとの座標を含む。
• 出力: アンサンブル全体にわたる残基の局所的な動的モチーフを表す、コードブック $\mathcal{C}$ からの離散トークン。
• 制約: トークナイザーは $P$ に依存せず（$P$-agnostic）、任意の正の基数を持つアンサンブルを受け入れ、フレーム順序に関して置換不変でなければならない。

2. 記述子設計

動的を捉えるため、ENSEMBITSはフレーム全体にわたってSE(3)不変記述子を計算する。著者らは、残基のフレームとその $K$ 個の空間的最近傍との間の SE(3) 変換として局所バックボーン環境を符号化する、ESM3 風の相対フレーム記述子（生産設定）を利用する。

• 動的モード: 決定的に、隣接する同一性（$N^p_r$）は、すべてのフレーム $p$ に対して独立して再計算される。これにより、記述子は固定された隣接アプローチでは見えない信号である軌跡に沿った接触の形成と切断を自然に符号化する。

3. アーキテクチャ

• セットエンコーダー: 残基ごとの記述子多重集合を単一の潜在ベクトル $z$ にマッピングする PerceiverIO 風のアーキテクチャ。これは、共有された要素ごとの MLP と学習可能なクエリートークンによるクロスアテンションを使用し、構成上置換不変性を保証する。
• 残差量子化器: $K=3$ レベルを持つ多段階 RVQ-VAE により、潜在 $z$ をコードブックインデックスのタプル $(c_1, \dots, c_K)$ に離散化する。最終トークンは、これらのレベルからの埋め込みの和である。
• デコーダーと損失: デコーダーは記述子の多重集合を再構成する。再構成損失は、固定された順序を必要とせずに置換不変性を保証するハンガリアンマッチングを用いて、予測された記述子とターゲット記述子を整合させる。

4. 単一フレームからトークンへの蒸留（SFTD）

動的データの希薄さに対処し、単一の静的構造での推論を可能にするため、著者らはSFTDを提案する。

• メカニズム: 訓練中、モデルは同一タンパク質からの 2 つのアンサンブル、すなわち完全なアンサンブル（$P_{max}$）と、ランダムにサンプリングされた部分アンサンブル（$P=1$ まで）を処理する。
• 目的: 部分アンサンブルの潜在表現は、完全なアンサンブルの潜在（完全なアンサンブルにはストップグラディエントを適用）へと引き寄せられる。
• 結果: これにより、トークナイザーはテスト時に予測された単一の構造から「動的トークン」を予測でき、アンサンブルレベルの情報を単一フレーム表現に効果的に蒸留する。

主要な貢献

1. 定式化とパイプライン: SE(3) 不変記述子と置換不変セットエンコーダーを利用した、タンパク質動力学のトークン化を多重集合からトークンへの圧縮問題として初めて端から端まで定式化したこと。
2. SFTD 目的: 部分アンサンブルの埋め込みを完全アンサンブルの対応物と整合させる新しい訓練戦略であり、単一の静的構造のみが利用可能な場合でも、モデルを動的認識型トークナイザーとして機能させる。
3. 実証的検証: ENSEMBITSがアンサンブルレベルの動力学を捉え、下流の機能タスクに効果的に転移することを示す包括的な検証。大幅に少ない事前学習データを使用しているにもかかわらず、静的トークナイザーを上回る性能をしばしば発揮する。

結果

1. 動力学表現（RMSF と ANOVA）

• RMSF 予測: ENSEMBITSは、二乗平均平方根変動（RMSF）予測ベンチマークで優位に立つ。
  • 多フレーム: ENSEMBITS（$P=full$）は、mdCATH-div において最強のベースライン（ProtProfileMD）を約 2.5 Spearman ポイント、MISATO スプリットにおいて約 11〜18 ポイント上回る。
  • 単一フレーム: $P=1$（蒸留済み）であっても、ENSEMBITSは RMSF 予測において、AminoAseed や ESM3structを含む他のすべての単一フレームトークナイザーを上回る。
• ANOVA 検定: 運動振幅（$s_1$）に関する一元配置分散分析により、ENSEMBITSトークンは局所運動振幅の分散の37.1%を説明することが明らかになった（$\eta^2 = 0.371$）。これは、次に良い多フレームベースライン（Vote_3Di、$\eta^2 \approx 0.128$）の約3倍、単一フレームの静的トークナイザーの約7倍である。これは、トークンが単に静的なフォールドクラスや位置ではなく、識別可能な局所動力学を符号化していることを確認する。

2. 下流の機能タスク

はるかに少ない事前学習データ（MD アンサンブルのみ、PDB スケールの逆フォールド事前学習なし）を使用しているにもかかわらず、ENSEMBITSは複数のタスクで静的トークナイザーと同等かそれ以上の性能を示す。

• 変異効果予測: PROTEINGYM ベンチマークにおいて、ESM2 と組み合わせた ENSEMBITS（$P=1$）は、ESM2 と組み合わされた他のすべての構造トークナイザーを上回り、ESM2 単独に対して6.9% の改善を達成した。
• 結合部位と親和性: ENSEMBITSは、結合部位予測（AUROC 0.750）およびすべてのスプリットにわたる結合親和性回帰で首位に立つ。特に、ESM3struct や AminoAseed を上回る結合部位検出性能を示しており、局所コンフォメーションの柔軟性が静的トークナイザーが見逃す界面残基にとって重要なシグナルであることを示唆している。
• EC および GO 予測: ENSEMBITSは、酵素委員会（EC）および遺伝子オントロジー（GO）タスクにおいて ESM3struct と競争力があるが、フォールドクラスと強く相関するタスクでは、逆フォールド事前学習に起因する可能性のある ESM3struct が優位性を保っている。

意義と主張

本論文は、ENSEMBITSを、静的構造予測からアンサンブル生成へと向かうタンパク質モデリングの進化において必要なステップとして位置づけている。

• ギャップの橋渡し: タンパク質動力学を言語モデリングと設計にもたらすために必要な離散語彙を提供する。
• データ効率: 表現が明示的に動力学向けに設計されていれば、ESM3 などのモデルが使用する膨大な配列/構造コーパスと比較して、比較的小さな MD コーパスから高品質な動的表現を学習できることを実証する。
• 実用性: SFTD メカニズムは、動的データの利用可能性というボトルネックを緩和し、学習された動力学の事前知識を活用しつつ、単一の静的構造（例：AlphaFold 予測）に対してモデルをクエリできるようにする。

著者らは、トークンの品質は基盤となる訓練軌跡の網羅性と忠実度によって制限されること、そして目標はすべてのタンパク質動力学を解決したと主張することではなく、動的情報を離散トークンに注入するための手法を導入することであることを控えめに指摘している。今後の課題には、より大規模なアンサンブルコーパスへのスケーリングと、これらのトークンをタンパク質言語モデルに直接統合して多フレーム構造を生成することが含まれる。