conversion, as shown in Alg. 2.

Here we briefly explain the derivation in Alg. 2. The rigid-body assumption
states that for each residue $\imath$ with $N_{\imath}$ atoms, the relative
coordinates $\mathbf{X}_{\imath,\mathrm{rel}}\in\mathbb{R}^{N_{\imath}\times
3}$ of all its atoms w.r.t its alpha-carbon at
$\mathbf{C}_{\imath}\in\mathbb{R}^{1\times 3}$ in a standard coordinate system
is fixed. Under this assumption, if the coarse-grained representation of this
residue is $(\mathbf{C}_{\imath},\mathbf{Q}_{\imath})$ where
$\mathbf{Q}_{\imath}$ is the orientation of the residue relative to the
standard coordinate system, then the absolute coordinates of these atoms are:

$\displaystyle\mathbf{X}_{\imath}=\mathbf{C}_{\imath}+\mathbf{X}_{\imath,\mathrm{rel}}\mathbf{Q}_{\imath},$
(71)

and if considering a protein with $I$ residues, the full-atom coordinates are
$\bar{\mathbf{R}}:=\\{\mathbf{C}_{\imath}+\mathbf{X}_{\imath,\mathrm{rel}}\mathbf{Q}_{\imath}\\}_{\imath=1}^{I}$.
So in this way, we can convert the full-atom energy function
$E(\bar{\mathbf{R}})=E(\\{\mathbf{X}_{\imath}\\}_{\imath=1}^{I})$ as a
function of the coarse-grained coordinates
$\mathbf{R}=(\mathbf{C},\mathbf{Q})$ using Eq. (71). The gradient w.r.t
$\mathbf{C}_{\imath}$ is then:

$\displaystyle\mathbf{g}_{\mathbf{C}_{\imath}}$
$\displaystyle:=\nabla_{\mathbf{C}_{\imath}}E=(\nabla_{\mathbf{C}_{\imath}}\mathbf{X}_{\imath})^{\top}\nabla_{\mathbf{X}_{\imath}}E=\sum_{a=1}^{N_{\imath}}\nabla_{\mathbf{X}_{\imath,a}}E,$
(72)

where
$(\nabla_{\mathbf{C}_{\imath}}\mathbf{X}_{\imath})_{a\mu,\nu}:=\frac{\partial\mathbf{X}_{\imath,a,\mu}}{\partial\mathbf{C}_{\imath,\nu}}$
is the Jacobian matrix (here $\mu,\nu\in\\{1,2,3\\}$ indices the spacial
dimension), and the last equality holds since the Jacobian is
$(\nabla_{\mathbf{C}_{\imath}}\mathbf{X}_{\imath})_{a\mu,\nu}=\delta_{\mu\nu}$
meaning that this matrix element is one if $\mu=\nu$ or it is zero.

As for the gradient w.r.t the orientation, the $\mathfrak{so}(3)$
representation denoted as $\mathbf{q}$ is finally required. The conversion
from $\mathbf{q}$ and $\mathbf{Q}$ is given by Eq. (51). Together with the
rigid-body assumption Eq. (71), the gradient is:

$\displaystyle\mathbf{g}_{\mathbf{q}_{\imath}}$
$\displaystyle:=\nabla_{\mathbf{q}_{\imath}}E=(\nabla_{\mathbf{q}_{\imath}}\mathbf{X}_{\imath})^{\top}\nabla_{\mathbf{X}_{\imath}}E=(\nabla_{\mathbf{Q}_{\imath}}\mathbf{X}_{\imath}\nabla_{\mathbf{q}_{\imath}}\mathbf{Q}_{\imath})^{\top}\nabla_{\mathbf{X}_{\imath}}E$
(73)
$\displaystyle=\sum_{a=1}^{N_{\imath}}\mathbf{X}_{\imath,\mathrm{rel},a}^{\top}\nabla_{\mathbf{q}_{\imath}}\mathbf{Q}_{\imath}\nabla_{\mathbf{X}_{\imath,a}}E,$
(74)

where the last term means
$(\mathbf{g}_{\mathbf{q}_{\imath}})_{\gamma}=\sum_{a=1}^{N_{\imath}}\mathbf{X}_{\imath,\mathrm{rel},a}^{\top}\frac{\partial\mathbf{Q}_{\imath}}{\partial\mathbf{q}_{\imath,\gamma}}\nabla_{\mathbf{X}_{\imath,a}}E$,
where $\gamma\in\\{1,2,3\\}$ indices one of the three dimensions of
$\mathbf{q}_{\imath}$, and
$\frac{\partial\mathbf{Q}_{\imath}}{\partial\mathbf{q}_{\imath,\gamma}}$ is
the $3\times 3$ matrix composed of the partial derivatives from Eq. (51). In
the equation, again $\nabla_{\mathbf{q}_{\imath}}\mathbf{X}_{\imath}$,
$\nabla_{\mathbf{Q}_{\imath}}\mathbf{X}_{\imath}$ and
$\nabla_{\mathbf{q}_{\imath}}\mathbf{Q}_{\imath}$ are Jacobian matrices, and
the second last equality holds due to the chain rule of differentiation. From
the rigid-body assumption Eq. (71),
$(\nabla_{\mathbf{Q}_{\imath}}\mathbf{X}_{\imath})_{a\nu^{\prime},\mu\nu}=\frac{\partial\mathbf{X}_{\imath,a,\nu^{\prime}}}{\partial\mathbf{Q}_{\imath,\mu,\nu}}=\mathbf{X}_{\imath,\mathrm{rel},a,\mu}\delta_{\nu\nu^{\prime}}$,
so:

$\displaystyle[(\nabla_{\mathbf{Q}_{\imath}}\mathbf{X}_{\imath}\nabla_{\mathbf{q}_{\imath}}\mathbf{Q}_{\imath})^{\top}\nabla_{\mathbf{X}_{\imath}}E]_{\gamma}=\sum_{a,\nu^{\prime},\mu,\nu}(\nabla_{\mathbf{Q}_{\imath}}\mathbf{X}_{\imath})_{a\nu^{\prime},\mu\nu}(\nabla_{\mathbf{q}_{\imath}}\mathbf{Q}_{\imath})_{\mu\nu,\gamma}(\nabla_{\mathbf{X}_{\imath}}E)_{a\nu^{\prime}}$
(75) $\displaystyle={}$
$\displaystyle\sum_{a,\mu,\nu}\mathbf{X}_{\imath,\mathrm{rel},a,\mu}(\nabla_{\mathbf{q}_{\imath}}\mathbf{Q}_{\imath})_{\mu\nu,\gamma}(\nabla_{\mathbf{X}_{\imath}}E)_{a\nu}=\sum_{a}\mathbf{X}_{\imath,\mathrm{rel},a}^{\top}\frac{\partial\mathbf{Q}_{\imath}}{\partial\mathbf{q}_{\imath,\gamma}}\nabla_{\mathbf{X}_{\imath},a}E,$
(76)

which gives the last equality.

Moreover, to avoid costly divergence evaluation, we use Hutchinson’s trace
estimator in Alg. 3 to handle
$\nabla(\nabla\cdot\mathbf{s}^{\theta}_{\mathcal{D},t}(\mathbf{R}_{t}))$
(recall that
$\mathbf{s}^{\theta}_{\mathcal{D},t}(\mathbf{R}_{t})=(\mathbf{s}^{(\mathbf{C}),\theta}_{\mathcal{D},t}(\mathbf{C}_{t},\mathbf{q}_{t}),\mathbf{s}^{(\mathbf{q}),\theta}_{\mathcal{D},t}(\mathbf{C}_{t},\mathbf{q}_{t}))$)
in Eqs. (57, 63) .

Optimizing PIDP from random initialization of the deep learning model is
extremely hard due to the complex landscapes of the training objectives in
Eqs. (57, 63) . Therefore, a good initialization of the model and some
training techniques are necessary to stabilize the optimization of PIDP. To
this end, before performing PIDP, we train DiG on the experimental structures
and use it as a more stable initialization for PIDP. See Supplementary Sec. C
for more details.

Algorithm 4 Protein score model data-based training (single step)

1:Score model $\mathbf{s}^{\theta}_{\mathcal{D},t}(\mathbf{R})$ to be trained,
randomly sampled system $\mathcal{D}$, a randomly sampled coarse-grained
protein structure
$\mathbf{R}_{\mathcal{D},0}^{(n)}=(\mathbf{C}_{\mathcal{D},0}^{(n)},\mathbf{q}_{\mathcal{D},0}^{(n)})$
from the MD simulation data for system $\mathcal{D}$, randomly sampled time
step $t\in[0,\tau]$.

2:Sample $(\mathbf{C}_{\mathcal{D},t}^{(n)},\mathbf{q}_{\mathcal{D},t}^{(n)})$
from $(\mathbf{C}_{\mathcal{D},0}^{(n)},\mathbf{q}_{\mathcal{D},0}^{(n)})$
using Eq. (64) and Eq. (59);

3:Evaluate the data-based loss $\mathcal{L}_{\textnormal{data}}$ as Eq. (60) +
Eq. (65);

4:Update the model parameter $\theta$ by performing an optimization step on
$\mathcal{L}_{\text{data}}$ with respect to $\theta$.

Next, we pick about 1000 protein complexes in PDBbind database wang2005pdbbind
, and simulate them with GROMACS, together with about 200 proteins in GPCRmd
rodriguez2020gpcrmd  dataset to perform DiG training with simulation data (See
more details in Supplementary Sec. D.1). We further train the score model pre-
trained by PIDP by minimizing a score matching loss from Eq. (8) that directly
supervises the score model with the empirical data distribution that
approximates the equilibrium distribution, using protein structures
$\mathbf{R}_{\mathcal{D},0}^{(n)}$ for each system $\mathcal{D}$. See
Supplementary Sec. B and C for more details on the model and training.

Algorithm 5 Protein structure sampling

1:A trained score model $\mathbf{s}^{\theta}_{\mathcal{D},t}(\mathbf{R})$,
target protein system $\mathcal{D}$.

2:Initialize random structure
$\mathbf{R}_{\tau}:=(\mathbf{C}_{\tau},\mathbf{q}_{\tau})$, where
$\mathbf{C}_{\tau}\sim\mathcal{N}(\mathbf{0},\sigma_{\tau}^{2}\mathbf{I})$,
and $\mathbf{q}_{\tau}\sim p_{\mathrm{Unif}}$ on $\mathfrak{so}(3)$ defined in
Eq. (55). $t_{N}:=\tau$.

3:for $i$ in $N,\cdots,1$ do

4: $t_{i-1}:=\frac{i-1}{N}\tau$;

5:
$\mathbf{C}_{t_{i-1}}:=\mathbf{C}_{t_{i}}+\frac{1}{2}(\sigma_{t_{i}}^{2}-\sigma_{t_{i-1}}^{2})\mathbf{s}^{(\mathbf{C}),\theta}_{\mathcal{D},t_{i}}(\mathbf{C}_{t_{i}},\mathbf{q}_{t_{i}})$
(c.f. Eq. (67));

6: Construct $\mathbf{Q}_{t_{i}}$ from $\mathbf{q}_{t_{i}}$, and
$\mathbf{s}^{(\mathbf{Q}),\theta}_{\mathcal{D},t_{i}}(\mathbf{C}_{t_{i}},\mathbf{q}_{t_{i}})$
from
$\mathbf{s}^{(\mathbf{q}),\theta}_{\mathcal{D},t_{i}}(\mathbf{C}_{t_{i}},\mathbf{q}_{t_{i}})$,
using Eq. (51);

7:
$\mathbf{Q}_{t_{i-1}}:=\operatorname{Exp}\Big{(}\frac{1}{2}(\sigma_{t_{i}}^{2}-\sigma_{t_{i-1}}^{2})\mathbf{s}^{(\mathbf{Q}),\theta}_{\mathcal{D},t_{i}}(\mathbf{C}_{t_{i}},\mathbf{q}_{t_{i}})\Big{)}\mathbf{Q}_{t_{i}}$
(c.f. Eq. (68));

8:end for

9:Return the sampled structure
$\mathbf{R}_{0}:=(\mathbf{C}_{0},\mathbf{q}_{0})$.

After training, amino acid sequences serve as input descriptors, denoted as
$\mathcal{D}$, for sampling protein conformations via DiG. During the sampling
procedure, an initial random structure is generated, which subsequently is
transformed into a physically plausible conformation, as shown in Alg. 5.

###  A.3 Ligand Structure Sampling around Binding Sites

In contrast to the coarse-grained representation employed in protein
conformation sampling, we train DiG of ligand structure sampling with all-atom
(except hydrogens) representations, which offer a more precise description of
atomic interactions between the binding site (pocket) and ligand during the
sampling of ligand-binding structures with proteins. However, the time
complexity and memory usage associated with attention layers in Transformer-
based architectures exhibit a quadratic increase with respect to the number of
input nodes. This becomes impractical when the atom count surpasses one
thousand. Consequently, we restrict our model to incorporate only the atoms of
the ligand and the protein atoms in close proximity to the pocket, using a
distance threshold. This threshold is set to $10\,\textnormal{\AA}$ for the
side length.

DiG defines a distribution over the vector space
$\bar{\mathbf{R}}:=(\bar{\mathbf{R}}_{\textrm{Rec}},\bar{\mathbf{R}}_{\textrm{Lig}})$,
where $\bar{\mathbf{R}}_{\textrm{Rec}}$ and $\bar{\mathbf{R}}_{\textrm{Lig}}$
are the absolute coordinates of the near-site receptor and ligand atoms,
respectively. Since the receptor atom coordinates may have different
distribution centers for different proteins while the diffusion process always
starts from a zero-centered Gaussian, we use the near-site receptor atom
coordinates $\bar{\mathbf{R}}_{\textrm{Rec}}^{\star}$ from the crystal
structure of the protein in the PDBBind database (wang2005pdbbind, ) as a
reference structure, and let the model predict the residue. This effectively
shifts the diffusion process to
$\mathbf{R}:=(\bar{\mathbf{R}}_{\textrm{Rec}}-\bar{\mathbf{R}}_{\textrm{Rec}}^{\star},\bar{\mathbf{R}}_{\textrm{Lig}})$.
DiG then generates the binding structures for the ligand and the protein
pocket by reversing the diffusion process, as shown in Eq. (5).

