[論文レビュー] Protein Discovery with Discrete Walk-Jump Sampling

Nathan C. Frey, Daniel Berenberg+11 more|arXiv (Cornell University)|Jun 8, 2023

Single-cell and spatial transcriptomics被引用数 4

ひとこと要約

本論文では、滑らか化されたエネルギーに基づくモデルとスコアベースのノイズ除去を組み合わせることで、効率的で高品質なタンパク質配列生成を可能にする、非自己回帰的生成モデルフレームワーク「離散ウォーク・ジャンプサンプリング（dWJS）」を提案する。単一のノイズレベルでの学習と、滑らか化されたデータ多様体上でランジュバンMCMCを用いることで、高速で多様なサンプリングが実現され、初回の実験で97–100%の発現・精製率を達成し、35%の抗体が結合親和性を向上させた。

ABSTRACT

We resolve difficulties in training and sampling from a discrete generative model by learning a smoothed energy function, sampling from the smoothed data manifold with Langevin Markov chain Monte Carlo (MCMC), and projecting back to the true data manifold with one-step denoising. Our Discrete Walk-Jump Sampling formalism combines the contrastive divergence training of an energy-based model and improved sample quality of a score-based model, while simplifying training and sampling by requiring only a single noise level. We evaluate the robustness of our approach on generative modeling of antibody proteins and introduce the distributional conformity score to benchmark protein generative models. By optimizing and sampling from our models for the proposed distributional conformity score, 97-100% of generated samples are successfully expressed and purified and 70% of functional designs show equal or improved binding affinity compared to known functional antibodies on the first attempt in a single round of laboratory experiments. We also report the first demonstration of long-run fast-mixing MCMC chains where diverse antibody protein classes are visited in a single MCMC chain.

研究の動機と目的

• タンパク質配列生成における離散的エネルギーに基づくモデル（EBM）の学習とサンプリングの課題に取り組む。
• 抗体設計における従来の自己回帰的および拡散ベースのモデルの脆さ、不安定性、非効率性を克服する。
• 再プレイバッファや冷却スケジューリングといった複雑な学習テクニックを必要とせず、単一のノイズレベルでのみ学習可能な、簡素化され一般化可能なフレームワークを開発する。
• 高いバイオフィジカル的妥当性と実験的成功率を備えた、新規で機能的な抗体のア・イニチオ探索を可能にする。
• タンパク質設計における生成モデルのベンチマークに用いる新しい評価指標「分布的整合性スコア（DCS）」を導入する。

提案手法

• 離散的データ多様体の滑らか化されたバージョン上でエネルギー関数を学習する新しいパラダイム「滑らか化離散的サンプリング（SDS）」を提案する。
• 滑らか化されたエネルギーに基づくモデルからのサンプリングにランジュバンMCMCを用い、データ分布の効率的探索を可能にする。
• 1ステップのノイズ除去を適用して、サンプルを真の離散的データ多様体に戻し、配列の妥当性を保持する。
• 学習をスコアベースモデルとエネルギーに基づくモデルに分離する：ノイズを含むデータ上でスコアモデルを学習し、滑らか化されたデータ上でEBMを学習する。
• 離散的データにおけるカーネル密度推定と経験ベイズを統合するためのニューラル経験ベイズ（NEB）形式を活用する。
• 一般的なEBM正則化技術（例：正則化項や受理採用サンプリング）を必要とせず、訓練安定性とサンプル品質を向上させる新しいアーキテクチャを設計する。

Figure 1: Selected samples from a single Markov chain Monte Carlo sampling run. Protein color corresponds to different antibody germlines (classes). Discrete walk-jump sampling exhibits fast mixing and explores diverse modes of the distribution in a single chain.

実験結果

リサーチクエスチョン

• RQ1 単一のノイズレベルに限定される非自己回帰的生成モデルは、自己回帰的および拡散ベースのベースラインを上回る機能的抗体配列の生成を達成できるか？
• RQ2 ウォーク・ジャンプサンプリングを用いた滑らか化離散的サンプリングは、1つのMCMC連鎖内で複数の抗体ゲノムクラスを高速に混合し、多様に探索できるか？
• RQ3 本手法は、再プレイバッファ、ℓ₂正則化、シミュレーテッドアニーリングといった一般的なEBM学習の補助技術の必要性を排除できるか？
• RQ4 分布的整合性スコア（DCS）は、タンパク質生成モデルの評価に信頼性があり、スカラーで評価可能な指標として機能するか？
• RQ5 dWJSは、1回の実験サイクルで、新規でバイオフィジカル的に妥当かつ機能的に改善された抗体をどの程度生成できるか？

主な発見

• 生成された抗体配列の97–100%が、in vitro実験の1ラウンドで正常に発現・精製された。
• 生成された機能的設計の35%が、既知の機能的抗体と同等または優れた結合親和性を示した。
• 長期間にわたる高速混合MCMC連鎖が達成され、1つのマルコフ連鎖内で多様な抗体タンパク質クラスをカバーした。
• 再プレイバッファ、ℓ₂正則化、受理採用サンプリングといった一般的なEBM学習安定化手法の必要性が排除された。
• dWJSは、サンプル品質および実験的成功率の両面で、自己回帰的大規模言語モデル、拡散モデル、スコアベースベースラインを上回った。
• 分布的整合性スコア（DCS）は、生成モデルを効果的にベンチマークし、実験的機能性と相関した。

Figure 2: Discrete walk-jump sampling. a The noising and denoising process is applied to antibody proteins. b Discrete inputs $x$ are smoothed with isotropic Gaussian noise, $\varepsilon\sim\mathcal{N}(0,\sigma^{2}I_{d})$ , to noisy inputs, $y=x+\varepsilon$ . A discrete energy-based model (dEBM) pa