技術要約：適応型タンパク質トークン化（Adaptive Protein Tokenization）

問題提起

現在のタンパク質構造のトークン化手法は、タンパク質配列に沿った空間的に隣接する残基から情報をプーリングすることによる**局所性（locality）**に依存しています。これらの手法は高い再構成精度を実現する一方で、著者らは以下の2つの決定的な限界を指摘しています。

1. 生成における誤差の蓄積： 局所的なトークン化に基づく生成モデルは、誤差の伝播に苦しみます。単一の誤ったサンプリングされたトークンが、予測構造の著しい逸脱を引き起こす可能性があり、その結果、離散的なアプローチは生成指標において連続的な拡散モデルと比較して性能が低下します。
2. スケーラビリティの低さ： トークンの数はタンパク質のサイズに対して線形にスケールします。これにより、巨大なタンパク質複合体のモデリングが計算コスト的に高価になり、タンパク質を効果的に圧縮する能力を制限します。
3. 表現力の制限： 局所的なトークナイザーは、ダウンストリームのタスクのために固定サイズの表現を生成する際、平均プーリング操作を必要とすることが多く、その結果、重要なグローバル情報を破棄したり、タンパク質のサイズに対してスケーリングが悪くなったりします。

手法

著者らは、**適応型タンパク質トークン化（Adaptive Protein Tokenization: APT）**を提案しています。これは、各トークンが特定の局所的な近傍に対応するのではなく、連続するトークンがグローバルな表現に対して段階的に詳細度を高めていく、グローバルなトークン化手法です。

アーキテクチャ

APTは、離散的なボトルネックを持つ**拡散オートエンコーダ（diffusion autoencoder）**として実装されています。

• 入力： 長さ $L$ の生の $C_\alpha$ 座標（ゼロ質量に正規化済み）のシーケンス。
• エンコーダ： 双方向アテンション・トランスフォーマーが、生の座標を潜在シーケンス $c \in \mathbb{R}^{L \times d}$ にマッピングします。
• 離散化： 潜在シーケンスは、レベル $(8, 5, 5, 5)$ の**有限スカラー量子化（Finite Scalar Quantization: FSQ）**を用いて離散化され、実効的なコードブックサイズは1,000となります。
• 適応メカニズム： 適応性を強制するため、モデルは**ネステッド・ドロップアウト（nested dropout）**を採用しています。学習中、上限カットオフ $U$ が $[1, \min(L, k_{max})]$ から一様にサンプリングされます。このカットオフを超えるトークンはドロップされます。これにより、モデルは最初の数個のトークンに重要なグローバル情報（低周波成分）を配置し、後続のトークンに高周波の詳細を配置するように促されます。
• デコーダ： **フローマッチング目的関数（flow-matching objective）**を用いて学習された拡散デコーダが、離散トークンを条件として原子座標を再構成します。デコーダは確率的等変性（stochastic equivariance）を利用し、対称性を明示的に強制するのではなく、学習によって獲得します。
• サイズ予測： タンパク質のサイズは、最初のトークンからクロスエントロピー損失を用いて回帰されます。これにより、モデルはタンパク質のサイズを条件付けシーケンスの長さから切り離すことができます。

学習と推論

• 学習： モデルは、Foldseekでクラスター化されたAFDBデータベースから得られた約473,000個の合成AlphaFold2予測を用いて学習されます。学習は2段階で行われます。(1) 確率的等変性を学習するためにランダムな回転を用いたオートエンコーダの学習、(2) APTトークンに対する自己回帰的（GPTスタイル）モデルの学習。
• 推論戦略：
  • テイル・ドロップアウト（Tail Dropout）： 生成中、モデルはトークン生成を早期に停止できます。トークンシーケンスの任意の接頭辞（prefix）は、有効な条件付け信号となります。末尾（高周波の詳細）は拡散デコーダに委ねられます。
  • エントロピーに基づく停止： 著者らは、表現の忠実度と生成能力のバランスを取るために、最適な停止点を決定するためのエントロピーに基づく停止基準（有限カットオフ、スプラインベースの最小エントロピー）を導入しています。
  • 分類器アニーリング（Classifier Annealing）： 真のデータ多様体から外れることによるアーティファクトを軽減するため、デコーディング中に条件付きガイダンス場と無条件ガイダンス場の間を補間する分類器アニーリングを採用しています。

主な貢献

1. グローバル適応型トークナイザー： 局所的な近傍ではなく、粗から密への階層構造としてグローバルな記述子を表すトークナイザーの導入。
2. 性能の検証： 本モデルが、生成、表現学習、および再構成タスクにおいて、最先端（SOTA）の局所的トークナイザー（例：DPLM2, ESM3, Kanzi）と同等またはそれ以上の性能を示すことを実証。
3. 推論時テクニック： サンプル複雑性を管理し、誤差への露出を軽減するためのエントロピーベースのサンプリングと分類器アニーリングの開発。これにより、デザイン可能性（designability）が向上しました。
4. ゼロショット応用： タンパク質のサイズを条件付けシーケンス長から切り離すことで、タンパク質縮小（protein shrinking）（構造を維持したまま残基数を減らす）や、推論時のスケーリングによる**アフィニティ成熟（affinity maturation）**などのゼロショット応用を可能にしました。

結果

• 再構成： CATH、CAMEO、およびAFDBテストセットにおいて、APTはSOTAの連続拡散モデルと同等のRMSDおよびTMスコアを達成しました。特に、わずか32〜64個のトークンを使用するだけでRMSDが2 Å未満となることは、タンパク質構造の圧縮可能性を裏付けています。
• 生成： APTトークンに対して学習された自己回帰モデルは、0.871の設計可能性（designability）（scRMSD < 2 Åとなるサンプルの割合）を達成しました。これは、DPLM2（0.486）やESM3といった離散拡散モデルを大幅に上回る数値であり、「Best-of-N」サンプリング手法を用いずに達成されました。
• 表現学習： CATH分類タスクにおいて、APTトークンに対する非線形（MLP）プロービングは、ESM3やDPLM2の表現に対するプロービングを上回りました。決定的なことに、APTは平均プーリングなしで固定サイズのグローバル表現を提供し、高度に圧縮された表現（16トークン）であっても、より大きなモデルを凌駕します。
• 応用：
  • タンパク質縮小： 拡散プロセスをより少ない残基数で条件付けることにより、グローバルおよびローカルな構造を維持したまま、タンパク質（ヘモグロビン、ベータバレルなど）の小型版を生成することに成功しました。
  • アフィニティ成熟： 潜在空間における報酬関数（例：iPAE）に導かれたビームサーチを用い、弱い結合体から強い結合体へと進化させる能力を実証しました。

重要性

本論文は、APTが生物学的タスクのスケーリングに対するタスク指向のアプローチを提供すると主張しています。局所的なトークン化からグローバルなトークン化へと移行することで、再構成の忠実度と生成能力の間のトレードオフを解決しています。平均プーリングによる情報損失なしにタンパク質を固定サイズのベクトルに圧縮できる能力は、巨大なタンパク質複合体のモデリングや、推論時のスケーリングを行うための実用的な経路を提供します。著者らは、この研究を、複数の長さスケールにわたって推論を行うことが可能な生成モデルへの一歩として位置づけています。

注：著者らは制限事項についても認めており、具体的には、APTトークンはグローバルなタスクには効果的であるが、モチーフのスキャフォールディングのようなローカルなタスクには適していないこと、また、現在のトークナイザーは特定の機能に不可欠な側鎖やアミノ酸組成を考慮していないことを挙げています。