To train DiG, we reuse the simulation data of protein complexes in PDBbind, as
detailed in Supplementary Sec. C. The data preprocessing and featurization
follows min2022predicting ; he2022masked . The atom representations are
further embedded into real-valued embedding vectors for a Graphormer
ying2021transformers . In this context, performing PIDP is not feasible due to
limitations imposed by the energy function. Specifically, DiG only considers
atoms surrounding the binding site, while conventional force fields
necessitate the inclusion of all atoms. As such, we sought to enhance ligand
sampling performance within the pocket by conducting a pre-training task
focused on binding structure prediction, utilizing the CrossDocked dataset
francoeur2020three . Further information regarding model architecture and
training can be found in Supplementary Sec. B and C.

###  A.4 Catalyst-Adsorbate Sampling

For catalyst-adsorbate sampling, DiG adopts the same input representation
strategy employed in the OC20 dataset chanussot2021open , employing the
descriptor $\mathcal{D}$ to characterize the system. Specifically, besides the
atom types $\mathcal{Z}$, also provided from the OC20 dataset are the absolute
coordinates $\bar{\mathbf{R}}^{\star}_{\textnormal{base}}$ for non-surface
catalyst atoms, $\bar{\mathbf{R}}^{\star}_{\mathrm{Cat}}$ for surface catalyst
atoms, and the initial absolute coordinates
$\bar{\mathbf{R}}^{\star}_{\mathrm{Ad}}$ for adsorbate atoms prior to
relaxation. Consequently, the system descriptor for this task was defined as
$\mathcal{D}:=(\mathcal{Z},\bar{\mathbf{R}}^{\star}_{\textnormal{base}},\bar{\mathbf{R}}^{\star}_{\mathrm{Cat}},\bar{\mathbf{R}}^{\star}_{\mathrm{Ad}})$.
The microscopic state of the system
$\bar{\mathbf{R}}:=(\bar{\mathbf{R}}_{\mathrm{Cat}},\bar{\mathbf{R}}_{\mathrm{Ad}})$
encompasses the absolute coordinates $\bar{\mathbf{R}}_{\mathrm{Cat}}$ of
surface catalyst atoms and $\bar{\mathbf{R}}_{\mathrm{Ad}}$ of the adsorbate
atoms. Similar to the ligand-receptor sampling case, to ease the prediction of
different distribution centers for different systems, we leverage the absolute
coordinates in the descriptor to define the diffusion-process variable as
relative coordinates, i.e.,
$\mathbf{R}:=(\mathbf{R}_{\mathrm{Cat}},\mathbf{R}_{\mathrm{Ad}})$ where
$\mathbf{R}_{\mathrm{Cat}}:=\bar{\mathbf{R}}_{\mathrm{Cat}}-\bar{\mathbf{R}}^{\star}_{\mathrm{Cat}}$
and
$\mathbf{R}_{\mathrm{Ad}}:=\bar{\mathbf{R}}_{\mathrm{Ad}}-\bar{\mathbf{R}}^{\star}_{\mathrm{Ad}}$,
whose distribution center is largely aligned across different systems.

During the reverse diffusion process, including the initial structure
$\bar{\mathbf{R}}^{\star}_{\mathrm{Ad}}$ into the model input is found crucial
for stable training. In the Graphormer model, in addition to the diffusion-
variable $\mathbf{R}$ in the input, the initial structure
$\bar{\mathbf{R}}_{\mathrm{Ad}}^{\star}$ is also encoded as an additional
structural attention bias term, as in shi2022benchmarking . More specifically,
the pairwise distance between atoms in
$\bar{\mathbf{R}}^{\star}_{\mathrm{Ad}}$ was calculated, which is then encoded
into a $K$-dimensional feature using $K$ radial basis function (RBF) kernels
with learnable means and variances. Likewise, the initial positions of the
atoms are encoded as extra node features and incorporated into the node
representation. By summing the pairwise features and aggregating them with the
node features, per-atom features are updated and projected into the atom
embedding dimension within the model. Further details regarding the structural
attention bias can be found in Supplementary Sec. B.3.

Catalyst systems in OC20 are inorganic and lack bond information between
atoms. However, adsorbates are organic, and the bonds between atoms in
adsorbates can benefit the model in generating more physically accurate
adsorbate structures. Consequently, we explicitly incorporated the 2D topology
of adsorbates, featuring bonds connecting the atoms, within the model. Bonds
are generated using the initial structure of the adsorbate and encoded in the
same manner as Graphormer ying2021transformers . The spatial encoding,
centrality encoding, and edge encoding of Graphormer are utilized alongside
the encodings derived from 3D information. It is important to note that, in
some instances, bonds in adsorbates may break upon adsorption to the catalyst
surface. Explicitly encoding bond information does not imply the enforcement
of bonded atoms to remain close in the sampled structures. Instead, the model
is allowed to learn when to separate two bonded atoms.

The training and sampling of DiG adhere to the general description in the main
text (Sec. 2). The training loss is based on Eq. (40), and Alg. 6 shows the
detailed training process. For sampling new structures using DiG, in
accordance with the training process, it is also conducted on the relative
coordinates w.r.t the initial structure. Alg. 7 outlines the sampling process.

Algorithm 6 Catalyst-adsorbate score model training (single step)

1:Noise-predicting model
$\mathbf{\epsilon}^{\theta}_{\mathcal{D},t}(\mathbf{R})$ to be trained,
randomly sampled time step $t\in[0,\tau]$, randomly sampled system with
descriptor
$\mathcal{D}=(\mathcal{Z},\bar{\mathbf{R}}^{\star}_{\textnormal{base}},\bar{\mathbf{R}}^{\star}_{\mathrm{Cat}},\bar{\mathbf{R}}^{\star}_{\mathrm{Ad}})$,
let
$\bar{\mathbf{R}}^{\star}:=\left(\bar{\mathbf{R}}^{\star}_{\mathrm{Cat}},\bar{\mathbf{R}}^{\star}_{\mathrm{Ad}}\right)$,
randomly sampled catalyst-adsorbent structure
$\bar{\mathbf{R}}_{\mathcal{D},0}^{(n)}=(\bar{\mathbf{R}}_{\mathcal{D},\mathrm{Cat},0}^{(n)},\bar{\mathbf{R}}_{\mathcal{D},\mathrm{Ad},0}^{(n)})$
from the MD simulation data for this system $\mathcal{D}$.

2:Let
$\mathbf{R}_{\mathcal{D},0}^{(n)}:=\bar{\mathbf{R}}_{\mathcal{D},0}^{(n)}-\bar{\mathbf{R}}^{\star}$;

3:Sample noise variable
$\mathbf{\epsilon}_{t}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ in the same
dimension as $\mathbf{R}_{\mathcal{D},0}^{(n)}$;

4:$\mathbf{R}^{(n)}_{\mathcal{D},t}:=\alpha_{t}\mathbf{R}_{\mathcal{D},0}^{(n)}+\sqrt{1-\alpha_{t}^{2}}\mathbf{\epsilon}_{t}$
(c.f. Eq. (38));

5:Evaluate the loss
$\left\lVert\mathbf{\epsilon}^{\theta}_{\mathcal{D},t}(\mathbf{R}^{(n)}_{\mathcal{D},t}+\bar{\mathbf{R}}^{\star})-\mathbf{\epsilon}_{t}\right\rVert^{2}$
(c.f. Eq. (40));

6:Update the model parameter $\theta$ by performing an optimization step on
the loss with respect to $\theta$.

Algorithm 7 Catalyst-adsorbate structure sampling

1:A trained noise-predicting model
$\mathbf{\epsilon}^{\theta}_{\mathcal{D},t}(\mathbf{R})$, the descriptor
$\mathcal{D}=(\mathcal{Z},\bar{\mathbf{R}}^{\star}_{\textnormal{base}},\bar{\mathbf{R}}^{\star}_{\mathrm{Cat}},\bar{\mathbf{R}}^{\star}_{\mathrm{Ad}})$
of the target system, let
$\bar{\mathbf{R}}^{\star}:=\left(\bar{\mathbf{R}}^{\star}_{\mathrm{Cat}},\bar{\mathbf{R}}^{\star}_{\mathrm{Ad}}\right)$.

2:Sample a noisy structure $\mathbf{R}_{\tau}\sim
p_{\textnormal{simple}}=\mathcal{N}(\mathbf{0},\mathbf{I})$;

3:for $i$ in $N,\cdots,1$ do

4: $t_{i-1}:=\frac{i-1}{N}\tau$;

5: Let $\beta_{i}:=\frac{\tau}{N}\beta_{t_{i}}$,
$\alpha_{i}:=\prod_{j=1}^{i}\sqrt{1-\beta_{j}}$;

6: Sample a noise variable
$\mathbf{\epsilon}_{t_{i}}\sim\mathcal{N}(\mathbf{0},\beta_{i}\mathbf{I})$ in
the same dimension as $\mathbf{R}_{t_{i}}$;

7:
$\mathbf{R}_{t_{i-1}}\\!:=\\!\frac{1}{\sqrt{1-\beta_{i}}}\Big{(}\\!\mathbf{R}_{t_{i}}-\frac{\beta_{i}}{\sqrt{1-\alpha_{i}^{2}}}\mathbf{\epsilon}^{\theta}_{\mathcal{D},t_{i}}(\mathbf{R}_{t_{i}}+\bar{\mathbf{R}}^{\star})\\!\Big{)}+\mathbf{\epsilon}_{t_{i}}$
(c.f. Eqs. (5, 39));

8:end for

9:Return $\mathbf{R}_{0}+\bar{\mathbf{R}}^{\star}$.

As catalysts are periodic systems along the $x$ and $y$ directions, DiG
expands the unit cell in these dimensions before feeding a system into the
model. Providing an exact descriptor of the infinitely repeated system in the
model is non-trivial; instead, DiG adopts a simple yet effective approach by
establishing a local cutoff for the infinitely repeating system. Specifically,
an atom outside the unit cell is included in the model only if its distance to
any atom inside the unit cell is within a threshold. A distance threshold of
$6\,\textnormal{\AA}$ is used in the experiments. In each layer of the
transformer model, the representation of a repeated atom outside the unit cell
is enforced to be identical to the representation of the corresponding atom in
the unit cell. More details about handling periodic boundary conditions can be
found in Supplementary Sec. B.5.

Lastly, to ensure a good initialization for the diffusion model, a model pre-
trained on the IS2RS task of OC20 is used to initialize the weights, except
for the time step embedding, which is dedicated to the diffusion task.

###  A.5 Property-Guided Structure Generation

For modeling carbon polymorphs, the structural variable needs to represent the
unit cell which defines a spatial period in the crystal, and the absolute
coordinates $\bar{\mathbf{X}}$ of the carbon atoms in the unit cell. The unit
cell is a parallelepiped to guarantee periodicity, so it can be determined by
the coordinates of 4 non-coplanar vertices, say $\bar{\mathbf{c}}_{0}$,
$\bar{\mathbf{c}}_{0}+\bar{\mathbf{l}}_{x}$,
$\bar{\mathbf{c}}_{0}+\bar{\mathbf{l}}_{y}$, and
$\bar{\mathbf{c}}_{0}+\bar{\mathbf{l}}_{z}$, where $\bar{\mathbf{c}}_{0}$ is
the origin of the unit cell, and
$\bar{\mathbf{L}}:=\\{\bar{\mathbf{l}}_{x},\bar{\mathbf{l}}_{y},\mathbf{l}_{z}\\}$
are known as the lattice vectors. (The locations of the other 4 vertices can
be determined as
$\bar{\mathbf{c}}_{0}+\bar{\mathbf{l}}_{x}+\bar{\mathbf{l}}_{y}$,
$\bar{\mathbf{c}}_{0}+\bar{\mathbf{l}}_{x}+\bar{\mathbf{l}}_{z}$,
$\bar{\mathbf{c}}_{0}+\bar{\mathbf{l}}_{y}+\bar{\mathbf{l}}_{z}$, and
$\bar{\mathbf{c}}_{0}+\bar{\mathbf{l}}_{x}+\bar{\mathbf{l}}_{y}+\bar{\mathbf{l}}_{z}$.)
In the structure representation, the origin $\bar{\mathbf{c}}_{0}$ is fixed as
$\mathbf{0}$, so the unit cell is fully determined by the lattice vectors
$\bar{\mathbf{L}}$.

Similar to the cases of protein-ligand sampling in Supplementary Sec. A.3 and
catalyst-adsorbate sampling in Supplementary Sec. A.4, we take the diffusion-
process variable $\mathbf{R}$ as relative coordinates to release the burden of
the diffusion model to also predict the distribution center for different
systems. For the lattice vectors, we introduce a reference lattice vector set
$\bar{\mathbf{L}}^{\star}$ taken as the mean lattice vector set on the
dataset, and diffuse the relative vectors
$\mathbf{L}:=\bar{\mathbf{L}}-\bar{\mathbf{L}}^{\star}$. For the carbon-atom
coordinates, we use their relative coordinates w.r.t the unit cell center
$\bar{\mathbf{z}}:=\frac{1}{2}(\bar{\mathbf{l}}_{x}+\bar{\mathbf{l}}_{y}+\bar{\mathbf{l}}_{z})$,
which means $\mathbf{X}:=\bar{\mathbf{X}}-\bar{\mathbf{z}}$. The diffusion-
process variable is then defined as $\mathbf{R}:=(\mathbf{X},\mathbf{L})$. The
diffusion-variable part of the input to the diffusion model still requires the
corresponding absolute coordinates $(\mathbf{X},\mathbf{L})$ of $\mathbf{R}$.
For the descriptor part of the input, in addition to the number of carbon
atoms $N_{\mathrm{atom}}$, we also include the reference lattice vectors
$\bar{\mathbf{L}}^{\star}$, which is an informative feature similar to the
discussion in Supplementary Sec. A.4.

The DiG model first learns the (unconditional) distribution of the structures
of carbon polymorphs from a dataset created by random structure search (RSS).
We take the noise-prediction form to learn the model (c.f. Eq. (39)). The
training process is detailed in Alg. 8. The DiG model is then asked to
generate structure samples conditioned on a given property value, specifically
a desired band gap value $c$ in our case. According to Eq. (12), this requires
a property predictor/classifier. For this, we use a GNN model M3GNet
chen2022universal , which provides the prediction for the band gap of a given
structure in absolute coordinates. To evaluate a probability, we convert the
regression task into a classification task by discretizing an inclusive range
of the band gap value into $K$ intervals of length $1.0$, represented by
$\mathcal{I}_{0}=[a_{0},b_{0}],\cdots,\mathcal{I}_{K-1}=[a_{K-1},b_{K-1}]$.
The property $c$ is then taken as the bin index. Using the predicted band gap
value from M3GNet, we define the probability of a given $c$ as:

