ProteinZero: Self-Improving Protein Generation via Online Reinforcement Learning

この論文は、ESMFold と独自に開発した ddG 予測器を報酬として用い、埋め込みレベルの多様性正則化を導入したオンライン強化学習フレームワーク「ProteinZero」を提案し、既存の最優秀手法を上回る設計成功率と多様性を実現しながら、単一 GPU ノードで 3 日以内に効率的に実行可能な自己改善型タンパク質生成モデルを構築したことを報告しています。

要約

論文「ProteinZero」のわかりやすい解説

～「失敗から学ぶ」AI が、タンパク質設計の天才になるまで～

この論文は、**「ProteinZero（プロテイン・ゼロ）」**という新しい AI の仕組みについて紹介しています。

簡単に言うと、これは**「自分で作ったタンパク質の設計図を、AI 自身が評価して、失敗から学びながら、どんどん進化させていく」**という画期的な方法です。

これまでの AI は「過去の成功例（教科書）」を暗記して設計をしていましたが、ProteinZero は**「試行錯誤しながら、自分自身で新しい教科書を作っていく」**ことができます。

---

1. 従来の方法の「壁」とは？

タンパク質（生体の部品）を設計する AI は、これまで「PDB（タンパク質のデータベース）」という、自然界に存在する成功例の教科書を見て学習していました。

• 問題点 1：教科書が少すぎる
  自然界にあるタンパク質は、宇宙の星の数ほどありますが、教科書に載っているのはその一部に過ぎません。AI は「教科書に載っていない新しいデザイン」を作ろうとすると、自信が持てません。
• 問題点 2：「正解」が「実用」ではない
  教科書は「自然に存在するもの」を覚えるのが得意ですが、私たちが欲しいのは「薬になる」「酵素として働く」という実用的な性能です。教科書通りに作っても、実際に使えない（不安定な）ものができてしまうことがありました。

2. ProteinZero の解決策：「オンライン・リインフォースメント学習」

ProteinZero は、AI が**「自分で設計図を作り、自分で評価し、自分で修正する」**というサイクルを回します。これを「オンライン強化学習」と呼びます。

① 評価係（レフェリー）の役割：「ESMFold」と「Fast-ddG」

AI が設計したタンパク質が「いいものか」を判断する必要があります。

• ESMFold（構造のレフェリー）： 設計図（アミノ酸の並び）から、実際に 3 次元の形ができるかを瞬時にシミュレーションします。
• Fast-ddG（安定性のレフェリー）： そのタンパク質が「壊れにくい（安定している）か」を計算します。

【重要なポイント】
従来の方法では、この評価に「物理シミュレーション」という重たい計算が必要で、1 回評価するのに数時間〜数日かかっていました。まるで、1 個の料理を作るのに、1 週間かけて味見をするようなものです。
しかし、ProteinZero は**「Fast-ddG」という「超高速な味見係」を使います。これにより、評価にかかる時間が「数秒」**に短縮されました。これによって、AI は 1 日で何万回も「設計→評価→修正」を繰り返せるようになったのです。

② 多様性の守り神：「埋め込み多様性正則化」

AI が学習を続けると、ある問題が起きます。
**「モード・クラッシュ（単調化）」**という現象です。

• 例え話： AI が「美味しい料理」を作るようになると、最初は「ラーメン、パスタ、カレー」などバラエティ豊かに作ります。しかし、ある時「ラーメンが一番高得点だ！」と気づくと、**「ラーメンしか作らなくなる」**という状態です。
• ProteinZero の対策： AI が「同じようなもの」ばかり作らないように、**「設計図の『雰囲気（埋め込み）』が似すぎないように」**というルールを追加しました。
  これにより、AI は「高得点」を狙いながらも、「ラーメンだけでなく、パスタや寿司も作れる」ようにバランスを保ちます。

3. 結果：驚異的な進化

ProteinZero は、既存の最高峰の AI（ProteinMPNN など）と比べて、以下の点で圧倒的に優れていました。

• 失敗率の激減： 設計したタンパク質が「形にならない」や「壊れる」という失敗が、36〜48% も減りました。
• 成功率 90% 超え： 10 個設計すれば、9 個以上が「安定して形になる」素晴らしい結果を出しました。
• 多様性の維持： 安定性が高まっても、バラエティ豊かなデザインを維持できました。

