# Generative Coarse-Graining of Molecular Conformations

Wujie Wang  Minkai Xu  Chen Cai  Benjamin Kurt Miller  Tess Smidt  Yusu Wang
Jian Tang  Rafael Gómez-Bombarelli

###### 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.

Molecular Dynamics, Coarse-Graining, Generative Models, Equivariance

##  1 Introduction

Defined as an operation to perform dimension reductions over continuous
distributions, coarse-graining (CG) was first proposed by Paul Ehrenfest and
Tatyana Afanasyeva in their study of molecular chaos (Ehrenfest & Ehrenfest,
1912). After nearly a century of theoretical developments, CG has become a
powerful technique to simplify many complicated problems in physics. Useful
applications of CG include the renormalization group techniques for critical
phenomena (Kadanoff, 1966; Wilson, 1971) and the usage of collective variables
to study reaction dynamics (Laio & Parrinello, 2002). In molecular
simulations, CG refers to the simplification of the particle representation by
lumping groups of original atoms $x$ into individual beads $X$ with a rule-
based CG mapping. CG molecular dynamics (MD) can significantly speed up
computational discovery over chemical spaces with simpler combination rules
(Marrink et al., 2007; Menichetti et al., 2019); for simulations of large
molecules with hundreds of thousands of atoms such as biological and
artificial polymers, CG MD allows accessing the long time scales of phenomena
like protein folding or polymer reptation (Levitt & Warshel, 1975; Kremer et
al., 1988). However, such acceleration comes at the cost of losing fine-
grained (FG) atomic details, which are important for studying properties and
interactions of atom-level structures or continuing simulations at FG scale.
How to accurately recover the FG structure $x$ from CG coordinates $X$ remains
a challenging problem.

Figure 1: An analogical illustration of the tackled generative coarse-graining
problem. Top: Image super-resolution recovers higher resolution
portraits111Portraits of P. Ehrenfest and T. Afanasyeva, taken from Dibner
Library of the History of Science and Technology.. Bottom: Generative CG
generates fine-grained molecular structures.

So far, several methods have been proposed for the backmapping problem based
on random projection (Wassenaar et al., 2014) or geometric rules (Shih et al.,
2007). However, these methods usually can only produce poor initial geometries
which require subsequent restraining molecular dynamics (Shih et al., 2007;
Brocos et al., 2012), incurring considerable computational cost and large
deviation from the original structures. These methods also require maintaining
system-specific fragment libraries that are only applicable to a predefined
mapping schema, preventing their usage from broader mapping choices and
resolutions. Furthermore, none of these methods are data-driven, and
geometries produced by these methods are biased toward predefined fragment
geometries. To produce data-informed backampping solutions, several machine
learning methods have recently been proposed (Wang & Gómez-Bombarelli, 2019;
An & Deshmukh, 2020) to use parameterized functions to deterministically
backmap. However, these methods also suffer from low-quality geometries and
have not been tested on complex molecular structures. We argue that the
problem is very challenging mainly because of the following: (C1)
Stochasticity of backmapping: Due to surjective nature of CG projections, many
different FG configurations can be mapped to the same CG conformation, hence
the reverse generative map is one-to-many. However, this vital consideration
on stochasticity has been overlooked in current deterministic methods (Brocos
et al., 2012; Wassenaar et al., 2014; Wang & Gómez-Bombarelli, 2019; An &
Deshmukh, 2020), making it hard to generate diverse structures that capture
the underlying FG distributions. (C2) Geometry consistency: The backmapped FG
coordinates should faithfully represent the underlying CG geometry, which
requires that they can be reduced back to the original CG structures.
Additionally, given the fact that CG transformation is $E(3)$ equivariant (see
details in Section 3.1), the backmapping function should also be $E(3)$
equivariant. A schematic illustration is given in Fig. 2. However, these
geometric consistency constraints have not been considered by existing
methods. Failure to satisfy these requirements leads to unrealistic FG
geometries that are not self-consistent. More discussions can be found in
Section 3.2. (C3) Generality w.r.t mapping protocols and resolutions: The
choice of dimensionality reduction degree for FG to CG mappings varies wildly
across CG modeling practices, depending on desired accuracy and simulation
speed (Marrink et al., 2007; Bereau & Kremer, 2015; Souza et al., 2021).
Therefore, it is desirable to have a general backmapping approach that is
compatible with arbitrary CG mapping choices, for better usability under
different CG simulation protocols.

Figure 2: A diagram illustrating the geometric constraints for the backmapping
function: 1) The backmapped coordinates should map back to the original CG
coordinates; 2) because the CG transformation is equivariant, the backmapping
transformation should also be equivariant.

To address the requirements and challenges mentioned above, we propose a novel
probabilistic model named Coarse-Graining Variational Auto-Encoder (CGVAE), a
principled data-driven framework to infer atomistic coordinates from coarse-
grained structures. We formulate the stochastic backmapping as a conditional
generative task, i.e., modeling the conditional distribution $p_{\theta}(x|X)$
of FG molecular structures $x$ conditioned on CG structures $X$. To model the
highly complex $p_{\theta}(x|X)$, we factorize the conditional distribution
$p_{\theta}(x|X)$ as a latent variable model and parameterize a generative map
for molecular geometries. This is achieved with $E(3)$ invariant encoding and
equivariant decoding operations, using the message passing framework (Gilmer
et al., 2017). Our model generates coordinates in a one-shot fashion,
requiring no predefined fragment geometry libraries to guide the backmapping.

Our contributions can be summarized as follows:

  * •

We provide a principled probabilistic formulation for the backmapping problem
of molecular conformations, and propose how to approximate this conditional
distribution with a latent variable model in CGVAE (C2).

  * •

We design a rigorous and expressive backmapping function that incorporates
equivariant convolutions to satisfy the geometry consistency criteria (C2).

  * •

CGVAE is CG mapping-agnostic (C3) and can be applied to diverse mapping
protocols used in MD simulations in practice.

  * •

To facilitate further developments in the field, we propose two benchmark
datasets and formulate a suite of metrics to evaluate the quality of generated
FG structures.

Experiments on proposed benchmarks show that our model can consistently
achieve superior performance compared with previous data-driven methods.

##  2 Related Works

Reverse algorithms from CG to FG usually require two steps: backmapping
projection followed by force field optimization (Guttenberg et al., 2013). The
projection typically relies on a deterministic linear backmapping operation
(Wassenaar et al., 2014; Rzepiela et al., 2010; Brocos et al., 2012) or a
parameterized deterministic function learned by supervised learning (Wang &
Gómez-Bombarelli, 2019; An & Deshmukh, 2020). However, these methods often
produce invalid structures because they lack equivariance and chemical rule
supervision, and thus rely on intensive force field equilibration of the
backmapped systems (Rzepiela et al., 2010; Wassenaar et al., 2014). Recently,
probabilistic backmapping (C1) has been incorporated to generate FG
coordinates, but its use is limited to condensed phase systems with voxelized
representation as model inputs (Stieffenhofer et al., 2020). To date, there
have not been works that considered C1-C3 all at once.

##  3 Preliminaries

###  3.1 CG Representations of Molecular Structures

We represent fine-grained (FG) systems as atomistic coordinates
$x=\\{x_{i}\\}_{i=1}^{n}\in\mathbb{R}^{n\times 3}$. Similarly, the coarse-
grained (CG) systems are represented by
$X=\\{X_{i}\\}_{i=1}^{N}\in\mathbb{R}^{N\times 3}$ where $N<n$. Let $[n]$ and
$[N]$ denote the set $\\{1,2,...,n\\}$ and $\\{1,2,...,N\\}$ respectively. CG
operation of molecular simulations can be defined as an assignment
$m:[n]\rightarrow[N]$, which maps each FG atom $i$ in $[n]$ to CG atom
$I\in[N]$, i.e., the CG node $I$ is composed from an ordered set of atoms
$C_{I}=(k\in[n]\mid m(k)=I)$. To ensure that particle mass and momentum are
defined consistently with CG (Noid et al., 2008), each atom $i$ can only
contribute to at most one $I$. This CG projection operation can be represented
as $X=Mx$. $M\in\mathbb{R}^{N\times n}$ with $M_{I,i}=\frac{w_{i}}{\sum_{j\in
C_{I}}w_{j}}$ if $i\in C_{I}$ and 0 otherwise. Here, $w_{i}$ is the projection
weight of atom $i$. It can be either the atomic mass $i$ or simply $1$, which
corresponds to placing $X_{I}$ at the center of mass or the center of geometry
of $C_{I}$. The normalization term in $M$ ensures that coarse coordinates are
the weighted averages of FG geometries.

###  3.2 Geometric Requirements of Backmapping

In this part, we discuss the symmetry properties of CG transformation to
identify the geometric requirements to design the backmapping decoder
function. First, we show that CG projection is $E(3)$ equivariant:

######  Property 3.1.

Let $f_{M}:\mathbb{R}^{n\times 3}\rightarrow\mathbb{R}^{N\times 3}$ be the
linear transformation $f_{M}(x):=Mx$. $f_{M}$ is $E(3)$ equivariant, i.e.,
$f_{M}(Qx+g)=Qf_{M}(x)+g$, where $Q$ is a $3\times 3$ orthogonal matrix, and
$g$ is a translational vector.222Following Satorras et al. (2021), here $Qx$
and $x+g$ are short-hands for $\\{Qx_{1},...,Qx_{n}\\}$ and
$\\{g+x_{1},...,g+x_{n}\\}$ . The same convention also applies to $QX$ and
$X+g$ at the coarse level.

The proof is provided in Section C.3. It states that the CG coordinates $X=Mx$
will always translate, rotate, and reflect in the same way as $x$, obeying
$E(3)$ equivariance. For a generic backmapping function
$\textnormal{Dec}:\mathbb{R}^{N\times 3}\rightarrow\mathbb{R}^{n\times 3}$, it
is natural to require that the backmapped coordinate $\widetilde{x}$ can also
be mapped back to $X$ for self-consistency, i.e. $M\textnormal{Dec}(X)=X$. As
a corollary, such consistency requirements implies that if $f_{M}$ is
equivariant, Dec should also be equivariant. (see Section C.4). Formally, we
consider the following requirements in our decoder design to ensure the
geometric consistency discussed above:

R1. $M\widetilde{x}=M\textnormal{Dec}(X)=X.$
R2. $\textnormal{Dec}(QX+g)$ = $Q\textnormal{Dec}(X)+g.$

To the best of our knowledge, these geometric requirements have not been
considered in previous works. In Section 4.3, we introduce how our approach
satisfies these requirements.

Figure 3: The overall framework of generative CG modeling. The encoder
$q_{\phi}(z|x,X)$ takes both FG and CG structures as inputs, and outputs
invariant latent distributions for each CG node via message passing and
pooling. Given sampled latent variables and the CG structure, the decoder
$p_{\theta}(x|X,z)$ learns to recover the FG structure through equivariant
convolutions. The whole model can be learned end-to-end by optimizing the KL
divergence of latent distributions and reconstruction error of generated FG
structures.

##  4 Method

In this section, we introduce our probabilistic framework of generating fine-
grained (FG) coordinates from coarse-grained (CG) coordinates. We first
present an overview of the conditional generation formalism in Section 4.1 and
describe the detailed model architecture in Sections 4.2 and 4.3. Then we show
the training and sampling protocol in Section 4.4. A notation table is
available in Table 1 in the appendix.

###  4.1 Generative Coarse-Graining Framework

The CG procedure defines a surjective mapping from the FG coordinates $x$ to
the CG coordinates $X$, and we study the inverse problem of recovering $x$
from $X$. This one-to-many problem is a generation task since the same CG
system can backmap to different FG conformations. We propose to learn a
parameterized conditional generative model $p_{\theta}(x|X)$ to approximate
the recovering function. We further incorporate the uncertainty of FG
coordinates into a continuous latent space $z$ and factorize the conditional
distribution as an integral over $z$, i.e., $p(x|X)=\int
p_{\theta}(x|X,z)p_{\psi}(z|X)dz$. This integral is known to be intractable
(Kingma & Welling, 2013), and we instead maximize its variational lower bound
with an approximate amortized posterior distribution $q_{\phi}(z|X,x)$ (Sohn
et al., 2015):

$\displaystyle\log p(x|X)\geq$
$\displaystyle\underbrace{\mathbb{E}_{q_{\phi}(z|x,X)}\log
p_{\theta}(x|X,z)}_{{\mathcal{L}}_{\text{recon.}}}$ (1)
$\displaystyle+\underbrace{\mathbb{E}_{q_{\phi}(z|x,X)}\log\frac{p_{\psi}(z|X)}{q_{\phi}(z|x,X)}}_{{\mathcal{L}}_{\text{reg.}}}$

where $p_{\theta}(x|X,z)$ is the decoder model to recover $x$ from $X$ and
$z$, and $q_{\phi}(z|x,X)$ and $p_{\psi}(z|X)$ are the encoder and prior model
respectively to model the uncertainties from CG reduction. The first term
${\mathcal{L}}_{\text{recon.}}$ is the expected reconstruction error of
generated FG structures, and the second term ${\mathcal{L}}_{\text{reg.}}$ can
be viewed as a regularization over the latent space. A schematic of the whole
framework is provided in Fig. 3. In the following sections, we elaborate on
how we estimate these distributions with graph neural networks.

###  4.2 CG Encoding and Prior

Overview of the data representation. The encoding process extracts information
from both the FG structure $\mathcal{G}_{\textnormal{FG}}$ and CG structure
$\mathcal{G}_{\textnormal{CG}}$. $\mathcal{G}_{\textnormal{FG}}$ represents
radius graphs with a distance cutoff $d_{\textnormal{cut}}$. Radius graph is a
widely adopted choice that can differentiate conformations with the same
molecular graph (Axelrod & Gomez-Bombarelli,
2020).$\mathcal{G}_{\textnormal{CG}}$ are induced graphs based on assignment
map $m$ as introduced in Section 3.1. See Appendix D for detailed descriptions
of these graphs in our model. In short, the nodes are labeled with atomic
types and the edges are labeled with distances.

Estimating the posterior distribution: $(X,x)\rightarrow z$. The encoder model
extracts CG-level invariant embeddings by three operations at each convolution
step: 1) message passing at FG level, 2) pooling operation that maps FG space
to CG space, and 3) message passing at CG level. After convolution updates,
the obtained CG embeddings are then used to parameterize the latent
distribution with CG latent variable $z\in\mathbb{R}^{N\times F}$ where $F$ is
the dimension of features. For each operation, we use a SchNet-like
architecture (Schütt et al., 2017). The first convolution process at FG level
can be described using message passing neural networks (MPNN) (Gilmer et al.,
2017):

$\small h_{i}^{t+1}=\textnormal{Update}\Big{(}\sum_{j\in
N(i)}\textnormal{Msg}(h_{i}^{t},h_{j}^{t},\textnormal{RBF}(d_{ij})),\;h_{i}^{t}\Big{)},$
(2)

where $N(i)$ denotes the neighbors of $i$ in $\mathcal{G}_{\textnormal{FG}}$
and $d_{ij}$ denotes the distance between atoms $i$ and $j$.
$\textnormal{RBF}(\cdot)$ is the radial basis transformation introduced in
Klicpera et al. (2020); Schütt et al. (2021), which transforms the distances
to high dimensional features. The initial node embeddings $h_{i}^{0}$ are
parameterized from atomic types. Since both atomic types and edge distances
are invariant features, the produced FG-wise features will also be invariant.
The FG embeddings are then pooled to CG level based on the CG map $m$:

$\small\widetilde{H}_{I}^{t}=\textnormal{Update}\Big{(}\sum_{i\in
C_{I}}\textnormal{Msg}\left(h_{i}^{t+1},H_{I}^{t},\textnormal{RBF}(d_{iI})\right),\;H_{I}^{t}\Big{)},$
(3)

where $\widetilde{H}_{I}^{t}\in\mathbb{R}^{F}$ is the pooled CG embeddings
from the FG graph and $d_{iI}$ is the distance between $x_{i}$ and $X_{I}$,
i.e., $d_{iI}=\|x_{i}-X_{I}\|_{2}$; $H_{I}^{0}$ is initialized as $0_{F}$.
This step pools messages from FG nodes to their assigned CG nodes, followed by
an updating procedure. The CG-level node features are further parameterized
with message passing at the CG level:

$\small H_{I}^{t+1}=\textnormal{Update}\Big{(}\sum_{J\in
N(I)}\textnormal{Msg}(\widetilde{H}_{I}^{t},\widetilde{H}_{J}^{t},\textnormal{RBF}(d_{IJ})),\;H_{I}^{t}\Big{)},$
(4)

where $N(I)$ denotes the neighbors of $I$ in $\mathcal{G}_{\textnormal{CG}}$,
and $d_{IJ}$ denotes the distance between CG beads $I$ and $J$.

The aforementioned three steps constitute the convolutional module for
extracting both CG and FG level information, with their detailed design
described in Section E.1. After several steps of convolutional updates in our
encoder, we can obtain the final embeddings
$\mathbf{Enc}_{\phi}(x,X)=H^{\phi}_{I}\in\mathbb{R}^{F}$ for each CG atom $I$.
Then two feedforward networks $\mu_{\phi}$ and $\sigma_{\phi}$ are applied to
parameterize the posterior distributions as diagonal multivariate Gaussians
$\mathcal{N}(z_{I}|\mu_{\phi}(H^{\phi}_{I}),\sigma_{\phi}(H^{\phi}_{I}))$.
Note that, $H^{\phi}$ are designed as invariant features and thus the
resulting likelihood is also invariant.

Estimating the prior distribution: $X\rightarrow z$. The prior model
$\mathbf{Prior}_{\psi}(X)$ only takes the CG structure
$\mathcal{G}_{\textnormal{CG}}$ as inputs to model the latent space. The
initial CG node features are set as the pooled sum of one-hot FG fingerprints
$H^{0}_{I}=\sum_{i\in C_{I}}h_{i}^{0}$. Then we employ the same networks as
defined in Eq. 4 to conduct graph convolutions over $\mathcal{G}_{CG}$ and
update the node embeddings. After obtaining the final CG embeddings
$\mathbf{Prior}_{\psi}(X)=H^{\psi}$, we also model $p_{\psi}(z|X)$ as the
products of CG-wise normal distributions
$\mathcal{N}(z_{I}|\mu_{\psi}(H^{\psi}_{I}),\sigma_{\psi}(H^{\psi}_{I})$ with
means and diagonal covariances parameterized from $H^{\psi}_{I}$ using two
separate neural networks $\mu_{\psi}$ and $\sigma_{\psi}$.

###  4.3 Multi-channel Equivariant Decoding

Our decoder design is inspired by recent advances in vector-based graph neural
nets (Schütt et al., 2021; Satorras et al., 2021; Jing et al., 2020; Batzner
et al., 2021). On the high level, the decoder $p_{\theta}(\widetilde{x}|X,z)$
maps the coarse geometry $X\in\mathbb{R}^{N\times 3}$ and invariant latent
variables $z\in\mathbb{R}^{N\times F}$ to produce
$\widetilde{x}\in\mathbb{R}^{n\times 3}$ as predictions of FG geometry $x$.
Our proposed decoder consists of three steps 1) Given $X$ and invariant
feature $z$ from the encoder, we will generate $F$ equivariant vector channel
for each bead, resulting in $V^{\theta}\in\mathbb{R}^{N\times F\times 3}$. 2)
As each bead corresponds to different number of atoms, not all channels will
be used. From vector channels $V^{\theta}$, we select a proper number of
channels for each bead to form ${\Delta\widetilde{x}}\in\mathbb{R}^{n\times
3}$ which is the prediction for the relative position $\Delta x$. 3) In the
last step, we go from relative position ${\Delta\widetilde{x}}$ to absolute
$\widetilde{x}$. This step has to satisfy R1. We elaborate on these three
steps next.

Equivariant convolution: $(X,z)\rightarrow V^{\theta}$. We use equivariant
convolutions based on inter-node vectors to predict ${\Delta\widetilde{x}}$.
Apart from invariant node and edge features, our decoder takes a set of edge-
wise equivariant vector features
$\\{\widehat{E}_{IJ}=\frac{X_{J}-X_{I}}{d_{IJ}}\mid(I,J)\in\mathcal{E}_{\textnormal{CG}}\\}$
as inputs, which form a basis set to construct $\Delta\widetilde{x}$.
Additionally, our model also include pseudoscalars/pseudovectors which contain
additional information of chirality. Unlike scalars/vectors, they undergo a
sign flip under reflection transformation (see Section E.2.1). The message
passing operations in the decoder update node-wise scalar features
$H\in\mathbb{R}^{F}$, pseudoscalar features $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}\in\mathbb{R}^{F}$, vector features $V\in\mathbb{R}^{F\times 3}$ , and
pseudovector features $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}\in\mathbb{R}^{F\times 3}$ separately. The convolutional updates in our
decoder are performed as follows:

$\displaystyle\Delta H_{I}^{t}=\sum_{J\in N(i)}$ $\displaystyle W_{1}\circ
H^{t}_{J}$ (5) $\displaystyle\Delta\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}_{I}^{t}=\sum_{J\in N(i)}$ $\displaystyle
V^{t}_{I}\cdot{\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}}^{t}_{J}$ $\displaystyle\Delta V_{I}^{t}=\sum_{J\in N(i)}$
$\displaystyle\Big{(}W_{2}\circ(V^{t}_{I}\times{\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}}^{t}_{J})+W_{3}\circ{\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}}^{t}_{I}\circ{\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}}^{t}_{J}$ $\displaystyle+W_{4}\widehat{E}_{IJ}+W_{5}\circ
V^{t}_{J}\Big{)}$ $\displaystyle\Delta\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}_{I}^{t}=\sum_{J\in N(i)}$
$\displaystyle\Big{(}W_{6}\circ(V^{t}_{I}\times
V^{t}_{J})+W_{7}\circ({\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}}^{t}_{I}\times{\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}}^{t}_{J})$ $\displaystyle+W_{8}\circ{\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}}^{t}_{J}+W_{9}\circ{\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}}^{t}\circ V^{t}_{J}\Big{)}$

where $\times$, $\circ$ and $\cdot$ are cross product, element-wise product
and dot product operations respectively. $W_{1}$-$W_{9}$ are parameterized
invariant edge-wise filters. Let
$L_{1}:\mathbb{R}^{F}\rightarrow\mathbb{R}^{F}$ and
$L_{2}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{F}$ be neural networks, then each
filter is implemented as as
$W\in\mathbb{R}^{F}=L_{1}(\textnormal{RBF}(d_{IJ}))\circ L_{2}(H_{J})$. The
message passing operations are followed by an equivariant update block
$\sigma(\cdot,\cdot)$ as proposed in PaiNN(Schütt et al., 2021). It is used as
an update block that linearly mixes vector channels and introduces non-linear
coupling between $H$ and $V$. For pseudoscalar and pseudovector
$\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{H}$
and $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$,
we simply apply a residual update to ensure that they flip sign under
reflection:

$\displaystyle H_{I}^{t+1},V_{I}^{t+1}$ $\displaystyle=\sigma(H_{I}^{t}+\Delta
H_{I}^{t},V_{I}^{t}+\Delta V_{I}^{t})$ (6) $\displaystyle\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}_{I}^{t+1},\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}_{I}^{t+1}$ $\displaystyle=\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}_{I}^{t}+\Delta\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}_{I}^{t},\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}_{I}^{t}+\Delta\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}_{I}^{t}$

Pseudoscalars and pseudovectors in our decoder come from cross product
updates. This introduces a richer basis set for the coordinate construction
especially for low $N$ cases. When $N=3$, the span of the vector basis will be
constrained in a plane, and cross product can overcome this limitation by
introducing a vector basis in the orthogonal directions, increasing the
expressiveness of the model.

The initial $H_{I}^{0}$ is obtained from $z\sim q_{\phi}(z|X,x)$ for training,
or $p_{\psi}(z|X)$ for sampling. We initialize other feature set with zero
vectors, i.e., $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}_{I}^{0}=0_{F}$, $V_{I}^{0}=0_{F\times 3}$, and $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}_{I}^{0}=0_{F\times 3}$. For 3-bead CG geometries which are not chiral,
we initialize $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}_{I}^{0}=1_{F}$ to break the symmetry (see Section F.5 for more
discussions). In this way, information in the pseudovector channel can be
propagated to vector channels to construct a basis set that spans 3D.

After several layers of convolution updates, the decoder produces the final
vector output $V^{\theta}\in\mathbb{R}^{N\times F\times 3}$. It is easy to
show that $V^{\theta}$ transforms in an $E(3)$ equivariant way because it is
updated strictly with vector operations, which satisfies R2 in Section 3.2.
More discussions on the equivariance properties of the decoder are available
at Section E.2 and Schütt et al. (2021).

Channel selection: $V^{\theta}\rightarrow{\Delta\widetilde{x}}$. In this part
we introduce how to obtain the relative position predictions
$\Delta\widetilde{x}\in\mathbb{R}^{n\times 3}$ from the equivariant vector
outputs $V^{\theta}\in\mathbb{R}^{N\times F\times 3}$. This requires a channel
selection procedure illustrated in Fig. 4A. As each atom $i$ is assigned to
$I=m(i)$, the corresponding prediction of position ${\Delta\widetilde{x}}_{i}$
is chosen to be the vector channel with the index of $i$ in $C_{I}$. This
requires an index retrieval function $\text{Index}(i,C_{I})$ given that
$C_{I}$ is ordered.333$\text{Index}(i,C_{I})$ retrieves the index of $i\in
C_{I}$. For example, $\text{Index}(1,(1,2,4))=0,\text{Index}(4,(1,2,4))=2$ The
prediction of ${\Delta\widetilde{x}}$ does not utilize all the $F>|C_{I}|$
vector channels at the final vector output $V^{\theta}$, while all the
channels participate in the convolutions before the read-out. We select a
proper number of vector channels from $V^{\theta}$ and concatenate them to
form the FG level prediction ${\Delta\widetilde{x}}$. More precisely,
${\Delta\widetilde{x}}=\bigoplus_{i}V^{\theta}_{m(i),\text{Index}(i,C_{I})}$
where $m(i)$ selects CG beads and $\textnormal{Index}(i,C_{I})$ selects vector
channel of bead $I$. This operation can be done conveniently in Pytorch
(Paszke et al., 2019). See implementation details in Section E.2.3.

Figure 4: A schematic diagram to describe how coordinates are reconstructed
from vector channels. A. The prediction of ${\Delta\widetilde{x}}_{i}$ are
compiled from selected vector channels of CG node $m(i)$ based on their index
in set $C_{I}$; this is followed by the concatenation of selected vectors to
compile
${\Delta\widetilde{x}}=\bigoplus_{i}V^{\theta}_{m(i),\text{Index}(i,C_{I})}$.
The obtained coordinates are further re-centered by subtracting
$M^{+}M{\Delta\widetilde{x}}$ so that R2 is satisfied. B. The CG coordinates
are “lifted” to the FG space via the lifting operator $M^{+}$ to get
$M^{+}X\in\mathbb{R}^{n\times 3}$ which are used as the geometric reference
for constructions.