$\displaystyle
q_{\mathcal{D}}(c\mid\bar{\mathbf{X}},\bar{\mathbf{L}}):=\frac{\exp\left\lparen-\left\lvert\mathrm{M3GNet}(\bar{\mathbf{X}},\bar{\mathbf{L}})-\frac{a_{c}+b_{c}}{2}\right\rvert\right\rparen}{\sum_{k=0}^{K-1}\exp\left\lparen-\left\lvert\mathrm{M3GNet}(\bar{\mathbf{X}},\bar{\mathbf{L}})-\frac{a_{k}+b_{k}}{2}\right\rvert\right\rparen}.$
(77)

This equation is used to construct the required conditional score
$\mathbf{s}^{\theta}_{\mathcal{D},t}(\bar{\mathbf{X}}_{t},\bar{\mathbf{L}}_{t}\mid
c)$ following Eq. (12). The value $\frac{a_{c}+b_{c}}{2}$ is the target band
gap for the interval $\mathcal{I}_{c}$.

Algorithm 8 Carbon structure score model training (single step)

1:Noise-predicting model
$\mathbf{\epsilon}^{\theta}_{\mathcal{D},t}(\bar{\mathbf{X}},\bar{\mathbf{L}})$
to be trained, descriptor
$\mathcal{D}=(N_{\mathrm{atom}},\bar{\mathbf{L}}^{\star})$, where
$N_{\mathrm{atom}}$ is the number of carbon atoms, and
$\bar{\mathbf{L}}^{\star}=\\{\bar{\mathbf{l}}^{\star}_{x},\bar{\mathbf{l}}^{\star}_{y},\bar{\mathbf{l}}^{\star}_{z}\\}$
is the reference lattice vectors; randomly sampled time step $t\in[0,\tau]$,
randomly sampled structure
$(\bar{\mathbf{X}}_{\mathcal{D},0}^{(n)},\bar{\mathbf{L}}_{\mathcal{D},0}^{(n)})$
from the dataset, where $\bar{\mathbf{X}}_{\mathcal{D},0}^{(n)}$ comprises
coordinates of carbon atoms, and
$\bar{\mathbf{L}}_{\mathcal{D},0}^{(n)}=\\{\bar{\mathbf{l}}_{\mathcal{D},0,x}^{(n)},\bar{\mathbf{l}}_{\mathcal{D},0,y}^{(n)},\bar{\mathbf{l}}_{\mathcal{D},0,z}^{(n)}\\}$
is the collection of lattice vectors.

2:Calculate the unit cell center
$\bar{\mathbf{z}}_{\mathcal{D},0}^{(n)}:=\frac{1}{2}(\bar{\mathbf{l}}_{\mathcal{D},0,x}^{(n)}+\bar{\mathbf{l}}_{\mathcal{D},0,y}^{(n)}+\bar{\mathbf{l}}_{\mathcal{D},0,z}^{(n)})$;

3:Let
$\mathbf{X}_{\mathcal{D},0}^{(n)}:=\bar{\mathbf{X}}_{\mathcal{D},0}^{(n)}-\bar{\mathbf{z}}_{\mathcal{D},0}^{(n)}$,
$\mathbf{L}_{\mathcal{D},0}^{(n)}:=\bar{\mathbf{L}}_{\mathcal{D},0}^{(n)}-\bar{\mathbf{L}}^{\star}$;

4:Sample the noise variable
$\mathbf{\epsilon}^{(\mathbf{X})}_{t}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{3N_{\mathrm{atom}}})$
and
$\mathbf{\epsilon}^{(\mathbf{L})}_{t}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{9})$;

5:Let
$\mathbf{X}_{\mathcal{D},t}^{(n)}:=\alpha_{t}\mathbf{X}_{\mathcal{D},0}^{(n)}+\sqrt{1-\alpha_{t}^{2}}\mathbf{\epsilon}^{(\mathbf{X})}_{t}$,
and
$\mathbf{L}_{\mathcal{D},t}^{(n)}:=\alpha_{t}\mathbf{L}_{\mathcal{D},0}^{(n)}+\sqrt{1-\alpha_{t}^{2}}\mathbf{\epsilon}^{(\mathbf{L})}_{t}$
(c.f. Eq. (38));

6:Let
$\bar{\mathbf{L}}_{\mathcal{D},t}^{(n)}:=\mathbf{L}_{\mathcal{D},t}^{(n)}+\bar{\mathbf{L}}^{\star}$
which is structured as
$\\{\bar{\mathbf{l}}_{\mathcal{D},t,x}^{(n)},\bar{\mathbf{l}}_{\mathcal{D},t,y}^{(n)},\bar{\mathbf{l}}_{\mathcal{D},t,z}^{(n)}\\}$;

7:Let
$\bar{\mathbf{z}}_{\mathcal{D},t}^{(n)}:=\frac{1}{2}(\bar{\mathbf{l}}_{\mathcal{D},t,x}^{(n)}+\bar{\mathbf{l}}_{\mathcal{D},t,y}^{(n)}+\bar{\mathbf{l}}_{\mathcal{D},t,z}^{(n)})$
and
$\bar{\mathbf{X}}_{\mathcal{D},t}^{(n)}:=\mathbf{X}_{\mathcal{D},t}^{(n)}+\bar{\mathbf{z}}_{\mathcal{D},t}^{(n)}$;

8:Evaluate the loss
$\left\lVert\mathbf{\epsilon}^{\theta}_{\mathcal{D},t}(\bar{\mathbf{X}}_{\mathcal{D},t}^{(n)},\bar{\mathbf{L}}_{\mathcal{D},t}^{(n)})-(\mathbf{\epsilon}^{(\mathbf{X})}_{t},\mathbf{\epsilon}^{(\mathbf{L})}_{t})\right\rVert^{2}$
(c.f. Eq. (40));

9:Update the model parameter $\theta$ by performing an optimization step on
the loss with respect to $\theta$.

Algorithm 9 Sampling for carbon structure inverse design

1:A trained noise-predicting model
$\mathbf{\epsilon}^{\theta}_{\mathcal{D},t}(\bar{\mathbf{X}},\bar{\mathbf{L}})$
which decomposes as
$(\mathbf{\epsilon}^{(\mathbf{X}),\theta}_{\mathcal{D},t},\mathbf{\epsilon}^{(\mathbf{L}),\theta}_{\mathcal{D},t})$
according to its output channels, descriptor
$\mathcal{D}=(N_{\mathrm{atom}},\bar{\mathbf{L}}^{\star})$ of the target
system, where $N_{\mathrm{atom}}$ is the number of carbon atoms, and
$\bar{\mathbf{L}}^{\star}=\\{\bar{\mathbf{l}}^{\star}_{x},\bar{\mathbf{l}}^{\star}_{y},\bar{\mathbf{l}}^{\star}_{z}\\}$
is the reference lattice vectors; a trained property classifier
$q_{\mathcal{D}}(c\mid\bar{\mathbf{X}},\bar{\mathbf{L}})$, the desired
property value $c$, guidance strength $\lambda_{\mathrm{guide}}$.

2:Sample a noisy initial structure
$\mathbf{X}_{\tau}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{3N_{\mathrm{atom}}})$,
$\mathbf{L}_{\tau}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{9})$;

3:for $i$ in $N,\cdots,1$ do

4: $t_{i-1}:=\frac{i-1}{N}\tau$;

5: Let $\bar{\mathbf{L}}_{t_{i}}:=\mathbf{L}_{t_{i}}+\bar{\mathbf{L}}^{\star}$
which is structured as
$\\{\bar{\mathbf{l}}_{t_{i},x},\bar{\mathbf{l}}_{t_{i},y},\bar{\mathbf{l}}_{t_{i},z}\\}$;

6: Calculate the unit cell center
$\bar{\mathbf{z}}_{t_{i}}:=\frac{1}{2}(\bar{\mathbf{l}}_{t_{i},x}+\bar{\mathbf{l}}_{t_{i},y}+\bar{\mathbf{l}}_{t_{i},z})$;

7: Let $\bar{\mathbf{X}}_{t_{i}}:=\mathbf{X}_{t_{i}}+\bar{\mathbf{z}}_{t_{i}}$
which is structured as
$\\{\bar{\mathbf{X}}_{t_{i},a}\\}_{a=1}^{N_{\mathrm{atom}}}$;

8: Let $\beta_{i}:=\frac{\tau}{N}\beta_{t_{i}}$,
$\alpha_{i}:=\prod_{j=1}^{i}\sqrt{1-\beta_{j}}$;

9:
$\mathbf{X}_{t_{i-1}}:=(2-\sqrt{1-\beta_{i}})\mathbf{X}_{t_{i}}-\frac{\beta_{i}}{2\sqrt{1-\alpha_{i}^{2}}}\mathbf{\epsilon}^{(\mathbf{X}),\theta}_{\mathcal{D},t_{i}}(\bar{\mathbf{X}}_{t_{i}},\bar{\mathbf{L}}_{t_{i}})+\lambda_{\mathrm{guide}}\frac{\beta_{i}}{2}\nabla_{\bar{\mathbf{X}}_{t_{i}}}\log
q_{\mathcal{D}}(c\mid\bar{\mathbf{X}}_{t_{i}},\bar{\mathbf{L}}_{t_{i}})$ (c.f.
Eqs. (80, 39, 12));

10:
$\mathbf{L}_{t_{i-1}}=(2-\sqrt{1-\beta_{i}})\mathbf{L}_{t_{i}}-\frac{\beta_{i}}{2\sqrt{1-\alpha_{i}^{2}}}\mathbf{\epsilon}^{(\mathbf{L}),\theta}_{\mathcal{D},t_{i}}(\bar{\mathbf{X}}_{t_{i}},\bar{\mathbf{L}}_{t_{i}})+\lambda_{\mathrm{guide}}\frac{\beta_{i}}{2}\big{(}\nabla_{\bar{\mathbf{L}}_{t_{i}}}\log
q_{\mathcal{D}}(c\mid\bar{\mathbf{X}}_{t_{i}},\bar{\mathbf{L}}_{t_{i}})+\frac{1}{2}\sum_{a=1}^{N_{\mathrm{atom}}}\nabla_{\bar{\mathbf{X}}_{t_{i},a}}\log
q_{\mathcal{D}}(c\mid\bar{\mathbf{X}}_{t_{i}},\bar{\mathbf{L}}_{t_{i}})\big{)}$
(c.f. Eqs. (80, 39, 12));

11:end for

12:Let $\bar{\mathbf{L}}_{0}:=\mathbf{L}_{0}+\bar{\mathbf{L}}^{\star}$ which
is structured as
$\\{\bar{\mathbf{l}}_{0,x},\bar{\mathbf{l}}_{0,y},\bar{\mathbf{l}}_{0,z}\\}$;

13:Calculate the unit cell center
$\bar{\mathbf{z}}_{0}:=\frac{1}{2}(\bar{\mathbf{l}}_{0,x}+\bar{\mathbf{l}}_{0,y}+\bar{\mathbf{l}}_{0,z})$;

14:Let $\bar{\mathbf{X}}_{0}:=\mathbf{X}_{0}+\bar{\mathbf{z}}_{0}$;

15:Return $(\bar{\mathbf{X}}_{0},\bar{\mathbf{L}}_{0})$.

The sampling process starts from a standard-Gaussian sample as $\mathbf{R}$,
and in each step $\mathbf{R}$ is converted to absolute coordinates using
$\bar{\mathbf{L}}^{\star}$ and the corresponding unit cell center for the
input to the model, and finally outputs the structure sample in absolute
coordinates. Similar to the case of protein structure sampling as explained at
the end of Supplementary Sec. A.2.1, for simulating the sampling process,
instead of simulating the SDE in the fashion of Eq. (5), it achieves better
results to simulate the equivalent ODE in Eq. (11). This is explained in
Supplementary Sec. A.1.3. Under discretization, Eq. (11) is written as:

$\displaystyle\mathbf{R}_{t_{i-1}}=\mathbf{R}_{t_{i}}+\frac{\beta_{i}}{2}\Big{\lparen}\mathbf{R}_{t_{i}}+\mathbf{s}^{\theta}_{\mathcal{D},t_{i}}(\bar{\mathbf{R}}_{t_{i}}\mid
c)\Big{\rparen},$ (78)

where $h$ is the discretization step size,
$\bar{\mathbf{R}}_{t}:=(\bar{\mathbf{X}}_{t},\bar{\mathbf{L}}_{t})$ is the
absolute coordinates corresponding to $\mathbf{R}_{t}$, and the plus sign is
due to the simulation is reversed in time. Also recall
$\beta_{i}:=h\beta_{t_{i}}$. Since $h$ hence $\beta_{i}$ is an infinitesimal,
the weight for the $\mathbf{R}_{t_{i}}$ term can be formulated as:
$1+\frac{\beta_{i}}{2}=2-(1-\frac{\beta_{i}}{2})=2-\sqrt{1-\beta_{i}}+o(h)$,
which gives an alternative up to $o(h)$ which is acceptable since the
disretization itself has $o(h)$ error. This then recovers the ODE simulation
in (song2021score, ). By leveraging Eq. (12), the simulation step becomes:

$\displaystyle\mathbf{R}_{t_{i-1}}=(2-\sqrt{1-\beta_{i}})\mathbf{R}_{t_{i}}+\frac{\beta_{i}}{2}\Big{\lparen}\mathbf{s}^{\theta}_{\mathcal{D},t_{i}}(\bar{\mathbf{R}}_{t_{i}})+\nabla_{\mathbf{R}_{t_{i}}}\log
q_{\mathcal{D}}(c\mid\bar{\mathbf{R}}_{t_{i}})\Big{\rparen},$ (79)

where $q_{\mathcal{D}}(c\mid\bar{\mathbf{R}}_{t_{i}})$ is given by Eq. (77).
Finally, by involing Eq. (39), we can express the generation process using the
noise-prediction model $\mathbf{\epsilon}^{\theta}_{\mathcal{D},t}$:

$\displaystyle\mathbf{R}_{t_{i-1}}=(2-\sqrt{1-\beta_{i}})\mathbf{R}_{t_{i}}-\frac{\beta_{i}}{2\sqrt{1-\alpha_{i}^{2}}}\mathbf{\epsilon}^{\theta}_{\mathcal{D},t_{i}}(\bar{\mathbf{R}}_{t_{i}})+\frac{\beta_{i}}{2}\nabla_{\mathbf{R}_{t_{i}}}\log
q_{\mathcal{D}}(c\mid\bar{\mathbf{R}}_{t_{i}}).$ (80)

Note that the input to the $q_{\mathcal{D}}(c\mid\bar{\mathbf{R}}_{t_{i}})$
model is absolute coordinates while the gradient is taken w.r.t the relative
coordinates $\mathbf{R}_{t_{i}}$. To conduct this conversion, note that the
conversion is $\bar{\mathbf{L}}=\mathbf{L}+\bar{\mathbf{L}}^{\star}$, and
$\bar{\mathbf{X}}=\mathbf{X}+\bar{\mathbf{z}}$ where
$\bar{\mathbf{z}}:=\frac{1}{2}(\bar{\mathbf{l}}_{x}+\bar{\mathbf{l}}_{y}+\bar{\mathbf{l}}_{z})$
with
$\bar{\mathbf{L}}=\\{\bar{\mathbf{l}}_{x},\bar{\mathbf{l}}_{y},\bar{\mathbf{l}}_{z}\\}$.
So $\nabla_{\mathbf{X}}\log
q_{\mathcal{D}}(c\mid\bar{\mathbf{X}},\bar{\mathbf{L}})=\nabla_{\bar{\mathbf{X}}}\log
q_{\mathcal{D}}(c\mid\bar{\mathbf{X}},\bar{\mathbf{L}})$, and
$\nabla_{\mathbf{L}}\log
q_{\mathcal{D}}(c\mid\bar{\mathbf{X}},\bar{\mathbf{L}})=\nabla_{\bar{\mathbf{L}}}\log
q_{\mathcal{D}}(c\mid\bar{\mathbf{X}},\bar{\mathbf{L}})+\frac{1}{2}\sum_{a=1}^{N_{\mathrm{atom}}}\nabla_{\bar{\mathbf{X}}_{a}}\log
q_{\mathcal{D}}(c\mid\bar{\mathbf{X}},\bar{\mathbf{L}})$. In practice, the
strength of the property guidance can be tuned for better performance, as is a
common practice in machine learning ho2022classifier . We hence introduce a
$\lambda_{\mathrm{guide}}$ parameter. Alg. 9 details the sampling process.

It is important to note that, in the conditional generation, the classifier
can be fully decoupled from the training of the diffusion model. As a result,
our model can serve to generate any demanded properties when provided with a
corresponding property prediction model. This flexibility allows the approach
to be adapted for various applications and desired properties, given an
appropriate predictive model.

###  A.6 Accelerating Inference

The inference procedure of DiG can be viewed as gradually removing the noise
from Gaussian random variables to sample clear conformations in an
approximated equilibrium distribution. Although it exhibits several orders of
magnitude speedup for the cases in Sec. 3 compared to traditional simulation
methods, the inference is still potentially expensive since it generally needs
to go over all time steps. Fortunately, the inference time can be further
saved with recently developed methods songdenoising ; baoanalytic ; ludpm ;
lu2022dpm ; karraselucidating . For example, Ref. (ludpm, ) shows that the
analytic solution of the diffusion ordinary differential equations (sampling
of DiG can be alternatively viewed as solving the corresponding diffusion
ordinary differential equations) can significantly accelerate the inference,
where high-quality samples can be drawn in around 10 steps, resulting in a
further 50 to 100 folds speedup. These recent advances demonstrate the
potential for even more efficient inference procedures in the context of DiG
and similar models. By combining these methods with the existing framework, it
becomes increasingly feasible to generate desired structures with specific
properties in a faster and more efficient manner.

###  A.7 Evaluation methods

##### Protein Conformation Sampling

As detailed in Sec. 3.1, this study aims to demonstrate the applicability of
using the DiG method to sample protein distributions. To this end, simulations
of two proteins from the SARS-CoV-2 virus are used to approximate their actual
distributions in equilibrium states. The molecular dynamics (MD) simulation
trajectories of the receptor-binding domain (RBD) of the spike protein and the
main protease are extracted from a public
dataset222https://covid.molssi.org/simulations/. For RBD, there are 2995
independent MD simulations with 1.8 ms of aggregate simulation time. For main
protease, the aggregated simulation time is 2.6 ms from 5688 independent
trajectories. Both sets of simulations are performed in the constant pressure
and temperature (NPT) conditions at 310 K and 1 atm pressure.

To facilitate the comparison between distributions obtained from MD
simulations and DiG generations, time-lagged independent component analysis
(TICA) is utilized to project the simulated structures onto a low-dimensional
manifold (perez2013identification, ; schwantes2013improvements, ). This
projection enables the visualization of the conformational distribution in the
low-dimension space, such as the one spanned by the two TICA coordinates (Fig.
3a). The TICA projection analysis is carried out as the following: first, the
backbone conformation is featurized with the cosine and sine values of
backbone torsion angles; then standard TICA projection is executed using
PyEmma (scherer_pyemma_2015, ), with lag times of 10 ns and 2 ns for RBD and
main protease respectively. We first parametrized the TICA transformation
matrix to MD simulation structures to obtain the probability distribution as
references, then the same transformation is applied to structures generated by
DiG. From MD simulation trajectories, about 1.8 million structures of each
system are used for TICA analysis. Furthermore, to reduce the influence of
disordered regions at protein termini, the terminal residues are excluded
during this TICA analysis and projection to the low-dimensional conformational
space. For the distribution comparison in the reduced 2D space, we focus on
the populated regions, which correspond to metastable states. The regions with
very low probability in the distribution map are not included for detailed
comparison, in order to focus on the functional relevant conformations.

For a further comparative analysis between DiG and atomistic MD simulation,
representative structures for both proteins are obtained from MD simulations
by clustering analysis. Initially, cluster centroid coordinates of all
dominant meta-stable clusters in the 2D TICA space (Fig. 3a) are estimated for
each protein. Following this, 1000 MD simulated structures near cluster
centroids are sampled for each structure cluster. The first 16 TICA components
of these structures are used as features to cluster the simulation structures
into 4 sub-clusters using the K-Means algorithm. Finally, the centroid
structure of the largest sub-cluster was extracted by MDTraj
(McGibbon2015MDTraj, ) and taken as the representative structure of the
respective meta-stable state.

In order to assess the degree of conformity of the distributions, we compute
quantitative metrics on the 2D TICA space. Particularly, we devise a $G\times
G$ grid to uniformly cover the populated regions of the 2D TICA space, wherein
the grid is labelled as positive if there exists at least one simulation
structure within the corresponding TICA range. This grid is referred to as the
“groundtruth” grid. With structures sampled by DiG or other sampling methods,
we can construct a similar “sampled” grid following the same procedure.

In the coverage analysis in Supplementary Sec. E, we employ four distinct
types of metrics, namely:

$\displaystyle\mathrm{Accuracy}$
$\displaystyle=\frac{\mathrm{TP}+\mathrm{TN}}{\mathrm{TP}+\mathrm{FP}+\mathrm{TN}+\mathrm{FN}},$
(81) $\displaystyle\mathrm{Precision}$
$\displaystyle=\frac{\mathrm{TP}}{\mathrm{TP}+\mathrm{FP}},$ (82)
$\displaystyle\mathrm{Recall}$
$\displaystyle=\frac{\mathrm{TP}}{\mathrm{TP}+\mathrm{FN}},$ (83)
$\displaystyle\mathrm{F1}$
$\displaystyle=\frac{2\times\mathrm{Precision}\times\mathrm{Recall}}{\mathrm{Precision}+\mathrm{Recall}}.$
(84)

In the above equations, the grid labelled as the “groundtruth” is considered
to be true, whereas the one labelled as “sampled” is considered to be
predicted. The values TP, FP, TN, and FN represent the values of true
positives, false positives, true negatives, and false negatives, respectively.

To quantitatively assess the similarity of structures to their reference
conformations, we employ two metrics: the template modeling score (TM-score)
(zhang2004scoring, ) and the root mean square deviation (RMSD). TM-score is a
normalized measure of structural similarity between two conformations, with a
score of 1 indicating a perfect match; and RMSD calculates the average
distance between the paired atoms of two optimally superimposed structures. In
our evaluations, we restrict our RMSD calculations to the alpha carbon atoms
in protein structure comparison, and for all non-hydrogen atoms for ligand
structure comparison. These metrics provide a quantitative means of gauging
the accuracy of our protein conformation samples relative to their
experimental counterparts.

The quality of DiG-generated structures was assessed using the TM-scores, by
comparing each generated structure against crystal structures (6M0J for RBD
and 6LU7 for main protease). For RBD, the mean value of TM-score is 0.84 and
all structures have TM-score larger than 0.8, indicating highly similar
structures; while in the case of main protease, the TM-score is more spread,
with about $94\%$ structures with TM-score $>$ 0.5. We adopt a criterion
suggested by (xu2010significant, ), using TM-score = 0.5 as a cutoff to remove
the structures that are dissimilar to the experimental model. For the
downstream analysis, such as structure distributions, only structures with TM-
score $>$ 0.5 are used to reduce noises from incorrect predictions.

##### Catalyst-Adsorbate Sampling

After training the model to find different adsorption configurations, we
traverse the initial positions of the adsorbate by shifting the original
initial structure of adsorbate in the dataset along the $x$ and $y$ vectors of
the unit cell. Specifically, we equally divide the unit cell in $x$ and $y$
dimensions into a grid of $15\times 15$ points. Without changing the
conformation and height in the $z$ dimension of the initial structure of the
adsorbate, we shift its coordinates in the $x$ and $y$ dimensions to match
these $15\times 15$ points, thus creating $15\times 15$ different initial
structures of the system. For each initial position, we use DiG to sample $10$
structures. The sampled structures are then verified by DFT relaxation with
VASP hafner2008ab . In VASP, we allow both the catalyst surface and the
adsorbate to move, which is consistent with our model and the dataset. The
structures generated by our model are close to the relaxed structures, as
described in Sec. 3. For some initial positions, we are able to find multiple
adsorption sites within the $10$ sampled structures. Fig. 5 shows such a case,
where two of the adsorption configurations are sampled from the same initial
structure. Note that the training dataset contains only very short MD
trajectories. The capability of our model is also limited by the dataset. With
longer MD trajectories that traverse more structures, our model should be able
to find more adsorption sites from an arbitrary initial structure. More
discussions can be found in Supplementary Sec. F.

The probability density map of DiG is generated with single-atom adsorbates.
To plot the map of probability density, we equally divide the unit cell along
$x$ and $y$ dimensions as mentioned previously, resulting in a $20\times 20$
grid. We place the adsorbate above each grid point, creating $20\times 20$
structures for each fixed height of the adsorbate atom. We traverse 10
different heights for the adsorbate atom, ranging from $0\,\textnormal{\AA}$
to $1\,\textnormal{\AA}$ from the surface of the catalyst. Atoms in the
catalyst surface are kept as the initial positions when density evaluation. In
total, we obtain $4000$ structures. We calculate the log-likelihood of our
model on these structures, given an initial structure where the adsorbate atom
is $2\,\textnormal{\AA}$ from the catalyst surface above the center of the
catalyst surface in the unit cell. The equation below summarizes the
calculation of the probability density map:

$\displaystyle p(x_{i},y_{j}\mid
x_{0},y_{0},z_{0})=\max_{k}{p_{\mathrm{model}}(x_{i},y_{j},z_{k}\mid
x_{0},y_{0},z_{0})},$ (85)

where
$i,j,k\in\\{1,\cdots,20\\}\times\\{1,\cdots,20\\}\times\\{1,\cdots,10\\}$,
$z_{0}=2\,\textnormal{\AA}$ is distance from the initial position of the atom
to the highest point of the catalyst surface, and $(x_{0},y_{0})$ is the
center of the catalyst surface in the unit cell. Finally, in the probability
density map, we plot the log values of the probabilities above. In the energy
map from VASP, for each $i,\,j$ we plot the negative of the energy of the
relaxed structure, with the $x_{i},y_{j}$, and the catalyst surface fixed.
Starting from an initial height $z_{k}=2\,\textnormal{\AA}$ from the catalyst
surface, we used DFT to relax along the $z$ dimension to obtain the minimum
energy upon the grid point $(x_{i},y_{i})$.

##### Property-Guided Structure Generation

To measure the ability of modeling the carbon polymorph structures, we use the
StructureMatcher in the Pymatgen ong2013python  package to calculate the ratio
that a sampled structure matches a structure in the training dataset. We use
hyperparameters stol=0.5, angle_tol=10 and ltol=0.3 for the StructureMatcher,
where stol is the tolerance for the displacement of atom positions, angle_tol
controls the difference in lattice vector angles between the matched
structures, and ltol is the tolerance for the difference in matched lengths of
lattice vectors. The get_rms_dist method is used to calculate the RMSD, which
is normalized by the average free length per atom.

##  Appendix B Model Details

###  B.1 Descriptors of Molecular Systems

We consider four types of molecular systems in previous sections: proteins,
protein-ligand, catalyst-adsorbate, and carbon polymorphs. A descriptor
$\mathcal{D}$ is used for each system that captures the relevant features of
the molecular structure and can be processed by DiG. The descriptor
$\mathcal{D}$ in the four systems is first processed into node representations
$\mathcal{V}$ describing the feature of each system-specific individual
element, and a pair representation $\mathcal{P}$ describing inter-node
features. Note that in some systems, the descriptor $\mathcal{D}$ also
contains structural features, which are treated as part of the node
representation $\mathcal{V}$. The $\\{\mathcal{V},\mathcal{P}\\}$
representation is the direct input from the descriptor part to the Graphormer
model, as illustrated in Fig. 1.