4. まとめ：なぜこれがすごいのか？

ProteinZero は、**「AI が自分で自分の先生になる」**ことを実現しました。

• 従来の AI： 過去の教科書（データ）を暗記して、似たようなものを作る。
• ProteinZero： 自分で新しい実験（設計）をして、その結果から学び、教科書のない領域（未知のタンパク質）でも活躍できる。

これは、**「試行錯誤のスピードを劇的に速め、AI が自律的に進化する」**ことを意味します。将来的には、新しい薬や、環境を浄化する酵素など、人間には思いつかないような革新的なタンパク質を、AI が次々と生み出せるようになるかもしれません。

一言で言えば：

「ProteinZero は、タンパク質設計の分野で、AI が『教科書』を捨てて、自らの『経験』から天才へと成長する道を開いた」
という画期的な研究です。

技術概要

ProteinZero: オンライン強化学習による自己改善型タンパク質生成モデルの技術的サマリー

本論文は、タンパク質設計における逆フォールディング（逆フォールディング：特定の3次元構造に折りたたまれるアミノ酸配列を生成するタスク）の課題を解決するため、ProteinZero と呼ばれる新しいオンライン強化学習（RL）フレームワークを提案しています。従来の教師あり学習モデルの限界を超え、計算コストを低く抑えながら、設計可能性、安定性、多様性を同時に最適化する自己改善型のシステムを実現しました。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細をまとめます。

---

1. 背景と課題 (Problem)

タンパク質逆フォールディングは、創薬や酵素設計の基盤となる重要なタスクですが、以下の課題に直面しています。

• データ依存性と探索範囲の限界: 既存の高性能モデル（ProteinMPNN, ESM-IF など）は、PDB（タンパク質構造データベース）の限られた配列 - 構造ペアデータで教師あり学習されています。これは自然に存在するタンパク質のわずかな部分しかカバーしておらず、新規な設計空間の探索が困難です。
• 目的関数の不一致: 教師あり学習の目的（既存配列の回復率最大化）は、実際の設計目標（高安定性、高設計可能性、多様な候補の生成）と一致していません。
• オンライン RL の実用障壁:
  • 計算コスト: 物理ベースの評価（FoldX, Rosetta）や高精度構造予測（AlphaFold2/3）は、1 回の評価に数分〜数時間を要し、オンライン RL での反復学習を非現実的にします。
  • モード崩壊（Mode Collapse）: 強化学習において、モデルが報酬を最大化する狭い解空間に収束し、多様な配列を生成できなくなる現象が頻発します。

2. 提案手法：ProteinZero (Methodology)

ProteinZero は、逆フォールディングモデルをオンライン RL で微調整し、自身の生成物から継続的に学習・改善するフレームワークです。

2.1 高速な代理報酬（Fast Proxy Rewards）

オンライン RL を実用的にするため、物理ベースや高精度モデルに依存しない高速な報酬パイプラインを構築しました。

1. 設計可能性報酬（Designability Reward）:
   • ESMFold を使用して生成配列の構造を予測し、US-Align を用いてターゲット構造との TM スコアを計算します。
   • AlphaFold2/3 の MSA 検索や反復処理を回避し、推論を高速化（25〜100 倍の高速化）しています。
2. 熱力学的安定性報酬（Thermal Stability Reward, $\Delta\Delta G$）:
   • 物理計算（FoldX など）の代わりに、逆フォールディングモデル自体のバックボーン条件付き尤度を利用します。
   • 無条件の配列事前分布（Unconditional Prior）で正規化し、野生型（Wild-type）を基準とした相対的な安定性（$\Delta\Delta G$）を推定します。
   • この手法は「Fast-ddG」として知られ、実験値との相関（PCC 0.60-0.62）を維持しつつ、FoldX に比べて 236〜760 倍高速です。

2.2 埋め込みレベルの多様性正則化 (Embedding-Level Diversity Regularization)

オンライン RL におけるモード崩壊を防ぐため、従来の配列レベル（ハミング距離など）ではなく、モデルの潜在空間（埋め込み空間）での多様性を正則化項として導入しました。