Compile predictions for FG coordinates:
${\Delta\widetilde{x}}\rightarrow\widetilde{x}$. As ${\Delta\widetilde{x}}$ is
the relative position, we need one more step to go from
${\Delta\widetilde{x}}$ to final FG coordinates $\widetilde{x}$. This is done
by setting
$\widetilde{x}:=M^{+}X+{\Delta\widetilde{x}}-M^{+}M{\Delta\widetilde{x}}$.
$M^{+}\in\mathbb{R}^{n\times N}$ is the lift operator with $M^{+}_{i,I}=1$ if
$I=m(i)$ and 0 otherwise (Section C.1). Similar constructions can also be
found in graph-coarsening methods (Loukas, 2019; Cai et al., 2020). As
described by Fig. 4B, $M^{+}$ maps the point sets space from
$\mathbb{R}^{N\times 3}$ to $\mathbb{R}^{n\times 3}$ by assigning or “lifting”
CG coordinates back to their contributing FG atoms in $C_{I}$. The term
$-M^{+}M{\Delta\widetilde{x}}$ is added to re-center ${\Delta\widetilde{x}}$
so that we can get the original $X$ back after CG projection in order to
satisfy R1. See Section E.2.4 a formal statement with proof.

###  4.4 Model Training and Sampling

Reconstruction loss calculation. We calculate the reconstruction loss
introduced in Eq. 1 as two invariant components:
${\mathcal{L}}_{\textnormal{MSD}}$ and ${\mathcal{L}}_{\textnormal{graph}}$.
The mean square distance loss
${\mathcal{L}}_{\textnormal{MSD}}=\frac{1}{n}\sum_{i=1}^{n}\|\widetilde{x}_{i}-x_{i}\|_{2}^{2}$
is a standard measurement to quantify differences between the generated
reference coordinates (Wang & Gómez-Bombarelli, 2019; Xu et al., 2021). To
supervise the model to produce valid distance geometries, we introduce another
loss term ${\mathcal{L}}_{\textnormal{graph}}$. Let $\mathcal{E}$ denote the
set of chemical bonds in the generated molecular graph. Following previous
works (Xu* et al., 2021), we also expand the set with some auxiliary multi-hop
edges to fix rotatable bonds. We also include all 2-hop neighbor pairs in
$\mathcal{E}$. Then the objective can be calculated as
${\mathcal{L}}_{\textnormal{graph}}=\frac{1}{|\mathcal{E}|}\sum_{(i,j)\in\mathcal{E}}(d(\widetilde{x}_{i},\widetilde{x}_{j})-d(x_{i},x_{j}))^{2}$,
which can be viewed as a supervision over the bond lengths. The total
reconstruction loss is computed as
${\mathcal{L}}_{\textnormal{recon.}}={\mathcal{L}}_{\textnormal{MSD}}+\gamma{\mathcal{L}}_{\textnormal{graph}}$,
where $\gamma$ is a weighting hyperparameter to balance the two loss
contributions.

Training and sampling. We adopt the training objective
${\mathcal{L}}={\mathcal{L}}_{\text{recon.}}+\beta{\mathcal{L}}_{\text{reg.}}$,
where $\beta$ is a hyperparameter to control the regularization strength
(Higgins et al., 2017). During training, the CG latent variable $z$ is sampled
from $q_{\phi}(z|x,X)$ by $z=\mu_{\phi}+\sigma_{\phi}\circ\epsilon$, where
$\epsilon\sim{\mathcal{N}}(0,I)$. Then both $z$ and $X$ are taken as inputs of
the decoding model to generate $\widetilde{x}$. For sampling, given the coarse
structure $X$, we first sample the invariant latent variable from the prior
model by $z\sim p_{\psi}(z|X)$, and then pass the sampled latent
representation $z$ and the coarse coordinates $X$ into the decoder to generate
FG coordinates $\widetilde{x}=\mathbf{Dec}_{\theta}(X,z)$.

##  5 Experiments

(a)

(b)

Figure 5: Left: Examples of CG and generated FG geometries of alanine
dipeptide (top) and chignolin (bottom) generated by CGVAE. The geometry
visualizations are created by PyMol (DeLano, ). Right: Benchmarks of
reconstructed and generated geometries for heavy-atom structures with
different resolution for alanine dipeptide (top) and chignolin (bottom).

In this section, we introduce our experimental settings and discuss both the
quantitative and qualitative results. We first introduce the benchmark
datasets and comprehensive metrics in Section 5.1 and Section 5.2, and then
visualize and analyze the results in Section 5.3.

###  5.1 Benchmarks

Datasets. We evaluate our models on 2 datasets that have been used as
benchmarks for data-driven CG modeling ((Wang & Gómez-Bombarelli, 2019; Wang
et al., 2019; Husic et al., 2020)): MD trajectories of alanine dipeptide and
chignolin. The details of the simulation protocol can be found in Husic et al.
(2020). Alanine dipeptide is a classical benchmark system for developing
enhanced sampling methods for molecular conformation, containing 22 atoms.
Chignolin is a 10-amino acid mini-protein with 175 atoms and features a
$\beta$-hairpin turn in its folded state. It is a more challenging structure
to evaluate our model’s capacity of decoding complex conformations. For
alanine dipeptide and chignolin, we randomly select 20,000 and 10,000
conformations respectively from the data trajectories and perform 5-fold
cross-validation to compute mean and standard errors of evaluation metrics.

Baselines. We compare our model to two baseline models proposed in recent
works. The first method is a learnable linear backmap matrix (Wang & Gómez-
Bombarelli, 2019). The method itself is $O(3)$ equivariant, and we fix the
translation freedom by re-centering the geometry to its geometric center for
each sample. This method is a data-driven and equivariant extension of the
approach described in Wassenaar et al. (2014), which uses linear and
deterministic (not learnable) projection followed by a random displacement.
The second baseline employs Multi-Layer Perceptrons (MLP) (An & Deshmukh,
2020) to learn the projection. The model transforms the CG coordinates input
and the FG coordinates output as 1D flattened arrays, shaped $3N$ and $3n$
respectively. The method uses neural networks to learn the mapping of the
flattened vectors, which is not equivariant. Note that, both baselines lack
the probabilistic formulation and are deterministic, which suffer from poor
sampling diversity.

CG mapping generations. Our method is compatible with arbitrary mapping
protocols. For training our CGVAE and baseline models, in addition to the FG
structures provided in the datasets, we also need a CG mapping to produce the
CG coordinates. In this paper, we use AutoGrain proposed in (Wang & Gómez-
Bombarelli, 2019), which learns the CG assignment via the objective function
that penalizes distant FG atoms being assigned to the same CG atoms. In
practice, we first train the AutoGrain model on each dataset, and then take
the learned assignment to compute the CG coordinates $X$ from FG coordinates
$x$. More details about CG mapping generation can be found in Section F.2.

(a)

(b)

Figure 6:  Left: 2D Ramachandran plots (top) and TICA plots (bottom) for
conformations of alanine dipeptide and chignolin respectively. The structures
for the plots are references in the dataset (first column) and samples
generated by CGVAE trained with different $N$ (the rest three columns).
Figures show that the generated samples recover the important meta-stable
states and the distributions agree well with the ground truth. It also shows
that the higher resolution model shows better resemblance with the ground
truth data. Right: Visualization of representative samples in the 2D plots.

###  5.2 Metrics

Ideally, the generated FG samples should match the original FG molecular
structures (quality) and consist of novel conformations (diversity). We
propose evaluation metrics for generation and reconstruction tasks that
measure model fitting capacity, sample quality, and sample diversity at
different CG resolutions $N$. More detailed descriptions of the proposed
evaluation metrics are in Section F.3. We evaluate the reconstructed and
sampled 3D geometries for both all-atom and heavy-atom (excluding hydrogen
atoms) structures because heavy-atom encodes more important structural
dynamics.

Reconstruction accuracy. The reconstruction task evaluates the model’s
capacity to encode and reconstruct coordinates. Given the reference and
reconstructed FG structures, we report their root mean squared distance
($\textnormal{RMSD}_{\textnormal{recon}}$). A lower
$\textnormal{RMSD}_{\textnormal{recon}}$ indicates the more powerful learning
capacity of the model for reconstructing FG geometries.

Sample quality. We evaluate the sample quality by measuring how well the
generated geometries preserve the original chemical bond graph. Specifically,
we report the graph edit distance ratio
$\lambda(\mathcal{G}_{\textnormal{gen}},\mathcal{G}_{\textnormal{mol}})$,
which measures the relative difference between the connectivity of the
generated geometries $\mathcal{G}_{\textnormal{gen}}$ and the ground truth
bond graph $\mathcal{G}_{\textnormal{mol}}$. For each generated conformation,
$\mathcal{G}_{\textnormal{gen}}$ is deduced from the coordinates by connecting
bonds between pair-wise atoms where the distances are within a threshold
defined by atomic covalent radius cutoff (see Appendix D). The lower
$\lambda(\mathcal{G}_{\textnormal{gen}},\mathcal{G}_{\textnormal{mol}})$ is,
the better the $\mathcal{G}_{\textnormal{gen}}$ resembles
$\mathcal{G}_{\textnormal{mol}}$, with $\lambda=0$ indicating
$\mathcal{G}_{\textnormal{mol}}$ and $\mathcal{G}_{\textnormal{gen}}$ are
isomorphic.

Sample diversity. This metric measures the model capacity to explore the large
conformation space. To evaluate sample diversity, we repeatedly sample
multiple $z\sim P(z|X)$ and generate an ensemble of geometries using
$\mathbf{Dec}_{\theta}(z,X)$. We compare chemically valid generated geometries
with the reference FG geometries using RMSD, and term this metric
$\textnormal{RMSD}_{\textnormal{gen}}$. Similar ideas have been adopted in
related conformation generation literature (Mansimov et al., 2019). A higher
value of $\textnormal{RMSD}_{\textnormal{gen}}$ indicates the better diversity
of the geometries generated by the model. Note that, since both baselines are
deterministic without the capacity to model sampling diversity, their
$\textnormal{RMSD}_{\textnormal{gen}}$ are zeros.

###  5.3 Results and Analysis

Quantitative results. We test our CGVAE models against baseline methods on
both datasets. A summary of all the metrics regarding the heavy-atom
geometries is reported in Fig. 5. Results for all-atom geometries are left in
the Section F.4. As shown in Fig. 5, CGVAE models can consistently achieve
better reconstruction accuracy, sample quality, and diversity than linear and
MLP models among all resolutions. The linear model ensures equivariance but
lacks the expressive capacity to encode complex conformations. The MLP model
enjoys higher expressiveness, but the lack of equivariance limits its
generalization to different roto-translational states. Combining equivariance
and MPNNs, CGVAEs can generate FG samples with high quality and diversity.
Moreover, we observe that higher-resolution CG representations lead to better
reconstruction accuracy, as channels for each CG node are responsible for
fewer atoms. Besides, the higher $\textnormal{RMSD}_{\textnormal{gen}}$
results of the lower resolution models indicate that they can generate more
diverse samples because the distribution of $z$ encodes a larger
conformational space. Surprisingly we notice that the complex geometries can
be well represented with a very small set of CG atoms. Our model works well
even for the extremely coarse order reduction of 22-to-3 and 175-to-6 for
alanine dipeptide and chignolin respectively. Such extent of CG reduction
means potential acceleration of all-atom simulations by one or two orders of
magnitude if a $\mathcal{O}(n\log n)$ algorithm such as Particle Mesh Ewald
(PME) is used (Darden et al., 1993).

Analyzing generated conformations. To gain a better qualitative understanding
of the generated samples, we project the generated samples onto lower-
dimensional reaction coordinates for visualization. The visualizations are
shown in Fig. 6. For alanine dipeptide samples, we visualize the histogram of
the generated samples on $\phi\times\psi$ which are two special central
dihedral angles. The resulting histogram is called the Ramachandran plot,
which is a classic visualization for analyzing peptide conformations
(Ramachandran et al., 1963). The distribution of these two dihedral angles is
an important observable because the heavy-atom dynamics of the central alanine
are completely described by $\phi$ and $\psi$ (Feig, 2008). As shown in the
plot, with only 3 CG atoms, the model can nearly recover all the important
meta-stable states, even though the dihedral angles can only be fully
described by at least 6 heavy atoms at the center of the backbone. For
analysis of generated chignolin samples, we perform Time-structure Independent
Components Analysis (TICA), a popular analysis method to project dynamics onto
slow degrees of freedom (Pérez-Hernández et al., 2013; Scherer et al., 2015).
It is a useful analysis to project high-dimensional complex conformation
variations onto a 2D space to identify meta-stable states of molecular
geometries. To conduct the analysis, we use pair distances for backbone atoms
(45 atoms, 990 distances) with a lag time of 10 ns for both the ground-truth
data and generated samples. We first parameterize the TICA transformation on
the simulation data by PyEmma (Scherer et al., 2015), and apply the
transformation on the generated geometry sets to produce a plot for
comparison. As shown in the figure, our TICA density plot made from generated
samples agrees well with the reference data, which demonstrates our model
accurately captures the distribution of conformational states.

##  6 Conclusion

In this work, we present the Coarse-Grained Variational Auto-Encoder (CGVAE),
a principled probabilistic model to perform stochastic reverse generation to
recover all-atom coordinates from coarse-grained molecular coordinates. Built
on equivariant operations, our model encodes all-atom information into
invariant latent variables which are sampled to generate coordinates
equivariantly. Comprehensive experiments demonstrate that CGVAE significantly
outperforms existing methods. For future work, we plan to extend our general
framework to the more complex condensed phase systems which demand accurate
backmapping at both intra-molecular and inter-molecular level. We hope that
our efforts further improve the predictive power of CG simulations for
accelerated chemical and materials discovery.

## Code Availability

A reference implementation can be found at
https://github.com/wwang2/CoarseGrainingVAE

## Acknowledgement

This work is supported by Toyota Research Institute and Google Faculty Award.

## References

  * An & Deshmukh (2020) An, Y. and Deshmukh, S. A.  Machine learning approach for accurate backmapping of coarse-grained models to all-atom models.  _Chemical Communications_ , 56(65):9312–9315, 2020.
  * Axelrod & Gomez-Bombarelli (2020) Axelrod, S. and Gomez-Bombarelli, R.  Molecular machine learning with conformer ensembles.  _arXiv preprint arXiv:2012.08452_ , 2020.
  * Batzner et al. (2021) Batzner, S., Musaelian, A., Sun, L., Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., and Kozinsky, B.  Se (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials.  _arXiv preprint arXiv:2101.03164_ , 2021.
  * Bereau & Kremer (2015) Bereau, T. and Kremer, K.  Automated parametrization of the coarse-grained martini force field for small organic molecules.  _Journal of chemical theory and computation_ , 11(6):2783–2791, 2015.
  * Bevilacqua et al. (2021) Bevilacqua, B., Frasca, F., Lim, D., Srinivasan, B., Cai, C., Balamurugan, G., Bronstein, M. M., and Maron, H.  Equivariant subgraph aggregation networks.  _arXiv preprint arXiv:2110.02910_ , 2021.
  * Brocos et al. (2012) Brocos, P., Mendoza-Espinosa, P., Castillo, R., Mas-Oliva, J., and Pineiro, Á.  Multiscale molecular dynamics simulations of micelles: coarse-grain for self-assembly and atomic resolution for finer details.  _Soft Matter_ , 8(34):9005–9014, 2012.
  * Cai et al. (2020) Cai, C., Wang, D., and Wang, Y.  Graph coarsening with neural networks.  In _International Conference on Learning Representations_ , 2020.
  * Cai et al. (2021) Cai, C., Vlassis, N., Magee, L., Ma, R., Xiong, Z., Bahmani, B., Wong, T.-F., Wang, Y., and Sun, W.  Equivariant geometric learning for digital rock physics: estimating formation factor and effective permeability tensors from morse graph.  _arXiv preprint arXiv:2104.05608_ , 2021.
  * Cohen et al. (2019) Cohen, T. S., Weiler, M., Kicanaoglu, B., and Welling, M.  Gauge equivariant convolutional networks and the icosahedral cnn.  _arXiv preprint arXiv:1902.04615_ , 2019.
  * Darden et al. (1993) Darden, T., York, D., and Pedersen, L.  Particle mesh ewald: An n log (n) method for ewald sums in large systems.  _The Journal of chemical physics_ , 98(12):10089–10092, 1993.
  * (11) DeLano, W.  Pymol.  URL http://www.pymol.org/pymol.
  * Ehrenfest & Ehrenfest (1912) Ehrenfest, P. and Ehrenfest, T.  _The Conceptual Foundations of the Statistical Approach in Mechanics_.  North-Holland Publishing Company, 1912.
  * Feig (2008) Feig, M.  Is alanine dipeptide a good model for representing the torsional preferences of protein backbones?  _Journal of Chemical Theory and Computation_ , 4(9):1555–1564, 2008.
  * Finzi et al. (2020) Finzi, M., Stanton, S., Izmailov, P., and Wilson, A. G.  Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data.  _arXiv preprint arXiv:2002.12880_ , 2020.
  * Fuchs et al. (2020) Fuchs, F. B., Worrall, D. E., Fischer, V., and Welling, M.  Se (3)-transformers: 3d roto-translation equivariant attention networks.  _arXiv preprint arXiv:2006.10503_ , 2020.
  * Gebauer et al. (2019) Gebauer, N. W., Gastegger, M., and Schütt, K. T.  Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules.  _arXiv preprint arXiv:1906.00957_ , 2019.
  * Geiger et al. (2020) Geiger, M., Smidt, T., M., A., Miller, B. K., Boomsma, W., Dice, B., Lapchevskyi, K., Weiler, M., Tyszkiewicz, M., Batzner, S., Uhrin, M., Frellsen, J., Jung, N., Sanborn, S., Rackers, J., and Bailey, M.  Euclidean neural networks: e3nn, 2020.  URL https://doi.org/10.5281/zenodo.5292912.
  * Gilmer et al. (2017) Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E.  Neural message passing for quantum chemistry.  In _International conference on machine learning_ , pp. 1263–1272. PMLR, 2017.
  * Guttenberg et al. (2013) Guttenberg, N., Dama, J. F., Saunders, M. G., Voth, G. A., Weare, J., and Dinner, A. R.  Minimizing memory as an objective for coarse-graining.  _The Journal of chemical physics_ , 138(9):094111, 2013.
  * Higgins et al. (2017) Higgins, I., Matthey, L., Pal, A., Burgess, C. P., Glorot, X., Botvinick, M. M., Mohamed, S., and Lerchner, A.  beta-vae: Learning basic visual concepts with a constrained variational framework.  In _ICLR_ , 2017.
  * Husic et al. (2020) Husic, B. E., Charron, N. E., Lemm, D., Wang, J., Pérez, A., Majewski, M., Krämer, A., Chen, Y., Olsson, S., de Fabritiis, G., et al.  Coarse graining molecular dynamics with graph neural networks.  _The Journal of Chemical Physics_ , 153(19):194101, 2020.
  * Jang et al. (2016) Jang, E., Gu, S., and Poole, B.  Categorical reparameterization with gumbel-softmax.  _arXiv preprint arXiv:1611.01144_ , 2016.
  * Jensen (2021) Jensen, J. H.  xyz2mol.  https://github.com/jensengroup/xyz2mol, 2021.
  * Jing et al. (2020) Jing, B., Eismann, S., Suriana, P., Townshend, R. J., and Dror, R.  Learning from protein structure with geometric vector perceptrons.  _arXiv preprint arXiv:2009.01411_ , 2020.
  * Kadanoff (1966) Kadanoff, L. P.  Scaling laws for ising models near $t_{c}$.  _Physics Physique Fizika_ , 2(6):263, 1966.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kingma & Welling (2013) Kingma, D. P. and Welling, M.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Kirkwood (1935) Kirkwood, J. G.  Statistical mechanics of fluid mixtures.  _The Journal of chemical physics_ , 3(5):300–313, 1935.
  * Klicpera et al. (2020) Klicpera, J., Groß, J., and Günnemann, S.  Directional message passing for molecular graphs.  _arXiv preprint arXiv:2003.03123_ , 2020.
  * Kremer et al. (1988) Kremer, K., Grest, G. S., and Carmesin, I.  Crossover from rouse to reptation dynamics: A molecular-dynamics simulation.  _Physical review letters_ , 61(5):566, 1988.
  * Laio & Parrinello (2002) Laio, A. and Parrinello, M.  Escaping free-energy minima.  _Proceedings of the National Academy of Sciences_ , 99(20):12562–12566, 2002.
  * Landrum (2013) Landrum, G.  Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling, 2013.
  * Levitt & Warshel (1975) Levitt, M. and Warshel, A.  Computer simulation of protein folding.  _Nature_ , 253(5494):694–698, 1975.
  * Loukas (2019) Loukas, A.  Graph reduction with spectral and cut guarantees.  _J. Mach. Learn. Res._ , 20(116):1–42, 2019.
  * Mansimov et al. (2019) Mansimov, E., Mahmood, O., Kang, S., and Cho, K.  Molecular geometry prediction using a deep generative graph neural network.  _Scientific reports_ , 9(1):1–13, 2019.
  * Maron et al. (2018) Maron, H., Ben-Hamu, H., Shamir, N., and Lipman, Y.  Invariant and equivariant graph networks.  _arXiv preprint arXiv:1812.09902_ , 2018.
  * Maron et al. (2020) Maron, H., Litany, O., Chechik, G., and Fetaya, E.  On learning sets of symmetric elements.  In _International Conference on Machine Learning_ , pp. 6734–6744. PMLR, 2020.
  * Marrink et al. (2007) Marrink, S. J., Risselada, H. J., Yefimov, S., Tieleman, D. P., and De Vries, A. H.  The martini force field: coarse grained model for biomolecular simulations.  _The journal of physical chemistry B_ , 111(27):7812–7824, 2007.
  * McGibbon et al. (2015) McGibbon, R. T., Beauchamp, K. A., Harrigan, M. P., Klein, C., Swails, J. M., Hernández, C. X., Schwantes, C. R., Wang, L.-P., Lane, T. J., and Pande, V. S.  Mdtraj: A modern open library for the analysis of molecular dynamics trajectories.  _Biophysical Journal_ , 109(8):1528 – 1532, 2015\.  doi: 10.1016/j.bpj.2015.08.015.
  * Menichetti et al. (2019) Menichetti, R., Kanekal, K. H., and Bereau, T.  Drug–membrane permeability across chemical space.  _ACS central science_ , 5(2):290–298, 2019.
  * Noid et al. (2008) Noid, W., Chu, J.-W., Ayton, G. S., Krishna, V., Izvekov, S., Voth, G. A., Das, A., and Andersen, H. C.  The multiscale coarse-graining method. i. a rigorous bridge between atomistic and coarse-grained models.  _The Journal of chemical physics_ , 128(24):244114, 2008.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S.  Pytorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems 32_ , pp. 8024–8035. Curran Associates, Inc., 2019.
  * Pérez-Hernández et al. (2013) Pérez-Hernández, G., Paul, F., Giorgino, T., De Fabritiis, G., and Noé, F.  Identification of slow molecular order parameters for markov model construction.  _The Journal of chemical physics_ , 139(1):07B604_1, 2013.
  * Qi et al. (2017) Qi, C. R., Su, H., Mo, K., and Guibas, L. J.  Pointnet: Deep learning on point sets for 3d classification and segmentation.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 652–660, 2017.
  * Ramachandran et al. (1963) Ramachandran, G., Ramakrishnan, C., and Sasisekharan, V.  Stereochemistry of polypeptide chain configurations.  _Journal of Molecular Biology_ , 7(1):95–99, 1963\.
  * Rzepiela et al. (2010) Rzepiela, A. J., Schäfer, L. V., Goga, N., Risselada, H. J., De Vries, A. H., and Marrink, S. J.  Reconstruction of atomistic details from coarse-grained structures.  _Journal of computational chemistry_ , 31(6):1333–1343, 2010.
  * Satorras et al. (2021) Satorras, V. G., Hoogeboom, E., and Welling, M.  E (n) equivariant graph neural networks.  _arXiv preprint arXiv:2102.09844_ , 2021.
  * Scherer et al. (2015) Scherer, M. K., Trendelkamp-Schroer, B., Paul, F., Pérez-Hernández, G., Hoffmann, M., Plattner, N., Wehmeyer, C., Prinz, J.-H., and Noé, F.  PyEMMA 2: A Software Package for Estimation, Validation, and Analysis of Markov Models.  _Journal of Chemical Theory and Computation_ , 11:5525–5542, October 2015.  ISSN 1549-9618.  doi: 10.1021/acs.jctc.5b00743.  URL http://dx.doi.org/10.1021/acs.jctc.5b00743.
  * Schütt et al. (2017) Schütt, K. T., Kindermans, P.-J., Sauceda, H. E., Chmiela, S., Tkatchenko, A., and Müller, K.-R.  Schnet: A continuous-filter convolutional neural network for modeling quantum interactions.  _arXiv preprint arXiv:1706.08566_ , 2017.
  * Schütt et al. (2021) Schütt, K. T., Unke, O. T., and Gastegger, M.  Equivariant message passing for the prediction of tensorial properties and molecular spectra.  _arXiv preprint arXiv:2102.03150_ , 2021.
  * Shi et al. (2021) Shi, C., Luo, S., Xu, M., and Tang, J.  Learning gradient fields for molecular conformation generation.  _arXiv preprint arXiv:2105.03902_ , 2021.
  * Shih et al. (2007) Shih, A. Y., Freddolino, P. L., Sligar, S. G., and Schulten, K.  Disassembly of nanodiscs with cholate.  _Nano letters_ , 7(6):1692–1696, 2007.
  * Sohn et al. (2015) Sohn, K., Lee, H., and Yan, X.  Learning structured output representation using deep conditional generative models.  _Advances in neural information processing systems_ , 28:3483–3491, 2015.
  * Souza et al. (2021) Souza, P. C., Alessandri, R., Barnoud, J., Thallmair, S., Faustino, I., Grünewald, F., Patmanidis, I., Abdizadeh, H., Bruininks, B. M., Wassenaar, T. A., et al.  Martini 3: a general purpose force field for coarse-grained molecular dynamics.  _Nature methods_ , 18(4):382–388, 2021.
  * Stieffenhofer et al. (2020) Stieffenhofer, M., Wand, M., and Bereau, T.  Adversarial reverse mapping of equilibrated condensed-phase molecular structures.  _Machine Learning: Science and Technology_ , 1(4):045014, 2020.
  * Thomas et al. (2018) Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., and Riley, P.  Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds.  _arXiv preprint arXiv:1802.08219_ , 2018.
  * Unke et al. (2021) Unke, O. T., Bogojeski, M., Gastegger, M., Geiger, M., Smidt, T., and Müller, K.-R.  Se (3)-equivariant prediction of molecular wavefunctions and electronic densities.  _arXiv preprint arXiv:2106.02347_ , 2021.
  * Wang et al. (2019) Wang, J., Olsson, S., Wehmeyer, C., Pérez, A., Charron, N. E., De Fabritiis, G., Noé, F., and Clementi, C.  Machine learning of coarse-grained molecular dynamics force fields.  _ACS central science_ , 5(5):755–767, 2019.
  * Wang & Gómez-Bombarelli (2019) Wang, W. and Gómez-Bombarelli, R.  Coarse-graining auto-encoders for molecular dynamics.  _npj Computational Materials_ , 5(1):1–9, 2019\.
  * Wassenaar et al. (2014) Wassenaar, T. A., Pluhackova, K., Bockmann, R. A., Marrink, S. J., and Tieleman, D. P.  Going backward: a flexible geometric approach to reverse transformation from coarse grained to atomistic models.  _Journal of chemical theory and computation_ , 10(2):676–690, 2014.
  * Weiler & Cesa (2019) Weiler, M. and Cesa, G.  General e (2)-equivariant steerable cnns.  In _Advances in Neural Information Processing Systems_ , pp. 14334–14345, 2019.
  * Weiler et al. (2018) Weiler, M., Geiger, M., Welling, M., Boomsma, W., and Cohen, T. S.  3d steerable cnns: Learning rotationally equivariant features in volumetric data.  In _Advances in Neural Information Processing Systems_ , pp. 10381–10392, 2018.
  * Wilson (1971) Wilson, K. G.  Renormalization group and critical phenomena. i. renormalization group and the kadanoff scaling picture.  _Physical review B_ , 4(9):3174, 1971.
  * Winkels & Cohen (2018) Winkels, M. and Cohen, T. S.  3d g-cnns for pulmonary nodule detection.  _arXiv preprint arXiv:1804.04656_ , 2018.
  * Worrall & Welling (2019) Worrall, D. and Welling, M.  Deep scale-spaces: Equivariance over scale.  In _Advances in Neural Information Processing Systems_ , pp. 7366–7378, 2019.
  * Xu* et al. (2021) Xu*, M., Luo*, S., Bengio, Y., Peng, J., and Tang, J.  Learning neural generative dynamics for molecular conformation generation.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=pAbm1qfheGk.
  * Xu et al. (2021) Xu, M., Wang, W., Luo, S., Shi, C., Bengio, Y., Gomez-Bombarelli, R., and Tang, J.  An end-to-end framework for molecular conformation generation via bilevel programming.  _arXiv preprint arXiv:2105.07246_ , 2021.
  * Zaheer et al. (2017) Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R., and Smola, A.  Deep sets.  _arXiv preprint arXiv:1703.06114_ , 2017.
  * Zhang et al. (2018) Zhang, L., Han, J., Wang, H., Car, R., and E, W.  Deepcg: Constructing coarse-grained models via deep neural networks.  _The Journal of chemical physics_ , 149(3):034101, 2018.
  * Zwanzig (2001) Zwanzig, R.  _Nonequilibrium statistical mechanics_.  Oxford university press, 2001.

Table 1: Summary of important notations. Symbol | Meaning
---|---
| Preliminaries
$x\in\mathbb{R}^{n\times 3}$ and $X\in\mathbb{R}^{N\times 3}$ | FG geometry and CG geometry
$i/I$ | index over atoms/beads in FG/CG geometry
$M\in\mathbb{R}^{N\times n}$ | CG projection operator
$M^{+}\in\mathbb{R}^{n\times N}$ |  lift operator. $M^{+}M=\Pi_{n}$, $MM^{+}=I_{N}$.
$Q\in\mathbb{R}^{3\times 3}$ |  3 $\times$ 3 orthogonal transformation $Qx=\\{Qx_{1},...Qx_{n}\\}$
$g\in\mathbb{R}^{3}$ |  translational vector. $x+g=\\{x_{1}+g,...,x_{n}+g\\}$
$m$ |  assignment map from $[n]$ to $[N]$
$w_{i}$ |  weights for atom $i$ for CG projection
$C_{I}$ | $C_{I}=(k\in[n]|m(k)=I)$
| Encoder and prior
$\mathbf{Enc}_{\phi}$ or $\mathbf{Enc}_{\phi}$(x,X) |  encoder in VAE formulation. In CGVAE, it is instantiated as $q_{\phi}(z|x,X)$
$\mathbf{Prior}_{\psi}$ or $\mathbf{Prior}_{\psi}(X)$ |  prior in VAE formulation. In CGVAE, it is instantiated as $p_{\psi}(z|X)$.
Msg | message function
Update | update module
$\textnormal{RBF}:\mathbb{R}\rightarrow\mathbb{R}^{K}$ |  radial basis function following Schütt et al. (2021); Klicpera et al. (2020)
$h^{t}_{i}/H^{t}_{I}\in\mathbb{R}^{F}$ |  FG/CG invariant embeddings at $t_{th}$ convolution step in $\mathbf{Enc}_{\phi}$ or $\mathbf{Prior}_{\psi}$ model
$\widetilde{H}_{I}^{t}$ |  CG node embeddings updates for node $I$ pooled from FG information in Eq. 4
$\Delta H_{I}^{t},\Delta V_{I}^{t},\Delta\macc@depth\char 1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{H}_{I}^{t},\Delta\macc@depth\char 1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}_{I}^{t}$ |  residual udpates for scalar, vector, pseudoscalar, and psuedovector at convolution depth $t$
$N(\cdot)$ | the neighbor function to obtain adjacent node in a graph
$F/K$ | dimension for node/edge features
$\mathcal{N}(\cdot|\cdot,\cdot)$ | multivariate gaussian distribution
$z\in\mathbb{R}^{N\times F}$ | latent representation on CG beads
$\epsilon$ | reparameterized Gaussian noise
$H^{\phi}\in\mathbb{R}^{N\times F}$ |  $H^{\phi}:=\mathbf{Enc}_{\phi}(x,X)=p_{\phi}(z|x,X)$ which is the product of CG-wise normal distribution.
$H^{\psi}\in\mathbb{R}^{N\times F}$ |  $H^{\psi}:=\mathbf{Prior}_{\psi}(X)=p_{\psi}(z|X)$ which is the product of CG-wise normal distribution.
| Decoder
$W_{1}-W_{9}$ |  parameters for equivariant update in Eq. 5
$L_{1},L_{2}$ |  MLPs used to parameterize $W$ in Eq. 5. $L_{1}:\mathbb{R}^{F}\rightarrow\mathbb{R}^{F}$, $L_{2}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{F}$
$\text{Index}(i,C_{I})$ |  index function. $\text{Index}(i,C_{I})$ retrieves the index of $i\in C_{I}$. For example, $\text{Index}(1,(1,2,4))=0$
$\textnormal{Dec}(\cdot)$ |  a generic decoder function : $\mathbb{R}^{N\times 3}\rightarrow\mathbb{R}^{n\times 3}$
$\mathbf{Dec}_{\theta}$ or $\mathbf{Dec}_{\theta}(X,z)$ |  decoder in VAE formulation. In CGVAE, it is instantiated as $p_{\theta}(x|X,z)$.
$\widetilde{x}$ |  $\mathbb{R}^{n\times 3}$, decoded FG geometry. $\widetilde{x}=p_{\theta}(x|X,z)$.
${\Delta\widetilde{x}}$/${\Delta\widetilde{x}_{i}}$ |  ${\Delta\widetilde{x}}:=\widetilde{x}-X$. ${\Delta\widetilde{x}}_{i}$ is $i_{th}$ coordinate of ${\Delta\widetilde{x}}$. Note that decoder in CGVAE generates ${\Delta\widetilde{x}}$.
$H$, $V$, $\macc@depth\char 1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{H}$, $\macc@depth\char 1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$ | scalar, vector, pseudoscalar, pseudovector features
$V^{\theta}\in\mathbb{R}^{N\times F\times 3}$ |  equivariant output of decoder, subsets of its vector channels are used as prediction for ${\Delta\widetilde{x}}$
$\sigma(\cdot,\cdot)$ |  equivariant update function proposed in Schütt et al. (2021)
| Misc
$d_{\textnormal{cut}}/D_{\textnormal{cut}}$ | cutoff distance to form a graph on FG/CG geometry
$\\{\mathcal{G},\mathcal{V},\mathcal{E}\\}$ | graphs, nodes, and edges
$d_{ij}/d_{IJ}\in\mathbb{R}$ | inter-node distances in FG/CG level.
$d_{iI}$ |  distance between $i_{th}$ FG atom and $I_{th}$ CG atom
$\widehat{E}_{IJ}=\frac{X_{J}-X_{I}}{d_{IJ}}$ | unit distance vector between I and J
| Experiments
$\mathcal{G}_{\textnormal{mol}}$ | molecular graphs based on chemical bond connectivities
$\mathcal{G}_{\textnormal{sample}}$ | deduced molecular graphs of samples generated by CGVAE
$\mathcal{E}$ |  edge set that includes 2-hop neighbors in $\mathcal{G}_{\textnormal{mol}}$
$\lambda(\cdot,\cdot)$ |  graph edit distance ratio to measure sample quality by comparing $\mathcal{G}_{\textnormal{mol}}$ and $\mathcal{G}_{\textnormal{sample}}$
$\textnormal{RMSD}_{\textnormal{recon}}$ | measure reconstruction loss; the lower, the better
$\textnormal{RMSD}_{\textnormal{gen}}$ | measure diversity of generated samples; the higher, the better.
| Operators
$\circ$ | pointwise product
$[\cdot,\cdot]$ | concatenation
$\times$ | cross product