For protein systems, we follow AlphaFold (jumper2021highly, ) and adopt a
coarse-grained representation that uses the position of the alpha-carbon atom
and the orientation of each residue (see Supplementary Sec. A.2.1). The node
representation $\mathcal{V}$ is a sequence of feature vectors that are
generated by the Evoformer module in (jumper2021highly, ), which takes as
input the amino acid sequence and the multiple-sequence alignment (MSA) of the
protein. The pair representation $\mathcal{P}$ is composed of two matrices:
the first represents all pairwise interactions between residues, which is also
produced by Evoformer, and the second represents the lengths of all residue
pairs on the amino acid sequence. The node and pair representations are
learnable embeddings in (jumper2021highly, ), but we hold them fixed in this
work to avoid additional computational cost from fine-tuning the Evoformer.
See more discussions about the limitation of fixing parameters of Evoformer
can be found in Supplementary Sec. F.

For protein-ligand binding systems, we use an all-atom representation that
includes atoms from both the protein and the ligand. The features are obtained
following the method in (min2022predicting, ). To be specific, the node
representation $\mathcal{V}$ for the protein part consists of the types of
atoms around the binding pocket and also the positions of these atoms
$\bar{\mathbf{R}}_{\textrm{Rec}}^{\star}$ in the crystal structure of the
protein, and the ligand part consists of a graph of the skeletal formula of
the ligand. The pair representation $\mathcal{P}$ is also composed of two
parts: one for the intra-molecular bonds and one for the inter-molecular
interactions. The intra-molecular bonds are represented by feature embeddings
of the chemical bonds between the atoms of the protein or the ligand, and the
inter-molecular interactions are represented by feature embeddings of
interactions like the hydrogen bonding between the protein and the ligand, as
defined in (min2022predicting, ).

For catalyst systems, we use an all-atomic representation that includes both
the catalyst surface and the adsorbates. The node representation $\mathcal{V}$
consists of the types of atoms of the catalyst and adsorbed molecules as well
as their positions
$\bar{\mathbf{R}}_{\mathrm{Cat}}^{\star},\bar{\mathbf{R}}_{\mathrm{Ad}}^{\star}$
in the initial structure from the OC20 dataset (chanussot2021open, ). The
features of the nodes are only embeddings of the atomic type and position,
following the method in (shi2022benchmarking, ). The pair representation
$\mathcal{P}$ is only defined for the adsorbates, which consists of feature
embeddings of the chemical bonds between the atoms of the adsorbates.

For carbon polymorphs, we use an all-atomic representation that only includes
carbon atoms. The node representation $\mathcal{V}$ consists of the initial
embedding of the carbon element. The only difference among the systems is the
number of carbon atoms. There is no pair representation $\mathcal{P}$ for the
carbon polymorphs, as the chemical bonds are not pre-defined and may change
during the diffusion process.

###  B.2 Backbone Architecture

The deep learning models used in DiG are extended by our previously proposed
Graphormer ying2021transformers ; shi2022benchmarking , which is a
Transformer-based graph neural network vaswani2017attention , and could
efficiently capture the topological information while keep the powerful
expressiveness from the Transformer architecture. The model is composed of a
few concatenated so-called attention layers and feed-forward layers. Each
attention layer takes the hidden node representation $\mathcal{H}$ as the
input tokens of the Transformer and uses the pair representation $\mathcal{P}$
as a learnable attention bias for the attention mechanism. Formally, given a
hidden node representation
$\mathcal{H}=\\{\mathbf{h}_{1},\cdots,\mathbf{h}_{I}\\}$, where $I$ is the
number of nodes, and a pair representation
$\mathcal{P}=\\{\mathbf{P}_{\imath\jmath}\\}_{\imath,\jmath=1}^{I}$, where
$\mathbf{P}_{\imath\jmath}$ is the learnable embedding of the edge features
from node $\imath$ to node $\jmath$, Graphormer computes the attention score
$\mathbf{A}_{\imath\jmath}$ from node $\imath$ to node $\jmath$ as:

$\displaystyle\mathbf{A}_{\imath\jmath}=\frac{(\mathbf{h}_{\imath}\mathbf{W}^{(Q)})(\mathbf{h}_{\jmath}\mathbf{W}^{(K)})^{\top}}{\sqrt{d}}+\mathbf{P}_{\imath\jmath}^{\top}\mathbf{w}^{(\mathcal{P})},$
(86)

where $\mathbf{W}^{(Q)},\mathbf{W}^{(K)}$ are the “query” and “key” linear
projections for the node representation, $\mathbf{w}^{(\mathcal{P})}$ is a
learnable weight vector for the pair representation, and $d$ is the dimension
of the query and key vectors. The hidden node representation $\mathcal{H}$ in
the first layer is $\mathcal{V}$. The attention bias term
$\mathbf{P}_{\imath\jmath}^{\top}\mathbf{w}^{(\mathcal{P})}$ enables the model
to learn the importance of the pair representation for the attention
mechanism. The attention score is then normalized by a softmax function over
all nodes and used to compute the attention output following the invariant
point attention mechanism in AlphaFold (jumper2021highly, ) for protein
systems, or standard Transformer architecture vaswani2017attention  for other
molecular systems.

###  B.3 Structural Attention Biases

Besides the descriptor $\mathcal{D}$ input of the molecular systems, the
Graphormer model also needs to process the geometric structure input
$\mathbf{R}$ and produce a physically finer structure. To more informatively
encode the geometric information during the diffusion process, we also
introduce a structural representation for the input $\mathbf{R}$, which is
used to help Graphormer to capture the spatial and rotational relationships
among the nodes and refine the noisy structures to more physically realistic
ones.

For all molecular systems in full-atom representation, we follow
(shi2022benchmarking, ) to encode the Euclidean distance $d_{\imath\jmath}$
between the positions of node $\imath$ and node $\jmath$ as a bias term
$b_{\phi}(d_{\imath\jmath})$, where $\phi$ is a learnable parameter. The
distance encoding bias is added to the attention score in Eq. (86) to modulate
the attention based on the distance between nodes. For protein systems that
use the coarse-grained representation $\mathbf{R}=(\mathbf{C},\mathbf{Q})$
(see Supplementary Sec. A.2.1), we adopt the invariant point attention
mechanism in (jumper2021highly, ) to construct the corresponding attention
score, which has been shown to be indispensable for capturing the local
rotational invariance feature of the protein structures.

###  B.4 Equivariant Graphormer

Algorithm 10 Equivariant Vector Prediction

1:Node 3D positions $\\{\mathbf{R}_{1},\cdots,\mathbf{R}_{I}\\}$, node
representation
$\mathbf{H}=[\mathbf{h}_{1}^{\top},\cdots,\mathbf{h}_{I}^{\top}]^{\top}\in\mathbb{R}^{I\times
C}$, attention bias $\mathbf{E}^{\mathrm{attn}}\in\mathbb{R}^{I\times I}$;
linear projector matrices
$\mathbf{W}^{(Q)},\mathbf{W}^{(K)}\in\mathbb{R}^{C\times d}$,
$\mathbf{W}^{(V)}\in\mathbb{R}^{C\times C}$,
$\mathbf{w}^{(F)}\in\mathbb{R}^{C}$;

2:Calculate the relative positions
$\mathbf{E}^{\mathrm{rel}}\in\mathbb{R}^{I\times I\times
3}:\mathbf{E}^{\mathrm{rel}}_{\imath\jmath}:=\mathbf{R}_{\imath}-\mathbf{R}_{\jmath}$;

3:$\mathbf{Q}=\mathbf{H}\mathbf{W}^{(Q)},\mathbf{K}=\mathbf{H}\mathbf{W}^{(K)},\mathbf{V}=\mathbf{H}\mathbf{W}^{(V)}$;

4:$\mathbf{A}=\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{d}}+\mathbf{E}^{\mathrm{attn}}$;

5:Calculate $\mathbf{F}\in\mathbb{R}^{I\times
3}:\mathbf{F}_{\imath}:=\sum_{\jmath=1}^{I}\mathrm{softmax}(\mathbf{A}_{\imath:})_{\jmath}({\mathbf{w}^{(F)}}^{\top}\mathbf{V}_{\jmath})\mathbf{E}^{\mathrm{rel}}_{\imath\jmath}$;

6:Return $\mathbf{F}$;

Due to the score model interpretation (gradient of log-density function), the
output of the Graphormer model is required to be equivariant w.r.t the
$\mathbf{R}$ input, which requires a proper design for processing the
$\mathbf{R}$ input. To ensure the rotational equivariance of the Graphormer
model, we add one equivariant attention layer shi2022benchmarking  as the 3D
vector output head, which produces geometric vectors that are equivariant to
any rotation transformations in 3D Euclidean space on the input, as detailed
in Alg. 10. Specifically, we first compute the attention matrix $\mathbf{A}$
and the transformed invariant features
$\mathbf{H}=[\mathbf{h}_{1}^{\top},\cdots,\mathbf{h}_{I}^{\top}]^{\top}$ as in
the previous layers, and then the vector output is obtained by attentively
aggregating the relative position information and the invariant features.
Since we only apply scalar multiplication and linear combination operations on
the vector features, the resulting vector $\mathbf{F}$ is naturally
equivariant to the $\mathrm{SO}(3)$ group of rotations. Moreover, $\mathbf{F}$
is translation-invariant because it only depends on the relative positions
between two nodes. This design strategy for processing vector features is
similar to those used in previous works on protein modeling jing2020learning
and quantum chemistry schutt2021equivariant .

###  B.5 Periodic Boundary Condition

In catalyst-adsorbate systems and carbon polymorphs, atoms in a 3D unit cell
are periodically repeated. Therefore, radius graphs with periodic boundary
conditions are constructed to represent the systems, where each atom in one
single cell (the centric cell) will connect with its neighboring atoms within
a pre-defined cutoff distance. Since atoms are periodically repeated, the same
atom in different cells may appear repeatedly in one graph as different nodes.
To avoid that the same atom has different node representations in the network,
typically a multi-graph will be constructed for message-passing neural
networks (MPNNs), where one node represents one atom, and multiple edges
between nodes represent the interactions with the same atom in different
cells. In this way, information on neighboring atoms will be aggregated by
MPNNs with multiple times through each edge.

Differently, message aggregation is done by attentively weighted sum on full
graph in Graphormer, and interactions between atoms are encoded into spatial
distance embeddings acting as attention bias. Multi-graph will lead to a
summation of multiple biases in the distance embedding space, which might be
projected to a new distance, and would not reflect multiple interactions with
the same atom in different cells. Therefore, to reflect the multiple
interactions correctly while enforce one representation for the same atom in
different cells, we use a cross-attention sub-layer to aggregate information
from all atoms in the radius graph into the atoms in the centric cell as shown
in Alg. 11.

Algorithm 11 Handling Periodic Boundary Condition

1:Atom positions in the centric unit cell
$\widetilde{\mathcal{X}}:=\left\\{\widetilde{\mathbf{x}}_{\imath}\right\\}_{\imath=1}^{I}$.

2:Lattice vectors $\mathbf{l}_{x},\mathbf{l}_{y},\mathbf{l}_{z}$. Cutoff
distance $D_{\text{cut}}$.

3:Atom representation
$\left\\{\mathbf{h}(\widetilde{\mathbf{x}_{\imath}})\mid\widetilde{\mathbf{x}}_{\imath}\in\widetilde{\mathcal{X}}\right\\}$
in current layer.

4:Atom representation
$\left\\{\mathbf{h}^{\prime}(\widetilde{\mathbf{x}_{\imath}})\mid\widetilde{\mathbf{x}}_{\imath}\in\widetilde{\mathcal{X}}\right\\}$
after attention.

5:${\mathcal{X}}\leftarrow\displaystyle\left\\{\widetilde{\mathbf{x}}+l\mathbf{l}_{x}+m\mathbf{l}_{y}+n\mathbf{l}_{z}\mid
l,m,n\in\mathbb{Z},\widetilde{\mathbf{x}}\in\widetilde{\mathcal{X}}\right\\}$;

6:$\mathcal{X}_{D}\leftarrow\displaystyle\left\\{\mathbf{x}\in\mathcal{X}\mid\exists\widetilde{\mathbf{x}}^{\prime}\in\widetilde{\mathcal{X}},\
\|\mathbf{x}-\widetilde{\mathbf{x}}^{\prime}\|\leqslant
D_{\text{cut}}\right\\}$;

7:$\forall\mathbf{x}\in\mathcal{X}_{D},\widetilde{\mathbf{x}}:=\mathbf{x}+l\mathbf{l}_{x}+m\mathbf{l}_{y}+n\mathbf{l}_{z}$,
s.t. $\widetilde{\mathbf{x}}\in\widetilde{\mathcal{X}},\ l,m,n\in\mathbb{Z}$;

8:$\mathbf{A}_{\imath\jmath}\leftarrow\frac{\left(\mathbf{h}(\widetilde{\mathbf{x}}_{\imath})\mathbf{W}^{(Q)}\right)\left(\mathbf{h}_{\jmath}(\widetilde{\mathbf{x}}_{\jmath})\mathbf{W}^{(K)}\right)^{\top}}{\sqrt{d}}+b_{\phi}\left(\left\|\mathbf{x}_{\imath}-\mathbf{x}_{\jmath}\right\|\right),\
\forall\mathbf{x}_{\imath},\mathbf{x}_{\jmath}\in\mathcal{X}_{D}$;

9:$\mathbf{h}^{\prime}(\widetilde{\mathbf{x}}_{\imath})\leftarrow\sum_{\jmath}\mathrm{softmax}(\mathbf{A}_{\imath:})_{\jmath}\left(\mathbf{h}(\widetilde{\mathbf{x}_{\jmath}})\mathbf{W}^{(V)}\right)$.

The basic idea is that learnable node embeddings are only assigned to atoms in
the centric cell, and the embeddings for the same atom appearing in
neighboring cells are its replicas. The attention bias encoded from pair-wise
atom distances is used to tell replicas of atoms in different cells.
Therefore, the representation of each node in the centric cell will be updated
by the correlation and interaction with all atoms in the radius graph.

In modeling the carbon polymorphs, with the lattice vectors, we can expand the
unit cell and handle periodic boundary conditions as described in Alg. 11.
However, we find that explicitly encoding the vertices in the unit cell as
tokens in the Graphormer encoder is very helpful for sampling physical
structures. Thus, each vertex of the unit cell is treated as an atom with a
special type in the model. In other words, for a structure with $n$ atoms in
the unit cell, the model will take $n+8$ tokens, and expand the tokens
according to the PBC handling method. Alg. 8 and 9 describe the details of the
training and sampling processes, respectively.

