SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Rui Jiao $^{1,2}$  Wenbing Huang $^{3,4*}$  Peijia Lin $^{5}$  Jiaqi Han $^{6}$  Pin Chen $^{5}$  Yutong Lu $^{5}$  Yang Liu $^{1,2*}$
$^{1}$ Dept. of Comp. Sci. & Tech., Institute for AI, Tsinghua University
$^{2}$ Institute for AIR, Tsinghua University
$^{3}$ Gaojing School of Artificial Intelligence, Renmin University of China
$^{4}$  Beijing Key Laboratory of Big Data Management and Analysis Methods, Beijing, China
5 National Supercomputer Center in Guangzhou,
School of Computer Science and Engineering, Sun Yat-sen University
6 Stanford University
Crystal Structure Prediction (CSP) is crucial in various scientific disciplines. While CSP can be addressed by employing currently-prevailing generative models (e.g. diffusion models), this task encounters unique challenges owing to the symmetric geometry of crystal structures—the invariance of translation, rotation, and periodicity. To incorporate the above symmetries, this paper proposes DiffCSP, a novel diffusion model to learn the structure distribution from stable crystals. To be specific, DiffCSP jointly generates the lattice and atom coordinates for each crystal by employing a periodic-E(3)-equivariant denoising model, to better model the crystal geometry. Notably, different from related equivariant generative approaches, DiffCSP leverages fractional coordinates other than Cartesian coordinates to represent crystals, remarkably promoting the diffusion and the generation process of atom positions. Extensive experiments verify that our DiffCSP significantly outperforms existing CSP methods, with a much lower computation cost in contrast to DFT-based methods. Moreover, the superiority of DiffCSP is also observed when it is extended for ab initio crystal generation. Code is available at https://github.com/jiaor17/DiffCSP.
Crystal Structure Prediction (CSP), which returns the stable 3D structure of a compound based solely on its composition, has been a goal in physical sciences since the 1950s [1]. As crystals are the foundation of various materials, estimating their structures in 3D space determines the physical and chemical properties that greatly influence the application to various academic and industrial sciences, such as the design of drugs, batteries, and catalysts [2]. Conventional methods towards CSP mostly apply the Density Functional Theory (DFT) [3] to compute the energy at each iteration, guided by optimization algorithms (such as random search [4], Bayesian optimization [5], etc.) to iteratively search for the stable state corresponding to the local minima of the energy surface [6].
The DFT-based approaches are computationally-intensive. Recent attention has been paid to deep generative models, which directly learn the distribution from the training data consisting of stable structures [7, 8]. More recently, diffusion models, a special kind of deep generative models are employed for crystal generation [9], encouraged by their better physical interpretability and enhanced performance than other generative models. Intuitively, by conducting diffusion on stable structures,
* Wenbing Huang and Yang Liu are corresponding authors.
the denoising process in diffusion models acts like a force field that drives the atom coordinates towards the energy local minimum and thus is able to increase stability. Indeed, the success of diffusion models is observed in broad scientific domains, including molecular conformation generation [10], protein structure prediction [11] and protein docking [12].
However, designing diffusion models for CSP is challenging. From the perspective of physics, any E(3) transformation, including translation, rotation, and reflection, of the crystal coordinates does not change the physical law and thus keeps the crystal distribution invariant. In other words, the generation process we design should yield E(3) invariant samples. Moreover, in contrast to other types of structures such as small molecules [13] and proteins [14], CSP exhibits unique challenges, mainly incurred by the periodicity of the atom arrangement in crystals. Figure 1 displays a crystal where the atoms in a unit cell are repeated infinitely in space. We identify such unique symmetry, jointly consisting of E(3) invariance and periodicity, as periodic  $E(3)$  invariance in this paper. Generating such type of structures requires not only modeling the distribution of the atom coordinates within every cell, but also inferring how their bases (a.k.a. lattice vectors) are placed in 3D space. Interestingly, as we will show in  $\S 4.1$ , such view offers a natural disentanglement for fulfilling the periodic E(3) invariance by separately enforcing constraints on fractional coordinates and lattice vectors, which permits a feasible implementation to encode the crystal symmetry.
In this work, we introduce DiffCSP, an equivariant diffusion method to address CSP. Considering the specificity of the crystal geometry, our DiffCSP jointly generates the lattice vectors and the fractional coordinates of all atoms, by employing a proposed denoising model that is theoretically proved to generate periodic-E(3)-invariant samples. A preferable characteristic of DiffCSP is that it leverages the fractional coordinate system (defined in § 3) other than the Cartesian system used in previous methods to represent crystals [9, 15], which encodes periodicity intrinsically. In particular, the fractional representation not only allows us to consider Wrapped Normal (WN) distribution [16] to better model the periodicity, but also facilitates the design of the denoising model via the Fourier transformation, compared to the traditional multi-graph encoder in crystal modeling [15].
CDVAE [9] is closely related with our paper. It adopts an equivariant Variational Auto-Encoder (VAE) based framework to learn the data distribution and then generates crystals in a score-matching-based diffusion process. However, CDVAE focuses mainly on ab initio crystal generation where the composition is also randomly sampled, which is distinct from the CSP task in this paper. Moreover, while CDVAE first predicts the lattice and then updates the coordinates with the fixed lattice, we jointly update the lattice and coordinates to better model the crystal geometry. Besides, CDVAE represents crystals by Cartesian coordinates upon multi-graph modeling, whereas our DiffCSP applies fractional coordinates without multi-graph modeling as mentioned above.
To sum up, our contributions are as follows:
- To the best of our knowledge, we are the first to apply equivariant diffusion-based methods to address CSP. The proposed DiffCSP is more insightful than current learning-based approaches as the periodic E(3) invariance has been delicately considered.
- DiffCSP conducts joint diffusion on lattices and fractional coordinates, which is capable of capturing the crystal geometry as a whole. Besides, the usage of fractional coordinates in place of Cartesian coordinates used in previous methods (e.g. CDVAE [9]) remarkably promotes the diffusion and the generation process of atom positions.
- We verify the efficacy of DiffCSP on the CSP task against learning-based and DFT-based methods, and sufficiently ablate each proposed component in DiffCSP. We further extend DiffCSP into ab initio generation and show its effectiveness against related methods.
Crystal Structure Prediction Traditional computation methods [4, 5, 17, 18] combine DFT [3] with optimization algorithms to search for local minima in the potential energy surface. However, DFT is computationally intensive, making it dilemmatic to balance efficiency and accuracy. With the improvement of crystal databases, machine-learning methods are applied as alternative energy predictors to DFT followed by optimization steps [19, 20, 21]. Apart from the predict-optimize paradigm, another line of approaches directly learns stable structures from data by deep generative models, which represents crystals by 3D voxels [7, 22, 23], distance matrices [8, 24, 25] or 3D
coordinates [26, 27, 28]. Unfortunately, these methods are unaware of the full symmetries of the crystal structure. CDVAE [9] has taken the required symmetries into account. However, as mentioned above, the initial version of CDVAE is for different task and utilizes different generation process.
Equivariant Graph Neural Networks Geometrically equivariant Graph Neural Networks (GNNs) that ensure E(3) symmetry are powerful tools to represent physical objects [29, 30, 31, 32, 33], and have showcased the superiority in modeling 3D structures [34, 35]. To further model the periodic materials, Xie and Grossman [15] propose the multi-graph edge construction to capture the periodicity by connecting the edges between adjacent lattices. Yan et al. [36] further introduce periodic pattern encoding into a Transformer-based backbone. In this work, we achieve the periodic invariance by introducing the Fourier transform on fractional coordinates.
Diffusion Generative Models Motivated by the non-equilibrium thermodynamics [37], diffusion models connect the data distribution with the prior distribution via forward and backward Markov chains [38], and have made remarkable progress in the field of image generation [39, 40]. Equipped with equivariant GNNs, diffusion models are capable of generating samples from the invariant distribution, which is desirable in conformation generation [10, 13], ab initio molecule design [41], protein generation [42], and so on. Recent works extend the diffusion models onto Riemann manifolds [43, 44], and enable the generation of periodic features like torsion angles [12, 16].
Representation of crystal structures A 3D crystal can be represented as the infinite periodic arrangement of atoms in 3D space, and the smallest repeating unit is called a unit cell, as shown in Figure 1. A unit cell can be defined by a triplet  $\mathcal{M} = (A,X,L)$ , where  $A = [a_{1},a_{2},\dots,a_{N}]\in \mathbb{R}^{h\times N}$  denotes the list of the one-hot representations of atom types,  $X = [x_{1},x_{2},\dots,x_{N}]\in \mathbb{R}^{3\times N}$  consists of Cartesian coordinates of the atoms, and  $L = [l_1,l_2,l_3]\in \mathbb{R}^{3\times 3}$  represents the lattice matrix containing three basic vectors to describe the periodicity of the crystal. The infinite periodic crystal structure is represented by
$$
\left\{\left(\boldsymbol {a} _ {i} ^ {\prime}, \boldsymbol {x} _ {i} ^ {\prime}\right) \mid \boldsymbol {a} _ {i} ^ {\prime} = \boldsymbol {a} _ {i}, \boldsymbol {x} _ {i} ^ {\prime} = \boldsymbol {x} _ {i} + \boldsymbol {L k}, \forall \boldsymbol {k} \in \mathbb {Z} ^ {3 \times 1} \right\}, \tag {1}
$$
where the  $j$ -th element of the integral vector  $\pmb{k}$  denotes the integral 3D translation in units of  $\pmb{l}_j$ .
Fractional coordinate system The Cartesian coordinate system  $X$  leverages three standard orthogonal bases as the coordinate axes. In crystallography, the fractional coordinate system is usually applied to reflect the periodicity of the crystal structure [26, 27, 28, 45], which utilizes the lattices  $(l_1, l_2, l_3)$  as the bases. In this way, a point represented by the fractional coordinate vector  $\pmb{f} = [f_1, f_2, f_3]^\top \in [0,1)^3$  corresponds to the Cartesian vector  $\pmb{x} = \sum_{i=1}^{3} f_i \pmb{l}_i$ . This paper employs the fractional coordinate system, and denotes the crystal by  $\mathcal{M} = (\pmb{A}, \pmb{F}, \pmb{L})$ , where the fractional coordinates of all atoms in a cell compose the matrix  $\pmb{F} \in [0,1)^{3 \times N}$ .
Task definition CSP predicts for each unit cell the lattice matrix  $\pmb{L}$  and the fractional matrix  $\pmb{F}$  given its chemical composition  $\pmb{A}$ , namely, learning the conditional distribution  $p(\pmb{L},\pmb{F}|\pmb{A})$ .
4 The Proposed Method: DiffCSP
This section first presents the symmetries of the crystal geometry, and then introduces the joint equivariant diffusion process on  $L$  and  $F$ , followed by the architecture of the denoising function.
4.1 Symmetries of Crystal Structure Distribution
While various generative models can be utilized to address CSP, this task encounters particular challenges, including constraints arising from symmetries of crystal structure distribution. Here, we consider the three types of symmetries in the distribution of  $p(\boldsymbol{L}, \boldsymbol{F} \mid \boldsymbol{A})$ : permutation invariance,  $O(3)$  invariance, and periodic translation invariance. Their detailed definitions are provided as follows.
Definition 1 (Permutation Invariance). For any permutation  $\pmb{P} \in \mathrm{S}_N$ ,  $p(\pmb{L},\pmb{F} \mid \pmb{A}) = p(\pmb{L},\pmb{F}\pmb{P} \mid \pmb{A}\pmb{P})$ , i.e., changing the order of atoms will not change the distribution.
Figure 2: Overview of DiffCSP. Given the composition  $\mathbf{A}$ , we denote the crystal, its lattice and fractional coordinate matrix at time  $t$  as  $\mathcal{M}_t$ ,  $L_t$  and  $F_t$ , respectively. The terms  $\epsilon_L$  and  $\epsilon_F$  are Gaussian noises,  $\hat{\epsilon}_L$  and  $\hat{\epsilon}_F$  are predicted by the denoising model  $\phi$ .
Definition 2 (O(3) Invariance). For any orthogonal transformation  $\pmb{Q} \in \mathbb{R}^{3 \times 3}$  satisfying  $\pmb{Q}^{\top} \pmb{Q} = \pmb{I}$ ,  $p(\pmb{QL}, \pmb{F} \mid \pmb{A}) = p(\pmb{L}, \pmb{F} \mid \pmb{A})$ , namely, any rotation/reflection of  $\pmb{L}$  keeps the distribution unchanged.
Definition 3 (Periodic Translation Invariance). For any translation  $\pmb{t} \in \mathbb{R}^{3 \times 1}$ ,  $p(\pmb{L}, w(\pmb{F} + \pmb{t}\pmb{1}^{\top}) \mid \pmb{A}) = p(\pmb{L}, \pmb{F} \mid \pmb{A})$ , where the function  $w(\pmb{F}) = \pmb{F} - [\pmb{F}] \in [0,1)^{3 \times N}$  returns the fractional part of each element in  $\pmb{F}$ , and  $\pmb{1} \in \mathbb{R}^{3 \times 1}$  is a vector with all elements set to one. It explains that any periodic translation of  $\pmb{F}$  will not change the distribution<sup>2</sup>.
The permutation invariance is tractably encapsulated by using GNNs as the backbone for generation [47]. We mainly focus on the other two kinds of invariance (see Figure 1), since GNNs are our default choices. For simplicity, we compactly term the  $O(3)$  invariance and periodic translation invariance as periodic  $E(3)$  invariance henceforth. Previous approaches (e.g. [9, 36]) adopt Cartesian coordinates  $\pmb{X}$  other than fractional coordinates  $\pmb{F}$ , hence their derived forms of the symmetry are different. Particularly, in Definition 2, the orthogonal transformation additionally acts on  $\pmb{X}$ ; in Definition 3, the periodic translation  $w(\pmb{F} + t\pmb{1}^{\top})$  becomes the translation along the lattice bases  $\pmb{X} + Lt\pmb{1}^{\top}$ ; besides,  $\pmb{X}$  should also maintain E(3) translation invariance, that is  $p(\pmb{L},\pmb{X} + t\pmb{1}^{\top}|\pmb{A}) = p(\pmb{L},\pmb{X}|\pmb{A})$ . With the help of the fractional system, the periodic E(3) invariance
Figure 1: (a)  $\rightarrow$  (b): The orthogonal transformation of the lattice vectors. (c)  $\rightarrow$  (d): The periodic translation of the fractional coordinates. Both cases do not change the structure.
is made tractable by fulfilling O(3) invariance w.r.t. the orthogonal transformations on  $L$  and periodic translation invariance w.r.t. the periodic translations on  $\pmb{F}$ , respectively. In this way, such approach, as detailed in the next section, facilitates the application of diffusion methods to the CSP task.
4.2 Joint Equivariant Diffusion
Our method DiffCSP addresses CSP by simultaneously diffusing the lattice  $\pmb{L}$  and the fractional coordinate matrix  $\pmb{F}$ . Given the atom composition  $\pmb{A}$ ,  $\mathcal{M}_t$  denotes the intermediate state of  $\pmb{L}$  and  $\pmb{F}$  at time step  $t$  ( $0 \leq t \leq T$ ). DiffCSP defines two Markov processes: the forward diffusion process gradually adds noise to  $\mathcal{M}_0$ , and the backward generation process iteratively samples from the prior distribution  $\mathcal{M}_T$  to recover the origin data  $\mathcal{M}_0$ .
2Previous works (e.g. [36]) further discuss the scaling invariance of a unit cell formed by periodic boundaries, allowing  $L\to \alpha L,\forall \alpha \in \mathbb{N}_{+}^{3}$  . In this paper, the scaling invariance is unnecessary since we apply the Niggli reduction [46] on the primitive cell as a canonical scale representation of the lattice vectors where we fix  $\alpha = (1,1,1)^{\top}$  . Additionally, periodic translation invariance in our paper is equivalent to the invariance of shifting periodic boundaries in [36]. We provide more discussions in Appendix A.4.
Joining the statements in § 4.1, the recovered distribution from  $\mathcal{M}_T$  should meet periodic E(3) invariance. Such requirement is satisfied if the prior distribution  $p(\mathcal{M}_T)$  is invariant and the Markov transition  $p(\mathcal{M}_{t-1} \mid \mathcal{M}_t)$  is equivariant, according to the diffusion-based generation literature [10]. Here, an equivariant transition is specified as  $p(g \cdot \mathcal{M}_{t-1} \mid g \cdot \mathcal{M}_t) = p(\mathcal{M}_{t-1} \mid \mathcal{M}_t)$  where  $g \cdot \mathcal{M}$  refers to any orthogonal/translational transformation  $g$  acts on  $\mathcal{M}$  in the way presented in Definitions 2-3. We separately explain the derivation details of  $\pmb{L}$  and  $\pmb{F}$  below. The detailed flowcharts are summarized in Algorithms 1 and 2 in Appendix B.3.
Diffusion on  $\pmb{L}$  Given that  $\pmb{L}$  is a continuous variable, we exploit Denoising Diffusion Probabilistic Model (DDPM) [38] to accomplish the generation. We define the forward process that progressively diffuses  $\pmb{L}_0$  towards the Normal prior  $p(\pmb{L}_T) = \mathcal{N}(0, \pmb{I})$  by  $q(\pmb{L}_t | \pmb{L}_{t-1})$  which can be devised as the probability conditional on the initial distribution:
$$
q \left(\boldsymbol {L} _ {t} \mid \boldsymbol {L} _ {0}\right) = \mathcal {N} \left(\boldsymbol {L} _ {t} \mid \sqrt {\bar {\alpha} _ {t}} \boldsymbol {L} _ {0}, (1 - \bar {\alpha} _ {t}) \boldsymbol {I}\right), \tag {2}
$$
where  $\beta_{t} \in (0,1)$  controls the variance, and  $\bar{\alpha}_{t} = \prod_{s=1}^{t} \alpha_{t} = \prod_{s=1}^{t} (1 - \beta_{t})$  is valued in accordance to the cosine scheduler [48].
The backward generation process is given by:
$$
p \left(\boldsymbol {L} _ {t - 1} \mid \mathcal {M} _ {t}\right) = \mathcal {N} \left(\boldsymbol {L} _ {t - 1} \mid \mu \left(\mathcal {M} _ {t}\right), \sigma^ {2} \left(\mathcal {M} _ {t}\right) \boldsymbol {I}\right), \tag {3}
$$
where  $\mu (\mathcal{M}_t) = \frac{1}{\sqrt{\alpha_t}}\Bigl (L_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\hat{\epsilon}_L(\mathcal{M}_t,t)\Bigr),\sigma^2 (\mathcal{M}_t) = \beta_t\frac{1 - \bar{\alpha}_{t - 1}}{1 - \bar{\alpha}_t}$ . The denoising term  $\hat{\epsilon}_{L}(\mathcal{M}_{t},t)\in \mathbb{R}^{3\times 3}$  is predicted by the model  $\phi (L_{t},F_{t},A,t)$ .
As the prior distribution  $p(\pmb{L}_T) = \mathcal{N}(0, \pmb{I})$  is already O(3)-invariant, we require the generation process in Eq. (3) to be O(3)-equivariant, which is formally stated below.
Proposition 1. The marginal distribution  $p(L_0)$  by Eq. (3) is  $O(3)$ -invariant if  $\hat{\epsilon}_{L}(\mathcal{M}_{t}, t)$  is  $O(3)$ -equivariant, namely  $\hat{\epsilon}_{L}(Q L_{t}, F_{t}, A, t) = Q \hat{\epsilon}_{L}(L_{t}, F_{t}, A, t), \forall Q^{\top} Q = I$ .
To train the denoising model  $\phi$ , we first sample  $\epsilon_{L} \sim \mathcal{N}(0, I)$  and reparameterize  $L_{t} = \sqrt{\bar{\alpha}_{t}} L_{0} + \sqrt{1 - \bar{\alpha}_{t}} \epsilon_{L}$  based on Eq. (2). The training objective is defined as the  $\ell_{2}$  loss between  $\epsilon_{L}$  and  $\hat{\epsilon}_{L}$ :
$$
\mathcal {L} _ {\boldsymbol {L}} = \mathbb {E} _ {\boldsymbol {\epsilon} _ {\boldsymbol {L}} \sim \mathcal {N} (0, I), t \sim \mathcal {U} (1, T)} [ \| \boldsymbol {\epsilon} _ {\boldsymbol {L}} - \hat {\boldsymbol {\epsilon}} _ {\boldsymbol {L}} (\mathcal {M} _ {t}, t) \| _ {2} ^ {2} ]. \tag {4}
$$
Diffusion on  $\pmb{F}$  The domain of fractional coordinates  $[0,1)^{3\times N}$  forms a quotient space  $\mathbb{R}^{3\times N} / \mathbb{Z}^{3\times N}$  induced by the crystal periodicity. It is not suitable to apply the above DDPM fashion to generate  $\pmb{F}$ , as the normal distribution used in DDPM is unable to model the cyclical and bounded domain of  $\pmb{F}$ . Instead, we leverage Score-Matching (SM) based framework [49, 50] along with Wrapped Normal (WN) distribution [43] to fit the specificity here. Note that WN distribution has been explored in generative models, such as molecular conformation generation [16].
During the forward process, we first sample each column of  $\epsilon_{\pmb{F}} \in \mathbb{R}^{3 \times N}$  from  $\mathcal{N}(0, I)$ , and then acquire  $\pmb{F}_t = w(\pmb{F}_0 + \sigma_t \epsilon_{\pmb{F}})$  where the truncation function  $w(\cdot)$  is already defined in Definition 3. This truncated sampling implies the WN transition:
$$
q \left(\boldsymbol {F} _ {t} \mid \boldsymbol {F} _ {0}\right) \propto \sum_ {\boldsymbol {Z} \in \mathbb {Z} ^ {3 \times N}} \exp \left(- \frac {\| \boldsymbol {F} _ {t} - \boldsymbol {F} _ {0} + \boldsymbol {Z} \| _ {F} ^ {2}}{2 \sigma_ {t} ^ {2}}\right). \tag {5}
$$
Basically, this process ensures the probability distribution over  $[z,z + 1)^{3\times N}$  for any integer  $z$  to be the same to keep the crystal periodicity. Here, the noise scale  $\sigma_{t}$  obeys the exponential scheduler:  $\sigma_0 = 0$  and  $\sigma_t = \sigma_1\left(\frac{\sigma_T}{\sigma_1}\right)^{\frac{t - 1}{T - 1}}$ , if  $t > 0$ . Desirably,  $q(\pmb {F}_t|\pmb {F}_0)$  is periodic translation equivariant, and approaches a uniform distribution  $\mathcal{U}(0,1)$  if  $\sigma_T$  is sufficiently large.
For the backward process, we first initialize  $\mathbf{F}_T$  from the uniform distribution  $\mathcal{U}(0,1)$ , which is periodic translation invariant. With the denoising term  $\hat{\epsilon}_{\mathbf{F}}$  predicted by  $\phi(\mathbf{L}_t, \mathbf{F}_t, \mathbf{A}, t)$ , we combine the ancestral predictor [38, 50] with the Langevin corrector [49] to sample  $\mathbf{F}_0$ . We immediately have:
Proposition 2. The marginal distribution  $p(\pmb{F}_0)$  is periodic translation invariant if  $\hat{\epsilon}_{\pmb{F}}(\mathcal{M}_t, t)$  is periodic translation invariant, namely  $\hat{\epsilon}_{\pmb{F}}(\pmb{L}_t, \pmb{F}_t, \pmb{A}, t) = \hat{\epsilon}_{\pmb{F}}(\pmb{L}_t, w(\pmb{F}_t + t\pmb{1}^\top), \pmb{A}, t), \forall t \in \mathbb{R}^3$ .
The training objective for score matching is:
$$
\mathcal {L} _ {\boldsymbol {F}} = \mathbb {E} _ {\boldsymbol {F} _ {t} \sim q (\boldsymbol {F} _ {t} | \boldsymbol {F} _ {0}), t \sim \mathcal {U} (1, T)} \left[ \lambda_ {t} \| \nabla_ {\boldsymbol {F} _ {t}} \log q (\boldsymbol {F} _ {t} | \boldsymbol {F} _ {0}) - \hat {\epsilon} _ {\boldsymbol {F}} (\mathcal {M} _ {t}, t) \| _ {2} ^ {2} \right],
$$
where  $\lambda_{t} = \mathbb{E}_{\pmb{F}_{t}}^{-1}\left[\| \nabla_{\pmb{F}_{t}}\log q(\pmb{F}_{t}|\pmb{F}_{0})\|_{2}^{2}\right]$  is approximated via Monte-Carlo sampling. More details are deferred to Appendix B.1.
Extension to ab initio crystal generation Although our method is proposed to address CSP where the composition  $\mathbf{A}$  is fixed, our method is able to be extended for the ab initio generation task by further generating  $\mathbf{A}$ . We achieve this by additionally optimizing the one-hot representation  $\mathbf{A}$  with a DDPM-based approach. We provide more details in Appendix G.
4.3 The Architecture of the Denoising Model
This subsection designs the denoising model  $\phi(L, F, A, t)$  that outputs  $\hat{\epsilon}_L$  and  $\hat{\epsilon}_F$  satisfying the properties stated in Proposition 1 and 2.
Let  $\pmb{H}^{(s)} = [h_1^{(s)},\dots ,h_N^{(s)}]$  denote the node representations of the  $s$ -th layer. The input feature is given by  $\pmb{h}_i^{(0)} = \rho (f_{\mathrm{atom}}(\pmb {a}_i),f_{\mathrm{pos}}(t))$ , where  $f_{\mathrm{atom}}$  and  $f_{\mathrm{pos}}$  are the atomic embedding and sinusoidal positional encoding [38, 51], respectively;  $\rho$  is a multi-layer perception (MLP).
Built upon EGNN [32], the  $s$ -th layer message-passing is unfolded as follows: