[Paper Review] Backdiff: a diffusion model for generalized transferable protein backmapping

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

Protein Structure and Dynamics4 citations

TL;DR

BackDiff proposes a conditional score-based diffusion model for generalized, transferable protein backmapping across diverse coarse-grained (CG) models and proteins. By leveraging self-supervised training to handle variable CG atoms and manifold constraint sampling for arbitrary CG auxiliary variables, BackDiff enables end-to-end, retraining-free backmapping with state-of-the-art accuracy and robust generalization on multiple CG models and protein types.

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.

Motivation & Objective

• To address the lack of generalization in existing data-driven backmapping methods across different coarse-grained (CG) models and protein types.
• To enable transferable backmapping without retraining by learning a universal all-atom distribution conditioned on CG structures.
• To handle variable CG atoms and arbitrary CG auxiliary variables through a self-supervised and manifold-constrained framework.
• To provide a plug-and-play solution for researchers using their own CG models without model fine-tuning.
• To improve sampling efficiency and geometric fidelity in backmapped all-atom structures, especially for nonlinear CG variables.

Proposed method

• BackDiff employs a conditional score-based diffusion model that learns the reverse denoising process to reconstruct all-atom configurations from CG inputs.
• It uses a self-supervised training strategy that randomly selects CG atoms in each training epoch to enhance transferability across different CG representations.
• Manifold constraint sampling is applied during inference, where CG auxiliary variables guide the reverse diffusion path to maintain consistency with the input CG model.
• The model is trained end-to-end on a diverse dataset of 92 proteins, enabling generalization across protein space and CG model types.
• The framework supports both fixed and transferable sampling modes, with the latter allowing zero-shot adaptation to new proteins and CG models.
• Geometric constraints are enforced via score-based denoising, ensuring that generated structures respect the input CG geometry and auxiliary variables.

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

Experimental results

Research questions

• RQ1 Can a single generative model achieve reliable backmapping across multiple, diverse coarse-grained protein models without retraining?
• RQ2 How does the model perform in transferring to new proteins and CG representations not seen during training?
• RQ3 To what extent does manifold constraint sampling improve fidelity when CG auxiliary variables vary across models?
• RQ4 How does BackDiff compare in accuracy and efficiency to existing state-of-the-art backmapping methods like GenZProt and Torsional Diffusion?
• RQ5 What are the limitations of diffusion-based backmapping in preserving bond angles and chirality, and how can they be mitigated?

Key findings

• BackDiff achieves the lowest root mean square deviation (RMSD) of 0.030(0.005) Å on the fixed CG setting and 0.030(0.006) Å on the transferable setting, significantly outperforming GenZProt and Torsional Diffusion.
• In multi-protein experiments on the UNRES model, BackDiff achieved an SCMSE of 0.147(0.072) Ų, the lowest among all compared methods, indicating superior geometric fidelity.
• On the MARTINI model, BackDiff achieved a mean RMSD of 0.028(0.005) Å, demonstrating high accuracy across diverse protein types and CG representations.
• The model shows strong generalization, maintaining low error across different proteins and CG models without retraining, including on previously unseen proteins.
• Despite high accuracy, BackDiff exhibits higher sampling latency (293 seconds per 100 samples) compared to VAE-based models like GenZProt (0.009 seconds), highlighting a trade-off in efficiency.
• The model struggles with bond angle accuracy due to nonlinear mappings from Cartesian coordinates, and does not explicitly preserve chirality, indicating areas for future improvement.

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