• 手法: バッチ内の生成配列の最終層デコーダ活性化を集約して固定次元の埋め込みベクトルを計算し、それらのコサイン類似度を基に多様性スコアを算出します。
• 効果: 意味的に多様な機能を持つ配列の生成を促しつつ、構造的一貫性を保ちます。この正則化項は報酬関数に直接組み込むのではなく、損失関数の別項として追加することで、トレーニングの安定性を確保しています。

2.3 最適化アルゴリズム

提案フレームワークは、以下の 2 つのオンライン RL アルゴリズムに実装されました。

• ProteinZeroRAFT: 報酬に基づいて生成候補をフィルタリングし、上位のものを教師データとして微調整する手法。
• ProteinZeroGRPO: グループ相対方策最適化（Group Relative Policy Optimization）を用い、クラッピングと KL 正則化を組み合わせて方策を直接最適化する手法。
• 両アルゴリズムとも、KL 発散（参照モデルからの乖離防止）と埋め込み多様性正則化を統合した目的関数で学習されます。

3. 主要な貢献 (Key Contributions)

1. ラベルなしでの自己改善: 手動でキュレーションされた好ましいデータセットに依存せず、モデル自身の出力から継続的に探索・改善するオンライン RL フレームワークを初めてタンパク質設計に適用。
2. 効率的な多目的最適化: 構造予測（ESMFold）と自己導出型 $\Delta\Delta G$ 予測を組み合わせ、計算的に実行可能な多目的オンライン RL 最適化を実現。
3. モード崩壊の防止: 配列空間ではなく埋め込み空間で動作する新しい多様性正則化を開発し、機能的一貫性を保ちながら多様性を維持。
4. 設計空間の解明: 異なる RL アルゴリズム（GRPO, RAFT, DPO など）と正則化戦略を比較し、タンパク質設計における最適な構成を特定。

4. 実験結果 (Results)

CATH-4.3 ベンチマーク（0-150 アミノ酸、150-300 アミノ酸の 2 つのカテゴリ）において、既存の最先端モデルと比較評価を行いました。

• 全体的な性能:
  • 成功率の向上: ProteinZero（特に GRPO 版）は、ProteinMPNN や ESM-IF、InstructPLM を上回る性能を示しました。
    • 0-150 残基：成功率 90.13%（ProteinMPNN は 81.95%）。
    • 150-300 残基：成功率 91.19%（ProteinMPNN は 84.67%）。
  • 失敗率の削減: 設計失敗率を 36-48% 削減しました。
• 多目的バランス:
  • 構造精度（TM スコア）、安定性（FoldX ddG）、配列回復率、多様性のすべての指標で改善が見られました。
  • 特に、DPO（オフライン RL）ベースの手法が多様性を犠牲にするのに対し、ProteinZero は多様性を維持・向上させながら成功率を高めることに成功しました。
• 独立した検証:
  • 学習時に使用した ESMFold だけでなく、AlphaFold3 や FoldX、Rosetta といった独立したオラクルでも同様の改善が確認され、学習された原理が一般化可能であることを示しました。
• 計算効率:
  • 1 つの 8 GPU ノードで、報酬計算とデータ生成を含めて3 日以内に RL 学習を完了可能です（従来の物理ベース RL は数ヶ月かかる可能性があります）。

5. 意義と展望 (Significance)

ProteinZero は、タンパク質設計の分野において以下の重要な転換点をもたらします。

• 教師あり学習の限界の突破: 既存の PDB データに依存せず、モデルが自身の生成物から学習することで、未知のタンパク質設計空間への探索を可能にします。
• 実用性の向上: 高速な代理報酬と埋め込み正則化により、オンライン RL の計算的・技術的障壁を解消し、実用的なタンパク質設計パイプラインへの統合を可能にしました。
• 多様性と性能の両立: 従来の RL 手法で課題だった「モード崩壊」を解決し、高品質かつ多様な設計候補を生成する能力を実証しました。

今後は、実験室での検証（ウェットラボ）を通じて、計算機上で予測された安定性や機能が実際に生物学的に機能するかを確認することが次のステップとなりますが、ProteinZero は AI 駆動型のタンパク質設計において、スケーラブルで自律的な進化を実現する強力な基盤技術として位置づけられます。