##  Appendix C Training Details

###  C.1 Protein Conformation Sampling

#### Training Pipeline and Dataset

Our training process for the protein system consists of three stages:
initialization, physics-informed diffusion pre-training (PIDP), and data-based
training using simulation data. In the first stage which aims to provide a
good initialization to stabilize PIDP training, we collect all experimental
structures from the Protein Data Bank (PDB) (wwpdb2019protein, ) before
December 25, 2020 as training data and employ Alg. 4 for model training. Such
experimental structures are widely used in structure prediction methods, in
which case a dataset is organized following the pattern
$(\mathcal{D},\mathbf{R}_{\mathcal{D}})$, where each amino-acid sequence
$\mathcal{D}$ is paired with one experimental structure
$\mathbf{R}_{\mathcal{D}}$. In contrast, to provide distributional information
during the training of DiG, we organize the structures to construct a physical
distribution dataset, in which each data point follows the pattern
$(\mathcal{D},\\{\mathbf{R}_{\mathcal{D}}^{(n)}\\}_{n=1}^{N_{\textnormal{data}}})$
where each amino acid sequence $\mathcal{D}$ is paired with a _set_ of
structures
$\\{\mathbf{R}_{\mathcal{D}}^{(n)}\\}_{n=1}^{N_{\textnormal{data}}}$. To
prepare this dataset, we adopt all the sequence identity clusters from PDB
obtained by MMSeqs2 steinegger2018clustering , and include all the available
experimental structures for each cluster. Following AlphaFold jumper2021highly
, we filter out structures from PDB that have a resolution worse than
$9\,\textnormal{\AA}$. This eliminates about 0.2% of structures. For proteins
longer than 256 amino acids, we divide them into segments of length no longer
than 256 amino acids. In each training step, we randomly draw clusters and
structures within each drawn cluster with equal probability. We would like to
remark that although this physical structure distribution is hard to verify to
obey the equilibrium distribution, it can still provide rich information about
the different modes of the equilibrium distribution.

In the second stage, to prepare the relevant structures for evaluating the
PIDP loss, we run short MD simulations for about 1000 proteins, for which the
details can be found in Supplementary Sec. D. We randomly pick 100 simulated
structures for each protein as the relevant structures. The training process
follows Alg. 1. In the final stage, we use a simulation dataset consisting of
the above simulation dataset, and 238 simulation trajectories randomly picked
from the GPCRmd dataset (rodriguez2020gpcrmd, ). The GPCRmd dataset contains
short simulations of various classes of G protein-coupled receptors (GPCRs).
The training process follows Alg. 4.

#### PIDP Training

As mentioned in the main text (Sec. 2), the structures
$\\{\mathbf{R}_{0}^{(m)}\\}_{m=1}^{M}$ for evaluating the PIDP loss are
ideally grid points (as in finite-element methods) spanning the structure
space, but this is unaffordable since the number of grid points increases
exponentially with the dimension of the space, which is typically exceedingly
high for molecular systems. But we only need to supervise the model on a low-
dimensional manifold of physically relevant structures (with low energy). Due
to the grid point nature, these structures do not have to follow the
equilibrium distribution, but only need to demonstrate the manifold. This then
enables a wide range of affordable methods to prepare such structures, such as
perturbation around experimentally observed structures (zhang2021prody, ), and
short MD simulation structures. In practice, we found protein structures
perturbed by (zhang2021prody, ) lead to overly large energy gradient which
hinders effective optimization and cannot be easily mitigated by e.g. gradient
clipping (see Supplementary Sec. F.2 for more details about the limitation of
energy function). For this reason, we adopt structure samples from short MD
simulation trajectories for PIDP training. This is much cheaper than
generating structures following equilibrium distribution by long enough MD
simulations. The short MD simulations provide structures that can be seen in a
physical process thus demonstrating the relevant manifold, and the information
of equilibrium distribution is provided by the energy function.

For evaluating the energy function (or its gradient, i.e., force field), we
use OpenMM to compute the full-atom force using the Amber force field as
$-\nabla E(\bar{\mathbf{R}})$, which also serves for calculating the coarse-
grained forces $-\nabla_{\mathbf{C}}E$ and $-\nabla_{\mathbf{q}}E$ through
Alg. 2. We use PDBFixer eastman2013openmm  to fix the input protein structure
files before processed by OpenMM to avoid potential failures.

We find that the range of magnitude of the calculated forces varies
drastically, which poses a significant challenge in optimizing the PIDP loss
in Eq. (7). We hence retain only the structures that have a force magnitude
within the smallest three orders of magnitude. To further address the
optimization challenge, we also modify the loss terms for matching the score
model at $t=0$ to the force field in Eqs. (57, 63) by only matching their
directions.

We also find that for loss terms for $t>0$ (or $i>0$) in the PIDP losses Eqs.
(57, 63) , different sampling time steps $t$ (or $i$) result in significant
differences in the scale of the loss, ranging from 0 to $1\times 10^{6}$. For
stable and effective training, we rescale and clip these losses. Specifically,
with $\ell^{(\mathbf{C})}_{t}$ and $\ell^{(\mathbf{Q})}_{t}$ denoting the PIDP
loss terms at time step $t$ of alpha-carbon coordinates and residue
orientations, we find $\ell^{(\mathbf{C})}_{t}$ and $\ell^{(\mathbf{Q})}_{t}$
increase exponentially with $t$, with exponents $\rho^{(\mathbf{C})}$ and
$\rho^{(\mathbf{Q})}$. To avoid an excessively large loss, we scale these
losses by:
$\tilde{\ell}^{(\mathbf{C})}_{t}:=\ell^{(\mathbf{C})}_{t}/(\rho^{(\mathbf{C})})^{1-\mathrm{clip}(t)}$,
and
$\tilde{\ell}^{(\mathbf{Q})}_{t}:=\ell^{(\mathbf{Q})}_{t}/(\rho^{(\mathbf{Q})})^{1-\mathrm{clip}(t)}$,
where $\mathrm{clip}(t):=\min(0.05\tau,t)$.

Furthermore, we find that a balanced combination of the losses in Eq. (32) and
Eq. (7) is essential for generating more physical structures after PIDP
training.

These training methods, implemented to ensure more stable optimization, may
compromise the accuracy of the distribution learned by the model. These
limitations will be discussed further in Supplementary Sec. F.

#### Hyperparameter Choices

In protein training, we discretize the diffusion-process time variable
$t\in[0,\tau]$ into $i\in\\{0,1,\cdots,N\\}$ discrete time steps. The noise
scales $\sigma_{i}$ in Eq. (52) and Eq. (61) (or in Algs. 1 and 4 for
training, and Alg. 5 for sampling) at the corresponding discretized time step
$t_{i}$ are taken in the form
$\sigma_{i}=(\sigma_{\min})^{(1-t_{i}/\tau)}(\sigma_{\max})^{t_{i}/\tau}$. The
parameters $\sigma^{(\mathbf{C})}_{\min}$ and $\sigma^{(\mathbf{C})}_{\max}$
for the diffusion process on alpha-carbon coordinates and
$\sigma^{(\mathbf{Q})}_{\min}$ and $\sigma^{(\mathbf{Q})}_{\max}$ for the
residue orientation are detailed in Supplementary Tab. 2, together with all
other hyperparameters. To strike a balance between computational efficiency
and performance, we train for at least one epoch at each stage and halt
training when the rate of loss reduction noticeably decelerates.

Table 2: Hyperparameters of protein model. Different hyperparameters are used in different stages. Hyperparameter | Initialization | PIDP | Data Training
---|---|---|---
Model depth | 12
Hidden dim (Single) | 768
Hidden dim (Pair) | 256
Hidden dim (Feed Forward) | 1024
Number of Heads | 32
Optimizer | Adam
Learning rate schedule | Inverse Square Root
Peak Learning rate | 1E-03 | 1E-05 | 1E-05
Dropout $p$ | 0.1 | 0.0 | 0.1
Adam $(\beta_{1},\beta_{1})$ | (0.9, 0.999) | (0.9, 0.999) | (0.9, 0.999)
Adam $\epsilon$ | 1E-06 | 1E-06 | 1E-06
Weight decay | 1E-02 | 1E-02 | 1E-02
Warmup ratio | 0.06 | 0.06 | 0.06
Batch size | 128 | 32 | 128
Diffusion steps $N$ | 500 | 500 | 500
$(\sigma^{(\mathbf{C})}_{\min},\sigma^{(\mathbf{C})}_{\max})$ | (0.1, 35) | (0.1, 35) | (0.1, 35)
$[\sigma^{(\mathbf{Q})}_{\min},\sigma^{(\mathbf{Q})}_{\max}]$ | (0.02, 1.65) | (0.02, 1.65) | (0.02, 1.65)
Boundary weight $\lambda_{1}$ | - | 5 | -
Hutchinson $N_{\mathrm{est}}$ | - | 20 | -

###  C.2 Ligand Structure Sampling around Binding Sites

Two-stage training was performed in ligand sampling. The first stage employs
CrossDocked (francoeur2020three, ) as a binding structure prediction task. The
CrossDocked dataset contains docked conformations of varying quality, so we
filter out all complexes whose RMSD between the docked and the experimental
crystal structures is larger than $2.5\,\textnormal{\AA}$. The second stage
employs the simulation data (Sec. D.1), where we set the threshold as
$6\,\textnormal{\AA}$, and 1 ns as the sampling stride. For data-based
training, we collect data from CrossDocked (francoeur2020three, ) and MD
simulations to explore the most concerned part in the conformational space.
The CrossDocked dataset contains docked conformations of varying quality, so
we filter out all complexes whose RMSD between the docked and the experimental
crystal structures is larger than $2.5\,\textnormal{\AA}$. For MD simulation
data, we set the threshold as $6\,\textnormal{\AA}$, and 1 ns as the sampling
stride. We conduct a quality screening on the simulation data, by filtering
out trajectories that there are ligand’s atoms around the protein laying
within $5\,\textnormal{\AA}$. After filtering, the MD simulation dataset
contains 1157 protein-ligand complex trajectories, and we split 80% for
training, 10% for evaluation, and 10% for testing. All hyperparameters are
listed in Supplementary Tab. 3.

Table 3: Hyperparameters of protein-ligand binding model used in CrossDocked datasset and MD dataset training. Hyperparameter | CrossDocked & MD Training
---|---
Model depth | 12
Hidden dim (Model) | 768
Hidden dim (Feed Forward) | 768
Number of Heads | 32
Optimizer | Adam
Peak learning rate | 0.0002
Warmup ratio | 0.06
Learning rate schedule | Linear Decay
Dropout $p$ | 0.1
Adam $(\beta_{1},\beta_{2})$ | (0.9, 0.98)
Adam $\epsilon$ | 1E-08
Weight decay | 0.0
Batch size | 64
Diffusion steps $N$ | 500
Diffusion $\beta$ schedule | Sigmoid
Diffusion $\beta$ start | 1E-07
Diffusion $\beta$ end | 0.02
EMA decay | 0.9999
EMA fp32 | true
Clip norm | 10.0

###  C.3 Catalyst-Adsorbate Sampling

Our model captures the distribution of the structure of an adsorbate on a
catalyst surface. DiG predicts the distribution conditioned on an initial
structure, which is the first frame in a relaxation trajectory provided in the
OC20 dataset chanussot2021open . To train DiG, we use the MD part of the OC20
dataset following Alg. 6. 20,000 systems of the MD dataset are separated for
validation. Before this data-based training, the Graphormer model is first
pretrained on the IS2RS task of OC20. Detailed hyperparameters for the IS2RS
pretraining are listed in Supplementary Tab. 4. In the data-based training, we
use a peak learning rate of $2\times 10^{-4}$, maximum number of epochs 300,
warm-up ratio $6\%$, a batch size of 64, and the number of diffusion steps
$N=5000$. The beta schedule follows a sigmoid form:

$\beta_{i}=\frac{1}{1+\exp{(12(0.5-i/N))}}(\beta_{\mathrm{end}}-\beta_{\mathrm{start}})+\beta_{\mathrm{start}},$
(87)

with $\beta_{\mathrm{start}}=1\times 10^{-7}$ and
$\beta_{\mathrm{end}}=2\times 10^{-3}$ where $i\in\\{0,\cdots,N\\}$ is the
diffusion time step. The training is stopped after 86 epochs. We use a cutoff
value of $6\,\textnormal{\AA}$ for PBC handling. Supplementary Tab. 4
summarizes detailed hyperparameters for training of DiG for catalyst-adsorbate
sampling.

Table 4: Hyperparameters of backbone model for pretraining on Open Catalyst IS2RS and training on Open Catalyst MD dataset. Hyper Parameter | IS2RS pretraining | MD training
---|---|---
Model depth | 12
Hidden dim (Model) | 768
Hidden dim (Feed Forward) | 768
Number of Heads | 32
Optimizer | Adam
Learning rate | 0.0002
Warmup ratio | 0.06
Learning rate schedule | Linear Decay
Dropout $p$ | 0.1
Adam $(\beta_{1},\beta_{2})$ | (0.9, 0.98)
Adam $\epsilon$ | 1E-08
Weight decay | 0.0
PBC Cutoff | 6.0
Batch size | 1024 | 64
Diffusion steps $N$ | - | 5000
Diffusion $\beta$ schedule | - | Sigmoid
Diffusion $\beta$ start | - | 1E-07
Diffusion $\beta$ end | - | 0.002

###  C.4 Property-Guided Structure Generation

For training the DiG on an unconditional distribution, we use $15,697$
structures of carbon polymorphs generated from ab initio random structural
search (RSS) at the DFT level (PBE/plane-wave basis, with an energy cutoff of
$520$ eV) with a range of number of atoms from 2 to 24, following the method
in lu2021computational . Only the relaxed structures, i.e., the final frames
in relaxation processes, are taken for training. We remark that for the
inverse design task, the distribution of the structures at local energy minima
does not follow an equilibrium distribution. However, the primary goal here is
to generate structure candidates with reasonable stability and targeted
properties. In this context, the relaxed structures from random structure
search can well represent the low energy structure manifold. The number of
training epochs is $50,000$, with a peak learning rate $2\times 10^{-4}$,
batch size $4096$, and a warm-up ratio $6\%$. For the reference lattice vector
set $\bar{\mathbf{L}}^{\star}$, we use the mean lattice vector set over the
dataset, which is close to three orthogonal vectors with a length of
$4\,\textnormal{\AA}$. As the initial structure in the catalyst-adsorbate
model, this reference lattice vector structure is also encoded into the model
with an additional attention bias term. Tab. 5 summarizes the hyperparameters
for training the diffusion model for property-guided structure generation. We
use a cutoff value of $20\,\textnormal{\AA}$ for PBC handling.