##  Appendix A Symbol Table

We summarize definitions of important notations in Table 1.

##  Appendix B Background and Related Works

###  B.1 MD Background

Molecular dynamics and CG. The dynamics and distribution of $x$ are governed
by a Hamiltonian $u(x)$. Under the ergodic hypothesis, the distribution of $x$
converges to the Boltzmann distribution: $p(x)\propto e^{-u(x)\beta}$ where
$\beta$ is the inverse temperature of the system heat reservoir. Given $m$ and
$p(x)$, the equilibrium distribution of $X$ can be derived by $\int
p(x)p(X|x)\>dx$ where $p(X|x)$ is the dirac delta function $\delta(X-M(x))$
which map $x$ deterministically onto the space of $X$. In statistical
mechanics, $-\frac{1}{\beta}\ln p(X)$ is also defined as the Potential of Mean
Force (PMF) which is a conservative free energy quantity to describe the
equilibrium distributions of coarse-grained variables. In this work we are
interested in solving the inverse problem of learning the generative map
$p(x|X)$ given some mapping $M(x)=X$.

CG simulation. The theory and methodology of parameterizing CG models have
been well established over decades (Kirkwood, 1935; Zwanzig, 2001; Noid et
al., 2008). Systematic parameterization and mapping protocols have been widely
applied to simulating large-scale biological molecules (Marrink et al., 2007;
Menichetti et al., 2019; Souza et al., 2021). Lately, there are also a handful
of works utilizing machine-learned potentials to simulate CG systems (Zhang et
al., 2018; Wang & Gómez-Bombarelli, 2019; Wang et al., 2019; Husic et al.,
2020).

###  B.2 More Related Works

Equivariant modeling of point sets. Efforts have been made to build
equivariant neural networks for different types of data and symmetry groups
(Zaheer et al., 2017; Qi et al., 2017; Maron et al., 2018; Thomas et al.,
2018; Weiler et al., 2018; Cohen et al., 2019; Worrall & Welling, 2019; Weiler
& Cesa, 2019; Cohen et al., 2019; Fuchs et al., 2020; Finzi et al., 2020;
Maron et al., 2020; Bevilacqua et al., 2021) . In the case of 3D roto-
translations for point clouds, tensor field network (TFN) (Thomas et al.,
2018) and its extension with attention mechanism $SE(3)$ transformer (Fuchs et
al., 2020) have been applied to various domain where modeling roto-translation
is crucial (Winkels & Cohen, 2018; Unke et al., 2021; Batzner et al., 2021;
Cai et al., 2021).

Generative model for molecular conformations. Many generative models have been
proposed to generate molecular conformation ensembles conditioned on molecular
graphs with modeling techniques like autoregressive models (Gebauer et al.,
2019), conditional normalizing flow (Xu* et al., 2021), bilevel programming
(Xu et al., 2021), and gradient-based optimizations (Shi et al., 2021).

##  Appendix C Mathematical Details

###  C.1 Lift Operator

######  Definition C.1.

