[论文解读] 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方法，dWJS实现了快速、多样化的采样，生成的抗体序列中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.

研究动机与目标

• 解决在蛋白质序列生成中训练和采样离散能量基准模型（EBMs）的挑战。
• 克服现有自回归模型与基于扩散的模型在抗体设计中存在脆弱性、不稳定性与低效性的问题。
• 开发一种简化且可泛化的框架，仅需单一噪声水平，避免使用如回放缓冲区或退火等复杂训练技巧。
• 实现从头发现具有高生物物理有效性与实验成功率的新型功能性抗体。
• 引入一种新评估指标——分布一致性得分（DCS），用于基准测试蛋白质设计中的生成模型。

提出的方法

• 提出平滑离散采样（SDS），一种新范式，其在离散数据流形的平滑版本上学习能量函数。
• 使用朗之万MCMC从平滑的能量基准模型中采样，实现对数据分布的高效探索。
• 应用一步去噪将采样结果投影回真实的离散数据流形，以保持序列的有效性。
• 将训练过程解耦为独立的基于得分的模型与基于能量的模型：在噪声数据上训练得分模型，在平滑数据上训练EBM。
• 利用神经经验贝叶斯（NEB）形式化统一离散数据的核密度估计与经验贝叶斯方法。
• 设计一种新型离散EBM架构，提升训练稳定性与采样质量，且无需依赖常见的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 通过行走-跳跃采样实现的平滑离散采样能否在单个MCMC链中实现快速混合并多样化探索多种抗体胚系类别？
• RQ3 所提出的方法是否可消除EBM训练中对复杂技巧（如回放缓冲区、ℓ₂惩罚项或模拟退火）的依赖？
• RQ4 分布一致性得分（DCS）能否作为评估蛋白质生成模型的可靠标量基准？
• RQ5 dWJS在单轮实验室测试中，能在多大程度上生成新颖、具有生物物理有效性且功能改进的抗体？

主要发现

• 97–100%的生成抗体序列在单轮体外实验中成功表达并纯化。
• 35%的生成功能性设计表现出与已知功能性抗体相当或更高的结合亲和力。
• 该方法实现了长运行时间下的快速混合MCMC链，在单个马尔可夫链中即可访问多样化的抗体蛋白类别。
• 该方法消除了对常见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