618ZXW

Natureサブジャーナルに掲載されました!筆頭著者が、タンパク質言語モデルの少数サンプル学習法を詳細に解説し、ウェット実験データの不足という問題を解決します。

「Meet AI4S」ライブストリームシリーズの第3回では、上海交通大学自然科学研究所および上海国家応用数学センターのポスドク研究員である周子怡博士をお招きしました。上海交通大学で周子怡博士の研究グループは、洪亮教授が率いており、AIタンパク質・創薬、分子生物物理学に焦点を当てています。同グループは目覚ましい成果を上げており、これまでに77本の論文を発表しており、その多くはNature誌に掲載されています。

このプレゼンテーションでは、周子怡博士が「タンパク質言語モデルのための少量学習法」に関するチームの最新の研究成果を共有し、AI 支援による指向性進化の新しいアイデアを探りました。

タンパク質言語モデル(PLM)の研究背景

タンパク質とタンパク質工学

タンパク質は生物学的機能の主要な担い手であり、生命活動の実行者です。天然アミノ酸、特にアンモニアは脱水縮合反応を起こしてタンパク質の残基配列を形成し、それが折り畳まれて三次構造を形成します。タンパク質中のアミノ酸の種類が変化することで、その構造と機能が変化することがあります。

天然のタンパク質は産業や医療のニーズを満たせないことが多いため、タンパク質工学はタンパク質を変異させることで、触媒活性、安定性、結合能力などの機能特性を強化することを目的としています。

タンパク質の機能特性の定量化は、一般的に適応度(Fitness)と呼ばれます。現在、タンパク質工学の主流は指向性進化です。この手法では、ランダム変異とハイスループット実験を用いて高い適応度を持つ変異体を見つけ出しますが、実験コストが高額です。そこで本日のテーマは、AI手法を用いて適応度を予測し、実験コストを削減する方法です。

PLMアーキテクチャ

ChatGPTに代表される言語モデルは非常に強力で、高品質なテキスト理解と生成が可能であることが知られています。これらのモデルは大量のテキストで事前学習されており、テキストの統計パターンを学習し、基本的な文法と文脈における単語の意味を習得します。では、タンパク質言語モデルを同様に大量のタンパク質配列で学習することは可能でしょうか?答えはイエスです。

タンパク質言語モデル(PLM)には、主に3つの機能があります。まず、PLMはタンパク質配列の共進化情報をモデル化し、残基間の相互依存性と進化的制約を学習します。これは、自然言語モデル(LM)がテキストの文法を学習するのと似ています。PLMはこの能力を利用して、どの変異が有害でどの変異が有益かを推定し、それによって変異の適応度を予測することができます。

第二に、適応度予測に加えて、PLM はタンパク質のベクトル表現を計算することもできます。これは構造予測やタンパク質マイニングに使用でき、微調整後は機能予測にも使用できます。

最後に、PLM は ChatGPT と同様に条件付きタンパク質生成を実行できるため、de novo タンパク質設計が可能になります。

PLMのアーキテクチャは自然言語学習(LM)のアーキテクチャに似ており、自己回帰モデルとマスクモデルで構成されています。どちらのモデルも、自己注意メカニズムと全結合層で構成されるTransformerネットワーク構造を採用しています。主な違いは、事前学習の目的にあります。

自己回帰モデルの事前学習目標は、次のアミノ酸を左から右へ順番に生成することです。一方、マスキングモデルの目標は、穴埋めテストのように、ランダムにマスクされたアミノ酸を再構築することです。自己回帰モデルは、各アミノ酸を予測する際に、既に左側に生成された配列のみに頼ることができるため、その注目度は一方向です。一方、マスキングモデルは予測時にマスクされた位置の両側にあるアミノ酸を見ることができるため、注目度は双方向です。

PLMにおける2つの注目の研究分野

