Crystal Structure Prediction by Joint Equivariant Diffusion
Authors: Rui Jiao
Wenbing Huang
Peijia Lin
Jiaqi Han
Pin Chen
Yutong Lu
Yang Liu
Venue: NeurIPS 2023 poster (2023)
Keywords: crystal structure prediction
equivariant graph neural networks
diffusion generative models
TL;DR: We propose DiffCSP, a diffusion framework to jointly generate the lattice and the fractional coordinates for the Crystal Structure Prediction (CSP) task.

Abstract: 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 remarkably outperforms existing CSP methods, with a much lower computation cost in contrast to DFT-based methods. Moreover, the superiority of DiffCSP is still observed when it is extended for ab initio crystal generation.

Crystal Structure Prediction by Joint Equivariant Diffusion
Rui Ji Wenbing Huang Peijia Lin Jiaqi Han Pin Chen Yutong Lu Yang Liu 1.2
Dept. of Comp. Sci. & Tech., Institute for AI, Tsinghua University Institute for AIR, Tsinghua University
GGollin School of Artificial Intelligence, Renmin University of China 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 Stanford University
Abstract
Crystal Structure Prediction (CSP) is crucial in various scientific disciplines. While CSP can addressed by employing surrently-prevailing generative models (e.g. diffusion models), this task encounters unique challenges owing to the symmet- ric 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, DifICSP jointly generates the lattice and atom coordinates for each crystal by employing a periodic-F(()--quivvarrr model the crystal geometry. Notably, different from related equivariant generative approaches, DiffCSP leverages fractional coordinates other than Cartesian coordi- nates represent crystals, remarkably promoting the diffusion and the generation process of atom positions. Extensive experiments verify that our DifICSP signifi- cantly 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 https: gi thub. com/jiaor 17
denoising model, to better
Introduction
Crystal Structure Prediction (CSP), which returns the stable 3D structure of compound based solely on its composition, has been goal in physical sciences since the 1950s (I]. As crystals the foundation of various materials, estimating their structures in 3D space determines the physical and chemical properties that greatly influence the application various academic and industrial sciences, such as the design of drugs, batteries, and catalysts 22 Conventional methods towards CSP mostly apply the Density Functional Theory (DFT) (3] compute the energy 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-intenn generative models, which directly learn the distribution from the training data consisting of stable structures 7 g More recently, diffusion models, a special kind of deep generative models are employed for crystal generation 99, encouraged by their better physical interpretability and enhanced performance than other generative models. Intuitively, by conducting diffusion on stable structures,
Recent attention has been paid to deep
"Wenbing Huang and Yang Liu are corresponding authors.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

the denoising process in diffusion models acts like force field that drives the atom coordinates to- wards 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 [TO] protein structure prediction
[1] 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 [ and proteins [14] CSP exhibits unique challenges, mainly incurred by the periodicity of the atom arrangement in crystals. Figure displays a crystal where the atoms in a unit cell are repeated infinitely in space. We identify such unique symmetry, jointly consisting of F(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 $ 1.11 such view offers natural disentanglement for fulfilling the periodic E(3) invariance by separately enforcing constraints on fractional coordinates and lattice vectors, which permits feasible implementation to encode the crystal symmetry.
In this work, we introduce DiffCSP, an equivariant diffusion method 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 proposed denoising model that is theoretically proved to generate periodic-E(3)) F(3))-nvariin 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 [[6] 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 (VAE) based framework to learn the data distribution and then generates crystals in a -aachinn based diffusion process. However, CDVAE focuses mainly on ab initio crystal generation where the composition is also randomly sampled, which 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 DifICSP applies fractional coordinates without multi- 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 barning-based approaches as the periodic E(3) invariance has been delicately considered.
DiffCSP conducts joint diffusion on lattices and fractional coordinates, which capable whole. Besides, the usage of fractional coordinates in place of Cartesian coordinates used in previous methods (e.g. CDVAE ( remarkably promotes the diffusion and the generation process of atom positions.
of capturing the crystal geometry
We verify the efficacy of DiffCSP on the CSP task against learning-based and DFT-based methods, and sufficiently ablate each proposed component DiffCSP. We further extend DiffCSP into ab initio generation and show its effectiveness against related methods.
2 Related Works
Crystal Structure Prediction Traditional computation methods 4
5
17 18 combine DFT B
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

---

Unfortunately, these methods are unaware of the full symmetries of the crystal structure. CDVAE 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.
coordinates 226 27 28
Equivariant Graph Neural Networks Geometrically equivariant Graph Neural Networks (GNNs) that ensure E(3) symmetry are powerful tools to represent physical objects (2 30
31
32 33], and
have showcased the superiority modeling 3D structures [34 35 To further model the periodic materials, Xie and Grossman [5] propose the multi-graph edge construction to capture the periodicity by connecting the edges between adjacent lattices. Yan et al. [6 further introduce periodic pattern encoding into a Transformer-bet 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 Markoy 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 [1O [3), ab initio molecule design [41],
protein generation 42,. and so on. Recent works extend the diffusion models onto Riemann manifolds 43 and enable the generation of periodic features like torsion angles 12 [6]
3 Preliminaries
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 ! A unit cell can be defined by a triplet M (A, X, L), where A [a1, a2 avl RaxN denotes the list of the one-hot representations of atom types, X [x 22...., xx 2 consists of Cartesian coordinates of the atoms, and L = [4, , 42 L3] E g represents the lattice matrix containing three basic vectors to describe the periodicity of the crystal. The infinite periodic crystal structure is represented by
a/)aa a;, xt. x, + Lk, Vk E Z3x1),
where the j-th element of the integral vector & denotes the integral 3D translation in units of lj.
Fractional coordinate system The Cartesian coordinate system X leverages three standard orthog- onal bases the coordinate axes. In crystallography, the fractional coordinate system is usually applied to reflect the periodicity of the crystal structure (26. 27 28, [5), which utilizes the lat- tices (4,, L3) as the bases. In this way, a point represented by the fractional coordinate vector f f= [fi f2, fs] E [0, 1) corresponds to the Cartesian vector x 3 file. This paper employs the fractional coordinate system, and denotes the crystal by M = [AFF. F. L), where the fractional coordinates of all atoms in a cell compose the matrix F E [0. 1) 3x
Task definition CSP predicts for each unit cell the lattice matrix L and the fractional matrix given its chemical composition A, namely, learning the conditional distribution pl L, F A).
4
The Proposed Method: DiffCSP
This section first presents the symmetries of the crystal geometry, and then introduces the joint equivaraint 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 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( F A): permutation invariance, 0(3) invariance, and periodic translation invariance. Their detailed definitions are provided as follows.
Definition 1 (Permutation Invariance). For any permutation P SN, p(L, F A) = p(L, FP AP), i.e., changing the order of atoms will not change the distribution.
3

---

Process Forward Difiuion Precess
Invariant Prior
Equivariant Process
Sample
Ey te MIA
Periodic
0 1-- 1) sxps
Mus
Equivaraze
(t- steps
Dencising
Mr
M,
Mai1
Compooilion
xF-3
F
Figure 2: Overview of DifCCSP Given the composition A, we denote the crystal, its lattice and fractional coordinate matrix time as Mr, L, and F;, respectively. The terms 6, and er are Gaussian noises, & and ép are predicted by the denoising model Ø.
Definition (0(3) Invariance). For any orthogonal transformation Q E R333 satisfying Q =Q=I. =1 p(QL, F A) p(L, F A), namely, any rotation/reflection of L keeps the distribution unchanged.
Definition 3 (Periodic Translation Invariance). For any translation t E R3x1, p(L, w(F + + =t1)) A) = p( F A), where the function w (F) = [F] E 1)3x returns the fructional purt of each element in F, and 1 E [R3x1 is vector with all elements set to one. It explains that any periodic translation of F will not change the distribution
The permutation invariance is tractably encapsulated by using GNNs as the backbone for generation 47). We mainly focus on the other two kinds of invari- ance (see Figure []. since GNNs our default choices. For simplicity, we compactly term the 0(3) invariance and periodic translation invariance as periodic E(3) invariance henceforth. Previous ap- proaches (e.g. [9] (36) adopt Cartesian coordinates X other than fractional coordinates F. hence their derived forms of the symmetry are different. Par- ticularly, in Definition _ the orthogonal transforma- tion additionally acts on X; in Definition 3 the pe- riodic translation w( F + +1 = becomes the transla- tion along the lattice bases X + LL1 besides, X should also maintain E(3) translation invariance, that is p( L( X + t1 IA) = P(L, X A)). With the help of the fractional system. the periodic E(3) invariance is made tractable by fulfilling O(3) invariance w.r.f. the orthogonal transformations on L and periodic translation invariance w.r.t. the periodic translations on F, respectively. In this way, such approach, as detailed in the next section, facilitates the application of diffusion methods to the CSP task.
(0)
(6)
11
50
(c)
(d)
Figure 1: (a)->(b): The orthogonal transfor- mation of the lattice vectors. (c)-xd): The periodic translation of the fractional coordi- nates. Both cases do not change the structure.
4.2 Joint Equivariant Diffusion
Our method DiffCSP addresses CSP by simultancously diffusing the lattice L and the fractional coordinate matrix F. Given the atom composition A, M, denotes the intermediate state of L and F at time step & (0 < t < T). DiffCSP defines two Markov processes: the forward diffusion process gradually adds noise to Ma, and the backward generation process iteratively samples from the prior distribution Mr to recover the origin data Ma-
Prrevoous works (e.g. 36 further discuss the scaling invariance of unit cell formed by periodic boundaries. allowing + L, Va N this paper. scaling invariance is unnecessary since we apply the Niggli reduction 466 on the primitive cell as a canonical scale representation of the lattice vectors where we fix 11, 1, 1) 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 Mr should meet periodic E(3)
invariance. Such requirement satisfied if the prior distribution p(Mr) invariant and the Markov
transition p(M,-1 M,) is equivariant, according to the diffusion- basoon generation literature [10]
Here, an equivariant transition specified as -M-1 | 9 MMe) = p(M{11 MMe) where
g .M refers to any orthogonal/rranslaaiona transformation acts on M in the way presented in
Definitions 3 We separately explain the derivation details of L and F below. The detailed
flowcharts are summarized in Algorithms and in Appendix B.3
Diffusion on Given that L is continuous variable, we exploit Denoising Diffusion Probabilistic Model (DDPM) (38) to accomplish the generation. We define the forward process that progressively diffuses Lu towards the Normal prior p(Lr) -O, I) by
which can be devised as the
probability conditional on the initial distribution:
q (L,)oo)) No ~ (L
(1 (1-a)11)
(2)
where 3, E (0, 1) controls the variance, and ā. = IT as = II (1 Be) is valued in accordance to the cosine scheduler 48).
The backward generation process is given by:
p(L,-1lMM) = -(L-)))))
(3)
where m(M)) yar
Jar (L. -i-ivia--- E (M t)) o2(M)) of) Be 1- The denoising term
VI-a.
Ey (M, E 12a33 is predicted by the model of L,, F1, A,t).
As the prior distribution p(Lr) = NO, I) already O(3)-invariant, we require the generation process in Eq. (3) to be -eui-equiv,r,ant, which is formally stated below.
Proposition 1. The marginal distribution Lo) by Eq. 3 is ((3)-invariant ê t) is 0(3)- equivariant, namely é (QL, F,, A,t) = Q&l (L., Fú, A. t). 2 QQ Q I.
To train the denoising model d, we first sample ér, (O, I) and reparameterize L, - varLo vI - â,Ey based on Eq. ( The training objective is defined as the bg loss between E7 and éL:
LL =
- (x, 1)
Diffusion on F The domain of fractional coordinates [0, 1)3xN forms quotient space 3 /23x induced by the crystal periodicity. It is not suitable to apply the above DDPM fashion to generate F, as the normal distribution used in DDPM is unable to model the cyclical and bounded domain of F. Instead, we leverage Score-Matching Matching (SM) based framework [49 50] along with Wrapped Normal (WN) distribution [3) 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 EF RR from (0,, 7), and then acquire F, = w( wFoo + OTEF) where the truncation function w--) is already defined in Definition 3 This truncated sampling implies the WN transition:
F. Fo Z1?
Fo) O
Σ exp Z
(5)
20
Basically, this process ensures the probability distribution over [z, + 1)3xN for any integer be the same to keep the crystal periodicity. Here, the noise scale obeys the exponential scheduler: ao 0 and o; a ( (gr) f if > 0. Desirably, q( (o) is periodic translation equivariant, and approaches uniform distribution (1, 1) if or sufficiently large.
For the backward process, we first initialize Fr from the uniform distribution 4((0, 1), which is periodic translation invariant. With the denoising term ε predicted by Φ( Le, F,, A, we combine the ancestral predictor (38 50 with the Langevin corrector [49] to sample Fo- We immediately have:
Proposition 2. The marginal distribution F(F)) is periodic translation invariant if é (M,,)) is periodic translation invariant, namely ĒF (L, F,, A, t) (L,,, w F, + t1 A, f), Vt E R².
5

---

The training objective for score matching is:
FF FF (o).T [Ael V., (F) - (
1)
where A EF I1V F, log g F. Fb) 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 A is fixed, our method able to be extended for the ab initio generation task by further generating A. We achieve this by additionally optimizing the one-hot representation A with DDPM-based approach. We provide more details in Appendix
4.3 The Architecture the Denoising Model
This subsection designs the denoising model @LL, A,t) that outputs & and & satisfying the properties stated in Proposition and
Let (*) [h (s) .... h N denote the node representations of the s-th layer. The input feature is
given by 11 p( a;), , (po (1)), where Jatom and fpe are the atomic embedding and sinusoidal positional encoding [38
51. respectively; o is multi-layer perception (MIP)..
Built upon EGNN 32]. the s-th layer message paaaggg--asinggggg is unfolded as follows:
Here Pm and are MLPs. The func- tion ę/r 1), - [-1, 1]3xK is Fourier Transformation of the relative fractional coordinate fj f.. Specif- ically, suppose the input to be f [2, Sa] then the e-th row and k- th column of the output is calculated by (f)(() k] = sin(2nmfa fe), if h 2m (even); and eher (f)lc, (odd). WFF extracts various frequencies of all relative fractional distances that are helpful for crystal structure modeling, and more importantly, WIFT is periodic translation invariant, namely, ( t) wff: t)) (fr - fi) for any translation t, which is proved in Appendix A.3
(8) Pm
(h
(e-1) h
(8-1) L ( f()),
(6)
(8) Σ
(8)
m;
mij
(7)
(1 = ni (.-1) + pp (h, (8-1)
m (1)).
(8)
cos cos(22mm.f))).m if = 2m +
After S layers of message passing conducted on the fully connected graph, the lattice noise é is acquired by a linear combination of L, with the weights given by the final layer:
N
)
é = Lip.
(9)
N
where fr. is an MLP with output shape as 3 x 3. The fractional coordinate score é is output by:
éFli, i] = ( (5)).
(10)
where é é defines the i-th column of éF, and yF is an MLP on the final representation.
We apply the inner product term L L in Eq. (6) to achieve invaainve,,, as (QL) (QL==1) for any orthogonal matrix Q E 3x33 This leads to O())-inarriacce of yr. Eq. 10). and we further luft-multiply L with yr to ensure the Dquivqquuv of é,. Therefore, the above formulation of the denoising model d L, F, A,t) 1) ensures the following property.
Proposition 3. The score é by Eq. 9 is 0(3) equivaran,,,,, and the score é from Eq. 10)) is periodic translation invariant. Hence, the generated distribution by DiffCSP is periodic E(3) invariant.
utilize Cartesian coordinates, and usually describe crystals with multi graph representation to encode the periodic structures. They create multiple edges to connect each pair of nodes where different edges refer to different integral cell translations. Here, we no longer require multi-graph representation, since we employ fractional coordinates that naturally encode periodicity and the Fourier transform v'rr in our message passing is already periodic translation invariant. We will ablate the benefit in Table 3
Comparison with multi-graph representation Previous methods 15 29

