【抗体配列生成】Prescient Design の抗体配列生成手法 dWJS を紹介

論文タイトル

Protein Discovery with Discrete Walk-Jump Sampling

出典

Protein Discovery with Discrete Walk-Jump Sampling

We resolve difficulties in training and sampling from a discrete generative model by learning a smoothed energy function, sampling from the smoothed data manifo...

要旨

望みの性質をもつタンパク質配列を効率的にサンプリングする手法、Discrete Walk-Jump Sampling (dWJS)を紹介した論文です。

解説など

Genentech で AI創薬を推進する Prescient Design からのレポートです。

本ブログではこれまで、Prescient Design の抗体設計に関わる技術をいくつか紹介してきました。

今回の記事で紹介するのは、Prescient Design の要素技術である Discrete Walk-Jump Sampling (dWJS)です。これは、コンディショナルにタンパク質のアミノ酸配列をサンプリングする手法です。

既存の配列生成モデルとして代表的な拡散モデルやエネルギーベースモデルには以下の課題がありました。

拡散モデル (Diffusion Models)
- 仕組み
  - 多段階のノイズ除去プロセスを通して、元の分布に近づくサンプルを生成する
- 課題
  - ノイズスケジュールを必要とし、学習とサンプリングが非常に遅い
  - 通常は連続値データに特化しており、離散系列への適用が難しい
エネルギーベースモデル (EBM)
- 仕組み
  - 入力データに対してエネルギーを定義し、各データが低いエネルギーを持つように学習
  - モデルからサンプリングした配列が高いエネルギーを持つように学習（negative phase)
  - Contrastive Divergence を使って学習し、Langevin MCMC でサンプリング
- 課題
  - サンプリングが局所モードに閉じ込められる（モード崩壊）
  - サンプリングが遅い
  - ノイズなしだと過学習しやすい

これに対して、dWJS は以下のような特徴があります。

入力に対してノイズを人為的に加え、スムージングされた分布に対して EBM を学習することで学習の発散を防ぐ
単一ステップのデノイジングでクリーンデータに戻れる
MCMC (Walk) と Denoiser (Jump) を独立に学習できるため、モデルの柔軟性が高い

dWJS は Walk と Jump の２つの要素で構成されています。Walk ではノイズを加えた配列 y 上の確率分布を Langevin MCMC を活用して探索します。Jump では最終的なノイズ付きデータ y を、スコア関数によりきれいな配列に変換（denoising) します。

dWJS の学習データはプロジェクトごとに異なります。基本的に、目的に応じて個別の抗体配列データを使って専用に学習させます。抗体配列は one-hot encoding で特徴付けされ、AHO numbering でアライメントされてモデルの学習に供されます。

筆者らは dWJS の性能評価目的に OAS の VH+VL ペア配列を学習させたモデルを構築しました。これは従来手法である IgLM や ESM2 に比べて、多様性・新規性・分布適合性の観点で優れていることが示されました。

またトラスツズマブの改変体設計の応用事例も紹介しています。Mason et al. の binder/non-binder ラベル付きの NGS データを学習したモデルを使って H3 配列のみを MCMC サンプリングで設計しました。結合活性をもつ配列の割合が70%と、LaMBO や AbDiffuser と比べて高い結果でした。一方で論文内で言及されていませんが、親抗体の KD がおよそ 5nM であることを考慮すると、ほとんどのクローンの親和性は低下しています。抗体の最適化には導入する変異数を抑制するなどの工夫が必要なようすです。

コードはこちら。