現在、PLMの研究ホットスポットは主に2つの方向に分類されます。1つ目は、検索拡張型PLMです。これは、現在のタンパク質多重配列アライメント(MSA)を学習または予測時に追加入力として用いることで、検索された情報を通じて予測性能を向上させるモデルです。例えば、MSA TransformerやTransceptionはこのタイプの典型的なモデルです。

第二に、マルチモーダルPLMがあります。これらのモデルは、タンパク質配列に加えて、タンパク質構造やその他の情報も追加入力として組み込むことで、表現能力を高めています。例えば、今年提出された私たちのグループのProSSTモデルは、タンパク質構造を構造トークン配列に定量化し、これをアミノ酸配列と共にTransformerモデルに入力することで、これら2種類の情報を別のアテンションメカニズムによって融合します。もう一つの例は、同時にリリースされたモデルESM-3です。このモデルは、アミノ酸の種類、完全な三次構造、三次構造トークン、二次構造、溶媒接触表面積(SASA)、タンパク質と残基の機能記述という7つの入力を含む、より豊富な情報セットを考慮しています。

教師なしおよび教師ありの適応度予測

次に、適応度予測問題についての議論を続けましょう。PLMはタンパク質配列の確率分布をモデル化できるため、ラベル付きデータを必要とせずに変異の適応度予測に直接使用できます。この手法はゼロショット予測または教師なし予測と呼ばれます。

具体的には、PLMは変異体と野生型の対数尤度比を計算することで変異をスコアリングします。自己回帰モデルの場合、配列の確率Pは各アミノ酸の生成確率の積です。変異スコアは、野生型のlogPから変異体のlogPを差し引くことで得られます。直感的に言えば、PLMは変異の確率が野生型と比較してどれだけ高いかを比較することで、変異の影響を評価する、経験的な評価方法です。

マスキングモデルでは、配列全体の確率を直接計算することはできません。しかし、特定の位置をマスクし、その位置におけるアミノ酸の確率分布を推定することは可能です。したがって、各変異位置について、マスク後の予測変異アミノ酸のlogPを野生型アミノ酸のlogPから差し引き、すべての位置における差を合計することで、変異スコアを得ることができます。

さらに、PLM はタンパク質配列のベクトル表現を提供するため、十分な実験データが利用可能な場合は微調整することができ、教師ありの適応度予測が可能になります。

具体的なアプローチとしては、PLMの最後の特徴層の後に適応度予測のための出力層(アテンションメカニズムや多層パーセプトロン(MLP)など)を追加し、適応度ラベルを全体的または部分的な学習に用いることが挙げられます。例えば、ECNetは大規模モデルの特徴量にMSA特徴量を追加し、LSTMを用いてそれらを融合することで教師あり学習を行います。当グループが昨年開発したSESNetモデルは、ESM-1b配列特徴量、ESM-IF構造特徴量、そしてMSA特徴量を統合し、教師あり適応度予測に利用しています。

FSFP法の紹介:PLMのための少数ショット学習法

適応度予測におけるスモールショット学習の重要性

FSFP法を紹介する前に、適応度予測における少数ショット学習の重要性を明確にする必要があります。教師なし学習法は学習にラベル付きデータを必要としませんが、ゼロショットスコアリングの精度は低いです。さらに、対数尤度比に基づくスコアリングはタンパク質の特定の自然法則しか反映できないため、タンパク質の非自然的な特性を効果的に予測することも困難です。

一方、教師あり学習法は正確である一方、PLMは多数のパラメータを必要とするため、パフォーマンスを大幅に向上させるには膨大な実験データを用いた学習が必要となります。教師あり学習モデルの評価には通常、既存の高スループットデータセットを8:2に分割しますが、学習セットの80%にはすでに数万点のデータポイントが含まれている場合があり、実際には取得コストが非常に高くなります。