---

Table 1: Results on stable structure prediction task.
of samples
Perov:
Carbon-24
MP-20
MPTS-S2
Match ratef
RMSEL
Match ratef RMSE
Match ratef
RMSEJ Match ratef RMSEL
20
29.22
0.2924
14.63
0.4041
8.73
0.2501
2.05
0.3329
RS
21
5,000
36.56
0.0886
14.63
0.4041
11.49
0.2822
2.68
0.3444
20
21.03
0.2830
0.44
0.3653
8.11
0.2402
2.05
0.3024
BO
5,000
55.09
0.2037
12.17
0,4089
12.68
0.2816
6.69
0.3444
20
20,90
0.0836
6.40
0.4204
4.05
0.1567
1.06
0,2339
PSO 21
5,000
21.88
0.0844
6.50
0.4211
4.35
0.1670
1.09
0.2390
P-cG-
48.22
0.4179
17.29
0.3846
15.39
0.3762
3.67
0,4115
SchNet 533
20
97.94
0.3463
55,91
0.3551
32.64
0.3018
12.96
0,3942
45.31
0.1138
17.09
0.2969
33.90
0.1045
5.34
0.2106
CDVAE 9
20
88.51
0.0464
88.37
0.2286
66.95
0.1026
20.79
0,2085
52.02
0.0760
17.54
0.2759
51.49
0.0631
12.19
0,1786
DifCCSP
98.60
0.0128
88.47
0.2192
77.93
0.0492
34.02
0.1749
MP-20
MPTS-52
DiaCSP
COVAR
P-oG-SchNet
Figure 3: Visualization of the predicted structures from different methods. We select the structure of the lowest RMSE over 20 candidates. We translate the same predicted atom by all methods the origin for better comparison. Our DifCCSP accurately delivers high quality structure predictions.
Experiments
In this section, we evaluate the efficacy of DiffCSP on a diverse range of tasks, by showing that can generate high-quality structures of different crystals in $ |5.1 with lower time comparing with DFT-based optimization method in $5.22 Ablations in $ 9.33 exhibit the necessity of each designed component. We further showcase the capability of DifICSP in the ab initio generation task in 5.4
5.1 Stable Structure Prediction
Dataset We conduct experiments on four datasets with distinct levels of difficulty. Perov-5 54 55
contains 18,928 perovskite materials with similar structures. Each structure has 5 atoms in unit cell. Carbon-24 [56] includes 10,153 carbon materials with 6~24 atoms cell. MP-20
selects 45,231 stable inorganic materials from Material Projects (57), which includes the majority of experimentally-gen generratd materials with at most 20 atoms unit cell. MPTS-52 is a more challenging extension of MP-20, consisting of 40,476 structures up to 52 atoms per cell, sorted according to the carliest published year in literature. For Perov-5, Carbon-24 and MP-20, we apply the 60-20-20 split in line with Xie et al. [9]. For MPTS-52, we split 27.380/5,000f8 training/vliidationnsstngg in chronological order.
for
Baselines We contrast two types of previous works. The first type follows the predict-optimize paradigm, which first trains a predictor of the target property and then utilizes certain optimization algorithms to search for optimal structures. Following Cheng et al. [2I]. we apply MEGNet [52] as the predictor of the formation energy. For the optimization algorithms, we choose Random Search (RS), Bayesian Optimization (BO), and Particle Swarm Optimization (PSO), all iterated over 5,000 steps. The second type based on deep generative models. We follow the modification in Xie
et al. [9). and leverage cG-SchNet (53) that utilizes SchNet (29 the backbone and additionally
7

---

considers the ground truth lattice initialization for encoding periodicity, yielding final model named P-cG-SchNet. Another baseline CDVAE (9] is VAE-based framework for pure crystal generation. by first predicting the lattice and the initial composition and then optimizing the atom types and coordinates via annealed Langevin dynamies [49]. To adapt CDVAF into the CSP task, we replace the original normal prior for generation with a parametric prior conditional on the encoding of the given composition. More details are provided in Appendix B.2
Evaluation metries Following the common practice 99. we evaluate by matching the predicted candidates with the ground-truth structure. Specifically, for each structure in the test set, we first generate k samples of the same composition and then identify the matching if at least one of the samples matches the ground truth structure, under the metric by the StructureMatcher class in pymatgen [58] with thresholds stol=0.5, angle tol=10, Itol=0.3. The Match rate is the proportion of the matched structures over the test set. RMSE is calculated between the ground truth and the best matching candidate, normalized by 2/VNN where V is the volume of the lattice, and averaged over the matched structures. For optimization methods, we select 20 structures of the lowest energy of all 5,000 structures from all iterations during testing as candidates. For generative baselines and our DifCCSP, we let k- = 1 and k = 20 for evaluation. We provide more details in Appendix B
C.
and
Results Table conveys the following observations. 1. The optimization methods encounter low Match rates, signifying the difficulty of locating the optimal structures within the vast search space. 2. In comparison other generative methods that construct structures atom by atom or predict the lattice and atom coordinates in two stages, our method demonstrates superior performance, highlighting the effectiveness of jointly refining the lattice and coordinates during generation. 3. All methods struggle with performance degradation as the number of atoms per cell increases, on the datasets from Perov-5 to MPTS-52. For example, the Match rates of the optimization methods are less than 10% in MPTS-52. Even so, our method consistently outperforms all other methods.
Visualization Figure 3 provides qualitative comparisons. DiffCSSP clearly makes the best predictions.
5.2 Comparison with DFT-based Methods
We further select 10 binary and 5 ternary compounds in MP-20 test- ing set and compare our model with USPEX [59]. a DFT-based software equipped with the evolutionary algo- rithm to search for stable structures. For our method, we sample 20 candidates for each compound following the setting in Table We select the model trained on MP-20 for inference, with a training duration of 5.2 hours. For USPEX, we apply 20 generations, 20 populations for each compound, and select the best sample in each generation, leading to 20 candidates as well. We summarize the Match rate over the 15 compounds, the Averaged RMSD over the matched structures, and the Averaged Inference Time to generate 20 candidates for each compound in Table [0] The detailed results for each compound are listed in Appendix F DiffCSP correctly predicts more structures with higher match rate, and more importantly, its time cost is much less than USPEX, allowing more potential for real applications.
Table 2: Overall results over the 15 selected compounds.
Match rate (%b)
Avg. RMSDJ
Avg. Time].
USPEX (59
53.33
0.0159
12.5h
DifCCSP
73.33
0.0172
10s
5.3 Ablation Studies
We ablate each component of DiffCSP in Table 3 and probe the following aspects. 1. To verify the necessity of jointly updating the lattice L and fractional coordinates F, we construct two variants that separate the joint optimization into two stages, denoted as L + F and F + L. Particu- larly, L - F applies two networks to learn the reverse processes po, (Lo.T- A, Fr, Lr) and Po, Fo.r-11 A, Fr, Lo). During inference, we first sample Lr, Fr from their prior distributions, acquiring Lo via po, and then Fo by po, based on Lg. F - L is similarly executed but with the gen- eration order of Lo and Fo exchanged. Results indicate that L - F performs better than the F - L. but both are inferior to the joint update in DifICSP, which endorses our design. We conjecture that the joint diffusion fashion enables L and F to update synergistically, which makes the generation process more tractable to learn and thus leads to better performance. 2. We explore the necessity of preserving the 0(3) invariance when generating L, which is ensured by the inner product L L in Eq. (6).
8

---

