[論文レビュー] Diffusion on language model encodings for protein sequence generation

Viacheslav Meshchaninov, П. В. Страшнов+5 more|arXiv (Cornell University)|Mar 6, 2024

Topic Modeling被引用数 7

ひとこと要約

DiMAはタンパク質言語モデルESM-2の埋め込み上で連続拡散を用いてアミノ酸配列を生成し、品質と多様性の面で自己回帰および離散拡散のベースラインを上回り、豊富なアブレーションと生物学的関連性分析を実施している。

ABSTRACT

Protein sequence design has seen significant advances through discrete diffusion and autoregressive approaches, yet the potential of continuous diffusion remains underexplored. Here, we present DiMA, a latent diffusion framework that operates on protein language model representations. Through systematic exploration of architectural choices and diffusion components, we develop a robust methodology that generalizes across multiple protein encoders ranging from 8M to 3B parameters. We demonstrate that our framework achieves consistently high performance across sequence-only (ESM-2, ESMc), dual-decodable (CHEAP), and multimodal (SaProt) representations using the same architecture and training approach. We extensively evaluate existing methods alongside DiMA using multiple metrics across two protein modalities, covering quality, diversity, novelty, and distribution matching of generated proteins. DiMA consistently produces novel, high-quality and diverse protein sequences and achieves strong results compared to baselines such as autoregressive, discrete diffusion and flow matching language models. The model demonstrates versatile functionality, supporting conditional generation tasks including protein family-generation, motif scaffolding and infilling, and fold-specific sequence design. This work provides a universal continuous diffusion framework for protein sequence generation, offering both architectural insights and practical applicability across various protein design scenarios. Code is released at \href{https://github.com/MeshchaninovViacheslav/DiMA}{GitHub}.

研究の動機と目的

• unconditionalなタンパク質配列生成を動機づけ、タンパク質宇宙全体の条件設計の基盤とする。
• DiMAを提案：pLM埋め込み上で動作する拡散モデルを用いてタンパク質配列を生成。
• 生成の品質、多様性、分布類似性、および配列と構造の生物学的関連性を評価。
• 自己条件付けやシャープな設計選択を特定するために自己条件付け、スキップ接続、時間条件付け、ESMエンコーダ、ノイズスケジュールなどのアーキテクチャと訓練の影響を比較・アブレーションを実施。

提案手法

• 事前学習済みのESM-2タンパク質言語モデルを用いてタンパク質配列をエンコードし潜在埋め込みを取得。
• 潜在空間で連続的なデノイジング拡散モデルを訓練し、汚損した埋め込みを再構成。
• 潜在埋め込みをアミノ酸配列へ戻すデコーダを使用。
• サンプリング時に前のz0予測を再利用する自己条件付けを導入し、訓練時には勾配停止を適用。
• 訓練で観測された経験的分布から配列長をサンプルし、デコード前に潜在ベクトルをデノormalize。

実験結果

リサーチクエスチョン

• RQ1 pLM潜在空間で拡散を行えば条件なしで高品質かつ多様なタンパク質配列を生成できるか。
• RQ2 DiMAは配列品質、多様性、分布類似性の観点から自己回帰および離散拡散ベースラインとどう比較されるか。
• RQ3 自己条件付け、スキップ、時間条件付け、ESMエンコーダ、ノイズスケジュールといった設計・訓練の選択が生成性能に与える影響は何か。
• RQ4 生成された配列は構造・機能・ドメイン注釈の観点で生物学的関連性を示すか。

主な発見

モデル	pLDDT (↑)	ESM-2 pppl (↓)	scPerplexity (↓)	TM-score (↑)	BLAST (↑)	FPD (↓)	MMD (↓)	OT (↓)
SwissProt Dataset	80.7	5.35	1.88	0.80	100	0.13	0.00	1.08
Random sequences	25.0	21.54	2.77	0.33	0	3.97	0.20	3.88
nanoGPT	61.0	8.18	2.04	0.63	43	1.24	0.03	2.53
EvoDiff-OADM	37.1	15.77	2.44	0.42	12	1.49	0.11	2.63
SeqDesign	43.1	11.89	2.35	0.41	17	3.53	0.19	5.12
proteinGAN	30.4	16.48	2.57	0.33	0	2.94	0.17	3.98
DiMA	80.8	5.20	1.80	0.85	68	0.41	0.01	1.41
w/o skip connections	77.0	5.84	1.87	0.82	61	0.48	0.02	1.51
w/o time layers	79.4	5.49	1.83	0.85	66	0.44	0.02	1.44
w/o ESM encoder	62.7	9.22	2.09	0.71	48	1.05	0.04	2.14
w/o self-conditioning	68.2	9.18	2.08	0.74	46	0.54	0.04	1.61
w linear schedule	77.0	6.29	1.89	0.82	58	0.50	0.02	1.51
w cosine schedule	54.1	10.86	2.16	0.60	34	0.97	0.06	2.02
AFDB Dataset	83.9	5.79	1.75	0.92	100	0.18	0.00	1.57
Random sequences	26.2	21.67	2.75	0.35	0	3.02	0.18	4.15
nanoGPT	68.5	8.21	1.94	0.77	40	0.62	0.02	1.99
DiMA	73.9	8.50	1.90	0.85	48	0.69	0.03	1.86
w/o self-conditioning	56.3	12.08	2.18	0.69	31	0.96	0.05	2.29
4.1 Ablation w/o self-conditioning	4.1 Baseline models	?	?	?	?	?	?	?

• DiMAはSwissProtおよびAFDBv4-90において複数の品質・多様性指標で自己回帰および離散拡散ベースラインを上回る。
• 自己条件付けとESM-2エンコーダの使用は性能にとって最も影響力のある設計選択の1つである。
• Simple Diffusion SD-10ノイズスケジュールはタンパク質潜在拡散において線形やコサインスケジュールよりも品質と多様性を向上させる。
• 生成された配列は生物学的関連性を示し、高いInterProSUPERFAMILY注釈率や妥当なIDR/二次構造プロファイルを持つ。
• DiMAは訓練データとの分布類似性をモダリティ間で維持し、SwissProtでFréchet ProtT5 Distanceおよび関連指標が良好で、AFDBv4-90では競合的な結果を示す。