この問題を解決するため、我々はPLMに適した少数ショット学習アプローチであるFSFP法を提案する。この手法は、少数の学習サンプル(数十個)のみを用いてPLMの適応度予測性能を大幅に向上させることができる。さらに、FSFP法は非常に柔軟性が高く、様々なPLMに適用可能である。

FSFP法:フィットネスレベルの順位付けを学習する

従来の教師あり学習法では、適応度予測を回帰問題として扱い、モデル出力と適応度ラベル間の平均二乗誤差(MSE)を計算することでモデルを最適化します。しかし、サンプルサイズが小さい場合、回帰モデルは過学習に陥りやすく、学習損失が急速に減少します。そこで私たちはアプローチを変更し、回帰を放棄してランキング学習に焦点を当てました。ランキング学習では、正確な数値的フィッティングを必要とせず、正確なランキングのみを必要とします。

この方法には2つの大きな利点があります。第一に、シーケンシング自体がタンパク質工学の基本要件と一致しており、変異の相対的な有効性を測定するだけで済むことです。第二に、シーケンシングは絶対値を予測するよりも簡単な作業です。

トレーニングの反復処理では、サンプリングしたミュータント群をラベルの逆順にソートし、これらのミュータントに対するモデルの予測に基づいてランキング損失(ListMLE)を計算します。モデルの予測ランキングが実際のランキングに近いほど、損失は小さくなります。各ミュータントに対するモデルのスコアリング関数fとして、対数尤度比に基づくゼロショットスコアリング関数を使用します。これは、ゼロショットスコアから開始し、トレーニングデータを使用して徐々にスコアを洗練させることで、モジュールを再初期化することなくパフォーマンスを向上させ、トレーニングの難易度を低減することを目的としています。

FSFP法: PLMパラメータを効率的に微調整する

PLMは通常数億個のパラメータを持つため、非常に少ないデータでモデル全体を微調整すると、必然的に過学習につながります。そこで、モデル内の学習可能なパラメータの数を制限する2つ目の手法であるLoRAを導入します。

LoRAは、Transformerの各ブロックの全結合層に、学習可能なランク分解行列のペアを挿入します。この際、事前学習済みのパラメータは変更しません。ランク分解行列は非常に小さいため、学習可能なパラメータの数は元の1.84%にまで削減できます。学習可能なパラメータの数が削減されたにもかかわらず、Transformerの各層は微調整されているため、モデルの学習能力は依然として保証されています。

FSFP法:メタ学習を適応度予測に適用する

過学習を回避するため、より優れた損失関数を使用するだけでなく、LoRA技術を用いて学習可能なパラメータの数を制限しました。しかし、それでもなお、少量の学習データに対して学習の反復回数を多くしすぎると、過学習のリスクが残ります。そのため、より少ない学習反復回数でモデル性能を迅速に向上させることを目指しました。このニーズに基づき、3つ目の手法であるメタ学習を採用しました。メタ学習の基本的な考え方は、まずモデルにいくつかの補助的なタスクで経験を積ませて初期モデルを取得し、その後、この初期モデルを用いて新しいタスクに迅速に適応するというものです。

下の画像は、メタ学習に基づく画像分類の例を示しています。目標タスクが馬を分類するモデルをトレーニングすることだと仮定しますが、ラベル付きの馬のデータは比較的少ないです。そこで、まず猫や犬の分類など、データが豊富な補助タスクを見つけ、これらの補助タスクに対してメタ学習アルゴリズムを用いてモデルをトレーニングします。これにより、モデルは新しいタスクの学習方法を学習し、メタ学習器を構築できます。次に、このメタ学習器を初期モデルとして使用し、少量のラベル付きの馬のデータを用いて数ステップのトレーニングを行うことで、馬の分類器を迅速に得ることができます。メタ学習の有効性は、補助タスクが目標タスクに十分近いかどうかにかかっていることは明らかです。

メタ学習は適応度予測シナリオにどのように適用できるでしょうか?私たちの主なタスクは、対象タンパク質の変異を適応度の観点からランク付けすることであり、トレーニング対象となるモデルはLoRAテクノロジーを用いたPLMです。

