Generative Coarse-Graining

In theoretical and computational science, order reduction or coarse-graining (CG) is a popular approach to describe complex systems with lower-dimensional representations that capture the key physics. In modeling the dynamics of large biological and materials systems at the molecular level, CG simulations unify selected groups of atoms into CG particles described by simpler equations of motions than the all-atom system. However, the restoration of fine-grained (FG) coordinates from CG representations is challenging due to the inherent information loss in the CG procedure. Traditional backmapping methods do not fully capture the geometrical constraints of the all-atom distribution and therefore tend to reconstruct unrealistic geometries. Inspired by the recent progress in deep generative models and equivariant neural nets, we propose a unified probabilistic formulation of backmapping, where the conditional likelihood of FG positions solely depends on the CG conformations. We design a conditional variational autoencoder that encodes the missing FG information into an invariant latent space and decodes back to FG geometries via convolutions over equivariant vector basis obtained from CG geometries. To evaluate both the diversity and quality of generated geometries, we propose three metrics and apply them to two molecular dynamics datasets: Alanine Dipeptide and Chignolin. Extensive experiments on these benchmarks demonstrate that our approach significantly outperforms all related baseline methods. Interestingly, the results suggest that most FG information can be robustly recovered even for ultra-low resolution CG models.