Table 5: Hyperparameters of diffusion model training for property-guided structure sampling. Hyperparameter | Data Training
---|---
Model depth | 12
Hidden dim (Model) | 768
Hidden dim (Feed Forward) | 768
Number of Heads | 32
Optimizer | Adam
Learning rate | 0.0002
Warmup ratio | 0.06
Learning rate schedule | Linear Decay
Dropout $p$ | 0.1
Adam $(\beta_{1},\beta_{2})$ | (0.9, 0.98)
Adam $\epsilon$ | 1E-08
Weight decay | 0.0
PBC Cutoff | 20.0
Batch size | 1024
Diffusion steps $N$ | 500
Diffusion $\beta$ schedule | Sigmoid
Diffusion $\beta$ start | 1E-07
Diffusion $\beta$ end | 0.02

##  Appendix D Molecular Simulation and Energy Evaluations

###  D.1 Molecular Dynamics Simulation for Protein-Ligand Complexes

We generate MD simulation data for complex systems selected from the PDBbind
v2020 su2018comparative  using an automatic pipeline called
protocolGromacs333https://github.com/tubiana/protocolGromacs. It utilizes
GROMACS van2005gromacs  as the backend engine with a common simulation setting
for all complexes, providing a capability of high-throughput MD simulations.
Specifically, this pipeline comprises four stages: preparation, minimization,
equilibration, and production simulations. In the system preparation stage, a
protein topology is generated with pdb2gmx with the amber99sb-ildn
lindorff2010improved  force field with the tip3p explicit water model; the
ligand parameter and topology are generated with acpype sousa2012acpype . For
cases with missing atoms/residues in the PDB files, PDBFixer eastman2013openmm
is applied to complete the molecules. Then, a cubic simulation box is used
with a minimum distance of 1.2 nm between the protein-ligand complex and the
box boundaries. Finally, a pre-equilibrated system of 216 water molecules is
repeated over the simulation box to provide the solvated environment. To
neutralize charged systems, appropriate ions ($\mathrm{Na}^{+}$ or
$\mathrm{Cl}^{-}$, depending on the net charge of the solute molecules) are
applied by replacing randomly selected water molecules. Once the simulation
box is prepared, an energy minimization process is carried out to remove the
atomic clashes and optimize the geometry of all molecules. In the
equilibration simulation stage, a thermostat is applied to heat the system
from 0 to 300K within 100 ps. The heated system is then further equilibrated
to 1 bar in an NPT ensemble for another 100 ps. During the equilibration
stage, the bonds for molecules are constrained. For the final production, the
leap-frog algorithm van1988leap  is used for integrating Newton’s equations of
motion and a Particle Mesh Ewald (PME) essmann1995smooth  method is used for
calculating long-range electrostatic interactions. The LINCS hess1997lincs
algorithm is adopted for resetting all bonds to their correct lengths after an
unconstrained update. Finally, the production is performed for 100 ns with an
integration time step of 2 fs.

Following the above protocol, we generate MD simulation trajectories for 1500
protein-ligand complexes. To facilitate model training, each 100-ns simulation
trajectory from the production run is divided into segments of length 1 ns,
resulting in 100 trajectory segments for each complex system. The protein-
ligand complex is mapped to the center of the primary simulation box by
applying periodic boundary conditions to ensure the integrity and connectivity
of the molecules.

###  D.2 Energy Evaluations in Physics-Informed Diffusion Pre-training

For physics-informed diffusion pre-training (PIDP), the energy and gradients
are evaluated using OpenMM eastman2017openmm  following the settings used in
Folding@home (zimmerman2021sars, ). The amber14sb force is used for the
ablation study in Supplementary Sec. E.3. The Generalized Born solvent model
mongan2004constant  is used for solvation energy calculation.

###  D.3 Density Functional Theory Computation

We use DFT for the grid search of adsorbate configurations on catalyst
surfaces, for verifying the adsorbate configuration distributions predicted by
DiG as shown in Fig. 5. They are carried out using VASP6.3 with setups
compatible with OC20 (chanussot2021open, ). Specifically, periodic boundary
conditions and projector-augmented wave pseudopotentials are adopted with
plane-wave electron kinetic energy cut-off of 350 eV. Generalized gradient
approximation and the revised Perdew-Burke-Ernzerhof (RPBE) functional are
employed (hammer1999improved, ; kresse1996efficient, ). The Monkhorst-Pack
grid is used to sample the reciprocal space. For the electronic degree of
freedom, the convergence criteria for self-consistent computations is set to
$1\times 10^{-3}$ eV/atom. For ionic degree of freedom, the relaxations are
carried out only on the atomic coordinates with the lattice being fixed.
Convergence is considered to be reached when the Hellmann-Feynman forces are
smaller than 0.02 eV/Å.

##  Appendix E Additional results

###  E.1 Protein Conformation Sampling

Table 6: Protein systems utilized in this paper. Reference structure 1 is denoted in cyan, while reference structure 2 is denoted in brown in Fig. 3b. Protein | Ref. 1 | Ref. 2 | TMscore
---|---|---|---
Adenylate Kinase | 4ake (chain A) | 1ake (chain A) | 0.6899
Lmrb | DEER-AF | 6t1z (chain A) | 0.7600
human B-Raf kinase | 6uan (chain A) | 3skc (chain A) | 0.9235
D-ribose | 3dri (chain A) | 1urp (chain A) | 0.7187

We present a list of proteins employed to demonstrate the efficacy of DiG in
generating multiple conformations in Fig. 3b in Table 6. The table provides a
detailed overview of the protein systems utilized in this study, including the
reference structures denoted in cyan and brown, respectively, and the
conformational differences between the two reference structures, measured in
TMscore.

For both RBD and main protease proteins, there are some structures in the PDB
dataset used for DiG model training. We project those structures onto the
reduced 2D space spanned by the first two TICA coordinates, and show the
results in Supplementary Fig. S8. Clearly, multiple states are observed in the
maps indicated by the diamonds, but these structures together only represent a
small fraction of the structure space. In contrast, the DiG-generated
structures show much broader overlaps with MD simulations in the structure
space. In the case of RBD, DiG even predicted a new region (lower right,
cluster-IV in the main text), which has no experimental determined structure.
This particular region was significantly sampled by MD simulation. We can
observe clear correspondence for the regions sampled by DiG and MD
simulations. The diverse structures reveal more information about the
functional states of proteins.

Besides the qualitative comparison of the distributions generated by DiG and
those sampled by MD simulations. We provide a detailed quantitative
measurement of the overlaps of explored regions by these two methods. Taking
the regions with MD simulation structures as references, the coverage analysis
is carried out following a standard binary classification approach (see
Supplementary Fig. S8). The distributions are first converted to masked binary
maps (divided to $50\times 50$ regions), then the two sets of binary maps (DiG
results vs. MD simulation results) are compared. The accuracy, precision,
recall, and F1-score are computed for each case, at different sampling data
sizes. There are 50,000 structures generated by DiG for each protein. We see
that the coverage (recall) increases as more structures are included for the
analysis, while the precision level is kept at high levels. In order to
compare with the MD simulation data, we took two approaches to draw samplings:
(1) take structures consecutively from simulation trajectories; (2) take
structures randomly from all simulation trajectories (i.i.d sampling). We
observed that the sampling coverage and efficiency of DiG are better than MD
simulations.

Figure S7: Protein structure distributions and sampling coverages. Figure S8:
a. Results for the RBD of SARS-CoV-2 spike protein. Experimentally determined
structures are mapped to the reduced 2D space, indicated using the purple
diamond symbols. On the top right panel, the same space is divided into 50x50
grids, which are classified into explored (blue) and unexplored (white) sub-
regions, depending on the presence of MD simulation structures. The accuracy,
precision, recall, and F1-score are shown as a function of sampling size (or
the ratio to the whole dataset). The cluster in the lower-right region has no
experimental structures, indicating a new state revealed DiG, which is
consistent with MD simulations. b. Results for the main protease following the
same analysis protocol and representations. In both plots a and b, blue star
symbols indicate the AlphaFold predicted structures in the 2D space, and red
circles show the cluster centers of MD simulation structures.

For two representative ligand-protein systems shown in Fig. 4b, structures
from 100 ns MD simulations are superposed to show the variation in ligand
structures (see Supplementary Fig. S9).

Figure S9: Ligand structures observed in MD simulations. a. For the case of
Tyk2 target, the binding pocket is deep and well confined, MD simulations show
highly similar ligand structures and binding poses. The RMSD values compared
to the crystal structure is small. b. the MD simulation results are shown for
the case of P38 protein. Similar to the DiG results, the ligand structure
exhibits larger variations compared to the cases of tyk2. For both proteins,
the same ligands shown in Figure 2 in the main text are used in the
simulations, and the simulation duration is 100 ns for both systems.

The ligand structure generations are carried out for a set of 16 proteins,
each paired with various numbers of ligands. In total, there are 409 ligand-
protein systems in this testing dataset. In Supplementary Fig. S10, five
proteins, each with four different ligands, are shown to illustrate the best
structures generated by DiG. The ligand binding poses and atomic structures
generated by DiG exhibit diversity and are correlated with the characteristics
of protein pockets.

Figure S10: Ligand structure generation in the binding pocket of target
proteins. The target names are indicated in the figure, with each row showing
different ligands with its best binding poses to the same target protein.
Here, best binding poses is defined as the most similar structure to the
experimental observations.

###  E.2 Catalyst-Adsorbate Sampling

Supplementary Fig. S11 shows top and front views of all adsorption
configurations generated by DiG, which covers all the adsorption
configurations found by DFT relaxation for this system by traversing initial
positions of the adsorbate. Note that for most systems, with a very short MD
trajectory as provided in the dataset, it can only cover one or two structures
close to these relaxed configurations. Thus, the result shows the ability of
DiG to generate to unseen systems from very short MD trajectories in the
training set.

Figure S11: Additional Results for catalyst surface adsorption. All adsorption
configurations found by DiG, with the configurations from model in color and
the configurations from grid search with DFT in white. The number of
adsorption sites is 6 in total. We divide all the sites into 2 groups in
(a)(c) and (b)(d), and show both the top view in (a)(b) and the front view in
(c)(d).

###  E.3 Ablation Study

We conduct an ablation study to investigate the effect of various components
in the training pipeline of DiG for protein systems. The training pipeline
consists of three stages: (i) initialization from experimental data, (ii)
physics-informed diffusion pre-training (PIDP), and (iii) training with
simulation data.

Figure S12: Ramachandran plots of the sampled structures at different stages
of the optimization process (initialization, PIDP, and simulation data
training) compared with the reference MD structures.

The main protease of SARS-CoV-2 is selected as a case study since it has long
simulation trajectories (2.6 ms) (zimmerman2021sars, ) and has been studied in
Section 3.1. We investigate the importance of each stage of the training
pipeline by evaluating the quality of structures sampled at different stages.
The quality is measured by comparing the torsion angle distribution of sampled
structures with that of the simulation trajectories. We use Ramachandran plots
to visualize the distribution of torsion angles of the 6 amino acids
corresponding to the first 10 TICA components of the simulation trajectories
(RAMACHANDRAN196395, ).

Supplementary Fig. S12 summarizes the results of the ablation study. We obtain
approximately 100,000 filtered results from each of the three stages. The blue
area represents the ground-truth simulation distribution, which encompasses a
relatively large and diverse range of torsion angles, reflecting the dynamic
and flexible nature of the protein. Supplementary Fig. S12a displays the
distribution of structures sampled by DiG initialized from experimental data,
i.e., the PDB protein structure dataset (wwpdb2019protein, ) (see
Supplementary Sec. C.1). The distribution is highly concentrated on a single
point, indicating that the sampled structures are only fit to experimental
structures and do not capture the protein’s dynamics. Supplementary Fig. S12b
illustrates the distribution of structures sampled by DiG after PIDP, which
improves the generated distribution towards the equilibrium distribution using
the energy function. We find that PIDP training may generate some failure
cases with very high RMSDs or low TMscores compared to the crystal structure.
Therefore, we filter out results with RMSD higher than $10\,\textnormal{\AA}$
and TMscore lower than 0.6.

However, PIDP alone is still insufficient to fully capture the equilibrium
distribution. Supplementary Fig. S12c displays the distribution of structures
sampled by DiG after further training with simulation data, where simulation
structures serve as direct signals to supervise DiG’s learning. The
distribution becomes even more similar to the equilibrium distribution,
demonstrating DiG’s ability to learn from simulation data and generate
realistic and diverse structures.

All sampled structures at various stages of the training pipeline are
physically reasonable and are similar to low energy values, as verified by the
structural quality metrics. The results show that DiG can effectively combine
information from both energy function and simulation data to produce high-
quality structures reflecting the protein systems’ equilibrium distribution.

###  E.4 Reproducibility

Figure S13: Reproducibility experiments. Taking the ligand structure
generation as an example, different DiG models were trained by varying
training parameters. The results of the two trained DiG models show high
similarity in terms of ligand structure differences compared to crystal
structures.

In this section, we investigate the reproducibility of DiG training, under
different initialization, and different random seeds which affect the order of
data batches fed into the model during the training. We conduct this
experiment on the protein-ligand systems and compare the distribution of the
generated ligand structures with respect to the crystal structures. Fig. S13
shows the histograms of RMSD statistics for ligand structures generated by
DiG. The left panel is identical to Fig. 4a, where the results are obtained
from the model checkpoint with the training pipeline described in
Supplementary Sec. A.3, which is used to generate all results for protein-
ligand systems in this paper. The initialization of this model is from pre-
training on the CrossDock dataset. The right panel shows the results from
another model checkpoint without any pre-training but a random initialization.
Other hyperparameters except the random seeds are kept the same in the two
training processes. We observe that the distributions from the two model
checkpoints are very similar, suggesting that DiG training is robustly
reproducible.