補助タスクの構築には2つの戦略を採用しました。1つ目は、既存のDMSデータベースから、標的タンパク質との類似性に基づいて類似タンパク質の変異実験データセットを検索し、最初の2つのデータセットを2つの補助タスクとして選択するというものです。このアプローチは、類似タンパク質は類似した物理的ランドスケープも共有するという仮定に基づいています。

2つ目の戦略は、MSAモデルを用いて標的タンパク質の候補変異をスコアリングし、擬似ラベルデータセットを作成することです。このデータセットは、3つ目の補助タスクとして機能します。MSAモデルが選択された理由は、その変異予測性能がPLMと概ね同等であるためです。本研究では、PLMの表現能力を最大限に活用するために、MSAをデータ拡張に用いることを目指しています。

使用するメタ学習アルゴリズムはMAMLです。その学習目標は、特定の補助タスクの学習データを用いてメタ学習器をkステップ微調整した後、テスト損失を最小化することです。これにより、kステップの微調整後に目標タスクが概ね収束します。

タンパク質適応度予測におけるFSFP法の性能評価

ベンチマークの作成

私たちのベンチマーク データは ProteinGym から取得されており、当初は 87 個の DMS データセットが含まれていましたが、現在は 217 個に更新されています。87個の DMS データセットに対応するタンパク質は、真核生物、原核生物、ヒト、ウイルスの 4 つのカテゴリに大まかに分類され、合計約 1,500 万の変異とそれに対応する適応度レベルをカバーしています。

各データセットについて、20、40、60、80、100個の単一点突然変異をランダムに選択し、小規模なトレーニングセットとして保存し、残りの突然変異をテストセットとして保存します。早期終了のための追加の検証セットは使用せず、代わりにトレーニングセットに対するクロスバリデーションによってトレーニングステップ数を推定することに注意してください。

前述のように、メタ学習には3つの補助タスクが必要です。そのうち2つは、標的タンパク質との類似性に基づいてDMSデータベースからタンパク質を取得することです。特定のデータセットで学習する場合、ProteinGymの残りのデータセット(ここではデータベースと仮定)からタンパク質を取得します。

右下の図に示すように、ProteinGymの各タンパク質をクエリとしてMMseqs2とFoldSeek検索を行い、最も類似したタンパク質の類似度分布を取得しました。最も類似したタンパク質の平均配列または構造類似度は約0.5であることがわかります。3つ目の補助タスクは、MSAモデルを用いた変異のスコアリングです。MSAに基づいて系統樹を構築し、系統樹上の各点における保存性を計算して変異をスコアリングするGEMMEモデルを選択しました。

評価指標には、適応度予測タスクで一般的に用いられるスピアマン/ピアソン係数とNDCGを使用しました。最終的な評価スコアは、87データセットの平均スコアです。

ESM-2におけるFSFPのアブレーション実験

下図に示すように、左側のグラフでは、x軸はトレーニングセットのサイズ、y軸はスピアマン係数を表しています。各線は異なるモデル構成に対応しています。一番上の線はFSFPモデルのフルバージョンを表し、2番目の線はメタ学習の3番目の補助タスクをMSAを使用せずに類似タンパク質のDMSデータに置き換えたことを示しています。MSA情報を削除した後、モデルのパフォーマンスが低下することがわかります。3番目の線はメタ学習を使用せず、ランキング学習とLoRAのみに依存しており、スピアマン係数がさらに低下していることを示しています。

緑の線は、以前 NBT で公開されたリッジ回帰モデルを表しています。これは、小規模なサンプル サイズに適した数少ないベースライン モデルの 1 つです。灰色の破線は ESM-2 のゼロ ショット スコアを表しています。下の 2 本の線は、従来の回帰法を使用して ESM-2 をトレーニングした結果を表しています。

