[论文解读] Backdiff: a diffusion model for generalized transferable protein backmapping

Yikai Liu, Ming Chen+1 more|arXiv (Cornell University)|Oct 3, 2023

Protein Structure and Dynamics被引用 4

一句话总结

BackDiff 提出了一种基于条件得分的扩散模型，用于在多种粗粒度（CG）模型和蛋白质之间实现通用且可迁移的蛋白质反向映射。通过利用自监督训练处理可变的CG原子，并采用流形约束采样以支持任意的CG辅助变量，BackDiff 实现了端到端、无需微调的反向映射，其在多种CG模型和蛋白质类型上均达到最先进水平的准确率和强大的泛化能力。

ABSTRACT

Coarse-grained (CG) models play a crucial role in the study of protein structures, protein thermodynamic properties, and protein conformation dynamics. Due to the information loss in the coarse-graining process, backmapping from CG to all-atom configurations is essential in many protein design and drug discovery applications when detailed atomic representations are needed for in-depth studies. Despite recent progress in data-driven backmapping approaches, devising a backmapping method that can be universally applied across various CG models and proteins remains unresolved. In this work, we propose BackDiff, a new generative model designed to achieve generalization and reliability in the protein backmapping problem. BackDiff leverages the conditional score-based diffusion model with geometric representations. Since different CG models can contain different coarse-grained sites which include selected atoms (CG atoms) and simple CG auxiliary functions of atomistic coordinates (CG auxiliary variables), we design a self-supervised training framework to adapt to different CG atoms, and constrain the diffusion sampling paths with arbitrary CG auxiliary variables as conditions. Our method facilitates end-to-end training and allows efficient sampling across different proteins and diverse CG models without the need for retraining. Comprehensive experiments over multiple popular CG models demonstrate BackDiff's superior performance to existing state-of-the-art approaches, and generalization and flexibility that these approaches cannot achieve. A pretrained BackDiff model can offer a convenient yet reliable plug-and-play solution for protein researchers, enabling them to investigate further from their own CG models.

研究动机与目标

• 解决现有数据驱动反向映射方法在不同粗粒度（CG）模型和蛋白质类型之间缺乏泛化能力的问题。
• 通过学习以CG结构为条件的通用全原子分布，实现在无需微调的情况下实现可迁移的反向映射。
• 通过自监督与流形约束框架，处理可变的CG原子和任意的CG辅助变量。
• 为使用自定义CG模型的研究人员提供即插即用的解决方案，无需模型微调。
• 提升反向映射全原子结构的采样效率与几何保真度，尤其针对非线性CG变量。

提出的方法

• BackDiff 采用条件得分驱动的扩散模型，学习逆去噪过程，以从CG输入重建全原子构型。
• 采用自监督训练策略，在每个训练周期中随机选择CG原子，以增强在不同CG表示之间的迁移能力。
• 在推理过程中应用流形约束采样，其中CG辅助变量引导逆向扩散路径，以保持与输入CG模型的一致性。
• 模型在包含92种蛋白质的多样化数据集上进行端到端训练，实现跨蛋白质空间和CG模型类型的泛化。
• 该框架支持固定采样和可迁移采样两种模式，后者可实现对新蛋白质和CG模型的零样本适应。
• 通过基于得分的去噪方法施加几何约束，确保生成的结构尊重输入CG的几何特征和辅助变量。

Figure 1: The sampling process of BackDiff. The reverse diffusion process gradually converts the noisy configuration into the plausible configuration, conditioned on CG atoms $\mathcal{R}_{\text{atm}}$ . In each diffusion step, the configuration is “corrected” with auxiliary variables, bond lengths

实验结果

研究问题

• RQ1 一个单一的生成模型是否能在不进行微调的情况下，可靠地实现对多个多样化粗粒度蛋白质模型的反向映射？
• RQ2 该模型在未在训练中见过的新蛋白质和CG表示上的迁移性能如何？
• RQ3 当CG辅助变量在不同模型间变化时，流形约束采样在多大程度上提升了结构保真度？
• RQ4 与现有最先进反向映射方法（如GenZProt和Torsional Diffusion）相比，BackDiff在准确率和效率方面表现如何？
• RQ5 基于扩散的反向映射在保持键角和手性方面存在哪些局限性，以及如何缓解这些问题？

主要发现

• 在固定CG设置下，BackDiff 的均方根偏差（RMSD）达到最低值 0.030(0.005) Å，在可迁移设置下为 0.030(0.006) Å，显著优于GenZProt和Torsional Diffusion。
• 在UNRES模型的多蛋白实验中，BackDiff 的SCMSE为0.147(0.072) Ų，为所有对比方法中最低，表明其具有更优的几何保真度。
• 在MARTINI模型上，BackDiff 的平均RMSD为0.028(0.005) Å，证明其在多种蛋白质类型和CG表示中均具有高精度。
• 该模型展现出强大的泛化能力，在无需微调的情况下，对不同蛋白质和CG模型均保持低误差，包括此前未见过的蛋白质。
• 尽管准确率高，BackDiff 的采样延迟较高（每100个样本293秒），高于VAE-based模型如GenZProt（0.009秒），凸显了其在效率上的权衡。
• 由于从笛卡尔坐标到非线性映射的转换，该模型在键角准确性方面表现不佳，且未显式保留手性，表明未来仍有改进空间。

Figure 2: Visualization of all-atom configurations sampled from different methods in multi-protein experiments backmapping from UNRES CG model.