潜在拡散を用いた生成のためのタンパク質折りたたみモデルの再利用

PLAIDは、タンパク質折り畳みモデルの潜在空間を学習することで、タンパク質の1次元配列と3次元構造を同時に生成するマルチモーダル生成モデルです。

2024年のノーベル賞がAlphaFold2に授与されたことは、生物学におけるAIの役割に対する認識の重要な瞬間を示しています。タンパク質折り畳みの次に来るものは何でしょうか？

PLAIDでは、タンパク質折り畳みモデルの潜在空間からサンプリングすることを学習し、新しいタンパク質を生成する手法を開発しました。これは、構成可能な機能や生物種のプロンプトを受け入れられ、構造データベースよりも2〜4桁大きい配列データベースで学習することができます。多くの従来のタンパク質構造生成モデルとは異なり、PLAIDはマルチモーダル共生成の問題設定、すなわち離散的な配列と連続的な全原子構造座標の両方を同時に生成する問題に取り組みます。

構造予測から実世界の薬剤設計へ

拡散モデルがタンパク質を生成する能力について有望であることが最近の研究で示されていますが、実世界の応用を非現実的にするような従来モデルの限界が依然として存在します。例えば：

全原子生成：既存の生成モデルの多くは、主鎖原子のみを生成します。全原子構造を生成し、側鎖原子を配置するためには、配列を知る必要があります。これは、離散的モダリティと連続的モダリティの同時生成を必要とするマルチモーダル生成問題を生み出します。

生物種特異性：ヒトでの使用を目的としたタンパク質医薬品は、ヒト免疫系によって破壊されるのを避けるために、ヒト化される必要があります。

制御仕様：創薬とそれを患者の手に渡すことは複雑なプロセスです。これらの複雑な制約をどのように指定すればよいでしょうか？例えば、生物学的问题が解決された後でも、錠剤がバイアルよりも輸送が容易だと判断し、溶解性に関する新たな制約を追加するかもしれません。

「有用な」タンパク質の生成

単にタンパク質を生成するよりも、生成を制御して有用なタンパク質を得ることの方が重要です。そのためのインターフェースはどのようなものになるでしょうか？

インスピレーションとして、構成可能なテキストプロンプトによる画像生成の制御方法を考えてみましょう（Liu et al., 2022の例）。

PLAIDでは、この制御仕様のインターフェースを模倣しています。最終目標はテキストインターフェースを介して完全に生成を制御することですが、ここでは概念実証として、機能と生物種という2つの軸に関する構成可能な制約を考察します：

機能-構造-配列の関連性の学習。PLAIDは、金属タンパク質でよく見られる四面体システイン-Fe2+/Fe3+配位パターンを学習しながら、高い配列レベルの多様性を維持します。

配列のみの学習データを用いた訓練

PLAIDモデルのもう一つの重要な側面は、生成モデルを訓練するために配列のみが必要であることです！生成モデルは、その学習データによって定義されるデータ分布を学習します。実験的な構造を得るよりも配列を得る方がはるかにコストが低いため、配列データベースは構造データベースよりもかなり大規模です。

より大きく広範なデータベースからの学習。タンパク質配列を取得するコストは、実験的に構造を決定するコストよりもはるかに低く、配列データベースは構造データベースよりも2〜4桁大きいです。

どのように機能するのか？

生成モデルを訓練して、配列データのみを使用して構造を生成できる理由は、タンパク質折り畳みモデルの潜在空間上で拡散モデルを学習するためです。その後、推論時に、有効なタンパク質のこの潜在空間からサンプリングした後、タンパク質折り畳みモデルの凍結された重みを使用して構造をデコードできます。ここでは、AlphaFold2モデルの後継であり、検索ステップをタンパク質言語モデルに置き換えたESMFoldを使用します。

我々の手法。訓練時には、埋め込みを得るために配列のみが必要です。推論時には、サンプリングされた埋め込みから配列と構造をデコードできます。❄️は凍結された重みを示します。