全体として、わずか20個の学習サンプルで、本手法はゼロショット・スピアマン法と比較して10ポイントの改善を達成し、すべてのモジュールがモデル性能にプラスの影響を与えています。右図は、学習サンプル数40個の87データセットにおける、ゼロショット法と比較した性能向上の分布を示しています。本手法はほとんどのデータセットでモデル性能を向上させ、一部のデータセットでは40ポイントを超える改善を示し、ベースラインよりも安定した性能を示していることがわかります。

メタ学習の有効性

メタ学習の目的は、PLMが少数の反復で目標タスクに迅速に収束できるようにすることです。これを説明するため、以下にいくつかの例を示します。

以下の3つのグラフは、3つのデータセットで40個の学習サンプルを用いて微調整された学習曲線を示しています。x軸は学習ステップ数、y軸はテストセットにおけるスピアマン係数を表しています。上部のオレンジと赤の線は、メタ学習を用いて学習されたモデルを表しています。前者は補助タスクの構築にMSAを使用し、後者は使用していません。黄色の線は、メタ学習を使用せず、ランキング学習とLoRAのみを使用したモデルを表しています。

ご覧のとおり、メタ学習によって学習されたモデルは、標的タンパク質に対するパフォーマンスをより迅速に向上させ、20ステップ以内に高いスコアを達成しています。また、微調整を行わない初期モデルでも良好なパフォーマンスを示す場合があります。これは、メタ学習によって効果的な初期モデルが得られることを示しています。一方、以下のMSEベースのモデルはパフォーマンスが低く、すぐに過学習に陥るため、ゼロショット法を上回ることは困難です。

FSFPをさまざまなPLMに適用した結果

我々は、3つの典型的なタンパク質モデリングメカニズム(PLM)、ESM-1v、ESM-2、そしてSaProtを選択しました。最初の2つのモデルはタンパク質配列情報のみを使用しますが、SaProtはタンパク質の三次構造トークンを組み込みます。

左の折れ線グラフは、異なるトレーニングセットサイズにおける単一点突然変異の予測におけるスピアマンスコアを示しています。同じ色は同じモデルを表し、異なる点の形状は異なるトレーニング方法を表しています。上の円はFSFP法、下の逆三角形はリッジ回帰、破線はモデルのゼロショット性能を表しています。紫色の線はGEMMEモデルを表しています。これはPLMではありませんが、リッジ回帰と組み合わせることができます。FSFP法は様々なPLMの性能を一貫して向上させることができ、リッジ回帰や対応するモデルのゼロショット性能をはるかに上回っていることがわかります。

2つ目の棒グラフは、異なるデータセットにおいて3つの戦略(ゼロショット、リッジ回帰、FSFP)を用いて得られた最高スコアの数を示しています。FSFPはほとんどのデータセットで最も優れたパフォーマンスを示しました。右側の2つのグラフは、11のマルチポイント変異データセットを用いたマルチポイント変異の予測におけるパフォーマンスを示しており、結論はシングルポイント変異の場合と同様です。ただし、リッジ回帰モデルはここでより大きな分散を示しており、データ分割の影響を受けやすいことを示しています。

次に、FSFPの外挿性能を評価しました。具体的には、トレーニングセットには見られない変異部位における予測性能を評価しました。この場合、テストセットははるかに小さく、トレーニングセットが大きくなるにつれて大きく変化するため、表のゼロショット性能は直線ではなくなりました。この設定は非常に困難であり、左側の単一変異に対するリッジ回帰の性能がゼロショット回帰の性能を上回ることはほとんどありませんが、FSFPは依然として一貫して性能を向上させていることがわかります。右側の複数変異の変異に対するテスト結果も、このトレーニング方法が優れた一般化能力を持っていることを示しています。

FSFPによるPhi29の修正

さらに、タンパク質修飾のケースにもFSFPを使用しました。標的タンパク質はDNAポリメラーゼであるPhi29であり、単一点変異によってTm値を向上させることが期待されました。