When we replace it with L and change the final
Table 3: Ablation studies on MP-20.
MG:
output ét, ( Ei ( in Eq. (9 to
Multi-Graph edge construction [S]. FT: Fourier-
break the equivariance, the model suffers foom extreme performance detriment. Only 1.66% structures are successfully matched, which ob- viously implies the importance of incorporating 0(3) equivariance. Furthermore, we introduce the chirality into the denoising model by adding sign( LI), the sign of the determinant of the lat- tice matrix, as an additional input Eq. [ The adapted model is SO(3) (iiiannaar,a but breaks the reflection symmetry and hence is NOT (3)- invariant. There is no essential performance change, indicating the chirality not quite cru- cial in distinguishing different crystal structures for the datasets used in this paper. 3. We further assess the importance of periodic translation in- variance from two perspectives. For the genera- tion process, generate F via the score-based model with the Wrapped Normal (WN) distribution. We replace this module with DDPM under standard Gaussian as q( (.F. MMM (F.l Fa, (1-aa) ()) similarly defined as Eq. (3). A lower match rate and higher RMSE are observed for this variant. For the model architecture, we adopt Fourier Transformation((FF) in Eq. (6) to capture periodicity. To investigate its effect, we replace (f - fi) with f, f., and the match rate drops from 51.49% to 29.15%. Both of the two observations verify the importance of retaining the periodic translation invariance. change the fully-connected graph into the multi-graph approach adopted in Xie and Grossman (15). The muuli--raph, approach decreases the match rate, since the multi-graphs constructed under different intermediate structures may differ vibrantly during generation, leading to substantially higher training difficulty and lower sampling stability. We will provide more discussions Appendix
Transformation.
Match rate (%)
RMSE
DiffCSP
$1.49
0.0631
w/o Joint Diffusion
L F+L
F
50.03
0.0921
36.73
0.0838
w/o 0(3) Equivariance
inner product w/ chirality
1.66
0,4002
49.68
0.0637
w/o Periodic Translation Invoriance
w/o WN w/o FT
34.09
0.2350
29.15
0.0926
MG Edge Construction
MG W/ FT MG w/o FT
25.85
0.1079
28.05
0.1314
We further
5.4 Ab Initio Crystal Generation
DiffCSP is extendable to ab initio crystal generation by further conducting discrete diffusion on atom types A. We contrast DifICSP against five generative methods following [9] FTCP (28). Cond-DFC- VAE 17). G-SchNet [60] and its periodic variant P-G-SchNet, and the orginal version of CDVAE [9. Specifically for our DiffCSP,
gather the statistics of the atom numbers from the training set, then sample the number based on the pre-computed distribution similar to Hoogeboom et al. 4T. which allows DiffCSP to generate structures of variable size. Following (9. we evaluate the generation performance in terms of there metrics: Validity, Coverage, and Property statistics, which respectively return the validity of the predicted crystals, the similarity between the test set and the generated samples, and the property calculation regarding density, formation energy, and the number of elements. The detalied definitions of the above metrics are provided in Appendix G
Results Table show that our method achieves comparable validity and coverage rate with previous methods, and significantly outperforms the baselines on the similarity of property statistics, which indicates the high reliability of the generated samples.
6 Discussions
Limitation 1. Composition generation. Our model yields slightly lower compositional validity in Table We provide more discussion in Appendix G and it is promising to propose more powerful generation methods on atom types. 2. Experimental evaluation. Further wet-lab experiments can better verify the effectiveness the model in real applications.
Conclusion In this work, we present DiffCSP, diffusion- based learning framework for crystal structure prediction, which particularly curated with the vital symmetries existing in crystals. The diffusion is highly flexible by jointly optimizing the lattice and fractional coordinates, where
Composition baaed metrics are meaningful for Carbon 24, as all structures composed of carbon.

---

Table 4: Results on ab initio generation task. The results of bascline methods are from Xie et al. 9
Validity
(
Coverage (%)
Property |
Data
Method
Struc.
Comp.
COV-R
COV-P
dp
dg
datem
Perov-5
FTCP
28
0.24
54.24
0.00
0.00
10.27
156.0
0.6297
Cond-DFC VAE 7
73.60
82.95
73.92
10.13
2.268
4.111
0.8373
G-SchNet
60
99.92
98.79
0.18
0.23
1.625
4.746
0.0368
P.G-SchNet 60
79.63
99.13
0.37
0.25
0.2755
1.388
0.4552
CDVAB
9
100.0
98.59
99.45
98,46
0.1258
0.0264
0.0628
DifICSP
100.0
98.85
99.74
98.27
0.1110
0.0263
0.0128
Carbon-24
FTCP
23
0.08
0.00
0.00
5.206
19.05
G-SchNet 60
99.94
0.00
0.00
0,9427
1.320
P-G-SchNet
60
48. 39
0.00
0.00
1,533
134.7
CDVAE
9
100.0
99.80
83.08
0.1407
0.2850
DifICSP
100.0
99.90
97.27
0.0805
0.0820
MP-20
FTCP 28
1,55
48.37
4.72
0.09
23.71
160.9
0.7363
G-SchNet
60
99.65
75.96
38.33
99.57
3.034
42.09
0.6411
P-G-SchNet 60
77.51
76.40
41.93
99.74
4.04
2.448
0.6234
CDVAE 9
100.0
86.70
99.15
99.49
0.6875
0.2778
1.432
DifTCSP
100.0
83.25
99.71
99.76
0,3502
0.1247
0,3398
the intermediate distributions are guaranteed to be invariant under necessary transformations. We demonstrate the efficacy of our approach on a wide range crystal datasets, verifying the strong applicability of DiffCSP towards predicting high-quality crystal structures.
7
Acknowledgement
This work was jointly supported by the following projects: the National Natural Science Founda- tion of China (61925601 & 62006137); Beijing Nova Program (20230484278); the Fundamental Research Funds for the Central Universities, and the Research Funds of Renmin University of China (23XNKJ19); Alibaba Damo Research Fund; CCF-Ant Research Fund (CCF-AFSG RF20220204); National Key R&D Program of China (122DD0117805)
References
[1] Gautam Desiraju. Cryptic crystallography. Nature materials, 122):77-79, 2002.
[2] Keith T Butler, Daniel Davies, Hugh Cartwright, Olexandr Isayev, and Aron Walsh. Machine learning for molecular and materials science. Nature, 559(7715):547-555, 2018.
[3] Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical review, 140(4A):A113,, 1965.
[4] Chris J Pickard and RJ Needs. Ab initio random structure searching. Journal of Physics: Condensed Matter, (3053300, 2011.
[5] Tomoki Yamashita, Nobuya Sato, Hiori Kino, Takashi Miyake, Koji Tsuda, and Tamio Oguchi. Crystal structure prediction accelerated by bayesian optimization. Physical Review Materials, (1):013003, 2018
[6] Artem Oganov, Chris J Pickard, Qiang Zhu, and Richard Needs. Structure prediction drives materials discovery. Nature Reviews Materials, 4(5):331-348, 2019.
[7] Callum J Court, Batuhan Yildirim, Apoorv Jain, and Jacqueline M Cole. 3-d inorganic crystal structure generation and property prediction via representation learning. Journal of chemical information and modeling, 60(10):4518 1111:44518-435, 2020.
[8] Wenhui Yang, Edirisuriya M Dilanga Siriwardane. Rongzhi Dong, Yuxin Li, and Jianjun Hu. Crystal structure prediction of materials with high symmetry using differential evolution. Journal of Physics: Condensed Matter, 33(45) 2021.
[9] Tian Xie, Xiang Fu, Octavian Eugen Ganea, Regina Barzilay, and Tommi S Jaakkola. Crystal diffusion variational autoencoder for periodic material generation. In International Conference on Learning Representations, 2021.
10

---

[10] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2021.
[11] Brian L. Trippe. Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones 3d for the motif- scaffolding problem. In The Eleventh International Conference on Learning Representations, 2023.
[12] Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi S. Jaakkola. Diff- dock: Diffusion steps, twists, and turns for molecular docking, In The Eleventh International Conference on Leurning Representations, 2023.
[13] Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In International Conference on Machine Learning, pages 9558-9568 PMLR, 2021.
[14] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A A Kohl, Andrew J Ballard, Andrew Cowie, Bernardino Romera-Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger. Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure predic- tion with AlphaFold. Nature, 596(7873) 2021. doi: 10.10381441586--22--583
[15] Tian Xie and Jeffrey C. Grossman. Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties. Phys. Rev. Lett., 120:14530,, Apr 2018. doi: 10. 103//hyssevL.ett. 220. 145301.
[16] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. arXiv preprint arXiv:2206.01729) 2022.
[17] Yanchao Wang, Jian Lv, Li Zhu, and Yanming Ma. Crystal structure prediction via particle- swarm optimization. Physical Review B, 82(9):0941 16, 2010.
[18] Yunwei Zhang, Hui Wang, Yanchao Wang. Lijun Zhang, and Yanming Ma. Computer-assisted inverse design of inorganic electrides. Physical Review X, 7(1 7(1):11001,, 2017.
[19] TL Jacobsen, MS Jørgensen, and B Hammer. On-the-fly fly machine learning of atomic potential in density functional theory structure optimization. Physical review letters, 120(2)022102, 2018.
[20] Exgeny Podryabinkin, Evgeny V Tikhonov, Alexander V Shapeev, and Artem R Oganov. Accelerating crystal structure prediction by machine-leamnin interatomic potentials with active learning. Physical Review B, 99(6):0641 14, 2019.
[21] Guanjian Cheng, Xin-Gao Gong, and Wan-Jian Yin. Crystal structure prediction by combining graph network and optimization algorithm. Nature communications, 13(1):1-8, 2022.
[22] Jordan Hoffmaan, Louis Maestrati, Yoshihide Sawada, Jian Tang, Jean Michel Sellier, and Yoshua Bengio. Data-driven approach to encoding and decoding 3-d crystal structures. arXiv preprint arXiv: 19w.:099999999999 2019.
Alan Aspuru-Guzik, and Yousung Jung. Inverse design of solid-state materials via a continuous representation. [23] Juhwan Noh, Jaehoon Kim, Helge S Stein, Benjamin Sanchez-Lengeling John M Gregoire,
Matter, 155):1370-1344 2019.
[24] Jianjun Hu, Wenhui Yang, and Edirisuriya M Dilanga Siriwardane. Distance matrix- crystal structure prediction using evolutionary algorithms. The Journal of Physical Chemistry A, 124651):10999-19919 2020.
[25] Jianjun Hu, Wenhui Yang, Rongzhi Dong, Yuxin Li, Xiang Li, Shaobo Li, and Edirisuriya MD Siriwardane. Contact map based crystal structure prediction using global optimization. Crys- IEngComm, 1765):175511776 2021.
11

---

[26] Asma Nouira, Nataliya Sokolovska, and Jean-Claude Crivello. Crystalgan: learning to discover crystallographic structures with generative adversarial networks. arXiv preprint arXiv: 1810. 1203, 2018.
[27] Sungwon Kim, Juhwan Noh, Geun Ho Gu, Alan Aspuru- Guzik, and Yousung Jung. Generative adversarial networks for crystal structure prediction. ACS central science, 6(8):1412-1400, 2020.
[28] Zekun Ren, Siyu Isaac Parker Tian, Juhwan Noh. Felipe Oviedo, Guangzong Xing, Jiali Li, Qiaohao Liang. Ruiming Zhu, Armin G. Aberle, Shijing Sun, Xiaonan Wang. Yi Liu, Qianxiao Li, Senthilnath Jayavelu, Kedar Hippalgaonkar, Yousung Jung, and Tonio Buonassisi. An invertible crystallographic representation for general inverse design of inorganic crystals with targeted properties. Matter, 2021. ISSN 2590-2385. doi: https://doi.org/ or..o//.////////. 1011100.1111111 m/aii//..tat.00001 11.032.
[29] Kristof T Schüit,, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722, 2018.
[30] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariaa neural networks for 3d point clouds. arXiv preprint vrXiv:/802.082//1 2018.
[31] Fabian Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling, Se(3)-transformmess: 3d roto- translation equivariant attention networks. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Floring Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-1 2020, virtual, 2020.
[32] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International Conference on Machine Learning, pages 9323-9332 PMLR, 2021.
[33] Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations, 2021.
[34] Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Mor- gane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi. Open catalyst 2020 (oc20) dataset and community challenges. ACS Catalysis, 2021. doi: 10. http//cccaaal.000452...
[35] Richard Tran, Janice Lan, Muhammed Shuaibi, Siddharth Goyal, Brandon M Wood, Ab- hishek Das, Javier Heras-Domingo, Adeesh Kolluru, Ammar Rizvi, Nima Shoghi, et al. The open catalyst 2022 (oc22) dataset and challenges for oxide electrocatalysis. arXiv preprint arXiv:2206.088 08911, 2022.
[36] Keqiang Yan, Yi Liu, Yuchao Lin, and Shuiwang Ji. Periodic graph transformers for crystal material property prediction. In The 36th Annual Conference on Neural Information Processing Systems, 2022.
[37] Jascha Sohl-Dicken Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. International Conference on Machine Learning, pages 2256-22622 2265. PMLR, 2015.
[38] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-688 6851, 2020.
[39] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models, 2021.
[40] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204006125, 2022.
12

---

[41] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning. pages 8867-8887. PMLR, 2022.
[42] Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigen- specific antibody design and optimization with diffusion-based generative models for protein structures. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[43] Valentin De Bortoli, Emile Mathieu, Michael John Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score- score-based generative modelling. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[44] Chin-Wa Huang, Milad Aghajohari, Jocy Bose, Prakash Panangaden, and Aaron Courville. Ric- mannian diffusion models. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[45] Detlef WM Hofmann and Joannis Apostolakis. Crystal structure prediction by data mining. Journal of Molecular Structure, 647(1-3):17-39, 2003.
[46] Ralf Grosse- Kunstleve,, Nicholas K Sauter, and Paul D Adams. Numerically stable algorithms for the computation of reduced unit cells. Acta Crystallographica Section A: Foundations of Crystallography, 60(1):1-6, 2004.
[47] Thomas N Kipf and Max Welling. Variational graph auto-encoders. arXiv preprint 11Xiv:1611.07300 2016.
[48] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162-8171. PMLR, 2021.
[49] Yang Song and Stefano Ermon. Improved techniques for training score- based generative models. Advances in neural information processing systems, 33:12438-12444 2020.
[50] Yang Song, Jascha DihhtDickst,in, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. bcore-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.132 131.-,66,1,1366, 2020.
[51] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[52] Chi Chen, Weike Ye, Yunxing Zuo, Chen Zheng, and Shyue Ping Ong. Graph networks as a universal machine learning framework for molecules and crystals. Chemistry of Materials, 31 (9):3564-355 3572, 2019.
[53] Niklas WA Gebauer, Michael Gastegger, Stefaan SP Hessmann, Klaus-Robert Müller, and Kristof T Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. Nature communications, 13(1):1-11, 2022.
[54] Ivano E Castelli, David D Landis, Kristian Thygesen, Soren Dahl, Ib Chorkendorff, Thomas Jaramillo, and Karsten W Jacobsen. New cubic perovskites for one-and two-photon water splitting using the computational materials repository. Energy & Environmental Science, 5(10): 9033-9043, 2012.
[55] Ivano E Castelli, Thomas Olsen, Soumendu Datta, David D Landis, Søren Dahl, Kristian S Thygesen, and Karsten W Jacobsen. Computational screening of perovskite metal oxides for optimal solar light capture. Energy & Environmental Science, 5(2):581455819, 2012.
[56] Chris J. Pickard. Airss data for carbon at 10gpa and the c+nth+o system at Igpa, 2020.
[57] Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, al. Commentary: The materials project: materials genome approach accelerating materials innovation. API. materials, 1(1):01 1002, 2013.
13

---

[58] Shyue Ping Ong, William Davidson Richards, Anubhav Jain, Geoffroy Hautier, Michael Kocher, Shreyas Cholia, Dan Gunter, Vincent L Chevrier, Kristin Persson, and Gerbrand Ceder. Python materials genomics (pymatgen): A robust, open-source python library for materials analysis. Computational Materials Science, 68:314-319, 2013.
[59] Colin W Glass, Artem R Oganov, and Nikolaus Hansen. Sxspex-evolutionaa crystal structure prediction, Computer physics communications, 175(11-12)771--22) 2006.
[60] Niklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry-udapted generation of 3d point sets for the targeted discovery of molecules. In Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buee E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 7566-7578. Curran Associates, Inc., 2019.
[61] Gerhard Kurz, Igor Gilitschenski, and Uwe D) Haneheck. Efficient evaluation of the probability density function of a wrapped normal distribution. In 2014 Sensor Data Fusion: Trends, Solutions, Applications (SDF), pages 1-5. IEEE, 2014.
[62] James Bergstra, Daniel Yamins, and David Cox. Making science of model search: Hyper- parameter optimization in hundreds of dimensions for vision architectures. In International conference on machine learning, pages 115-123. PMLR, 2013.
[63] Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and uncertainty-awart directional message passing for non-equilibrium molecules. In Machine Learning for Molecules Workshop, NeurlPS, 2020.
[64] Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. In Conference on Neural Information Processing Systems (NeurIPS), 2021.
[65] Nils ER Zimmermann and Anubhav Jain. Local structure order parameters and site fingerprints for quantification of coordination environment and crystal structure similarity. RSC advances, 10(10):606366081 2020.
[66] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 3748-8763. PMLR, 2021.
[67] Peter E Blöchl. Projector augmented-wer wave method. Physical review B, 50(22)):7795, 1994.
[68] Georg Kresse and Jürgen Furthmüller. Efficiency of ab-initio total energy calculations for metals and semiconductors using a plane-wave wave basis set. Computational materials science, 6(1):15-50, 1996.
[69] John Perdew, Kieron Burke, and Matthias Ernzerhof. Generalized gradient approximation made simple. Physical review letters, 77(1 77885538 1996.
[70] Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems, 34: 12454-12465. 2021.
[71] Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang. Volkan Cevher, and Pas- cal Frossard. Digress: Discrete denoising diffusion for graph generation. The Eleventh International Conference on Learning Representutions, 2023.
[72] Daniel Davies, Keith Butler, Adam Jackson, Jonathan M Skelton, Kazuki Morita, and Aron Walsh. Smact: Semiconducting materials by analogy and chemical theory. Journal of Open Source Software, 13381 2019.
[73] Logan Ward, Ankit Agrawal, Alok Choudhary, and Christopher Wolverton. A general- machine learning framework for predicting properties of inorganic materials. npj Computational Materials, 2(1):1-7,2015
14

---

[74] Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu. EGSDE: Unpaired image-to-imas -inaage translation via energy-guides stochastic differential equations. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[75] Fan Bao, Min Zhao, Zhongkai Hao, Peiyao Li, Chongxuan Li, and Jun Zhu. Equivariant energy-guided SDE for inverse molecular design. In The Eleventh International Conference on Learning Representations, 2023.
15

---

Contents
Introduction
2
Related Works
Preliminaries
4
The Proposed Method: DifCCSP
4.1
Symmetries of Crystal Structure Distribution
Joint Equivariant Diffusion
The Architecture of the Denoising Model
Experiments
Stable Structure Prediction
5.2
Comparison with DFT-based Methods
8
5.3
Ablation Studies
5.4
Ab Initio Crystal Generation
6
Discussions
Acknowledgement
10
Theoretical Analysis
18
Proof of Proposition
18
Proof of Proposition
19
A.3 Proof of Proposition 3
20
Discussion on Periodic Translation Invariance
21
Implementation Details
22
B.
Approximation of the Wrapped Normal Distribution
22
Adaptation of CDVAE
................................
22
B.3
Algorithms for Training and Sampling
23
B.4
Hlpee--arameterr and Training Details
24
Exploring the Effects Sampling and Candidate Ranking
25
Impact of Sampling Numbers
25
C.2
Diversity..
25
Ranking among Multiple Candidates
25
D Impact of Noise Schedulers
26
Learning Curves of Different Variants
26
Comparison with DFT-based Methods
26
16

---

Implementation Details
26
F22
Results
28
Extension to More General Tasks
28
Overview
28
G.2
Ab Initio Generation
28
G.3 Properly Optimization
30
H
Computational Cost for Inference
32
Error Bars
32
More Visualizations
32
17

---

A Theoretical Analysis
A.1 Proof of Proposition
We first introduce the following definition describe the equivariance and invariance from the perspective of distributions.
Definition 4. We call distribution p(æ)) is G-invarian for any transformation g in the group G, p(g. 5) = p(r), and a conditional distribution p(æ|c) is G-equivarians p(g zlg c) = p(x(c), Vg E G.
then provide and prove the following lemma capture the symmetry of the generation process.
Lemma 1 (Xu et al. [O]). Consider the generation Markov process p(ro) If the prior distribution p(ær) is G-invariant and the Markov transitions ((4) < S T are G-equivariant, the marginal distribution p(ro) is also G-invariant.
p(xT) J (((0or-1/24dddaa
Proof. For any G, we have
plg 2o) plg aT) / pla zoT-1ly xn dr
/ I r1-1-g x drir
pla xT)
tm1
-((rr) /
II plg 22-1lg x4 drrr
p((rr) /
II p (x doir
=p(or) / ((
= -p(0))-
Hence, the marginal distribution p(æo) is G-invariant.
The proposition Proposition T is rewritten and proved as follows.
Proposition 1. The marginal distribution pl Lo) by Eq. 3 is 0(3)-invariant if t) is 013)- equivariant, namely é (QL, F,, A, = Qē2 Lt F,, t) QQQ=1 =I
Proof. Consider the transition probability in Eq. we have
p(Ln-1 Fr. A) -N bLe Fú, A,1)),
where 0, Jar b
of B. 1- for simplicity, and éz (Me, t) is completed as VI-aa*
3.
êy (L, F., A,t)
As the denoising term éL, A,t) is 00)-eqquvariiant, have QL, F,, t) for any orthogonal transformation Q 3x3 QQQ=I.
For the variable o(L,oI), we have QL NIQL -
That is,
N(L o2) =(QLLQI,
(11)
18

---

For the transition probability o(L:-1 F,, A), we have
p(QLi-1QQL, F, A) = (N( (QL, biê A,t)), of)
-N(Q 44 (QL, b.QEI Pt, A, A,
(0(3)-equivariant Ez)
- (e b, (L,, Fb, A, A,)) . o(()
([, deL (L., F., A.t)), o.)
(Eq. (1
p(L, Lt, F,, A).
As the transition 0(3) -qquivaraat and the prior distribution (0, II is (3) invariant, we prove that the the marginal distribution p(Lo) is (133)iiiiiaan based on lemma
A.2 Proof of Proposition? 2
Let (M, o² o denote the wrapped normal distribution with mean μa, variance o² and period 1. We first provide the following lemma.
Lemma 3. If the denoising term é (L, A, is periodic translation invariant, and the transition probabilty can be formulated as p(Fi-i|Lr, Fr, A) (F. F, + uérpLr Fr, A, where ue, are functions of t, the transition is periodic translation equivariant.
Proof. As the denoising term ép Lt, Ft, 1, is periodic translation invariant (for short PTT), have é (FFF t1 A,t) = êF(L, Fr, A, t), for any translation t E R²
For the wrapping function w--), we have
w(a + b) = w(((() + b), Ya, R.R.
(12)
For wrapped normal distribution (u, with mean μa, variance o² and period 1, and for any W, k" E Z, we have
(x + W' (u + k²) - k)3
Nu + Rl + k², 0²)
exp
VnT k==00
20
(r-
m²)
exp
(m= k - k' kk)
2
2n²
Na (x)p,,)))) 02)
Let k' 0,K" w(() - μ. we directly have
(o -0) (N)(
(13)
For any R, we have
Nulr 1/k 1,02)=
V2To Σ exp
+
(u+ +
k)2
202
k=
- 11 - k)
exp
V
202
k=-
Let W/' w(x + -xx +)" w(H (μ + t), have
((((x + o²) = No + t/ + t, 0,02) = (z)
(14)
19

---

For the transition probability p(F:-1 L1, Fi, A), we have
w/ (FF [
1) A)
t1 Pe + t1 uer(LL, 1/ Fr + +t11), A,t),
-Nu((o(FF-11 w0 Fi + +1 ttt1) uur (L, Ft, A, u)
(PTI Ep)
-Nu(wo(Ft-)) 11 (
w: +F 1 + uuér(L L,, Ft, A, 11) .)?)
(Eq. [3
=Nu(w (F.[ t1 w/
F; u,ēF( Fi, A,t) + +t) r
(Eq. 22
-Nro(F,-1 F. un e Li, Fú, (.0)
(Eq. (14)
-P(Fi-1| Fr, F).
The transition probability of the fractional coordinates during the Predictor-Corrector sampling can be formulated as
p O(F-11 L1 F2 A) (P F L,, F,,
Pp(F, Lt F. A) No F Fi + (o ,ं EF- (Le, Fr. A,t), o--100 - I),
of
F A) No. (F. Ft + T(-1 ér ér (1-11, F A,t 1) 0t- I),
σ1
01
where PP, PC the transitions of the predictor and corrector. According to lemma 3 both of the transitions are periodic translation equivariant. Therefore, the transition p(F.-1 L, F. A)) is periodic translation equivariant. As the prior distribution (100, 1) is periodic translation invariant, finally prove that the marginal distribution p( Fo) is periodic translation invariant based on lemma
A.3 Proof of Proposition
We rewrite proposition aas follows.
Proposition 3. The score é by Eq. ( is 0 -quivvrra,, and the score ér from Eq. (70) is periodic translation invariant. Hence, the generated distribution by DiffCSP is periodic E(3) invariant.
Proof. We first prove the orthogonal invariance of the inner product term transformation E R3x3 Q Q I, we have
L. For any orthogonal
= Q"
For the Fourier Transformation, consider i is even, have
VFr(w((f t) w(f. t))(c, k] =sin (2rm(u(f3e) te) w(fic te))
sin 2mm Si,e) - 2mm 2 Si,c) (w(Jec+l) w(Sec 6o))
(i,e)
=,rT(ff - [c)
Similar results can be acquired as is odd. Therefore, we have (PF(ww(f (((((()()))) = VFr
fa), Vt R3, i.e., the Fourier Transformation thr is periodic translation invariant. According to
the above, the message passing layers defined in Eq. (6) (8) is periodic E(3) invariant. Hence, we can directly prove that the coordinate denoising term EF periodic translation invariant. Let
é F, A,t) = ( Σ h. ((9)). For the lattice denoising term é,, = Lé,, we have
én ( F, A,
F, A,
QLE F, FL, A,t)
= A. 333
3333 QQQ -I
20

---

Translation space.
Periodic vanslation on a unit cell.
Figure 4: An example of periodic translation invariance. From the view of unit cell, the atoms translated across the right boundary will brought back to the left side.
Above all, é -quivvaaaar and é is periodic translation invariant. According to proposition 2 the generated distribution by DifCCSP in Algorithm 2 is periodic E(3) invariant.
A.4 Discussion on Periodic Translation Invariance
In Definition 3] we define the periodic translation invariance combination translation invariance and periodicity. To see this, we illustrate an additional example in Figure From a global view, when we translate all atom coordinates from left right, the crystal structure remains unchanged, which indicates translation invariance. At the same time, from the view of unit cell, the atom translated across the right boundary will be brought back the left side owing to periodicity. Therefore, for convenience, we define the joint effect of translation invariance and periodicity as periodic translation invariance.
Previous works 36 have shown that shifting the periodic boundaries will not change the crystal structure. In this section, we further show that such periodic boundary shifting is equivalent the periodic translation defined in Definition 3
Consider two origin points P1, P2 33x1 and the lattice matrix L, the constructed unit cells by P1, P2 can be represented as M1 = (A1, F1, L) and M2 = (A2, F2, L), where Fú, F2 E R33
fractional coordinates and
{(a', x1, = 01
P1 Lf1, + Lk, vk 2 21,1 01,
(15)
={(a',, - |a., = 2 P2 + Lf2.5 + Lk, Vk E záx1),
(16)
which means that the unit cells formed by different origin points actually represent the same infinite crystal structures (36). We further construct bijection M1 + M2 mapping each atom in Mi to the corresponding atom in unit cell M2. For the pair (a1,), f1,)) E M1, (a2 fe.) E M2, we have T(a1, f1,)) = (az.j. f2.5) iff 3k; E 23x1 S.T.
a 02.1
pi + Lfi, + Lk; = P2 + Lf2.1
After proper transformation, we have
f2 = f1; + L-1(pi -p2) + k,
(17)
As f1J, f2, E [0, 1)3x1, we have
k; - -1fii 1- (p: P2)1 f2.1 w f1.1 + L (p: p2))
which means shifting the periodic boundaries by changing the origin point into is equivalent to a periodic translation F2 = w w((F F1 + L-i(p) - -P2)1T)
21

---

B Implementation Details
B.1 Approximation of the Wrapped Normal Distribution
The Probability Density Function (PDF) of the wrapped normal distribution (O, (0, of)
is
(x (x- -:)3
(x of) = VTT
exp
20
where x E [0, 1). Because the above series is convergent, is reasonable to approximate the infinite
summation finite truncated summation [6]] as
(xr k)²
(2;
exp
2 k=-n
20
And the logarithmic gradient of f can be formulated as
(- (ar - 1:)2 -k)2
V log (w; 02)
Vz log
exp
v 2TTT
20
(x 10-kk k)?
V log
exp
20
n2
12
k= (k. - I) exp (- (x=)))²
20
Ek=-n exp (
(2-k) k3
20
To estimate A = - (0. [| log. No (x 13). we first sample m points from Now (0, of). and the expectation is approximated as
V log Surn (1;; 0, 0,
k= (ki 2.) exp (
(3re R)²²
20
m
of Σラ ビ exp
(xr k3
20
For implementation, select n 10 and 10000.
B.2 Adaptation of CDVAE
As illustrated in Figure 5] the original CDVAE g mainly consists of three parts: (1) 3D encoder encode the structure into the latent variable 23D, (2) property predictor to predict the lattice L, the number of nodes the unit cell N, and the proportion of each element in the composition c, (3) a 3D decoder to generate the structure from Z3D, L. N, c via the Score Matching with Langevin Dynamics (SMLD, Song and Ermon (49 method. The training objective composed of the loss functions on the three parts, iee. the KL divergence between the encoded distribution and the standard normal distribution LKL. the aggregated prediction loss LAGG and the denoising loss on the decoder LDEC- Formally, we have
LORI LAGG LDEC + 3DKL (M (IN(O.))
We formulate LxL (DDx(N(Hstt I)I I)) for better comparison with the adapted method. 3 is the
to balance the scale of the KL divergence and other loss functions.
To adapt the CDVAE framework to the CSP task, we apply two main changes. Firstly, for the encoder side, take the composition as the condition, we apply an additional ID prior encoder to encode the composition set into latent distribution A D. oiD!) and minimize the KI. divergence between the 3D and distribution. The training objective is modified into
LADA LAGG + LDEC + +3DrI (o.Nn//)nn
22

---

H
preftet
Lace
Encoder
(c)
Pao 6a0
3D Decoder
Foar
Axu
X(O.T)
Training
I
Inference
pradict
(8)
Original Franework
Adapted Franework
predict
fass
(e)
Encoder
30 Decoder
Losc
Axi
Encoder
(d)
1D Encoder
732
3D Decoder
Figure 5: Overview original (a.b) and adapted (c,d) CDVAE. The key adaptations lie two points. (1) We introduce additional ID prior encoder to fit the latent distribution of the given composition. (2) We initialize the generation procedure the 3D decoder with the ground truth composition and keep the atom types unchanged to ensure the generated structure conforms to the given composition.
During the inference procedure, the composition is given, the latent variable 2 is sampled from N ((HLD. I). For implementation, we apply a Transformer [51] without positional encoding as the ID encoder to ensure the permutation invariance. Secondly, for the generation procedure, apply the ground truth composition for initialization and keep the atom types unchanged during the Langevin dynamies ensure the generated structure conforms the given composition.
B.3 Algorithms for Training and Sampling
Algorithm summarizes the forward diffusion process as well as the training of the denoising model 6, while Algorithm gegorithm@jllusttee iillstraate the backward sampling process. They can maintain the symmetries à is delicately constructed. Notably. We apply the predictor-correst correcto sampler [50] to sample Fg. In Algorithm Line refers to the predictor while Lines 9-10 correspond corrector.
Algorithm Training Procedure of DifTCSP
Input: lattice matrix Lo. atom types fractional coordinates Fo. denoising model d, and the number of sampling steps I'.
2: Sample ~ -N(0, 1).eF ~ N(0, I) and ~ u(1,T).
3: L. + Lo 11- dLEL
4: F, + +
5: Êg t + F., A,t)
6: LI. + ||E1. - ên.1
7: LF + V log q (F Fo) éF
8: Minimize Lt. LF
23

---

Algorithm 2 Sampling Procedure of DifICSP
1: Input: atom types A, denoising model &, number of sampling steps T, step size of Langevin dynamics
2: Sample Lr ~ (0, Fr 1).
for t T.......
4: Sample eL, CF; 6x ~ ()
5: éL, ê + Φ Lt, Fi, A,t).
6: L--1 t yp
L
[I - vI-ai êt.) ) +
8: 1-
T-a EL.
- EF)
7:
F wFF: + (o? - 12-1
11
8:
_ t 6(b1-1: F: A.A. 11)
9:
de 16-00-1101
10: Fi-1 t wl 1 + d + vd(t'z)
11:
end for
12: Return Lo, Fo.
B.4 hyper-parameters and Training Details
We acquire the origin datasets from CDVAR
9 and MPTS-52 (57 We utilize the codebases from for baseline implementations.
GN-OA 21 eG-SchNet (53 and CDVAR
For the optimization methods, we apply the MEGNet [52] with 3 layers, 32 hidden states as property predictor. The model is trained for 1000 epochs with Adam optimizer with learning rate 1 x 10 10-3 As for the optimization algorithms, we apply RS, PSO, and BO according to Cheng et al. [21
For RS and BO, We employ random search and TPE-based BO as implemented in Hyperopt 662
Specifically, we choose observation quantile ry as 0.25 and the number of initial random points as 200 for BO. For PSO, we used opi and choose the momentum parameter as 0.8, the cognitive as 0.5, the social parameters as 0.5 and the size of population as 20.
For P-cG-SchNet, we apply the SchNet [29] with 9 layers, 128 hidden states as the backbone model. The model is trained for 500 epochs on each dataset with an Adam optimizer with initial learning rate 1 1x 10 and Plateau scheduler with decaying factor 0.5 and patience of 10 epochs. We select the element proportion and the number of atoms in a unit cell as conditions for the CSP task. For CDVAE, we apply the DimeNet++ [63] with 4 layers, 256 hidden states as the encoder and the GemNet-T [64] with 3 layers. 128 hidden states as the decoder. We further apply a Transformer
[51 model with
layers, 128 hidden states as the additional prior encoder as proposed in Appendix B.2 The model is
trained for 3500, 4000, 1000, 1000 epochs for Perov-5, 24r MP-20 and MPTS-52 respectively with an Adam optimizer with initial learning rate 1 10 and Plateau scheduler with decaying factor 0.6 and a patience of 30 epochs. For our DiffCSP, we utilize the setting of 4 layer, 256 hidden states for Perov-5 and 6 layer, 512 hidden states for other datasets. The dimension of the Fourier embedding is set to l- = 256. We apply the cosine scheduler with & = 0.008 to control the variance of the DDPM process on L,, and an exponential scheduler with 0 0.005, oT 0.5 to control the noise scale of the score matching process on F.. The diffusion step is set to T = 1000. Our model is trained for 3500, 4000, 1000, 1000 epochs for Perov-5, Carbon- 24, MP-20 and MPTS-52 with the same optimizer and learning rate scheduler as CDVAE. For the step size 7 in Langevin dynamics for the structure prediction task, we apply y = 5 x 10-1 for Perov-5, x 10 5 for MP-20 and MPTS-52, and for Carbon-24, we apply y 5 x 10 to predict one sample and samples. For the ab initio generation and optimization task on Perov-5, Carbon-24 and MP-20. we apply y 1 x 10-6, x 10-5, 5--,, 5 x 10-6, respectively. All models are trained on GeForce RTX 3090 GPU.
5 x 10 for multiple
httppp//gtthh...coomm//aaaaaaaap split
24

---

Exploring the Effects of Sampling and Candidate Ranking
C.1 Impact Sampling Numbers
Figure 6 illustrates the impact of sampling numbers on the match rate. The match rate of all methods increases when sampling more candidates, and DiftCSP outperforms the baselines methods under the arbitrary number of samples.
102
Farzu3
100
Carben-24
BificoP
92
100
03
23
DHCSP
Foc Seh
COWAE
8
10
12
24
28
Numbeer Sam
Nunobe of samppae
103
HP 20
100
NPT'
F-sd Soh Mer
SoWiet
10a 000 PN
10
32
Number
Figure 6: Comparison on different number of samples.
C.2 Diversity
We further evaluate the diversity by yield- ing the CrystalNN (65] fingerprint of each generated structure, calculating the L2- distances among all pairs in the 20 samples of each composition, collating the mean and max value of the distances, and finally averaging the results from all testing candi- dates. We list the diversity of CDVAE and DiffCSP on Perov-5 and MP-20 in Table [
Table 5: Comparison on the diversity.
Perov-5
MP-20
Mean
Max
Mean
Max
CDVAE
0.3249
0.7316
0.3129
0.6979
DiffCSP
0.3860
0.8911
0.2030
0.5292
C.3 Ranking among Multiple Candidates
In $5.1 we match each candidate with the ground truth structure pick up the best sample. However, in real CSP scenarios, the ground truth structure is not available, necessitating a confidence model to rank the generated candidates. To address this, we develop three types of confidence mod- els to score each sample for ranking from different perspectives: Energy Predictor (EP). Since lower formation energy typi- cally leads to more stable structures, we directly train a predictor using energy bels and apply the negative of the predicted
Table 6: Results different confidence models. Oracle means applying the negative RMSD against the ground truth as the ranking score.
Match rate (%)
RMSE 1
EP
51.96
0.0589
MD (d=0.1)
60.30
0.0357
MD (d=0.3)
60.13
0.0382
MD (d=0.5)
59.20
0.0469
CS
58.81
0.0443
Oracle
77.93
0.0492
25

---

energy as the confidence score. Match Discriminator (MD). Inspired by Diffdock [2], we first generate five samples for each composition in the training/validation set using DiffCSP and calculate their RMSDS with the ground truth. We then train a binary classifier to predict whether the sample matches the ground truth with an RMSD below a threshold d. The predicted probability serves as the score. Contrastive Scorer (CS). Drawing inspiration from CLIP [66], we train contrastive model between a ID and 3D model to align the corresponding compositions and structures. The inner product of the ID and 3D models is used as the score. We select the Top-1 result among the 20 candidates ranked by each confidence model on the MP-20 dataset, as shown in Table The results indicate that MD and CS perform relatively better, but there remains a gap between the heuristic ranking models and the oracle ranker. Designing powerful ranking models is an essential problem, which we leave for future research.
D Impact of Noise Schedulers
We explore the noise schedulers from three perspectives and list the results in Table
Table 7: CSP results on different schedulers. DiffCSP utilizes cosine scheduler on and mma 0.5 on F.
and [8] 1. For lattices, we originally use
Match rate (%) 1 RMSE +
the cosine scheduler with = 0.008, and we
change it into the linear and sigmoid sched- ulers with 3 0.0001, 3 = 0.02. We find that the linear scheduler yields compa- rable results, while the sigmoid scheduler hinders the performance. 2. For fractional coordinates,
DiffCSP
51.49
0.0631
Schedulers of
Linear Sigmoid
50.06
0.0590
45.24
0.0664
the exponential sched- uler with mmnn = 0.005, mea 0.5, and we change the value of maz into 0.1 and 1.0. Results show that the small-O vari- ant performs obviously worse, only suf- ficiently large mar could approximate the prior uniform distribution. We visualize the PDF curves in Figure 7 for better un- derstanding. 3. For atom types, duct similar experiments as lattices, and the results indicate that the original cosine scheduler performs better. In conclusion, we suggest applying the proposed noise schedulers.
Schedulers of F
T max = 0.1
32.56
0.0913
T max = 1.0
47.89
0.0675
Table 8: Ab initio generation results on different type schedulers. DifICSP utilizes cosine scheduler on A.
Validity
Coverage
Struct.(%)
Comp.(%)
Recall
Precision
DifCCSP
100.00
83.25
99.71
99.76
Linear Sigmoid
99.70
79.78
98.29
99.48
99.88
81.59
99.33
99.55
E
Learning Curves of Different Variants
We plot the curves of training and validation loss of different variants proposed in $5.3 in Figure with the following observations. The multi-graph methods struggle with higher training and validation loss, the edges constructed under different disturbed lattices vary significantly, complicating the training procedure. 2. The Fourier transformation, expanding the relative coordinates and maintaining the periodic translation invariance, helps the model converge faster at the beginning of the training procedure. 3. The variant utilizing the fully connected graph without the Fourier transformation (named "DiffCSP w/o FT in Figure ()) encounters obvious overfitting the periodic translation invariance violated, highlighting the necessity to leverage the desired invariance into the model.
F
Comparison with DFT-based Methods
F1
Implementation Details
We utilize USPEX [59]. a DFT-based software equipped with the evolutionary algorithm to search for stable structures. We use 20 populations in each generation, and end up with 20 generations for each compound. We set 60% of the lowest-enthalpy structures allowed to produce the next generation through heredity (50%), lattice mutation (10%), and atomic permutation (20%). Moreover, two
26

---

4.0
0=0.1
0=0.2
3.5
0=0.3
0=0.5
3.0
0=1.0
2.5
PDF(x)
2.0
1.5
1.0
0.5
0.0
0.0
0.2
0.4
0.6
0.8
1.0
X
Figure 7: PDF curves of the wrapped normal distribution with periodic as and different noise scales. It can be find that 0.5 is practically large enough to approximate the prior uniform distribution.
WC
16
MOW/
MS FT
DICSPP WT
12
1.2
10
08
Spock
Epoch
Figure 8: Learning curves of different variants proposed in $ 5.3 MG and FT denote multi-graph edge construction Fourier transformation, respectively.
lowest-enthalpy structures are allowed to survive into the next generation. The structural relaxations are calculated by the frozen-core all-electron projector augmented wave (PAW) method [67] as implemented in the Vienna ab initio simulation package (VASP) [68]. Each sample's 's calculation is performed using one node with cores (Intel(R) Xeon(R) CPU ES-2692 @ 2.20GHz), while
populations within the same generation are concurrently computed across 20 nodes in parallel. The
exchange-correlation energy is treated within the generalized gradient approximation (GGA), using the Perdew-Bal Burke-- -er---Buukeretttte Ernzerhof (PBE) function 669
27

---

F.2 Results
We select 10 binary and 5 ternary compounds MP-20 testing set and compare our model with USPEX. For our method, we sample 20 candidates for each compound following the setting in T For USPEX, we apply 20 generations, 20 populations for each compound, and select the best sample in each generation, leading to 20 candidates as well. We list the minimum RMSD) of each compound in 9 and additionally summarize the match rate over the compounds, the averaged RMSD over the matched structures, and the averaged inference time to generate 20 candidates for each compound in Table I0 The results show that DifICSP correctly predicts more structures with higher match rate and significantly lower time cost.
Table 9: The minimum RMSD of 20 candidates of USPEX and DiffCSP, "N/A" means none of the candidates match with the ground truth.
Binary
CoySb2
SraO1
AlAg4
YMgs
Cr,Si
USPEX
0.0008 0.0005
0.0121
N/A
0.0000
N/A
DiffCSP
0.0133
0.0229
0.0003
0.0066
Binary
Pd,
AgoO2
Co
BayCd,
BigFs
USPEX
0.0184
0.0079
0.0052
N/A
N/A
DifCCSP
0.0264
N/A
N/A
0.0028
N/A
Temary
KZnF3
CryCuOs
Sig(CN2)4
Hg2S2Os
USPEX
0.0123 0.0006
N/A
N/A
N/A
0.0705
DiffCSP
0.0482
0.0203
N/A
0.0473
Table 10: Overall results over the 15 selected compounds.
Match Rate (%)
Avg. RMSDJ
Avg. Time.
USPEX DiffCSP
53.33 73.33
0.0159
12.5h
0.0172
10s
G Extension to More General Tasks
G.1
Overview
Our method mainly focuses on CSP, aiming at predicting the stable structures from the fixed com- position A. We first enable the generation on atom types and extend to ab initio generation task G.2, and then adopt the energy guidance for property optimization
$ G.3 Figure 9 illustrated the differences and connections of CSP, ab initio generation, and property optimization. Besides, CDVAE g proposes reconstruction task specific for VAE-based models, which first encodes the ground truth structure, and require the model to recover the input structure. As our method follows diffusion-based framework instead of VAE, it is unsuitable to conduct the reconstruction task of our method.
G.2 Ab Initio Generation
Apart from L and F, the ab initio generation task additionally requires the generation on A. As the atom types can be viewed as either N discrete variables in h classes, or the one-hot representation A E R continuous space. We apply two lines of diffusion processes for the type generation, which are detailedly shown as follows.
Multinomial Diffusion for Discrete Generation Regarding the atom types as discrete features, we apply the multinomial diffusion with the following forward diffusion process [70] 42 7T].
qAA, Ao) = (á + (1 - āt) h 1 1,
(18)
28

---

Crystal Structure Prediction (CSP)
(b) Ab Initio Generation
Property Optimization
(d) Reconstruction
Figure 9: Different tasks for crystal generation. Our approach mainly focuses on the CSP task (a), and capable to extend into the ab initio generation task (b) by further generating the composition, and the property optimization task (c) via introducing the guidance on the target property. The reconstruction task (d) in Xie et al. 9 is specific for VAE, which is unnecessary for our diffusion-based method.
where 1 R In
Nx1 are vectors with all elements setting to one, Ao is the one-hot representation of the origin composition, and the function C(-) samples the multinomial distribution with the conditional probability and returns the one-hot representation of A,.
The corresponding backward generation process defined as
p(A,-I = 1(A,-1
10.// Σ Ock);
(19)
k=1
where
0. (a. (1) 1,15
(aréa(Mo,))) + (1-ār) 1 1,1))
(20)
h
h
and A RAxN is predicted by the denoising model. We further find that specific for 1, Ao = ar 1)) works better.
The training objective for multinomial diffusion is
L Adiicret
(L(g(A, |440(
(21)
One-hot Diffusion for Continuous Generation Another approach is to simply consider the com- position A as continuous variable in real space RRx,,, which enables the application of standard DDPM-based method [41]. Similar to Eq. (2)-44) the forward diffusion process is defined as
((A.AA)) = N Lelva, (1 (1-)))))
(22)
And the backward generation process defined as
p(A:-1 AA (M(,)))
(23)
where (A(M) = 1 A. vー (M( 03 (MM) = B, -- The denoising term
Vas
é ( 1) E RaxN predicted by the model.
The training objective for one-hot diffusion is
A continuns = EE EEA. - (M,,,
(24)
The entire objective for training the joint diffusion model L, F. A combined
LM = Xr + XFLF + AACA
(25)
29

---

We select ΛY = = = for the CSP task as A is fixed during the generation process, and dz, = Ap = 1, 20 for the ab initio generation task to balance the scale of each loss component. Specifically, we do not optimize LA on the Carbon-24 24 dataset, as all atoms in this dataset are carbon.
Sample Structures with Arbitrary Numbers of Atoms As the number of atoms in unit cell (i.e. N) is unchanged during the generation process, we first sample N according to the distribution of N in the training set, which similar Hoogeboom et al. [4]]. The sampled distribution p(M) can be more concisely described as p(M, N) p(N) (( N). The former term p(N) ) is sampled from pre-coepeee data distribution, and the latter conditional distribution NM modeled by
DiffCSP.
Evaulation Metrics The results are evaluated from three perspectives. Validity: We consider both the structural validity and the compositional validity. The structural valid rate is calculated as the percentage of the generated structures with all pairwise distances larger than 0.5A and the generated composition is valid if the entire charge is neutral as determined by SMACT [22]. Coverage: It measures the structural and compositional similarity between the testing set St and the generated sam- ples S.. Specifically, letting ds(M1. M22), dc(M1. denote the L2 distances of the CrystalNN structural fingerprints [65] and the normalized Magpie compositional fingerprints (73), the COVer- age Recall (COV-R) is determined as COV-R = [.. [M:]MM S, 3M₁ Sg (M) <
5g. dc(Ms. < &cll where ōs, dc are pre-defined thresholds. The COVerage Precision (COV-P) defined similarly by swapping Sg. St. Property statisties: We calculate three kinds of Wasserstein distances between the generated and testing structures, in terms of density, formation energy, and the number of elements [9]. denoted as dp. dE and delem. individually. The validity and coverage metrics are calculated on 10,000 generated samples, and the property metrics are evaluated on a subset with 1,000 samples passing the validity test.
Results We denote the abovementioned discrete and continuous generation methods as DiffCSP-D, DiffCSP-C, respectively. The results of the two variants and the strongest baseline CDVAE on MP-20 are provided in Table [T] We observe that DiffCSP-D yields slightly lower validity and coverage rates than DiffCSP-C. Moreover, DiffCSP-D tends generate structures with more types of elements, which is far from the data distribution. Hence, we select DifiCSP-C for the other experiments in DiffCSP). We further find that DiffCSP-C supports the property optimization task in the next section. Besides, both variants have lower composition validity than CDVAE, implying that more powerful methods for composition generation are required. As our paper mainly focuses on the CSP task, we leave this for future studies.
Table (abbeeeaabbbevatee
G.3 Property Optimization
On top of DiffCSP-C, we further equip our method with energy guidance 74 75 for property optimization. Specifically we train a time-dependent property predictor E(L,, F,, A,, f) with the same message passing blocks as Eq. (6)- And the final prediction is acquired by the final layer as
E PE Σ h ())
(26)
And the gradients w.r.t F, A additionally guide the generation process. As the inner product term L L is 0(3) invariant, and the Fourier transformation term (fr - fi) is periodic translation invariant, the predicted energy E is periodic E(3) invariant. That is,
E(QL, w + t1 Ar. E(L,, A.,
(27)
Table 11: Results on MP-20 ab initio generation task.
Validity (%) t
Coverage (%) t COV-R
Property de
Data
Method
Struc.
Comp.
COV-P
dp
delem
MP-20
CDVAE 9
100.0
86.70
99.15
99.49
0.6875
0,2778
1.432
DinCCP-D
99.70
83.11
99.68
99.53
0.1730
0.1366
0.9959
DinCSP-C
100.0
83.25
99.71
99.76
0.3502
0.1247
0.3398
30

---

Algorithm 3 Energy-Guided Sampling Procedure of DiffCSP(-C)
1: Input: denoising model 6, energy predictor E, step size of Langevin dynamics guidance magnitude input structure before optimization M (L, F. A), number of sampling steps 7", maximum number of sampling steps
2: if T'= then
3: Sample LT ~ (O,, ~ U(0, 1), AT ~ ~N(O,)) I).
4: else
5: Sample LT (L) ~ (Frr F), Ar (Ar (A) according to Eq.
and 22
end
7:
fort- ...... do
8:
Sample EL, EF; EF I)
9: ÉL,
oi Lt. Fi, A.,
10: Acquire VLE, E, V A from E( , Fi, At,
11: 1--1
val b
8ßt- 1-6 F v
I
v1-u,
I-ac
12:
A t vai (A: ア-
2 _
éA)
=8, : 1-
T-
V +
1-01-1 I-ai
EA.
2-] -03-1 VF
13:
F w(F: F. (ot o(-)
of
d1
_ + (I+ P Ac-1,t -1).
15:
d 116-0--1/01
16:
Fi-1 t wi(k (F,, d,ēp + 2d(
17:
end for
18:
Return Lo, Fo, argamax(XAA
Table 12: Results on property optimization task.
The results of baselines are from Xie al. 9
Perov-5 SR10
Carbon-24
MP-20 SR10
Method
SR5
SR15
SR5
SR10
SR15
SR5
SR15
FTCP
0.06
0.11
0.16
0.00
0.00
0.00
0.02
0.04
0.05
Cond-DFC-VAL VAE
0.55
0.64
0.69
CDVAE
0.52
0.65
0.79
0.00
0.06
0.06
0.78
0.86
0.90
DifCCSP
1.00
1.00
1.00
0.50
0.69
0.69
0.82
0.98
1.00
Taking gradient both sides w.r.l. L, F, A, respectively, we have
Q Vi,E(L' w (FF + +1 t), = EL, F,, Ar,t), V F F, A,,t) =/=((F++11)) E(, F,, Ar. VA, F(QL,, w| (F, F, + +11 AA,,) VA, F,, A.,t)
(28)
(29)
(30)
which implies that the gradient to L, is 0(3) cquivariant, and the gradient to F, and A, is periodic E(3) invariant. Such symmetries maintain that the introduction of energy guidance does not violate the periodic E(3) invariance of the marginal distribution.
The detailed algorithm for energy- -grided sampling is provided in Algorithm 3 We find that s 50 works well on the three datasets. We evaluate the performance of the energy-guided model with the same metrics as Xie et al. 9 We sample 100 structures from the testing set for optimization. For each structure, we apply T 1, 000 and T' = 100, 200, ... 1, 000, leading to 10 optimized structures. We use the same independent property predictor as in Xie et al. [ to select the best one from the 10 structures. We calculate the success rate (SR) as the percentage of the 100 optimized structures achieving 5, 10, and 15 percentiles the property distribution. We select the formation energy per atom the target property, and provide the results on Perov-5, Carbon-24 and MP-20 in Table 12 showcasing the notable superiority of DiffCSP over the baseline methods.
Aside from the Carbon- 24 dataset, the composition is flexible in the above experiments. We also attempt to follow the CSP setting and optimize the crystal structures for lower energy based on the fixed composition. We visualize eight cases in Figure T0]and calculate the energies of the structures before and after optimization by VASP (8). Results show that our method is capable to search novel structures with lower energies compared to existing ones.
31

---

Orignal Strctures
Optimized Siructures
Energy:
Goergy: 41.11.
Energs 22.12
Energy Gol
Energy 120 Ste
Energy: 44003
Energy: 72
Energy 73
Figure 10: Visualization of 8 pairs of structures before and after optimization.
Table 13: GPU hours for yielding 20 candidates over the testing set.
Table 14: Results on MP-20 with different
inference steps.
Perov-5
Carbon-24
MP-20
MPTS-52
Steps
Match rate (%)T
RMSD
DifCCSP
1,000
51.49
0.0631
P-cG-SchNet
2.1
3.0
10.0
22.5
5,000
$2.95
0.0541
CDVAE
21.9
9.8
95.2
178.0
CDVAR
1,000
30.71
0.1288
DifCCSP
1.5
3.2
18.4
34.0
5,000
0.1045
H
Computational Cost for Inference
We provide the GPU hours (GeForce RTX 3090) for different generative methods predict 20 candidates on the 4 datasets. Table 13 demonstrates the diffusion- models (CDVAE and Yet, the computation overhead of our method still
our DiffCSP) are slower than acceptable given its clearly better performance than P-cG-SchNet. Additionally, our DiffCSP is much faster than CDVAE across all datasets mainly due to the fewer generation steps. CDVAE requires 5,000 steps for each generation, whereas our approach only requires 1,000 steps. We further compare the performance of CDVAE and DiffCSP with 1,000 and 5,000 generation steps on MP-20 Table [4 Our findings indicate that both models exhibit improved performance with increased number of steps. Notably, DifCCSP with 1,000 steps outperforms CDVAE with 5,000 steps.
I Error Bars
We provide single run to generate 20 candidates in Table We apply two more inferences each generative method on Perov-5 and MP-20. Table 15 shows similar results as Table
Table 15: Results on Perov-5 and MP-20 with error bars.
Perov-5
MP-20
# of samples Match rate (%)t
RMSEJ
Match (%)t RMSEL
1
47.64
0.4170-10.0006
15.59-10.41
0.3747-10.0000
P-eG-Schnet
20
97.92+10.02
0.3464=000004
32.70=0.12
0.3020::0.0002
1
45.31
0. 0.1123=+0.002
33.93+2.115
0. 0.1009::00011
CDVAE
20
88.20-+0.26
0.0473++0.0077
67.20-10.23
0.1012++0.0016
1
52.35.10.26
0.077811000030
51.89
0.0611110000
DiffCSP
20
98.58.10.02
0.012910.00003
77.85.10.23
0.049311.0001
J More Visualizations
In this section, we first present additional visualizations of the predicted structures from DiffCSP and other generative methods in Figure [1 In line with Figure our DiffCSP provides more accurate predictions compared with the baseline methods. Figure illustrates generated structures on Perov-5, Carbon-24 24 and MP-20. The visualization shows the capability of DifICSP to generate diverse
32

---

structures. We further visualize the generation process of 5 structures from MP-20 in Figure 13 We find that the generated structure Mo is periodically translated from the ground truth structure. indicating that the marginal distribution p(Ma) follows the desired periodic translation invariance. We provide the detailed generation process the Supplementary Materials.
Perov-5
Carbon-24
MP-20
MPTS-52
Ground Truth
DifCCSP
CDVAE
P-cG-SchNel
Figure 11: Additional visualizations of the predicted structures from different methods. We translate the same atom to the origin for better visualization and comparison.
Perov-5
Carbon-24
MP-20
Figure 12: Visualization of the generated structures by DiffCSP.
33

---

== 1000
t= 750
500
250
t= 0
Translated
Ground Truth
Figure 13: Visualization of the generation process MP-20. The column "Translated" means translating the same atom in the generated structure Mg to the origin as the ground truth for better comparison.
34


Reviews:
[{"review_id": "3zj0s8uKOv", "reviewer": "Reviewer_1HBn", "summary": "The paper studies crystal structure prediction, an important topic in the field of AI for scientific discovery. Diffusion models are designed considering necessary symmetries for crystal study, with diffusing lattice vectors and fractional coordinates for generating lattices and atoms, respectively. Experiments show the effectiveness of the proposed method.", "questions": "See Weaknesses. An additional one is-- how's the efficiency compared with CDVAE in empirical studies?", "limitations": "N/A", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper presents novel and shining ideas, e.g., considers fractional coordinates rather than Cartesian coordinates. This eliminates the need of constructing multi-edge graphs and naturally guarantees periodic invariance. The theory for processing fractional coordinates in the diffusion model looks rigid and solid.\n\n2. The joint learning fashion for lattices and coordinates should be the right way to go in the field.\n\n3. Experiments are solid and performance is impressive.", "weaknesses": "1. One major concern is-- the technical contribution is not clear or strong enough. Diffusion models for 3D scientific data ate not new. The authors may need to clearly state the significant and unique design of the proposed diffusion model to distinguish it from other diffusion models.\n\n2. Another concern is-- necessary symmetries are not completely considered. For example, similar to CDVAE, translation invariance (not periodic translation invariance) may not be considered.\n\n3. A key foundation is-- the consideration of fractional coordinates can achieve periodic invariance, which is roughly correct. However, if the three lattice vectors are not defined in a deterministic way, this may not be true.\n\n4. There are some issues in the organization. For example, much content in 4.1 overlaps with ref [36], and this content should be put in background/preliminary and credits should be given properly.\n\nOverall, this is a neat work. I would raise my scores if my concerns are fully addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful to the reviewer for the constructive comments. We answer the reviewer’s questions as follows.\n\n> **Q1: The technical contribution is not clear or strong enough. Diffusion models for 3D scientific data are not new. The authors may need to clearly state the significant and unique design of the proposed diffusion model to distinguish it from other diffusion models.**\n\nSorry for the insufficient clarity regarding our technical contribution. We sincerely thank the reviewer and will further clarify our contributions with the statements below:\n\nEven diffusion models for GENERAL 3D molecules [41] are not new, this paper focuses on crystals which exihibit an extra form of symmetry, namely, periodicity. As a consequence, our diffusion models should not only generate the 3D coordinates but also the lattice matrix. This leads to the unique design of our model:\n- **Joint Diffusion Framework**: Our model performs joint diffusion on both lattice structures and atom coordinates, and these two processes are coupled and promoted with each other (Figure 2). Conversely, prior crystal diffusion frameworks such as CDVAE[36] focus solely on coordinate diffusion within a fixed region bounded by the predicted lattice. \n- **Periodicity-Aware Design**: While previous models for crystals[36] employ Cartesian coordinates, we depart from this convention by utilizing fractional coordinates, which naturally fit the periodicity consideration. Given that the domain of fractional coordinates forms a quotient space, we leverage the Score-Matching (SM) based framework along with Wrapped Normal distribution (Line 182-203), other than the commonly-used DDPM paradigm.\n- **Novel Denoising Model**: In Section 4.3, we design a novel denoising model to fulfill the above diffusion process. Firstly, Fourier Transformation of the relative fractional coordinate is adopted to better encode periodicity (Eq.6). Secondly, the output layer is specifically designed to jointly predict the denoising terms of the entire lattice (Eq.9) and each fractional coordinate (Eq.10). Lastly, the proposed denoising model is novel compared to traditional multi-graph representation, and it is proved to be periodic E(3) equivariance (Proposition 3).\n\nThe above statements will be added to the revised paper.\n\n> **Q2: Necessary symmetries are not completely considered. For example, similar to CDVAE, translation invariance (not periodic translation invariance) may not be considered.**\n\nThank you for the comment. Actually, for crystals, our definition of periodic translation invariance already implies translation invariance. To see this, we provide an example in Figure R1 (General Response). From a global view, when we translate all atom coordinates from left to right, the crystal structure remains unchanged, which indicates translation invariance. At the same time, from the view of a unit cell, the atom translated across the right boundary will be brought back to the left side owing to periodicity ($L$ is fixed), which indicates periodic invariance. Therefore, for convenience, we define the joint effect of translation invariance and periodicity as periodic translation invariance. We will further clarify this point in the paper. \n\n> **Q3: The consideration of fractional coordinates can achieve periodic invariance, which is roughly correct. However, if the three lattice vectors are not defined in a deterministic way, this may not be true.**\n\nWe are unsure if we understand the reviewer's question (please point it out if not). Our point is that the definition of periodic invariance is independent of whether the lattice vectors are deterministic or not. When defining periodic invariance, we mean given a lattice matrix, the crystal distribution remains the same under any global translation of atom coordinates. This can be checked in the equation in Definition 3, where both $L$ and $F$ are fixed and only the translation $t$ varies.  \n\n\n> **Q4: There are some issues in the organization. For example, much content in 4.1 overlaps with ref [36], and this content should be put in background/preliminary and credits should be given properly.**\n\nThanks for pointing it out. We are willing to move Section 4.1 to \"Background/Preliminary\", and will further emphasize the connections and differences with [36]. Indeed we have discussed the difference between our definitions and those presented in [36] in Appendix A.4, which will be modified into the main paper.\n\n> **Q5: How's the efficiency compared with CDVAE in empirical studies?**\n\nThank you for the comment. We have indeed compared the computational cost of each generative method in Appendix G and Table 9. We provide the GPU hours (GeForce RTX 3090) for different generative methods to predict 20 candidates on the 4 datasets. Our method is more efficient mainly because it requires fewer diffusion steps for generation. Particularly, to achieve the accuracies reported in Table 1, CDVAE requires 5,000 steps for each generation, whereas our approach only requires 1,000 steps. We bring the results here for convenience.\n\n|             | Perov-5 | Carbon-24 | MP-20 | MPTS-52 |\n|-------------|---------|-----------|-------|---------|\n| CDVAE       | 21.9    | 9.8       | 95.2  | 178.0   |\n| DiffCSP     | 1.5     | 3.2       | 18.4  | 34.0    |"}, {"review_id": "sdosjqNTFd", "reviewer": "Reviewer_eNoa", "summary": "This paper addresses the crystal structure prediction problem, one of the important problems in this field.\nThe authors propose a diffusion model generating both lattices and fractional coordinates given chemical composition.\nThe proposed denoising model satisfies O(3) equivariance and periodic translation invariance.\nThe experimental results show that the proposed method outperforms the conventional methods in generating existing crystals.", "questions": "Is it reasonable to compare the inference time of the proposed method with the computation time of USPEX? Should the training time be included in the CSP time? How long does the training take?", "limitations": "Yes, the limitation is discussed in Section 6.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is well-organized and well-written.\n* The representation of crystals is reasonable.\n* The empirical result is strong, and the baseline methods are appropriate for the CSP task.\n* The results compared with USPEX, a strong GA-based method,  attract the CSP community's interest.\n* Ablation studies support the design choice of the proposed method.", "weaknesses": "The study does not report any new crystal structures. It will increase the impact of this paper if the proposed method can find a new crystal that achieves lower energies via DFT than that of the known crystals of a given multi-component system.\n\nA minor comment\n\n\"$\\psi_L$ in Eq. (10)\" in the main text of section 4.3 seems to be a typo in Eq. (9).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful to the reviewer for the constructive comments. We answer the reviewer’s questions as follows.\n\n> **Q1: The study does not report any new crystal structures. It will increase the impact of this paper if the proposed method can find a new crystal that achieves lower energies via DFT than that of the known crystals of a given multi-component system.**\n\nNice suggestion! In fact, we have equipped our method with a guidance model to generate structures with lower energies. Please refer to Appendix F.3 for more details. Here according to the reviewer's suggestion, we fix the composition of the initial crystal and only optimize its structures for lower energy. Results are visualized in Figure R2 (General Response), where the energies of the structures before and after optimization by VASP are shown. It reads that our method does search novel structures with lower energies compared to existing ones. \n\n\n> **Q2: Is it reasonable to compare the inference time of the proposed method with the computation time of USPEX? Should the training time be included in the CSP time? How long does the training take?**\n\nThanks for pointing it out. In our analysis of inference time, we aim to focus on the efficiency of generating samples once the model has been trained and is ready for deployment. This specific aspect is crucial when considering the practical application of a method in real-world scenarios, where users are primarily concerned with obtaining results in a timely manner. While the training time is indeed an essential consideration in understanding the overall computational cost of a method, it represents a one-time investment that occurs during the model development phase. Once the model is trained, its inference time becomes a recurring factor every time predictions are needed.\n\nMoreover, the training time for our method is approximately 5.2 hours on MP-20, and we do acknowledge that this is an important consideration for users when implementing our model. Even taking the training time into consideration, our method still showcases more efficiency against USPEX (the total time is 5.2h + 10s vs 12.5h)."}, {"review_id": "o3wVQ1MYy9", "reviewer": "Reviewer_pLZt", "summary": "DiffCSP is a novel diffusion model for predicting crystal structures conditioned on the crystal's atom composition; it can also be extended to an initio crystal generation (sampling both atom composition and positions). Crystal modeling is challenging due to its periodic symmetry. DiffCSP proposes to solve periodicity constraints by diffusing on fractional coordinates and jointly generating the lattice and atom coordinates. Fractional coordinate diffusion is achieved with torsional diffusion that constraints each intermediate state to stay within a bounding box, i.e. a fixed unit cell. When compared with existing methods, DiffCSP is state-of-the-art across DFT and ML-based methods in predicting the crystal structures on three datasets -- Perov-5, Carbon-24, and MP-20 -- while requiring less computation. It is also state-of-the-art on ab initio generation for the same datasets.", "questions": "Concrete questions pertaining to my weaknesses:\n\n- If there's a single right answer of the structure for a given atom composition then is it right to generate many possible structures? The author evaluate DiffCSP and other generative models by matching the samples to the ground truth structure to pick the best sample. But how would one select the best sample if the ground truth was not available?\n\n- What is the rationale behind gaussian diffusion for the lattice parameters? Are magnitudes of the lattice bounded? If it is unbounded then it seems strange to use gaussian diffusion.\n\n- Why is DiffCSP chosen to be E(3) equivariant but not SE(3) equivariant like in CDVAE? Is this a limitation?\n\n- Was the noising schedule of the atom types, lattice, and coordinates important? Did it make a practical difference?\n\nOther questions:\n\n- Could Table 1 include mean and standard deviation for RMSE? If it doesn't fit then in the appendix would be fine. It would be helpful to know the spread of the metrics when multiple samples are used.\n\n- How diverse are the crystals when multiple samples are drawn?\n\n\n\nTypoes/suggestions:\n\n- Line 33. Reference 11 does protein structure generation, not prediction.\n\n- It would be helpful to include mean and standard deviation for RMSE in Table 1. If it doesn't fit then in the appendix would be fine. It would be helpful to know the spread of the metrics when multiple samples are used.\n\n- It is odd to frame diffusion on the fractional coordinates as a DDPM when it is really a VE-SDE following the SDE formulation. I believe it will cause less confusion if some clarification is given to the original source of the torsional diffusion formulation as Riemannian score-matching [1].\n\n[1] https://arxiv.org/abs/2202.02763\n", "limitations": "- A potential limitation is SE(3) vs. E(3) equivariance. The authors do not discuss this. I have mentioned this already in the questions.\n\n- The authors do not discuss choice of diffusion hyperarameters which I believe would have lots of practical use for researchers who would build on top of this work.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- Using fractional coordinates to handle periodicity is supported by several prior studies on crystal structures. It is encouraging to see DiffCSP explore this representation in diffusion models. The authors do a in-depth comparison with CDVAE [1] which alternatively explored a multi-graph construction for handling periodicity.\n\n- DiffCSP studies joint diffusion over all three components of crystal structures: atom types, lattice matrix, and coordinates. The comparisons and benchmarks will be useful to the broader community.\n\n- The denoising model, a modified EGNN, to achieve the periodic E(3) symmetry has architectural choices that are well motivated. Each component is ablated to demonstrate their individual importances.\n\n- DiffCSP achieves state-of-the-art in crystal structure prediction and ab initio crystal generation.\n\n\n[1] https://arxiv.org/abs/2110.06197\n[2] https://arxiv.org/pdf/2102.09844.pdf", "weaknesses": "- The authors claim DiffCSP is the first generative model for crystal structure prediction (CSP). I am curious to know why a generative viewpoint of CSP is beneficial. Discussion of why diffusion is important for this task would help, i.e. a discussion about uncertainty in the problem and how it manifests.\n\n- The benefit of diffusing over the lattice could use more explanation. The ablation seems to show there is marginal benefit of diffusing the lattice. It is also not clear to me if gaussian diffusion is the right choice for the lattice vectors which can also be parameterized with 6 unique, rotation invariant parameters using the Niggli algorithm [1].\n\n- DiffCSP is E(3) and periodicity equivaraint however CDVAE is SE(3) and periodicity equivariant. Perhaps I missed it but I do not see discussion on the subtley between E(3) and SE(3). It would be important to note this difference and include it in the discussion.\n\n- A major consideration when performing joint diffusion is the noise schedules used for each diffusion process. A discussion of whether this was important practically would be helpful in the main text.\n\n \n\n[1] https://pubmed.ncbi.nlm.nih.gov/14691322/", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful to the reviewer for the detailed comments. We answer the reviewer’s questions as follows.\n\nFor weaknesses:\n\n> **W1: Why a generative viewpoint of CSP is beneficial.**\n\nThank you for your question. The generative viewpoint in CSP holds substantial benefits, particularly in addressing inherent uncertainties within the problem and capturing the underlying distribution of possible crystal structures. One of the fundamental challenges in CSP is the non-uniqueness of crystal structures, even when the composition is given and fixed. A classic example is carbon, which can form various allotropes, each with distinct structures (diamond, graphite, etc.). DiffCSP acknowledges this inherent diversity and leverages diffusion processes to explore the rich space of possible structures. \n\n> **W2: The benefit of diffusing over the lattice could use more explanation.**\n\nThank you for reminding. To address your concern, we provide a variant model named DiffCSP-L, which utilizes a Transformer-based predictor to directly predict the 6 invariant parameters from a given composition. Table R1 (General Response) displays the results on MP-20, highlighting the diffusion model's superior performance over the predictor-based model.  We conjecture that the distribution of crystals depends on both the lattice and the atom coordinates within each unit cell. By diffusing both, the generative model is able to simultaneously adjust their values, leading to more accurate generation. By contrast, the predictor-based variant conducts a two-stage generation paradigm, and the lattice is no longer adjustable during the generation of the atom coordinates. We will add the above explanation to the revised version.\n\n> **W3: SE(3) v.s. E(3) invariance.**\n\nGreat suggestion! We attempt to introduce the chirality into the denoising model by adapting the message-passing block in Eq. (6) into\n$$\n\\textbf{m}\\_{ij}^{(s)}=\\varphi\\_m(\\textbf{h}\\_{i}^{(s-1)}, \\textbf{h}\\_{j}^{(s-1)}, \\textbf{L}^\\top\\textbf{L}, \\text{sign}(|\\textbf{L}|), \\psi_{FT}(\\textbf{f}\\_{j} - \\textbf{f}\\_{i})),\n$$\nwhere sign$(|\\textbf{L}|)$ is the sign of the determinant of the lattice matrix, which is SO(3)-invariant, but breaks the reflection symmetry and hence is NOT E(3)-invariant. The results on MP-20 are listed in Table R1 (General Response). There is no essential difference between the performance of the SO(3) model and the original E(3) model, which is probably because the chirality is not so crucial in distinguishing different crystal structures for the datasets used in this paper. \n\n> **W4: The noise schedules used for each diffusion process.**\n\nThanks for your suggestion. We explore the noise schedulers from 3 perspectives and list the results in Table R2 and R3 (General Response). 1) For lattices, we originally use the cosine scheduler with $s=0.008$, and we change it into the linear and sigmoid schedulers with $\\beta_1=0.0001, \\beta_T=0.02$. We find that the linear scheduler yields comparable results, while the sigmoid scheduler hinders the performance. 2) For coordinates, we use the exponential scheduler with $\\sigma_{min}=0.005, \\sigma_{max}=0.5$, and we change $\\sigma_{max}$ into 0.1 and 1.0. Results show that the small-$\\sigma$ variant performs obviously worse, as only sufficiently large $\\sigma$ could approximate the prior uniform distribution. We visualize the PDF curves in Figure R3 for better understanding. 3) For types, we conduct similar experiments as lattices, and the results indicate that the original cosine scheduler performs better. In conclusion, we suggest applying the proposed noise schedulers, and we will provide the above details in the appendix.\n\nFor questions (Note that we have merged some repeated questions due to space limitation):\n\n> **Q1: How would one select the best sample if the ground truth was not available?**\n\nNice question! If the ground truth is unavailable, ranking the generated samples is indeed an important problem. To this end, we here design 3 kinds of confidence models to score each sample for ranking, named Energy Predictor(EP), Match Discriminator(MD), and Contrastive Scorer(CS). See General Response for more details.\n\nWe select the Top-1 result among the 20 candidates by each confidence model on MP-20 in Table R4 (General Response). It shows that MD and CS perform relatively better.  We agree that structure ranking is crucial and are willing to add the above discussion to the revised paper. \n\n> **Q2: What is the rationale behind gaussian diffusion for the lattice parameters?**\n\nThank you for this insightful question. The magnitudes of the lattice lengths are within [3.29, 5.71], [2.40, 16.67], [2.25, 46.75], and [0.98, 130.46] on Perov-5, Carbon-24, MP-20, and MPTS-52, respectively. In this sense, the magnitudes are practically bounded, and we apply gaussion diffusion for the lattice vectors since their values are real and continue. \n\n> **Q3: Could Table 1 include mean and standard deviation for RMSE?**\n\nYes, we did provide the error bars in Table 11 of Appendix H. We will further clarify this point in the main paper. \n\n> **Q4: How diverse are the crystals when multiple samples are drawn?**\n\nGood question! We further evaluate the diversity by yielding the CrystalNN fingerprint of each generated structure, calculating the L2-distances among all pairs in the 20 samples of each composition, collating the mean and max value of the distances, and finally averaging the results from all testing candidates. The results are listed in Table R5 (General Response). Clear diversity is observed for both CDVAE and our DiffCSP.\n\nFor typoes & discussions:\n\nThanks for your suggestions. We are glad to modify the typoes and provide more discussions on Riemann score-matching. The error bars were provided in Table 11 of Appendix H.\n\nFor limitations:\n\nThanks for your kind suggestions. We have explored the influence of SO(3) invariance (W3) and noise schedulers (W4). We will add these discussions to the revised paper."}, {"review_id": "CEXULB5pXl", "reviewer": "Reviewer_pv4b", "summary": "DiffCSP is proposed in this paper for the task of crystal structure prediction task, which generates the stable crystal structures of given compositions. The diffusion process of DIffCSP considers both lattice matrix and fractional coordinates of atoms in the unit cell. Experimental results demonstrate the effectiveness and efficiency of this method. ", "questions": "- Maybe a typo: In line 277, it is '20 structures of the lowest energy or' or of?\n\nSuggestions:\n- How about rename the Permutation Invariance, O(3) Invariance, etc, to Distribution Permutation Invariance, Distribution O(3) Invariance? The latter ones seem to be more accurate.", "limitations": "They listed the limitations in the main paper.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The proposed method shows good performances compared to previous methods in the task of Crystal Structure Prediction.\n- The problem setting is new to ML, and interesting.\n- Writing is clear with informative figures, formulas, and proofs.\n- Experiments are comprehensive with additional adjustments to change the current framework to the generative model setting and compare with previous generative models. Ablation studies are provided.", "weaknesses": "- Limited novelty: the proposed method is straightforward by further considering lattice matrix diffusion; the novelty is somehow limited. The way they diffuse lattice matrix is similar to previous works for molecule conformer generation. And the way they diffuse fractional atom coordinates is similar to CDVAE.\n- Other distribution invariant constraints are not considered: (1) The distribution of crystal structures may need to be SO(3) invariant instead of O(3) invariant, given the fact that chiral crystals (similar cases in molecules) are different from each other. (2) Additionally, given the diffusion process proposed, it may not have the ability to deal with the distribution periodic invariance in the following case: given a lattice matrix L = [l_1, l_2, l_3] \\in R^{3*3}, the crystal structure remain the same when you change it to L = [l_1, l_2 + l_3, l_3 - l_2], and many other cases when you choose a different set of periodic patterns to describe the same infinite crystal structure. \n- Experiments about CSP: adapting CDVAE to conditional generative model and comparing with CSP model is reasonable but not that fair. But they adjust their model to the generative setting and compare with CDVAE, so it may not be a big issue.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful to the reviewer for the constructive comments. We answer the reviewer’s questions as follows.\n\n\n> **W1: Limited novelty.**\n\nThere could be some misunderstandings here. We admit that the way we diffuse lattice matrix is similar to previous works for molecule conformer generation, but we disagree that the way we diffuse fractional atom coordinates is similar to CDVAE. Different from CDVAE that uses a score-matching method on Cartesian coordinates, we utilize fractional coordinates in our diffusion process. As also raised by other reviewers, the usage of fractional coordinates eliminates the need to construct multi-edge graphs (which is used in CDVAE) and naturally guarantees periodic invariance. Given that the domain of fractional coordinates forms a quotient space, we leverage the Score-Matching (SM) based framework along with Wrapped Normal (WN) distribution (Line 182-203).\n\nMoreover, in contrast to CDVAE which focuses solely on coordinate diffusion within a fixed region bounded by the predicted lattice, our model performs joint diffusion on both lattice structures and atom coordinates, and these two processes are coupled and promoted with each other (Figure 2). \n\nBesides, in Section 4.3, we design a novel denoising model to fulfill the joint diffusion process. Firstly, Fourier Transformation of the relative fractional coordinate is adopted to better encode periodicity (Eq.6). Secondly, the output layer is specifically designed to jointly predict the denoising terms of the entire lattice (Eq.9) and each fractional coordinate (Eq.10). Lastly, the proposed denoising model is novel compared to traditional multi-graph representation, and it is proved to be periodic E(3) equivariance (Proposition 3).\n\nOverall, we highlight the novelty of our paper in terms of the three aspects: the diffusion of fractional coordinates via an SM-based framework along with WN distribution, the joint diffusing process, and the novel denoising model. \n\n> **W2.1: Other distribution invariant constraints are not considered.**\n\nThank you for your valuable suggestions.  We appreciate your insight into the importance of considering SO(3) invariance. In response to this concern, we conduct an ablation study by additionally introducing the chirality into the denoising model. Specifically, we adapt the message-passing block in Eq. (6) into\n$$\n\\textbf{m}\\_{ij}^{(s)}=\\varphi\\_m(\\textbf{h}\\_{i}^{(s-1)}, \\textbf{h}\\_{j}^{(s-1)}, \\textbf{L}^\\top\\textbf{L}, \\text{sign}(|\\textbf{L}|), \\psi\\_{FT}(\\textbf{f}\\_{j} - \\textbf{f}\\_{i})),\n$$\nwhere sign$(|\\textbf{L}|)$ is the sign of the determinant of the lattice matrix, which is SO(3)-invariant, but breaks the reflection symmetry and hence is NOT E(3)-invariant. The results of the adapted model on MP-20 are listed in the following table. There is no essential difference between the performance of the SO(3) model and the original E(3) model, which is probably because the chirality is not so crucial in distinguishing different crystal structures for the datasets used in this paper. However, we are willing to add the discussion here to the revised paper. \n\n|                | Match Rate (%)$\\uparrow$ |  RMSD$\\downarrow$  |\n| ---------------| :------------: | :----: |\n| DiffCSP        |     51.49      | 0.0631 |\n| DiffCSP-SO(3)  |     49.68      | 0.0637 |\n\n> **W2.2: Additionally, given the diffusion process proposed, it may not have the ability to deal with the distribution periodic invariance in the following case...**\n\nThank you for this nice point. Actually, this invariance is caused by an infinite number of selections for the lattice matrix, to reflect the periodicity of crystal structures. To avoid the exploration of this kind of invariance, we follow the setting from CDVAE and apply the Niggli reduction [46] on the primitive cell as a canonical representation of the lattice vectors, which was mentioned in the footnote of Page 4. We will further clarify this point in the paper. \n\n\n> **W3: Experiments about CSP.**\n\nWe understand your concern. Hence, we conducted two kinds of protocols in our experiments, one adapting CDVAE to the settings of the CSP task for fair comparisons with our method (Table 1), the other one extending our method to the setting of CDVAE for ab initio generation (Table 4). Indeed, we discussed the relationship and distinction among various crystal-related tasks in Appendix F.1 and Figure 7. \n\n\n\nWe will also fix the raised typos, and further emphasize that the notions of the Permutation Invariance, O(3) Invariance, etc., are defined in terms of distribution. "}]

Author rebuttal:
We sincerely thank all reviewers and ACs for their time and efforts on reviewing the paper. We are glad that the reviewers recognized the contributions of our paper, which we briefly summarize as follows.

- **Novelty.** "The paper presents novel and shining ideas."(1HBn)"The problem setting is new to ML, and interesting."(pv4b)
- **Presentation.** "The paper is well-organized and well-written."(eNoa)"Writing is clear with informative figures, formulas, and proofs."(pv4b)
- **Experiments.** "Experiments are solid and performance is impressive."(1HBn)"The empirical result is strong, and the baseline methods are appropriate for the CSP task."(eNoa)"The comparisons and benchmarks will be useful to the broader community."(pLZt)"Experiments are comprehensive with additional adjustments to change the current framework to the generative model setting and compare with previous generative models."(pv4b)

We also appreciate the reviewers for their insightful comments, and provide additional visualizations and experment results in the supplementary PDF file for more details. We summarize the extra contents as follows.

- **Figure R1 (to Reviewer 1HBn)** illustrates the connection between the translation invariance in previous works and the periodic translation invariance defined in our paper. The periodic translation invariance already considers the translation invariance.
- **Figure R2 (to Reviewer eNoa)** visualizes 8 pairs of original and optimized structures, indicating that DiffCSP is capable to obtain structures with lower energies from existing ones.
- **Figure R3 (to Reviewer pLZt)** compares the PDF curves of wrapped normal distributions with period 1 and diffenent variances. It can be find that $\sigma=0.5$ is practically large enough to approximate the prior uniform distribution.
- **Table R1 (to Reviewer pLZt and pv4b)** conducts two extra ablation studies. The first explores the necessity of lattice diffusion against a predictor-based counterpart, and the second compares the performances of SO(3) and O(3) equivariant backbones.
- **Table R2 and R3 (to Reviewer pLZt)** evaluate the influence of noise schedulers of the diffusion processes on lattices, coordinates and atom types.
- **Table R4 (to Reviewer pLZt)** compares different confidence models for ranking candidate structures. Specifically, we design 3 kinds of confidence models:

1) **Energy Predictor(EP).** We directly train a predictor of the energy labels, and apply the opposite of the predicted energy as the confidence score.
2) **Match Discriminator(MD).** Inspired by Diffdock[1], we first generate 5 samples for each composition in the train/validation set via DiffCSP, and calculate their RMSDs with the ground truth. Finally, we train a binary classifier to predict whether the sample matches the ground truth with RMSD less than a threshold $d$. The predicted probability is used as the score.
3) **Contrastive Scorer(CS).** Inspired by CLIP[2], we train a contrastive model between a 1D and a 3D model to align the corresponding compositions and structures, and the inner product of the 1D and 3D model is used as the score.

- **Table R5 (to Reviewer pLZt)** compares the diversity of generated structures of CDVAE and DiffCSP.

[1] Corso, Gabriele, et al. "DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking." International Conference on Learning Representations (ICLR 2023). 2023.
[2] Radford, Alec, et al. "Learning transferable visual models from natural language supervision." International conference on machine learning. PMLR, 2021. similar tothan other schedulers

Meta-review: This paper studies crystal structure prediction which is an important application of AI for scientific discovery. The newly proposed diffusion model is equipped with nice ideas (fractional coordinates and consideration symmetries) and the experimental evaluation is concrete. This paper is likely to influence more researchers to study the problem of neural crystal structure prediction. 

The reviewers are all positive about the paper, appreciating the importance of the problem, empirical setup, and novelty of the fractional coordinates. I vote for acceptance of this work.
Decision: Accept (poster)