Towards Deep Generative Backmapping of Coarse-Grained Molecular Systems

Coarse-graining (CG) is a widely used technique for simplifying complex molecular systems in molecular dynamics simulations. It involves mapping fine-grained (FG) atoms to several CG beads, which enables the study of larger systems and longer simulation timescales. However, this simplification results in the loss of atomistic details. To address this, a technique reverse to CG named backmapping is designed to recover the FG structures from CG beads. Traditional methods of backmapping yield poor reconstructions and often necessitate computationally costly refinement processes. Recently proposed deep learning approaches have shown promise in directly generating atomistic structures with high accuracy. One of the most advanced models deploys a conditional variational auto-encoder (VAE) to encode the FG uncertainties into an invariant latent space and decode them back to FG geometries via equivariant convolutions. Motivated by the recent advancements of diffusion models, we explore an alternative approach to capture the multimodal distribution of FG structures. In this study, we present CGDiff, an equivariant CG-conditional diffusion model for generative backmapping. It models the underlying distribution by condensing the FG information into geometric diffusion processes. Conditioned on CG beads, our model progressively refines the coordinates of FG atoms through equivariant graph neural networks. We evaluate the performance of CGDiff against baselines by conducting experiments on two public datasets. To further promote research in this field, we additionally provide a demo interactive web service for generative backmapping, available at https://huggingface.co/spaces/jsli47/mlbackmap.