このようにして、事前学習されたタンパク質折り畳みモデルの重みに含まれる構造的理解情報を、タンパク質設計タスクに利用できます。これは、ロボティクスにおける視覚-言語-行動（VLA）モデルが、インターネット規模のデータで訓練された視覚-言語モデル（VLM）に含まれる事前情報を利用して知覚、推論、理解情報を提供する方法と類似しています。

タンパク質折り畳みモデルの潜在空間の圧縮

この方法を直接適用する際の小さな問題点は、ESMFoldの潜在空間、実際には多くのトランスフォーマーベースモデルの潜在空間が、多くの正則化を必要とすることです。この空間はまた非常に大きいため、この埋め込みを学習することは、高解像度画像合成へのマッピングになってしまいます。

この問題に対処するため、我々はまたCHEAP（Compressed Hourglass Embedding Adaptations of Proteins）を提案します。これは、タンパク質配列と構造の結合埋め込みのための圧縮モデルを学習するものです。

潜在空間の調査。（A）各チャネルの平均値を可視化すると、一部のチャネルが「大規模な活性化」を示します。（B）中央値（灰色）と比較して上位3つの活性化を調べ始めると、これが多くの層で起こっていることがわかります。（C）大規模な活性化は、他のトランスフォーマーベースモデルでも観察されています。

この潜在空間は実際には高度に圧縮可能であることがわかりました。我々が扱っている基本モデルをよりよく理解するために、少しメカニズム的解釈可能性を行うことで、全原子タンパク質生成モデルを作成することができました。

この研究ではタンパク質配列と構造の生成の場合を検討していますが、この方法は、より豊富なモダリティからより乏しいモダリティへの予測器が存在する任意のモダリティに対して、マルチモーダル生成を実行するように適応させることができます。タンパク質の配列から構造への予測器が、より複雑なシステム（例えば、AlphaFold3は核酸や分子リガンドと複合体を形成したタンパク質も予測可能）に取り組み始めているため、同じ方法を使用してより複雑なシステム上でマルチモーダル生成を実行することは容易に想像できます。我々の方法を拡張するための共同研究、またはウェットラボで我々の方法をテストすることに興味がある方は、ぜひご連絡ください！

我々の論文があなたの研究に役立った場合は、PLAIDとCHEAPについて以下のBibTeXの使用をご検討ください：

@article{lu2024generating, title={Generating All-Atom Protein Structure from Sequence-Only Training Data}, author={Lu, Amy X and Yan, Wilson and Robinson, Sarah A and Yang, Kevin K and Gligorijevic, Vladimir and Cho, Kyunghyun and Bonneau, Richard and Abbeel, Pieter and Frey, Nathan}, journal={bioRxiv}, pages={2024--12}, year={2024}, publisher={Cold Spring Harbor Laboratory} }

@article{lu2024tokenized, title={Tokenized and Continuous Embedding Compressions of Protein Sequence and Structure}, author={Lu, Amy X and Yan, Wilson and Yang, Kevin K and Gligorijevic, Vladimir and Cho, Kyunghyun and Abbeel, Pieter and Bonneau, Richard and Frey, Nathan}, journal={bioRxiv}, pages={2024--08}, year={2024}, publisher={Cold Spring Harbor Laboratory} }

プレプリント（PLAID, CHEAP）とコードベース（PLAID, CHEAP）もご覧いただけます。

ボーナス：タンパク質生成の楽しみ！

PLAIDによる追加の機能プロンプト生成。

PLAIDによる無条件生成。

膜貫通タンパク質は、脂肪酸層に埋め込まれたコア部分に疎水性残基を持ちます。これらは、膜貫通タンパク質のキーワードでPLAIDにプロンプトを与えると一貫して観察されます。

機能プロンプトに基づく活性部位再現の追加例。