A lift operator $M^{+}$ of size $n\times N$ that maps the point sets space
from $\mathbb{R}^{N\times 3}$ to $\mathbb{R}^{n\times 3}$ satisfies the
following property. $M^{+}_{i,I}=\left\\{\begin{array}[]{ll}1&\text{ if
}I=m\left(i\right)\\\ 0&\text{ otherwise }\end{array}\right.$

In the case of $w_{i}=1$ ($w_{i}$ is the projection weight defined in Section
3.1) for all $i$, the $M^{+}$ is the Moore-Penrose pseudoinverse of $M$ that
has the following properties (Loukas, 2019; Cai et al., 2020).

######  Property C.1.

$MM^{+}=I_{N}.$

###### Proof.

First note that $M^{+}_{i,I}=\mathrm{I}_{I\in m(i)}$ where $\mathrm{I}$ is the
indicator function. Expanding the definition, we have

$\sum_{i\in[n]}M_{I,i}M^{+}_{i,J}=\sum_{i\in[n]}\frac{w_{i}}{\sum_{j\in
C_{I}}w_{j}}\mathrm{I}_{J\in m(i)}=\sum_{i\in[n]}\frac{w_{i}\mathrm{I}_{I\in
m(i)}}{\sum_{j\in C_{I}}w_{j}}\delta_{IJ}=\frac{\sum_{i\in
C_{I}}w_{i}}{\sum_{j\in C_{I}}w_{j}}\delta_{IJ}=\delta_{IJ}$

∎

###  C.2 A Toy Example

Consider a cluster map $m$ that atoms $\\{1,2,3,4\\}$ into 2 groups
$\\{\\{1,2,3\\},\\{4\\}\\}$. Assuming the mapping is constructed as geometric
centers, i.e., the mapping weights are equal for each CG particle, then
$M=\begin{bmatrix}1/3&1/3&1/3&0\\\ 0&0&0&1\end{bmatrix}$ and
$M^{+}=\begin{bmatrix}1&0\\\ 1&0\\\ 1&0\\\ 0&1\end{bmatrix}$. It naturally
follows that $MM^{+}=\begin{bmatrix}1&0\\\ 0&1\end{bmatrix}$ and
$M^{+}M=\Pi_{n}=\begin{bmatrix}1/3&1/3&1/3&0\\\ 1/3&1/3&1/3&0\\\
1/3&1/3&1/3&0\\\ 0&0&0&1\\\ \end{bmatrix}.$

Note that the normalization condition $\sum_{i}M[I,i]=1$ is essential for $M$
to respect translational equivariance because CG needs to preserve the norm of
the translation vector $g$.

Note that the construction of CG projections do _not_ require that
$M_{I,i}\neq 0\;\textnormal{ }\forall i\in C_{I}=\\{k\in[n]|m(k)=I\\}$ because
$w_{i}$ can take arbitrary non-negative values. Intuitively, this means that
some atoms do not contribute to the construction of CG coordinates, but they
are assigned to a bead via $m$ in $C_{I}$. This can sometimes be seen in some
coarse-grained simulations in practice: for instance protein simulations are
often run with $\alpha$ carbons only. In the case where $w_{i}=0$ for some
$i$, $M^{+}$ the operator defined is no longer the pseudo-inverse of $M$. For
example, if we assign 0 weights to two of the nodes in $M$ induced by $m$ is
$M=\begin{bmatrix}0&1&0&0\\\ 0&0&0&1\end{bmatrix}$. Based on Definition C.1,
because $M^{+}$ is constructed from $m$, $M^{+}$ remains the same but is not
the pseudo-inverse of $M$ in this case.

###  C.3 Proof of Property 3.1

See 3.1

###### Proof.

Without loss of generality, we prove $f_{M}(Qx+g)=Qf_{M}(x)+g$ holds for any
bead $I$.

$[f_{M}(Qx+g)]_{I}=\sum_{i}M_{I,i}(Qx_{i}+g)=\sum_{i}M_{I,i}(Qx_{i})+\sum_{i}M_{I,i}g=QX_{I}+\sum_{i}M_{I,i}g=QX_{I}+g=[Qf_{M}(x)+g]_{I}$

The first equality follows from the definition of $f_{M}$. The second equality
follows from the linearity of $M$. The third equality definition of $Q$ and
$M$. The fourth equality follows from the definition of $M$. The last equality
follows from the definition of $f_{M}$. ∎

###  C.4 On the geometric Requirements for Backmapping

Given the fact that $M$ is $E(3)$ equivariant, we can show that Dec has to be
equivariant for R1 to be true.

######  Proposition C.1.

If $M\textnormal{Dec}(X)=X$ and $M$ is $E(3)$ equivariant, Dec is also $E(3)$
equivariant, i.e. Dec(QX + g) = Q Dec(X) + g.

###### Proof.

We show $\textnormal{Dec}(Qx+g)=Q\textnormal{Dec}(X)+g$ by proof of
contradiction. Assume $\textnormal{Dec}(QX+g)\neq Q\textnormal{Dec}(X)+g$.
Given $M\textnormal{Dec}(X)=X$ and $M$ is $E(3)$ equivariant, we have

$M\textnormal{Dec}(QX+g)\neq
MQ\textnormal{Dec}(X)+Mg=MQ\textnormal{Dec}(X)+g=QM\textnormal{Dec}(X)+g=QX+g$
(7)

The first inequality follows from the assumption. The second equality follows
from the definition of $M$. The third equality follows from the property of
$M$ and $Q$. The last equality follows from the property of Dec. We reaches a
contradiction that $M\textnormal{Dec}(QX+g)\neq QX+g$. Therefore we conclude
that the assumption has to be false and
$\textnormal{Dec}(QX+g)=Q\textnormal{Dec}(X)+g$.

∎

###  C.5 ELBO Derivation

Following the probabilistic graphical model in Fig. 7, we carry out the
standard derivation of the Evidence Lower Bound (ELBO) described in Eq. 1 for
our conditional VAE here (Sohn et al., 2015):

$\displaystyle\log p(x|X)$
$\displaystyle=\log\mathbb{E}_{q_{\phi}(z|x,X)}\frac{p(x|X)}{q_{\phi}(z|x,X)}$
$\displaystyle\geq\mathbb{E}_{q_{\phi}(z|x,X)}\log\frac{p(x|X)}{q_{\phi}(z|x,X)}$
$\displaystyle=\mathbb{E}_{q_{\phi}(z|x,X)}\log\frac{p_{\theta}(x|X,z)p_{\psi}(z|X)}{q_{\phi}(z|x,X)}$
$\displaystyle=\underbrace{\mathbb{E}_{q_{\phi}(z|x,X)}\log
p_{\theta}(x|X,z)}_{{\mathcal{L}}_{\text{recon.}}}+\underbrace{\mathbb{E}_{q_{\phi}(z|x,X)}\log\frac{p_{\psi}(z|X)}{q_{\phi}(z|x,X)}}_{{\mathcal{L}}_{\text{reg.}}}$

Figure 7: The probabilistic graphical model for the generation process (solid
line) and inference process (dashed line) of the conditional geometric reverse
mapping.

##  Appendix D Graph Representations

Table 2: Summary of graph data inputs used in CGVAE. | $\mathcal{G}_{\textnormal{FG}}$ | $\mathcal{G}_{\textnormal{CG}}$
---|---|---
Coordinate | $x\in\mathbb{R}^{n\times 3}$ | $X\in\mathbb{R}^{N\times 3}$
Node | $\mathcal{V}_{\textnormal{FG}}$ | $\mathcal{V}_{\textnormal{CG}}$
Edge | $\mathcal{E}_{\textnormal{FG}}=\\{(i,j)\mid d(x_{i},x_{j})<d_{\textnormal{cut}},\forall\;i,j\in\mathcal{V}_{\textnormal{FG}}\\}$ | $\mathcal{E}_{\textnormal{CG}}=\\{(I,J)\mid(C_{I}\times C_{J})\cap\mathcal{E}_{\textnormal{FG}}\neq\ \emptyset,\forall I,J\in\mathcal{V}_{\textnormal{CG}}\\}$
Node feature |  invariant features $h\in\mathbb{R}^{n\times F}$ |  invariant features $H\in\mathbb{R}^{N\times F}$, equivariant features $V\in\mathbb{R}^{N\times F\times 3}$
Edge feature | $d_{ij}\in\mathbb{R}$ |  $d_{IJ}\in\mathbb{R}$, $\widehat{E}_{IJ}=\frac{X_{I}-X_{J}}{d_{IJ}}\in\mathbb{R}^{3}$

Here we describe the graph data structure we used for CG and FG graphs for
geometric convolutions. The detailed description is summarized in Table 2.

FG 3D graph $\mathcal{G}_{\textnormal{FG}}$. The geometric graph at the FG
level is represented as an undirected
$\mathcal{G}_{\textnormal{FG}}=(\mathcal{V}_{\textnormal{FG}},\mathcal{E}_{\textnormal{FG}})$,
where $\mathcal{V}_{\textnormal{FG}}=[n]$ is the node set. The edge set is
determined by a predefined distance cutoff $d_{\textnormal{cut}}$ that
$\mathcal{E}_{FG}=\\{(i,j)\mid
d(x_{i},x_{j})<d_{\textnormal{cut}},\forall\;i,j\in\mathcal{V}\\}$, where
$d(a,b)=\|a-b\|_{2}$ is the euclidean distance between $a$ and $b$. Following
standard featurization procedure in machine learning techniques for molecular
graph, the initial node features $h\in\mathbb{R}^{F}$ is simply one-hot
encoding of atomic types. Following (Schütt et al., 2021), we parameterize
edge features using $\textnormal{RBF}(d_{ij})$ to expand $d_{ij}$ with
coefficients of $K$ sine radial basis:
$\sin(\frac{n\pi}{d_{\textnormal{cut}}}d_{ij})/d_{ij}$ with a cosine filter as
a continuous cutoff. $n$ ranges from 1 to $K$. Both $K$ and
$d_{\textnormal{cut}}$ are hyper-parameters of the model.

CG graph. Given $m$, we introduce a corresponding undirected CG graph
$\mathcal{G}_{\textnormal{CG}}$ with $\mathcal{V}_{CG}=[N]$. There is an edge
between node $I$ and $J$ if there exists at least an edge between sub-nodes in
$C_{I}$ and $C_{J}$: $\mathcal{E}_{\textnormal{CG}}=\\{(I,J)\mid(C_{I}\times
C_{J})\cap\mathcal{E}_{\textnormal{FG}}\neq\ \emptyset,\forall
I,J\in\mathcal{V}_{\textnormal{CG}}\\}$. Here, $C_{I}\times C_{J}$ indicates
the Cartesian product between $C_{I}$ and $C_{J}$. The initial node feature
$H_{I}$ is parameterized using the pooling operation described in Section 4.2.
The distances for CG point sets are similarly featurized with sine radial
basis functions, with a cutoff distance at $D_{\textnormal{cut}}$. For the
equivariant decoding described in Section 4.3, CG graph also carries a set of
equivariant features parameterized from linear combinations of
$\widehat{E}_{IJ}$.

FG bond graph. A molecule can have many conformations (stereo-isomers) which
share the same bond graph structure. For classical MD simulations where there
are no chemical reactions, the chemical bond graph does not change. The
molecular graph
$\mathcal{G}_{\textnormal{mol}}=(\mathcal{V}_{\textnormal{mol}},\mathcal{E}_{\textnormal{mol}})$
can be deduced from chemical rules and can be obtained from open-sourced
chemoinformatic analysis packages like RDKIT (Landrum, 2013).

Determine bond graph from FG geometries. For evaluation of quality of
generated samples, we compute the bond graphs of generated samples and compare
them with the ground truth chemical bond graph. The chemical graphs are
essentially radius graphs but with pairwise distance cutoffs determined by the
atom types of the pairs involved. The distance cutoffs for each pair of atoms
are determined from the covalent radii of each atoms. We obtain covalent
radius of each atom from RDKIT using rdkit.Chem.PeriodicTable.GetRCovalent.
Let $r_{\textnormal{covalent}}(i)$ indicates the covalent radius of atom $i$,
the bond graph of generated geometries can be deduced:
$\\{(i,j)\mid\|x_{i}-x_{j}\|_{2}<1.3\times(r_{\textnormal{covalent}}(i)+r_{\textnormal{covalent}}(j)\forall
i,j\in\mathcal{V}_{\textnormal{FG}}\\})$. The scaling factor of 1.3 is chosen
based on the xyz2mol public repository (Jensen, 2021).

##  Appendix E Model Details

###  E.1 Details of Encoder and Prior Networks

Encoder. We introduce the detailed design of our encoder. Recall that $K$ is
the dimension of invariant edge features, and $F$ is the dimension of node
embedding. We define
$L_{3},L^{\prime}_{3},L^{\prime\prime}_{3},L^{\prime\prime\prime}_{3}:\mathbb{R}^{K}\rightarrow\mathbb{R}^{F}$
and
$L_{4},L^{\prime}_{4},L^{\prime\prime}_{4},L^{\prime\prime\prime}_{4}:\mathbb{R}^{F}\rightarrow\mathbb{R}^{F}$
be layer-wise MLPs with swish function as activation. The detailed form of FG
MPNN is:

$\displaystyle h_{i}^{t+1}$ $\displaystyle=\textnormal{Update}\left(\sum_{j\in
N(i)}\textnormal{Msg}(h_{i}^{t},h_{j}^{t},\textnormal{RBF}(d_{ij})),\;h_{i}^{t}\right)=h_{i}^{t}+\sum_{j\in
N(i)}L_{3}(\textnormal{RBF}(d_{ij}))\circ L_{4}(h_{i}^{t}\circ h_{j}^{t})$

The detailed form of the pooling operation to map information from FG to CG
level is:

$\displaystyle\widetilde{H}_{I}^{t}$
$\displaystyle=\textnormal{Update}\left(\sum_{i\in
C_{I=m(i)}}\textnormal{Msg}(h_{i}^{t},H_{I}^{t},\textnormal{RBF}(d_{iI})),\;H_{I}^{t}\right)=H_{I}^{t}+\sum_{i\in
C_{I}}L^{\prime}_{3}(\textnormal{RBF}(d_{iI}))\circ L^{\prime}_{4}(h_{i}^{t})$

Similar to MPNNs at the FG resolution, the CG-level MPNN is:

$\displaystyle H_{I}^{t+1}$ $\displaystyle=\textnormal{Update}\left(\sum_{J\in
N(I)}\textnormal{Msg}(\widetilde{H}_{I}^{t},\widetilde{H}_{J}^{t},\textnormal{RBF}(d_{IJ})),\;H_{I}^{t}\right)=H_{I}^{t}+\sum_{J\in
N(I)}L^{\prime\prime}_{3}(\textnormal{RBF}(d_{IJ}))\circ
L^{\prime\prime}_{4}(\widetilde{H}_{I}^{t}\circ\widetilde{H}_{J}^{t})$

This form of MPNN is used for CG information propagation for both
$\mathbf{Enc}_{\phi}$ and $\mathbf{Prior}_{\psi}$.

Prior. For $\mathbf{Prior}_{\psi}$, the detailed depth-wise form of the
message passing operation is :

$\displaystyle H_{I}^{t+1}$ $\displaystyle=\textnormal{Update}\left(\sum_{J\in
N(I)}\textnormal{Msg}(H_{I}^{t},H_{J}^{t},\textnormal{RBF}(d_{IJ})),\;H_{I}^{t}\right)=H_{I}^{t}+\sum_{J\in
N(I)}L^{\prime\prime\prime}_{3}(\textnormal{RBF}(d_{IJ}))\circ
L^{\prime\prime\prime}_{4}(H_{I}^{t}\circ\widetilde{H}_{J}^{t})$

###  E.2 Details of Decoder

In this subsection, We first give details about Eq. 5 and Eq. 6 in decoder and
discuss the its equivariance property. We then discuss the geometric
consistency of the decoder. Lastly, we give the implementation of channel
selection operations.

####  E.2.1 Brief Introduction to Pseudoscalars and Pseudovectors

Pseudoscalars and pseudovectors are like scalars and vectors, but undergoes an
additional sign flip under reflection. We use $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{}$
to indicate a quantity is pseudoscalar or psuedovector. Let
$A,B,C\in\mathbb{R}^{3}$ be normal vectors. pseudoscalar $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{H}$
and pseudovector $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$
can be constructed as follows: $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}=A\cdot(B\times C),\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}=A\times B$. Note that cross product obeys the following identity under
matrix transformations $(MA)\times(MB)=\det(M)(A\times B)$. Let $P$ be a
$3\times 3$ reflection transformation with $\det(P)=-1$. It is easy to see
that $(PA\times PB)=-P(A\times B)$. Along with the simple property of $PA\cdot
PB=P(A\cdot B)$, we have $P\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{H}$
and $P\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$