実験手順は以下の通りである。まず、ESM-1vを用いて飽和した単一点変異のゼロショットスコアリングを行い、スコア上位20個の変異をウェット実験に用いてTm値を測定する。次に、これらの20個の実験データをトレーニングセットとして用い、FSFPを用いてESM-1vをトレーニングする。トレーニング済みのモデルを用いて再び飽和した単一点変異のスコアリングを行い、スコア上位20個の変異を再度テストに用いる。

右図は、2回の実験におけるTm分布の比較を示しています。第1ラウンドでは20個の変異のうち7個が陽性でしたが、第2ラウンドでは12個に増加し、平均Tmは1度向上しました。第2ラウンドで発見された陽性変異のうち、9個は新規でした。陽性率と平均Tmは向上しましたが、残念ながら最高Tmは向上しませんでした。これは、第2ラウンドで得られた最高Tm変異が、第1ラウンドで既に存在していたものであったためです。しかし、より多くの陽性単一点変異が得られたため、これらの部位を組み合わせて高点変異実験を実施することで、Tmをさらに向上させることが可能になります。

FSFP手法の概要と今後の研究展望

FSFPは、PLM(プロトコル変異効果予測)のための少数ショット学習戦略であり、少数(数十個)のラベル付き学習サンプルを用いて変異効果を予測するPLMの性能を大幅に向上させることができ、様々なPLMに柔軟に適用できます。実験により、FSFPの設計が合理的であることが示されています。

  • シーケンス学習は、タンパク質工学における変異シーケンスの基本要件を満たし、トレーニングの難易度を軽減します。
  • LoRA は、PLM 内のトレーニング可能なパラメータの数を制御することで過剰適合のリスクを軽減します。
  • メタ学習はモデルに適切な初期パラメータを提供し、モデルをターゲットタスクに素早く転送できるようにします。

最後に、AI支援による指向性進化の将来的な方向性について議論しましょう。AI支援による指向性進化の一般的なプロセスは、一連の初期突然変異から始まり、ウェット実験を通じてそれらの適応度ラベルを取得し、実験フィードバックから得られたラベル付きデータを用いて機械学習モデルを学習します。そして、モデルの予測に基づいて次にテストする突然変異を選択し、これを繰り返します。

FSFPは主に、各実験反復における小規模サンプルトレーニングの問題に対処し、モデルの予測精度を向上させます。しかし、次のラウンドでテストする変異、つまり次のラウンドで追加される新しいトレーニングサンプルを効果的に選択する方法についてはまだ議論していません。前述のPhi29タンパク質改変の例では、モデルスコアが最も高い上位20個の変異を直接選択しました。しかし、複数ラウンドの反復シナリオでは、貪欲な選択戦略は必ずしも最良のアプローチではありません。局所最適解に陥りやすいためです。したがって、探索と活用のバランスを見つける必要があります。

実際、ラベル付けのためのテストサンプルを反復的に選択し、トレーニングデータを徐々に拡張するプロセスは能動学習の問題であり、タンパク質工学の分野では既に研究が進展しています。例えば、指向性進化の第一人者であるフランシス・H・アーノルドは、論文「能動学習支援による指向性進化」で関連する問題を探求しました。
論文の宛先:
https://www.biorxiv.org/content/10.1101/2024.07.27.605457v1.full.pdf

不確実性定量化技術を用いることで、各変異体に対するモデルのスコアの不確実性を評価することができます。この不確実性に基づいて、テストサンプルの選択戦略をより多様なものにすることができます。一般的に用いられる戦略はUCB法です。これは、モデル予測の不確実性が最も高い変異体サンプルを次のラベリングラウンドに選択する、つまり予測分散が最も大きいサンプルを優先するものです。これは人間の学習プロセスに似ています。十分な知識が不足している場合や、特定の知識点について不確実性がある場合、学習を強化することに重点を置きます。