##  Appendix F Limitations

###  F.1 Limitations on Data Quantity and Quality

One of the major challenges and limitations of DiG is the scarcity of data for
training and evaluating the deep learning models for equilibrium distribution
prediction. The ground truth data of equilibrium distribution of different
molecular systems are not easily available, as they require massive
computational resources and time to generate by molecular dynamics simulation
or other methods. Therefore, we only have access to very little data for some
molecular systems, and the data quality and quantity may not be sufficient to
support the learning and generalization of DiG.

For example, for the catalyst systems, we use the Open Catalyst dataset
(chanussot2021open, ), which contains DFT-based molecular dynamics simulations
of catalyst-adsorbate systems for only 80 or 320 femtoseconds. However, this
simulation time is too short to capture the dynamics and transitions of the
systems, and the structures may not move significantly from their initial
positions. Thus, we need to traverse initial positions of the adsorbate to
find out all the adsorption configurations within a unit cell. With longer MD
trajectories, the model should be able to sample more adsorption
configurations from a single initial position of adsorbate. Moreover, the
adsorption configurations in the MD trajectories, which start from the relaxed
configurations, tend to have low energies, and high-energy configurations are
rare in the dataset. Thus, the density for high-energy configurations
estimated by the learned DiG may be inaccurate.

For property-guided structure generation, we use a dataset of 15,697 carbon
crystal structures (lucarbondataset, ) to train the model, which is also very
limited compared to the huge space of possible carbon polymorphs. The
generative model trained on such data may not be able to recover all the
structures in the dataset, let alone generalize to unseen carbon polymorphs.
For example, our conditionally generated structures (including 2, 4, 6, 8
carbon atoms per unit cell) only match $88.33\%$ of the structures in the
dataset with the same numbers of atoms, using the StructureMatcher from
Pymatgen.

For the protein conformation sampling, we collect MD simulations of about only
1000 proteins, each with 100 nanoseconds of simulation time. However, this
amount of data may not be enough to cover the diversity and complexity of
different protein structures and functions, especially for large and complex
proteins that may have longer time scales and more energy barriers for
conformational changes. Furthermore, the desired equilibrium distribution for
protein model training is not well represented by the available data. Although
some simulated data have sufficient length to approximate the equilibrium
distribution, they are too scarce to support the neural network models with
robust generalization and practical accuracy. Therefore, we resort to using a
larger number of experimental structures, such as the PDB dataset, as the
initial training data for the models. However, these data influence the final
distribution learned by the models. We observe that these data lead to more
accurate structures and better generalization, but also to a more concentrated
learned distribution. Besides the issue of simulation length, the accuracy of
simulation may also be problematic. We find that some structures simulated by
molecular dynamics in some systems deviate too much from the experimental
structures in terms of structural accuracy. Moreover, in some systems
molecular dynamics is highly sensitive to the initial state, and different
initial states can result in different distributions in practice. These
factors compromise the use of simulated data as approximations of the
equilibrium distribution and cause the model to learn inaccurate
distributions.

Similar issues also exist in protein-ligand training, where the simulation
time does not warrant equilibrium distributions. The simulation trajectories
used for ligand-structure model training are limited to 100 ns. Yet, we
observed ligand dissociation from the pocket in some of the systems, which
were excluded from model training. Moreover, our training set only covers a
small fraction of systems. In the CrossDocked and our self-generated MD
simulation datasets, there are about only 1000 unique proteins, which may
affect the generalization of our model.

The data scarcity also affects the evaluation of DiG, as we do not have enough
ground truth data of equilibrium distribution to compare with the predictions
of DiG. Therefore, we have to rely on indirect metrics, such as the energy
function, the structural quality metrics, or properties, to measure the
quality and diversity of the generated structures. However, these metrics may
not fully capture the accuracy and reliability of the equilibrium distribution
prediction, and may have some limitations or biases themselves. For example,
the structural quality metrics may not account for the dynamic and stochastic
nature of the molecular structures, and may have some dependencies on the
reference structures or the alignment methods. The property prediction may not
be sensitive to the subtle changes or variations of the molecular structures,
and may have some noise or uncertainty in the measurements or the models.

Therefore, we acknowledge that data scarcity is a serious limitation for DiG,
and we hope that more and better data of equilibrium distribution of molecular
systems can be generated and shared in the future, to enable more robust and
reliable learning and evaluation of DiG and other equilibrium distribution
prediction methods.

Data limitations also affect the training of the property predicting model
$q_{\mathcal{D}}(c\mid\mathbf{R})$, which guides the structure generation
based on properties. In the sampling process for a desired property value $c$
as described by Eqs. (5, 12) , the structures in early stages (i.e., large $i$
or $t$) are nearly random noise, for which the predictor model
$q_{\mathcal{D}}(c\mid\mathbf{R})$ are not typically trained on, making its
contribution $\nabla_{\mathbf{R}}\log q_{\mathcal{D}}(c\mid\mathbf{R})$ less
controlled. This may not be as harmful as it appears, since even with an
oracle property predictor, what the early stage of sampling does is still
refining the random structure to physically reasonable ones in which the
predictor model contribution $\nabla_{\mathbf{R}}\log
q_{\mathcal{D}}(c\mid\mathbf{R})$ does not dominate. The effect of this
process also largely aligns with the desired property since for which a random
structure is unlikely to achieve. However, we do observe that in some cases
small perturbance to the structure causes significant changes in the band gap
prediction using the M3GNet predictor, which makes sampling structures with
demanded band gap more challenging. This also adds to the evidence of
overfitting of the predictor to the limited stable structures. We also find
that the model produces more unphysical structures that violate the geometric
or energetic constraints when conditional generation. For example, a
conditionally generated structure close to a graphite may not have perfect
bond angles of exactly $120^{\circ}$. The band gap predictor model is trained
on stable carbon polymorphs. But structures in the denoising process can be
quite unstable. Thus the gradients from the predictor used to guide the
denoising process may be inaccurate and does not always lead to physical
structures. Adding guidance from an energy prediction model in the conditional
generation process may guide to more physical structures. Training the
property predictor with more abundant carbon polymorphs can also improve the
quality of conditionally generated structures. Thus, it would still be helpful
to generate labeled data on more noisy structures and train the property
predictor model on them, so that the model could take effect earlier in the
sampling process to better guide the structure to the desired property.

###  F.2 Limitations on Energy Function

In applications involving proteins, we adopt the common choice of a coarse-
grained representation for proteins to reduce the dimensionality of the
problem while maintaining most of the structural features. This nevertheless
incurs challenges from the energy function (equivalently, force field) side:
we have to convert a full-atom force field to the coarse-grained level. This
is required in the PIDP training as shown in Eqs. (57, 63) , while employing
an established coarse-grained force field is neither suitable since it is
unnecessarily the coarse-grained version of the full-atom force field used in
the simulation to generate the dataset. Such a conversion is conducted in Alg.
2, but this is not precise for coarse-graining for statistical use.
Specifically, if denoting the invertible transformed full-atom coordinate
$\bar{\mathbf{R}}$ and energy function as
$(\mathbf{R}_{\mathrm{CG}},\mathbf{R}_{\mathrm{FG}})$ and
$E(\mathbf{R}_{\mathrm{CG}},\mathbf{R}_{\mathrm{FG}})$ where
$\mathbf{R}_{\mathrm{CG}}$ denotes the coarse-grained coordinates
($(\mathbf{C},\mathbf{q})$ in Alg. 2) and $\mathbf{R}_{\mathrm{FG}}$ the fine-
grained details ($\mathbf{X}$ excluding $\mathbf{C}$ in Alg. 2), the required
coarse-grained energy (equilibrium free energy) under temperature $T$ would
be:

$\displaystyle
E_{\mathrm{CG}}(\mathbf{R}_{\mathrm{CG}})=-k_{\mathrm{B}}T\log\int\exp\left\\{-\frac{E(\mathbf{R}_{\mathrm{CG}},\mathbf{R}_{\mathrm{FG}})}{k_{\mathrm{B}}T}\right\\}\,\mathrm{d}\mathbf{R}_{\mathrm{FG}}.$
(88)

In practice, this integral is hard to evaluate, and is a long-standing problem
in statistical mechanics and Bayesian statistics. Even in the case of Alg. 2
where we have access to the full-atom coordinates of a query structure, the
estimation is still an approximation. To see this, the gradient (negative
force) of Eq. (88) can be written as:

$\displaystyle\nabla_{\mathbf{R}_{\mathrm{CG}}}E_{\mathrm{CG}}(\mathbf{R}_{\mathrm{CG}})=\mathbb{E}_{p_{T}(\mathbf{R}_{\mathrm{FG}}\mid\mathbf{R}_{\mathrm{CG}})}[\nabla_{\mathbf{R}_{\mathrm{CG}}}E(\mathbf{R}_{\mathrm{CG}},\mathbf{R}_{\mathrm{FG}})],$
(89) $\displaystyle
p_{T}(\mathbf{R}_{\mathrm{FG}}\mid\mathbf{R}_{\mathrm{CG}}):=\frac{\exp\left\\{-\frac{E(\mathbf{R}_{\mathrm{CG}},\mathbf{R}_{\mathrm{FG}})}{k_{\mathrm{B}}T}\right\\}}{\int\exp\left\\{-\frac{E(\mathbf{R}_{\mathrm{CG}},\mathbf{R}_{\mathrm{FG}})}{k_{\mathrm{B}}T}\right\\}\,\mathrm{d}\mathbf{R}_{\mathrm{FG}}},$
(90)

so in principle, the coarse-grained gradient is an average of the full-atom
gradient over samples from
$p_{T}(\mathbf{R}_{\mathrm{FG}}\mid\mathbf{R}_{\mathrm{CG}})$. Under this
perspective, the rigid body assumption that Alg. 2 is based on can be
understood as assuming
$p_{T}(\mathbf{R}_{\mathrm{FG}}\mid\mathbf{R}_{\mathrm{CG}})$ only
concentrates on one value of $\mathbf{R}_{\mathrm{FG}}$ (i.e., a Dirac delta
distribution), meaning $\mathbf{R}_{\mathrm{FG}}$ can be uniquely determined
from the given $\mathbf{R}_{\mathrm{CG}}$. In the algorithm, this
$\mathbf{R}_{\mathrm{FG}}$ is provided from the corresponding full-atom
coordinates. This is a good approximation if the true
$p_{T}(\mathbf{R}_{\mathrm{FG}}\mid\mathbf{R}_{\mathrm{CG}})$ distribution
indeed concentrates at the determined $\mathbf{R}_{\mathrm{FG}}$ value;
otherwise (e.g., there are very flexible residue or in relatively high
temperature), a more precise coarse-graining method for the energy function is
required (e.g. (koehler2023flow, )).

Moreover, in PIDP training Eq. (7), although the samples
$\\{\mathbf{R}_{\mathcal{D},0}^{(m)}\\}_{m=1}^{M}$ for evaluating the loss can
be taken as any that are relevant to the problem in principle, overly loosely
chosen structures may cause numerical difficulties as the corresponding
gradient energy gradient would be too large. This is the limiting issue from
using normal mode perturbed structures hence we have to resort to MD
structures. A possible approach to mitigate this limitation is using a
“milder” energy function, which does not increase its value as steeply on off-
equilibrium structures. For PIDP training, the energy function only needs to
indicate a very small probability, and it does not matter much how small it
is, as all small values almost equally indicate a vacuum.

###  F.3 Limitations on Model Architecture and Scale

Another crucial limitation of DiG is the model restriction, which is resulted
from the compromise between the model capacity and the required computational
resource. The model capacity determines the expressiveness and generalization
ability of the deep learning models for equilibrium distribution prediction,
while the computational resource determines the availability and speed of the
training and inference processes. In this work, we have to face the constraint
of the computational resource and make some choices that may affect the
performance of DiG.

For example, for the protein systems, we use a 12-layer Graphormer with about
80M learnable parameters, which is relatively small considering the complexity
and diversity of protein structures and distributions. The model capacity of
DiG may not be enough to capture the intricate and high-dimensional energy
landscapes and distributions of protein systems, and this can be evidenced by
the structural quality of the generated protein structures. We observe that
smaller models are easily outperformed by larger models. For example, a
4-layer Graphormer with about 10M learnable parameters only produces a median
TM-socre zhang2004scoring  of 0.46 on the PDB validation dataset, while a
12-layer Graphormer can easily reach more than 0.8.

Besides, we fix the parameters of the pre-trained Evoformer module in
AlphaFold, which is used to extract the features from the protein sequence and
the MSA. Evoformer is a powerful and sophisticated module that can encode rich
and informative features for protein structure prediction, but it is also
computationally expensive and complex, and hasn’t been fine-tuned during the
training of DiG for predicting equilibrium distribution. This may lead to a
significant performance drop since the frozen parameters of Evoformer restrict
the expressiveness of DiG very much. In Supplementary Fig. S8a and S8b, the
high-density regions of MD simulation are perfectly aligned with both known
structures and predicted structures by AlphaFold, but there is an observable
shift of the high-density regions generated by DiG. We suspect that the shift
is due to the limitation of model capacity. Specifically, a fixed Evoformer
implemented in DiG does not perform as well as a learnable Evoformer that is
used in AlphaFold. A similar observation is that, although for RBD protein,
both AlphaFold and DiG can generate high-quality structures with TMscores
$>0.8$ in all cases, their performances are different for the case of main
protease, where AlphaFold could generate high-quality structures, but
structures with TMscores $>0.8$ generated by DiG only accounted for about
$6.8\%$ of all generated structures (The average TMscore of all generated
structures of main protease by DiG is about 0.64). This performance gap
between AlphaFold and DiG can be possibly caused by the fixed Evoformer. If
so, the performance of DiG could be significantly improved if we can fine-tune
the Evoformer module with the data and objective of DiG.

Moreover, in this work, we mainly focus on algorithm development, but not on
model architecture development. We mainly use the existing deep learning
architectures, such as Graphormer and Evoformer. More advanced and specialized
architectures that can better exploit the 3D conformational information and
the physics principles of molecular systems may improve the performance and
efficiency of DiG.

Therefore, we acknowledge that the model architecture restriction is a serious
limitation for DiG in the current implementation, and this will be resolved in
the future with enhanced capacity of advanced models.