$\displaystyle P\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{H}$
$\displaystyle=PA\cdot(PB\times PC)=PA\cdot(-P(B\times C))=-A\cdot(B\times
C)=-\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{H}$
$\displaystyle P\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$
$\displaystyle=PA\times PB=-P(A\times B)=-\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$

Next, we list different ways to generate
scalars/vectors/pseudoscalars/pseudovectors. These are used as the basis in
Eq. 5. The details of these constructions for arbitrary order tensors can be
found in the documentation of the e3nn package (Geiger et al., 2020).

  * •

scalar: i) $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}\circ\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{H}$
ii) $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}\cdot\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$

  * •

vector: i) $V\times\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$
ii) $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}\circ\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$

  * •

pseudoscalar: i) $V\cdot\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$
ii) $H\circ\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{H}$

  * •

pseudovector: i) $V\times V$ ii) $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{V}\times\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a 111{V}$
iii) $\macc@depth\char
1\relax\frozen@everymath{\macc@group}\macc@set@skewchar\macc@nested@a
111{H}\circ V$

The output of Eq. 5 are scalars, pseudoscalars, vectors, and pseudovectors. As
the construction of each equation in Eq. 5 ensures that the “type” of output
matches with the “type” of the basis, we immediately get the equivariance
guarantee for the outputs of Eq. 5.

####  E.2.2 Details of Equivariant Update Layers

In this subsection, we provide details of equivariant update layers in Eq. 6
$H_{I}^{t+1},V_{I}^{t+1}=\sigma(H_{I}^{t}+\Delta H_{I}^{t},V_{I}^{t}+\Delta
V_{I}^{t})$. Eq. 6 updates an invariant scalar feature
$H_{I}\in\mathbb{R}^{F}$ as the first argument and an equivariant vector
feature $V_{I}\in\mathbb{R}^{F\times 3}$ as the second argument. Under the
action of rotation $Q$, $\sigma$ has to satisfy the following constraint:
$[\sigma(H_{I},QV_{I})]_{1}=[\sigma(H_{I},V_{I})]_{1}$ and
$[\sigma(H_{I},QV_{I})]_{2}=Q[\sigma(H_{I},V_{I})]_{2}$. We next introduce the
detailed forms of $\sigma$.

Let
$L_{5},L^{\prime}_{5},L^{\prime\prime}_{5}:\mathbb{R}^{2F}\rightarrow\mathbb{R}^{F}$
be MLPs and $W_{\alpha},W_{\beta}\in\mathbb{R}^{F\times F}$ be two learnable
matrices. Denote the invariant input of $\sigma(\cdot,\cdot)$ as
${H^{\prime}_{I}}^{t+1}:=H_{I}^{t}+\Delta{H_{I}}\in\mathbb{R}^{F}$ and the
equivariant input as ${V^{\prime}_{I}}^{t+1}:=V_{I}^{t}+\Delta
V_{I}\in\mathbb{R}^{F\times 3}$. The detailed form of $\sigma$ to the update
$H^{t+1}$ and $V^{t+1}$ is

$\displaystyle H_{I}^{t+1}$
$\displaystyle={H^{\prime}_{I}}^{t+1}+L_{5}(\left[{H^{\prime}_{I}}^{t+1},\left\|W_{\beta}{V^{\prime}_{I}}^{t+1}\right\|\right])+L^{\prime}_{5}(\left[{H^{\prime}_{I}}^{t+1},\left\|W_{\beta}{V^{\prime}_{I}}^{t+1}\right\|\right])\langle
W_{\alpha}{V^{\prime}_{I}}^{t+1},W_{\beta}{V^{\prime}_{I}}^{t+1}\rangle$ (8)
$\displaystyle V_{I}^{t+1}$
$\displaystyle={V^{\prime}_{I}}^{t+1}+L^{\prime\prime}_{5}(\left[H^{\prime
t+1}_{I},\left\|W_{\beta}{V^{\prime}_{I}}^{t+1}\right\|\right]){V^{\prime}_{I}}^{t+1}$

Here $\|\cdot\|:\mathbb{R}^{F\times 3}\rightarrow\mathbb{R}^{F}$ computes the
norm for each feature channel. $\langle\cdot,\cdot\rangle:\mathbb{R}^{F\times
3}\times\mathbb{R}^{F\times 3}\rightarrow\mathbb{R}^{F}$ is the feature-wise
dot product, and $[\cdot,\cdot]:\mathbb{R}^{F\times
d_{1}}\times\mathbb{R}^{F\times
d_{2}}\rightarrow\mathbb{R}^{F\times(d_{1}+d_{2})}$ indicates the
concatenation along the second feature channel. It is easy to show that the
update function produces an invariant $H_{I}^{t+1}$ and an equivariant
$V_{I}^{t+1}$ because 1) dot product and Euclidean norms are invariant under
rotation and reflection, and 2) linear combinations of vectors are
equivariant. More discussions can be found in Schütt et al. (2021).

####  E.2.3 Implementation of Channel Selection Operations

The final prediction of ${\Delta\widetilde{x}}$ are compiled from vector
channels in $V_{I,\text{Index}(i,C_{I})}$. This can be done efficiently by
generating pairs of $(I,\text{Index}(i,C_{I}))$ used as indices array for
advanced indexing in Pytorch. A PyTorch implementation is provided here:

⬇

1def channel_select(V, m):

2 # m: n x 1 CG map for each FG node

3 # V : N x F x 3 vector channels on CG node

4 channel_idx = torch.zeros_like( m)

5 for cg_type in torch.unique(m):

6 cg_filter = m == cg_type

7 # find size of C_I

8 k = cg_filter.sum().item() # find size of C_I

9 # construct (I, Index(i, C_I))

10 channel_idx[cg_filter] = torch.LongTensor(list(range(k)))

11 dx = V[m, channel_idx]

12 return dx

####  E.2.4 Ensuring Geometry Self-consistency

Just as the lift map $M^{+}$ satisfies $MM^{+}X=X$ , the learnt decoder
$\mathbf{Dec}_{\theta}$ also needs be compatible with pre-defined surjective
CG map $M$, i.e. $M(\mathbf{Dec}_{\theta}(X,z))=X$.

######  Proposition E.1.

$\mathbf{Dec}_{\theta}(X,z)=M^{+}X+\Delta\widetilde{x}-M^{+}M\Delta\widetilde{x}$
satisfies $M(\mathbf{Dec}_{\theta}(X,z))=X$.

###### Proof.

From linearity of $M$ and the property that $MM^{+}=I$, we have

$\displaystyle
M(\mathbf{Dec}_{\theta}(X,z))=MM^{+}X+M\Delta\widetilde{x}-MM^{+}M(\Delta\widetilde{x})=X+M\Delta\widetilde{x}-M\Delta\widetilde{x}=X$

∎

##  Appendix F Experiments Details

###  F.1 Baselines

Equivariant linear model. (Wang & Gómez-Bombarelli, 2019) We construct the
first model following (Wang & Gómez-Bombarelli, 2019), the reconstructed
coordinates are obtained from linear combinations of the coarse coordinates
with learnable coefficients $D_{i,I}$.
$\widetilde{x}_{i,m}:=\sum_{I}D_{i,I}X_{I,m}$ where $m\in[3]$ indicates the
Cartesian index of the coordinates. To incorporate translational equivariance
which the model fails to incorporate, we recenter the FG structure and the
mapped CG structure to its geometric center for training and testing of the
model.

Table 3: Hyper-parameters for the MLP and equivariant linear baseline models. We used the same set of parameters for both alanine dipeptide and chignolin datasets.  | Alanine Dipeptide | Chignolin
---|---|---
${\mathcal{L}}_{\textnormal{graph}}$ weight $\gamma$ | 5 | 5
Training epochs | 500 | 500
Learning rate | 0.0001 | 0.0001
Batch size | 32 | 32
Order for multi-hop graph | 2 | 2

Multi-Layer Perceptrons.(An & Deshmukh, 2020) We use an MLP regressor as
proposed by (An & Deshmukh, 2020) to map between CG space to FG space. After
flattening the coordinate arrays, the MLP maps $\mathbb{R}^{3N}$ to
$\mathbb{R}^{3n}$. The model does not incorporate $E(3)$ equivariance. We use
MLPs with 2 hidden layers with hidden dimensions of size $3n$ with ReLU as the
activation function.

Table 4: hyperparameters used for CGVAE | Alanine Dipeptide | Chignolin
---|---|---
| Dataset size
---
20000 | 10000
| Edge feature dimension $K$
---
8 | 10
|  ${\mathcal{L}}_{\textnormal{graph}}$ weight. $\gamma$
---
25.0 | 50.0
| Encoder conv depth $T^{\textnormal{enc}}$
---
4 | 2
| Prior conv depth $T^{\textnormal{prior}}$
---
4 | 2
| Decoder conv depth $T^{\textnormal{dec}}$
---
5 | 9
| $d_{\textnormal{cut}}$
---
8.5 | 12.0
| $D_{\textnormal{cut}}$
---
9.5 | 25.0
Node embedding dimension $F$ | 600 | 600
| Batch size
---
32 | 2
| Learning rate
---
8e-5 | 1e-4
| Activation function
---
swish | swish
| Training epochs
---
500 | 100
Order for multi-hop graph | 2 | 2
$\beta$ regularization strength of KL divergence | 0.05 | 0.05

###  F.2 Details for CGVAE

Parameterize mappings with Auto-Grain. Our model works with arbitrary mapping
protocols that satisfy the requirement described in Section 3.1. In practice,
we find mappings that assign close atoms together perform better. We
parameterize our CG mapping using Auto-Grain proposed in (Wang & Gómez-
Bombarelli, 2019) which readers can refer for details of the method. Auto-
Grain parameterizes the discrete assignment matrix
$C_{I,i}\in\mathbb{R}^{N\times n}$ in a differentiable way using Gumbel-
Softmax transformation (Jang et al., 2016) with a tunable fictitious
temperature $\tau$ for different level of discreteness. The resulting CG
transformation using the center of geometries is
$M_{I,i}=\frac{C_{I,i}}{\sum_{i=1}^{n}C_{i,I}}$. A decoding matrix is used to
equivariantly reconstruct $x$ based on $Mx$ with minimization of the
reconstruction MSD loss. To encourage assignment of closed-together atoms into
the same CG bead, we introduce a new geometric loss :
${\mathcal{L}}_{\textnormal{geo}}=\frac{1}{n}\sum_{i=1}^{n}\sum_{k=1}^{3}((C^{T}Mx)_{i,k}-x_{i,k})^{2}$
which penalizes mappings that groups atoms that are far away from each other.
$C^{T}$ can be thought as the lifting operator of the soft-parameterized $M$,
satisfying $MC^{T}=I_{N}$. The total loss for Auto-Grain is
${\mathcal{L}}_{MSD}+a{\mathcal{L}}_{\textnormal{geo}}$ where $a$ re-weights
the two loss contributions. We choose $a=0.25$ in our experiments. For each
experiment, we randomly select 1000 samples to train Auto-Grain to obtain a
mapping. We train our model for 1500 epochs, with $\tau$ starting at 1.0 and
gradually decreasing its value by 0.001 till it reaches a value of 0.025. The
optimization is stochastic, and the obtained mapping will be slightly
different, but compatible with CGVAE, for different parameter initializations.
We find mapping obtained this way generate higher-quality samples than
backbone-based CG protocol (introduced below), especially at low resolution
for Chignolin. We present results obtained from CGVAE using the two mapping
protocols in Fig. 8.

Hyperparameters for CGVAE. We use the hyperparameter set described in Table 4
for the CGVAE experiments we presented in the main text. The same set of
hyperparameters are used for all choices of $N$. For optimization, we use the
Adam optimizer (Kingma & Ba, 2014). We additionally apply a learning rate
scheduler to adaptively reduce learning rate when convergence is reached. We
use torch.optim.lr_scheduler.ReduceLROnPlateau implemented in Pytorch (Paszke
et al., 2019) with patience=15 and factor=0.3.

Figure 8: Comparing CGVAE results for alanine dipeptide (top) and chignolin
(bottom) using two different mapping protocols. In general, the results from
the two mappings are comparable, yet Auto-Grain yields better sample qualities
for low-resolution Chignolin representations.

CG based on backbones. Performing CG transformation based on backbones is an
alternative mapping strategy that groups neighboring atoms together. We
randomly partition the linear backbone chain of the protein into $N$ segments.
We then assign the side chain atoms to the closest backbone segments based on
the center of geometry of the backbone segments. We infer the backbone
structure of the protein using the MDTraj package (McGibbon et al., 2015).
After mapping is obtained, we compute $X$ as the center of geometry.

###  F.3 Evaluation Metrics

Reconstruction accuracy. For the reconstruction task, we are interested in the
difference between the original reference atomistic structure $x$ and the
reconstructed structure $\widetilde{x}$. This informs how well a model can
reconstruct samples. For the training of CGVAEs, we remove the stochastic
reparameterization to train a deterministic auto-encoders. This choice is to
align our evaluation with our deterministic baselines. The discrepancy between
$x$ and $\widetilde{x}$ can be evaluated with root mean square distance
$\textnormal{RMSD}_{\textnormal{recon}}=\sqrt{\frac{\sum_{i=1}^{n}(x_{i}-\widetilde{x}_{i})^{2}}{n}}.$

Sample qualities. The sampling task is performed based on the information of
CG coordinates only. The invariant features $z\sim p_{\psi}(z|X)$ is then used
to generate coordinates with the decoding function
$\mathbf{Dec}_{\theta}(X,z)$. We are primarily interested in the quality of
the generated geometries by evaluating how well the generated geometry
ensembles preserve the bond graph $\mathcal{G}_{\textnormal{gen}}$ compared
with the graph $\mathcal{G}_{\textnormal{mol.}}$ induced by the reference
structure in the data. We infer $\mathcal{G}_{\textnormal{gen}}$ and
$\mathcal{G}_{\textnormal{mol.}}$ with the protocol described in Appendix D.
We quantify the similarity of the two generated graphs by computing the
minimum graph edit distance. Because the generated geometries preserve the
number of nodes and their identities, we simply need to compute the number of
edge addition and removal operations required between
$\mathcal{G}_{\textnormal{gen}}$ and $\mathcal{G}_{\textnormal{mol.}}$ and we
denote $\mathcal{E}_{\textnormal{gen}}$ and $\mathcal{E}_{\textnormal{mol.}}$
as their edge sets respectively. Because the atom orders of the reference
graph and the generated graphs are the same, the graph edit distance is the
size of the symmetric difference ($\Delta$), i.e, the union without the
intersection of the two edge sets. We define the quantify
$\lambda(\mathcal{E}_{\textnormal{gen}},\mathcal{E}_{\textnormal{mol.}})$, the
graph edit distance normalized by the size of the edge set of the original
graph $\mathcal{E}_{\textnormal{mol.}}$.

$\lambda(\mathcal{E}_{\textnormal{gen}},\mathcal{E}_{\textnormal{mol}})=\frac{\textnormal{GED}(\mathcal{E}_{\textnormal{gen}},\mathcal{E}_{\textnormal{mol}})}{|\mathcal{E}_{\textnormal{mol}}|}=\frac{\mathcal{E}_{\textnormal{gen}}\Delta\mathcal{E}_{\textnormal{mol}}}{|\mathcal{E}_{\textnormal{mol}}|}$

where $\Delta$ denotes the set difference. If $\lambda=0$, it means that the
molecular graph of the generated geometries is identical to the reference
molecular graph. We evaluate $\lambda$ for graphs for both the all-atom
systems and heavy-atom only systems. We make the distinction because heavy-
atoms encode most of the important dynamics of the system. Also, substructures
that involve hydrogen are highly fluctuating, make the model (especially for
models with few $N$) hard to capture its distribution fully. Additionally, we
also compute the ratio of valid graphs of the generated samples. Such ratio
tends to be low for Chignolin because of its more complex structures. For
sample diversity metrics, we only calculate on structures with valid graphs.
These values can be found in Table 5 and Table 6.

Sample diversities. We are also interested in the diversity of sampled
geometries for our proposed generative model. Because $M$ is a surjective map,
each CG configuration corresponds to multiple possible configurations. The
more diverse the sampled geometries given $X$, the more information that $z$
has learned for the missing information encoded in $z$. For this purpose, we
want to evaluate RMSD compared with the reference geometry to quantify how
well the model can generate samples that are not in the reference set. The
evaluation is only performed for generated samples with valid molecular
graphs. The diversity RMSD is defined as:
$\textnormal{RMSD}_{\textnormal{gen}}=\sqrt{\frac{\sum_{i}^{n}\int
p_{\psi}(z|X)(x_{i}-\mathbf{Dec}(X,z))^{2}}{n}}.$ Based on the discussion
above, a higher value of $\textnormal{RMSD}_{\textnormal{gen}}$ indicates
better diversities of generated samples. To compute
$\textnormal{RMSD}_{\textnormal{gen}}$, we sample 32 values of $z$ for each
$X$ in the test set.

###  F.4 Tables and Sample Geometries

We show tabulated benchmark results in Table 5 for alanine dipeptide and
chignolin datasets respectively for both heavy-atom and all-atom geometries.
For $\textnormal{RMSD}_{\textnormal{gen}}$, the higher the value the better
the model is in generating diverse samples. To complement the results of heavy
atoms shown in Fig. 5 for heavy atoms, we also include plots for all-atom
reconstruction and sampling in Fig. 9. We also compare generated chignolin
samples ($N=6$) from different methods in Fig. 10.

Figure 9:  Benchmarks for all-atom reconstruction and sampled geometries at different resolutions for alanine dipeptide (top) and chignolin (bottom). $\textnormal{RMSD}_{\textnormal{gen}}$ is not available for our deterministic baselines because the generation process is deterministic.  Figure 10: Comparison between generated chignolin samples ($N=6$) from CGVAE and backmapped samples from baseline methods. Table 5: Alanine dipeptide benchmarks Metric | Model | $N$
---|---|---
3 | 5 | 8
Heavy atom | $\textbf{RMSD}_{\textnormal{recon}}$ | linear | $1.211\pm 0.075$ | $0.277\pm 0.032$ | $0.086\pm 0.009$
| MLP | $0.439\pm 0.003$ | $0.288\pm 0.004$ | $0.269\pm 0.009$
| CGVAE | $0.108\pm 0.050$ | $0.044\pm 0.005$ | $0.031\pm 0.001$
$\lambda$ | linear | $0.081\pm 0.045$ | $0.001\pm 0.001$ | $0.000\pm 0.000$
| MLP | $0.011\pm 0.001$ | $0.001\pm 0.000$ | $0.001\pm 0.000$
| CGVAE | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
$\textbf{RMSD}_{\textnormal{gen}}$ | linear | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
| MLP | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
| CGVAE | $0.326\pm 0.033$ | $0.087\pm 0.012$ | $0.051\pm 0.003$
| Valid graph ratio | linear | $0.561\pm 0.153$ | $0.990\pm 0.009$ | $1.000\pm 0.000$
|  | MLP | $0.908\pm 0.011$ | $0.992\pm 0.001$ | $0.993\pm 0.002$
|  | CGVAE | $1.000\pm 0.000$ | $1.000\pm 0.000$ | $1.000\pm 0.000$
All atoms | $\textbf{RMSD}_{\textnormal{recon}}$ | linear | $1.448\pm 0.071$ | $0.937\pm 0.005$ | $0.895\pm 0.006$
| MLP | $0.986\pm 0.002$ | $0.949\pm 0.003$ | $0.931\pm 0.008$
| CGVAE | $0.962\pm 0.040$ | $0.923\pm 0.002$ | $0.818\pm 0.028$
$\lambda$ | linear | $0.075\pm 0.037$ | $0.027\pm 0.003$ | $0.019\pm 0.005$
| MLP | $0.015\pm 0.001$ | $0.011\pm 0.001$ | $0.012\pm 0.002$
| CGVAE | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
$\textbf{RMSD}_{\textnormal{gen}}$ | linear | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
| MLP | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
| CGVAE | $1.000\pm 0.041$ | $0.924\pm 0.001$ | $0.772\pm 0.054$
| Valid graph ratio | linear | $0.230\pm 0.151$ | $0.667\pm 0.019$ | $0.802\pm 0.048$
|  | MLP | $0.732\pm 0.015$ | $0.808\pm 0.023$ | $0.790\pm 0.038$
|  | CGVAE | $0.990\pm 0.010$ | $0.998\pm 0.001$ | $0.999\pm 0.000$
Table 6: Chignolin benchmarks Metric | Model | N
---|---|---
6 | 10 | 20 | 30
heavy atom | $\textbf{RMSD}_{\textnormal{recon}}$ | linear | $2.723\pm 0.055$ | $1.845\pm 0.013$ | $1.077\pm 0.036$ | $0.702\pm 0.011$
MLP | $2.558\pm 0.026$ | $1.894\pm 0.013$ | $1.318\pm 0.030$ | $1.388\pm 0.024$
CGVAE | $1.188\pm 0.010$ | $0.897\pm 0.007$ | $0.658\pm 0.020$ | $0.445\pm 0.019$
$\lambda$ | linear | $1.124\pm 0.127$ | $0.411\pm 0.028$ | $0.085\pm 0.011$ | $0.021\pm 0.002$
MLP | $0.611\pm 0.009$ | $0.331\pm 0.008$ | $0.125\pm 0.008$ | $0.146\pm 0.010$
CGVAE | $0.027\pm 0.001$ | $0.006\pm 0.001$ | $0.002\pm 0.000$ | $0.001\pm 0.000$
$\textbf{RMSD}_{\textnormal{gen}}$ | linear | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
MLP | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
CGVAE | $1.642\pm 0.034$ | $1.231\pm 0.038$ | $0.762\pm 0.028$ | $0.502\pm 0.033$
valid graph ratio | linear | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.002\pm 0.002$ | $0.082\pm 0.022$
MLP | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
CGVAE | $0.415\pm 0.032$ | $0.716\pm 0.034$ | $0.843\pm 0.015$ | $0.955\pm 0.012$
all atom | $\textbf{RMSD}_{\textnormal{recon}}$ | linear | $3.008\pm 0.040$ | $2.163\pm 0.013$ | $1.382\pm 0.050$ | $0.902\pm 0.007$
MLP | $2.844\pm 0.026$ | $2.175\pm 0.020$ | $1.592\pm 0.031$ | $1.659\pm 0.027$
CGVAE | $1.574\pm 0.009$ | $1.353\pm 0.012$ | $1.081\pm 0.033$ | $0.795\pm 0.029$
$\lambda$ | linear | $1.160\pm 0.119$ | $0.368\pm 0.018$ | $0.052\pm 0.007$ | $0.031\pm 0.002$
MLP | $0.588\pm 0.006$ | $0.286\pm 0.005$ | $0.081\pm 0.005$ | $0.098\pm 0.011$
CGVAE | $0.020\pm 0.001$ | $0.010\pm 0.001$ | $0.008\pm 0.002$ | $0.010\pm 0.000$
$\textbf{RMSD}_{\textnormal{gen}}$ | linear | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
MLP | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
CGVAE | $1.914\pm 0.103$ | $1.586\pm 0.027$ | $1.175\pm 0.036$ | $0.877\pm 0.049$
valid graph ratio | linear | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
MLP | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$ | $0.000\pm 0.000$
CGVAE | $0.072\pm 0.027$ | $0.133\pm 0.024$ | $0.278\pm 0.099$ | $0.132\pm 0.011$

###  F.5 On Pseudoscalar Initialization for $N=3$

As discussed in the main text, for $N=3$, the vector basis sets are confined
in a plane, so it is required to initialize pseudoscalar update with non-zero
values to break the reflection symmetry. As it is shown in Fig. 11, the
decoder with pseudoscalar initialized with $0_{F}$ produces coordinates that
are confined in a plane, while a non-zero pseudoscalar initialization produces
proper geometries that span the 3D space. For the case $N>3$, it is less
likely for CG beads to be in the same plane. However, when it is needed, one
can always initialize pseudoscalar as non-zeros to break the symmetry even for
larger $N$ cases.

Figure 11: Comparing zero and non-zero pseudoscalar initializations for CGVAE
at the resolution of $N=3$ for alanine dipeptide.