[Paper Review] Chemically Transferable Generative Backmapping of Coarse-Grained Proteins

Soojung Yang, Rafael Gómez‐Bombarelli|arXiv (Cornell University)|Mar 2, 2023

Machine Learning in Materials Science12 citations

TL;DR

The paper introduces GenZProt, a transferable, chemistry-aware backmapping model that reconstructs all-atom protein structures from alpha-carbon coarse-grained representations using an internal-coordinate, SE(3)-equivariant VAE with physics-informed losses.

ABSTRACT

Coarse-graining (CG) accelerates molecular simulations of protein dynamics by simulating sets of atoms as singular beads. Backmapping is the opposite operation of bringing lost atomistic details back from the CG representation. While machine learning (ML) has produced accurate and efficient CG simulations of proteins, fast and reliable backmapping remains a challenge. Rule-based methods produce poor all-atom geometries, needing computationally costly refinement through additional simulations. Recently proposed ML approaches outperform traditional baselines but are not transferable between proteins and sometimes generate unphysical atom placements with steric clashes and implausible torsion angles. This work addresses both issues to build a fast, transferable, and reliable generative backmapping tool for CG protein representations. We achieve generalization and reliability through a combined set of innovations: representation based on internal coordinates; an equivariant encoder/prior; a custom loss function that helps ensure local structure, global structure, and physical constraints; and expert curation of high-quality out-of-equilibrium protein data for training. Our results pave the way for out-of-the-box backmapping of coarse-grained simulations for arbitrary proteins.

Motivation & Objective

• Motivate fast, reliable backmapping of coarse-grained protein simulations to recover atomistic detail.
• Achieve chemical transferability by training on diverse experimental ensembles from PED.
• Preserve topology and physical plausibility via internal-coordinate generation and physics-informed losses.
• Demonstrate out-of-the-box applicability to arbitrary proteins and complex protein–IDP systems.

Proposed method

• Use a VAE framework to model p(x|X), where X is the CG structure and x the all-atom structure.
• Represent structures in internal coordinates (Z-matrix) to preserve topology and enable rule-based Cartesian reconstruction.
• Employ an SE(3)-equivariant encoder/prior with multi-level graph message passing (atom-atom, atom-residue, residue-residue).
• Decode with a Z-matrix based, invariant decoder that constrains local geometry (bond lengths/angles) and allows torsion flexibility.
• Incorporate physics-inspired losses: L_bond, L_angle, L_torsion, L_xyz, and L_steric, combined as L_recon = γL_local + δL_torsion + ηL_xyz + ζL_steric; train with ELBO optimization.

Figure 1: Overview. We aim to build a transferable and reliable backmapping tool for proteins. Our method builds on a VAE framework (Wang et al., 2022 ) . We train the VAE model on the protein structural ensemble data curated from PED. Our model can be characterized with three components : internal

Experimental results

Research questions

• RQ1 Can a generative backmapping model learn atomistic reconstructions that generalize across diverse protein chemistries?
• RQ2 Does internal-coordinate decoding with an equivariant encoder improve topology preservation and reduce steric clashes compared to Cartesian decoders?
• RQ3 What is the impact of physics-informed loss terms on reconstruction quality and validity (sterics, bonds, angles, torsions)?
• RQ4 Is a transferable model trained on PED ensembles capable of accurate backmapping for unseen proteins and protein–IDP complexes?

Key findings

• GenZProt (m1) consistently achieves the best RMSD, GED, and steric clash metrics across test proteins compared to ablated variants.
• Equivariant encoder/prior with an internal-coordinate Z-matrix decoder outperforms invariant counterparts and Cartesian decoders for large proteins.
• Training on diverse PED ensembles yields a transferable model that generalizes beyond single-chemistry training data (PED00151 alone).
• Physics-informed losses, especially the xyz and steric terms, are crucial for high-quality reconstructions and reducing steric clashes.
• Qualitative analyses show reconstructed and sampled structures preserve topology and long-range interactions with limited steric issues; hydrogen-bonding contacts are reasonably recovered.

Figure 2: Internal coordinate-based reconstruction. (a) Backbone atoms $N_{i},C_{i}$ are placed using adjacent three $C_{\alpha}$ as anchors. (b) Side chain atoms are placed using adjacent three atoms within the same residue.