[Paper Review] Generative Coarse-Graining of Molecular Conformations

Wujie Wang, Minkai Xu+6 more|arXiv (Cornell University)|Jan 28, 2022

Machine Learning in Materials ScienceMaterials Science23 citations

TL;DR

CGVAE introduces a probabilistic, equivariant framework to backmap coarse-grained molecular structures to fine-grained conformations, capturing FG uncertainties and geometric consistency.

ABSTRACT

Coarse-graining (CG) of molecular simulations simplifies the particle representation by grouping selected atoms into pseudo-beads and drastically accelerates simulation. However, such CG procedure induces information losses, which makes accurate backmapping, i.e., restoring fine-grained (FG) coordinates from CG coordinates, a long-standing challenge. Inspired by the recent progress in generative models and equivariant networks, we propose a novel model that rigorously embeds the vital probabilistic nature and geometric consistency requirements of the backmapping transformation. Our model encodes the FG uncertainties into an invariant latent space and decodes them back to FG geometries via equivariant convolutions. To standardize the evaluation of this domain, we provide three comprehensive benchmarks based on molecular dynamics trajectories. Experiments show that our approach always recovers more realistic structures and outperforms existing data-driven methods with a significant margin.

Motivation & Objective

• Motivate accurate backmapping from coarse-grained (CG) to fine-grained (FG) molecular coordinates under information loss.
• Formulate backmapping as a conditional generative task to capture FG uncertainty.
• Develop an equivariant, geometry-consistent decoder that respects E(3) symmetries.
• Provide CG-agnostic benchmarks and metrics to standardize evaluation in the CG backmapping domain.

Proposed method

• Propose a probabilistic CG-to-FG model (CGVAE) that learns p(x|X) via latent variable z with an encoder qϕ(z|x,X) and decoder pθ(x|X,z).
• Use an invariant FG-to-CG encoder with message passing and pooling to obtain CG embeddings parameterizing z.
• Model the prior pψ(z|X) from CG structure alone to enable sampling without FG data.
• Implement a multi-channel equivariant decoder that uses vector and pseudoscalar/pseudovector features and E(3) equivariant convolutions to predict FG coordinates.
• Enforce geometric consistency with R1 and R2: M x̃ = X and Dec(QX+g) = Q Dec(X)+g, ensuring backmapping preserves CG geometry and equivariance.
• Train with a reconstruction loss (MSD plus graph-based loss) and a KL divergence regularizer; sample by drawing z from the prior and decoding.

Experimental results

Research questions

• RQ1 Can a probabilistic, equivariant backmapping model recover realistic FG geometries from CG inputs while capturing FG uncertainty?
• RQ2 How can backmapping be made independent of specific CG mappings while maintaining geometric consistency?
• RQ3 Do proposed benchmarks and metrics reliably evaluate reconstruction quality and diversity across CG resolutions?

Key findings

• CGVAE consistently achieves superior reconstruction accuracy, sample quality, and diversity compared with baseline deterministic methods.
• Higher-resolution CG representations yield better FG reconstruction accuracy due to fewer atoms per CG bead.
• Lower-resolution CG representations exhibit greater sample diversity (larger FG conformational space captured by the latent z).
• The model is CG-mapping-agnostic and applicable to diverse CG mapping protocols used in MD simulations.
• Equivariant decoding with vector, pseudoscalar, and cross-product-based updates enhances expressiveness for constructing FG coordinates.