Title: Crystal Structure Prediction by Joint Equivariant Diffusion
Venue: NeurIPS 2023
Accepted: yes
Keywords: crystal structure prediction

equivariant graph neural networks

diffusion generative models
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

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 significantly outperforms existing CSP methods, with a much lower computation cost in contrast to DFT-based methods. Moreover, the superiority of DiffCSP is also observed when it is extended for ab initio crystal generation. Code is available at https://github.com/jiaor17/DiffCSP .

1 Introduction

Crystal Structure Prediction (CSP), which returns the stable 3D structure of a compound based solely on its composition, has been a goal in physical sciences since the 1950s [1]. As crystals are the foundation of various materials, estimating their structures in 3D space determines the physical and chemical properties that greatly influence the application to various academic and industrial sciences, such as the design of drugs, batteries, and catalysts [2]. Conventional methods towards CSP mostly apply the Density Functional Theory (DFT) [3] to compute the energy at each iteration, guided by optimization algorithms (such as random search [4], Bayesian optimization [5], etc.) to iteratively search for the stable state corresponding to the local minima of the energy surface [6].

The DFT-based approaches are computationally-intensive. Recent attention has been paid to deep generative models, which directly learn the distribution from the training data consisting of stable structures [7, 8]. More recently, diffusion models, a special kind of deep generative models are employed for crystal generation [9], encouraged by their better physical interpretability and enhanced performance than other generative models. Intuitively, by conducting diffusion on stable structures,

the denoising process in diffusion models acts like a force field that drives the atom coordinates towards the energy local minimum and thus is able to increase stability. Indeed, the success of diffusion models is observed in broad scientific domains, including molecular conformation generation [10], protein structure prediction [11] and protein docking [12].

However, designing diffusion models for CSP is challenging. From the perspective of physics, any E(3) transformation, including translation, rotation, and reflection, of the crystal coordinates does not change the physical law and thus keeps the crystal distribution invariant. In other words, the generation process we design should yield E(3) invariant samples. Moreover, in contrast to other types of structures such as small molecules [13] and proteins [14], CSP exhibits unique challenges, mainly incurred by the periodicity of the atom arrangement in crystals. Figure 1 displays a crystal where the atoms in a unit cell are repeated infinitely in space. We identify such unique symmetry, jointly consisting of E(3) invariance and periodicity, as periodic E(3) invariance in this paper. Generating such type of structures requires not only modeling the distribution of the atom coordinates within every cell, but also inferring how their bases ( a.k.a. lattice vectors) are placed in 3D space. Interestingly, as we will show in § 4.1, such view offers a natural disentanglement for fulfilling the periodic E(3) invariance by separately enforcing constraints on fractional coordinates and lattice vectors, which permits a feasible implementation to encode the crystal symmetry.

In this work, we introduce DiffCSP, an equivariant diffusion method to address CSP. Considering the specificity of the crystal geometry, our DiffCSP jointly generates the lattice vectors and the fractional coordinates of all atoms, by employing a proposed denoising model that is theoretically proved to generate periodic-E(3)-invariant samples. A preferable characteristic of DiffCSP is that it leverages the fractional coordinate system (defined in § 3) other than the Cartesian system used in previous methods to represent crystals [9, 15], which encodes periodicity intrinsically. In particular, the fractional representation not only allows us to consider Wrapped Normal (WN) distribution [16] to better model the periodicity, but also facilitates the design of the denoising model via the Fourier transformation, compared to the traditional multi-graph encoder in crystal modeling [15].

CDVAE [9] is closely related with our paper. It adopts an equivariant Variational Auto-Encoder (VAE) based framework to learn the data distribution and then generates crystals in a score-matchingbased diffusion process. However, CDVAE focuses mainly on ab initio crystal generation where the composition is also randomly sampled, which is distinct from the CSP task in this paper. Moreover, while CDVAE first predicts the lattice and then updates the coordinates with the fixed lattice, we jointly update the lattice and coordinates to better model the crystal geometry. Besides, CDVAE represents crystals by Cartesian coordinates upon multi-graph modeling, whereas our DiffCSP applies fractional coordinates without multi-graph modeling as mentioned above.

To sum up, our contributions are as follows:

To the best of our knowledge, we are the first to apply equivariant diffusion-based methods to address CSP. The proposed DiffCSP is more insightful than current learning-based approaches as the periodic E(3) invariance has been delicately considered.

DiffCSP conducts joint diffusion on lattices and fractional coordinates, which is capable of capturing the crystal geometry as a whole. Besides, the usage of fractional coordinates in place of Cartesian coordinates used in previous methods ( e.g. CDVAE [9]) remarkably promotes the diffusion and the generation process of atom positions.

We verify the efficacy of DiffCSP on the CSP task against learning-based and DFT-based methods, and sufficiently ablate each proposed component in DiffCSP. We further extend DiffCSP into ab initio generation and show its effectiveness against related methods.

2 Related Works

Crystal Structure Prediction Traditional computation methods [4, 5, 17, 18] combine DFT [3] with optimization algorithms to search for local minima in the potential energy surface. However, DFT is computationally intensive, making it dilemmatic to balance efficiency and accuracy. With the improvement of crystal databases, machine-learning methods are applied as alternative energy predictors to DFT followed by optimization steps [19, 20, 21]. Apart from the predict-optimize paradigm, another line of approaches directly learns stable structures from data by deep generative models, which represents crystals by 3D voxels [7, 22, 23], distance matrices [8, 24, 25] or 3D

coordinates [26, 27, 28]. Unfortunately, these methods are unaware of the full symmetries of the crystal structure. CDV AE [9] has taken the required symmetries into account. However, as mentioned above, the initial version of CDV AE is for different task and utilizes different generation process.

Equivariant Graph Neural Networks Geometrically equivariant Graph Neural Networks (GNNs) that ensure E(3) symmetry are powerful tools to represent physical objects [29, 30, 31, 32, 33], and have showcased the superiority in modeling 3D structures [34, 35]. To further model the periodic materials, Xie and Grossman [15] propose the multi-graph edge construction to capture the periodicity by connecting the edges between adjacent lattices. Yan et al. [36] further introduce periodic pattern encoding into a Transformer-based backbone. In this work, we achieve the periodic invariance by introducing the Fourier transform on fractional coordinates.

Diffusion Generative Models Motivated by the non-equilibrium thermodynamics [37], diffusion models connect the data distribution with the prior distribution via forward and backward Markov chains [38], and have made remarkable progress in the field of image generation [39, 40]. Equipped with equivariant GNNs, diffusion models are capable of generating samples from the invariant distribution, which is desirable in conformation generation [10, 13], ab initio molecule design [41], protein generation [42], and so on. Recent works extend the diffusion models onto Riemann manifolds [43, 44], and enable the generation of periodic features like torsion angles [12, 16].

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 1. A unit cell can be defined by a triplet M = ( A , X , L ) , where A = [ a 1 , a 2 , ..., a N ] ∈ R h × N denotes the list of the one-hot representations of atom types, X = [ x 1 , x 2 , ..., x N ] ∈ R 3 × N consists of Cartesian coordinates of the atoms, and L = [ l 1 , l 2 , l 3 ] ∈ R 3 × 3 represents the lattice matrix containing three basic vectors to describe the periodicity of the crystal. The infinite periodic crystal structure is represented by

where the j -th element of the integral vector k denotes the integral 3D translation in units of l j .

Fractional coordinate system The Cartesian coordinate system X leverages three standard orthogonal bases as the coordinate axes. In crystallography, the fractional coordinate system is usually applied to reflect the periodicity of the crystal structure [26, 27, 28, 45], which utilizes the lattices ( l 1 , l 2 , l 3 ) as the bases. In this way, a point represented by the fractional coordinate vector f = [ f 1 , f 2 , f 3 ] ⊤ ∈ [0 , 1) 3 corresponds to the Cartesian vector x = ∑ 3 i =1 f i l i . This paper employs the fractional coordinate system, and denotes the crystal by M = ( A , F , L ) , where the fractional coordinates of all atoms in a cell compose the matrix F ∈ [0 , 1) 3 × N .

Task definition CSP predicts for each unit cell the lattice matrix L and the fractional matrix F given its chemical composition A , namely, learning the conditional distribution p ( 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 to address CSP, this task encounters particular challenges, including constraints arising from symmetries of crystal structure distribution. Here, we consider the three types of symmetries in the distribution of p ( L , F | A ) : permutation invariance, O (3) invariance, and periodic translation invariance. Their detailed definitions are provided as follows.

Definition 1 (Permutation Invariance) . For any permutation P ∈ S N , p ( L , F | A ) = p ( L , FP | AP ) , i.e., changing the order of atoms will not change the distribution.

Composition

𝑨

Initialize

Invariant Prior

𝑇 - 𝑡

steps

Joint Equivariant Markov Process

Periodic E(3)

Equivariant

Denoising

Model

𝜙

t

-

1

steps

Sample

Forward Diffusion Process

. . .

. . .

. . .

. . .

EL

LT ~ N(0, I)

Lt

=

Mt,t

MT

Mt

Mt-1

Mo

↑

EF

w(-)

Fr ~U(0,1)

Ft

w(-)

Ft-1

w(Ft1)

Fo

EF

Definition 2 (O(3) Invariance) . For any orthogonal transformation Q ∈ R 3 × 3 satisfying Q ⊤ Q = I , 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 ∈ R 3 × 1 , p ( L , w ( F + t 1 ⊤ ) | A ) = p ( L , F | A ) , where the function w ( F ) = F -⌊ F ⌋ ∈ [0 , 1) 3 × N returns the fractional part of each element in F , and 1 ∈ R 3 × 1 is a vector with all elements set to one. It explains that any periodic translation of F will not change the distribution 2 .

The permutation invariance is tractably encapsulated by using GNNs as the backbone for generation [47]. We mainly focus on the other two kinds of invariance (see Figure 1), since GNNs are our default choices. For simplicity, we compactly term the O (3) invariance and periodic translation invariance as periodic E(3) invariance henceforth. Previous approaches ( e.g. [9, 36]) adopt Cartesian coordinates X other than fractional coordinates F , hence their derived forms of the symmetry are different. Particularly, in Definition 2, the orthogonal transformation additionally acts on X ; in Definition 3, the periodic translation w ( F + t 1 ⊤ ) becomes the translation along the lattice bases X + Lt 1 ⊤ ; besides, X should also maintain E(3) translation invariance, that is p ( L , X + t 1 ⊤ | A ) = p ( L , X | A ) . With the help of the fractional system, the periodic E(3) invariance

f

Q

Ql2

f2

(b)

l

f

f2

l2

(a)

t

w(fi + t)

w(f2+ t)

(d)

l

f

f2

l2

(c)

is made tractable by fulfilling O(3) invariance w.r.t. the orthogonal transformations on L and periodic translation invariance w.r.t. the periodic translations on F , respectively. In this way, such approach, as detailed in the next section, facilitates the application of diffusion methods to the CSP task.

4.2 Joint Equivariant Diffusion

Our method DiffCSP addresses CSP by simultaneously diffusing the lattice L and the fractional coordinate matrix F . Given the atom composition A , M t denotes the intermediate state of L and F at time step t (0 ≤ t ≤ T ) . DiffCSP defines two Markov processes: the forward diffusion process gradually adds noise to M 0 , and the backward generation process iteratively samples from the prior distribution M T to recover the origin data M 0 .

Joining the statements in § 4.1, the recovered distribution from M T should meet periodic E(3) invariance. Such requirement is satisfied if the prior distribution p ( M T ) is invariant and the Markov transition p ( M t -1 |M t ) is equivariant, according to the diffusion-based generation literature [10]. Here, an equivariant transition is specified as p ( g ·M t -1 | g ·M t ) = p ( M t -1 | M t ) where g · M refers to any orthogonal/translational transformation g acts on M in the way presented in Definitions 2-3. We separately explain the derivation details of L and F below. The detailed flowcharts are summarized in Algorithms 1 and 2 in Appendix B.3.

Diffusion on L Given that L is a continuous variable, we exploit Denoising Diffusion Probabilistic Model (DDPM) [38] to accomplish the generation. We define the forward process that progressively diffuses L 0 towards the Normal prior p ( L T ) = N (0 , I ) by q ( L t | L t -1 ) which can be devised as the probability conditional on the initial distribution:

where β t ∈ (0 , 1) controls the variance, and ¯ α t = ∏ t s =1 α t = ∏ t s =1 (1 -β t ) is valued in accordance to the cosine scheduler [48].

The backward generation process is given by:

where µ ( M t ) = 1 √ α t ( L t -β t √ 1 -¯ α t ˆ ϵ L ( M t , t ) ) , σ 2 ( M t )= β t 1 -¯ α t -1 1 -¯ α t . The denoising term ˆ ϵ L ( M t , t ) ∈ R 3 × 3 is predicted by the model ϕ ( L t , F t , A , t ) .

As the prior distribution p ( L T ) = N (0 , I ) is already O(3)-invariant, we require the generation process in Eq. (3) to be O(3)-equivariant, which is formally stated below.

Proposition 1. The marginal distribution p ( L 0 ) by Eq. (3) is O(3)-invariant if ˆ ϵ L ( M t , t ) is O(3)equivariant, namely ˆ ϵ L ( QL t , F t , A , t ) = Q ˆ ϵ L ( L t , F t , A , t ) , ∀ Q ⊤ Q = I .

To train the denoising model ϕ , we first sample ϵ L ∼ N (0 , I ) and reparameterize L t = √ ¯ α t L 0 + √ 1 -¯ α t ϵ L based on Eq. (2). The training objective is defined as the ℓ 2 loss between ϵ L and ˆ ϵ L :

Diffusion on F The domain of fractional coordinates [0 , 1) 3 × N forms a quotient space R 3 × N / Z 3 × N 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 (SM) based framework [49, 50] along with Wrapped Normal (WN) distribution [43] to fit the specificity here. Note that WN distribution has been explored in generative models, such as molecular conformation generation [16].

During the forward process, we first sample each column of ϵ F ∈ R 3 × N from N (0 , I ) , and then acquire F t = w ( F 0 + σ t ϵ F ) where the truncation function w ( · ) is already defined in Definition 3. This truncated sampling implies the WN transition:

Basically, this process ensures the probability distribution over [ z, z +1) 3 × N for any integer z to be the same to keep the crystal periodicity. Here, the noise scale σ t obeys the exponential scheduler: σ 0 = 0 and σ t = σ 1 ( σ T σ 1 ) t -1 T -1 , if t > 0 . Desirably, q ( F t | F 0 ) is periodic translation equivariant, and approaches a uniform distribution U (0 , 1) if σ T is sufficiently large.

For the backward process, we first initialize F T from the uniform distribution U (0 , 1) , which is periodic translation invariant. With the denoising term ˆ ϵ F predicted by ϕ ( L t , F t , A , t ) , we combine the ancestral predictor [38, 50] with the Langevin corrector [49] to sample F 0 . We immediately have:

Proposition 2. The marginal distribution p ( F 0 ) is periodic translation invariant if ˆ ϵ F ( M t , t ) is periodic translation invariant, namely ˆ ϵ F ( L t , F t , A , t ) = ˆ ϵ F ( L t , w ( F t + t 1 ⊤ ) , A , t ) , ∀ t ∈ R 3 .

The training objective for score matching is:

where λ t = E -1 F t [ ∥∇ F t log q ( F t | F 0 ) ∥ 2 2 ] 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 is 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 a DDPM-based approach. We provide more details in Appendix G.

4.3 The Architecture of the Denoising Model

This subsection designs the denoising model ϕ ( L , F , A , t ) that outputs ˆ ϵ L and ˆ ϵ F satisfying the properties stated in Proposition 1 and 2.

Let H ( s ) = [ h ( s ) 1 , · · · , h ( s ) N ] denote the node representations of the s -th layer. The input feature is given by h (0) i = ρ ( f atom ( a i ) , f pos ( t )) , where f atom and f pos are the atomic embedding and sinusoidal positional encoding [38, 51], respectively; ρ is a multi-layer perception (MLP).

Built upon EGNN [32], the s -th layer message-passing is unfolded as follows:

Here φ m and φ h are MLPs. The function ψ FT : ( -1 , 1) 3 → [ -1 , 1] 3 × K is Fourier Transformation of the relative fractional coordinate f j -f i . Specifically, suppose the input to be f = [ f 1 , f 2 , f 3 ] ⊤ , then the c -th row and k -th column of the output is calculated by

ψ FT ( f )[ c, k ] = sin(2 πmf c ) , if k = 2 m (even); and ψ FT ( f )[ c, k ] = cos(2 πmf c ) , if k = 2 m + 1 (odd). ψ FT extracts various frequencies of all relative fractional distances that are helpful for crystal structure modeling, and more importantly, ψ FT is periodic translation invariant, namely, ψ FT ( w ( f j + t ) -w ( f i + t )) = ψ FT ( f j -f i ) for any translation t , which is proved in Appendix A.3.

After S layers of message passing conducted on the fully connected graph, the lattice noise ˆ ϵ L is acquired by a linear combination of L , with the weights given by the final layer:

where φ L is an MLP with output shape as 3 × 3 . The fractional coordinate score ˆ ϵ F is output by:

where ˆ ϵ F [: , i ] defines the i -th column of ˆ ϵ F , and φ F is an MLP on the final representation.

Weapply the inner product term L ⊤ L in Eq. (6) to achieve O(3)-invariance, as ( QL ) ⊤ ( QL ) = L ⊤ L for any orthogonal matrix Q ∈ R 3 × 3 . This leads to the O(3)-invariance of φ L in Eq. (10), and we further left-multiply L with φ L to ensure the O(3)-equivariance of ˆ ϵ L . Therefore, the above formulation of the denoising model ϕ ( L , F , A , t ) ensures the following property.

Proposition 3. The score ˆ ϵ L by Eq. (9) is O(3)-equivariant, and the score ˆ ϵ F from Eq. (10) is periodic translation invariant. Hence, the generated distribution by DiffCSP is periodic E(3) invariant.

Comparison with multi-graph representation Previous methods [9, 15, 29, 52] 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 ψ FT in our message passing is already periodic translation invariant. We will ablate the benefit in Table 3.

# of

Perov-5

Carbon-24

MP-20

MPTS-52

samples

Match rate

↑

RMSE

↓

Match rate

↑

RMSE

↓

Match rate

↑

RMSE

↓

Match rate

↑

RMSE

↓

RS [21]

20

29.22

0.2924

14.63

0.4041

8.73

0.2501

2.05

0.3329

5,000

36.56

0.0886

14.63

0.4041

11.49

0.2822

2.68

0.3444

BO [21]

20

21.03

0.2830

0.44

0.3653

8.11

0.2402

2.05

0.3024

5,000

55.09

0.2037

12.17

0.4089

12.68

0.2816

6.69

0.3444

PSO [21]

20

20.90

0.0836

6.40

0.4204

4.05

0.1567

1.06

0.2339

5,000

21.88

0.0844

6.50

0.4211

4.35

0.1670

1.09

0.2390

P-cG-

1

48.22

0.4179

17.29

0.3846

15.39

0.3762

3.67

0.4115

SchNet [53]

20

97.94

0.3463

55.91

0.3551

32.64

0.3018

12.96

0.3942

CDVAE [9]

1

45.31

0.1138

17.09

0.2969

33.90

0.1045

5.34

0.2106

20

88.51

0.0464

88.37

0.2286

66.95

0.1026

20.79

0.2085

DiffCSP

1

52.02

0.0760

17.54

0.2759

51.49

0.0631

12.19

0.1786

20

98.60

0.0128

88.47

0.2192

77.93

0.0492

34.02

0.1749

Ground Truth

DiffCSP

CDVAE

P-cG-SchNet

Perov-5

MP-20

MPTS-52

Carbon-24

5 Experiments

In this section, we evaluate the efficacy of DiffCSP on a diverse range of tasks, by showing that it can generate high-quality structures of different crystals in § 5.1, with lower time cost comparing with DFT-based optimization method in § 5.2. Ablations in § 5.3 exhibit the necessity of each designed component. We further showcase the capability of DiffCSP 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 a unit cell. Carbon-24 [56] includes 10,153 carbon materials with 6 ∼ 24 atoms in a cell. MP-20 [57] selects 45,231 stable inorganic materials from Material Projects [57], which includes the majority of experimentally-generated materials with at most 20 atoms in a 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 earliest 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,000/8,096 for training/validation/testing in chronological order.

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. [21], 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 is based on deep generative models. We follow the modification in Xie et al. [9], and leverage cG-SchNet [53] that utilizes SchNet [29] as the backbone and additionally

considers the ground-truth lattice initialization for encoding periodicity, yielding a final model named P-cG-SchNet . Another baseline CDVAE [9] is a V AE-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 dynamics [49]. To adapt CDVAE 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 metrics Following the common practice [9], 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, ltol=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 3 √ V/N 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 DiffCSP, we let k = 1 and k = 20 for evaluation. We provide more details in Appendix B, C.1 and I.

Results Table 1 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 to 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.DiffCSP clearly makes the best predictions.

5.2 Comparison with DFT-based Methods

We further select 10 binary and 5 ternary compounds in MP-20 testing set and compare our model with USPEX [59], a DFT-based software equipped with the evolutionary algorithm to search for stable structures.

For our method, we sample 20 candidates for each compound following the setting in Table 1. 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 10. 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.

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 . Particularly, L → F applies two networks to learn the reverse processes p θ 1 ( L 0: T -1 | A , F T , L T ) and p θ 2 ( F 0: T -1 | A , F T , L 0 ) . During inference, we first sample L T , F T from their prior distributions, acquiring L 0 via p θ 1 , and then F 0 by p θ 2 based on L 0 . F → L is similarly executed but with the generation order of L 0 and F 0 exchanged. Results indicate that L → F performs better than the F → L , but both are inferior to the joint update in DiffCSP, 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 O(3) invariance when generating L , which is ensured by the inner product L ⊤ L in Eq. (6).

When we replace it with L and change the final output as ˆ ϵ L = φ L ( 1 N ∑ N i =1 h ( S ) i ) in Eq. (9) to break the equivariance, the model suffers from extreme performance detriment. Only 1.66% structures are successfully matched, which obviously implies the importance of incorporating O(3) equivariance. Furthermore, we introduce the chirality into the denoising model by adding sign ( | L | ) , the sign of the determinant of the lattice matrix, as an additional input in Eq. 6. The adapted model is SO (3) -invariant, but breaks the reflection symmetry and hence is NOT O (3) -invariant. There is no essential performance change, indicating the chirality is not quite crucial in distinguishing different crystal structures for the datasets used in this paper. 3. We further assess the importance of periodic translation invariance from two perspectives. For the generation process, we 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 t |M 0 ) = N ( F t | √ ¯ α t F 0 , (1 -¯ α t ) I ) 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(FT) in Eq. (6) to capture periodicity. To investigate its effect, we replace ψ FT ( f j -f i ) with f j -f i , 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. 4. We further change the fully-connected graph into the multi-graph approach adopted in Xie and Grossman [15]. The multi-graph 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 in Appendix E.

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 DiffCSP against five generative methods following [9]: FTCP [28], Cond-DFC-VAE [7], G-SchNet [60] and its periodic variant P-G-SchNet , and the orginal version of CDVAE [9]. Specifically for our DiffCSP, we 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. [41], 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 4 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 4. 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 of the model in real applications.

Conclusion In this work, we present DiffCSP, a diffusion-based learning framework for crystal structure prediction, which is particularly curated with the vital symmetries existing in crystals. The diffusion is highly flexible by jointly optimizing the lattice and fractional coordinates, where

the intermediate distributions are guaranteed to be invariant under necessary transformations. We demonstrate the efficacy of our approach on a wide range of crystal datasets, verifying the strong applicability of DiffCSP towards predicting high-quality crystal structures.

References

Gautam R Desiraju. Cryptic crystallography. Nature materials , 1(2):77-79, 2002.

Keith T Butler, Daniel W Davies, Hugh Cartwright, Olexandr Isayev, and Aron Walsh. Machine learning for molecular and materials science. Nature , 559(7715):547-555, 2018.

Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical review , 140(4A):A1133, 1965.

Chris J Pickard and RJ Needs. Ab initio random structure searching. Journal of Physics: Condensed Matter , 23(5):053201, 2011.

Tomoki Yamashita, Nobuya Sato, Hiori Kino, Takashi Miyake, Koji Tsuda, and Tamio Oguchi. Crystal structure prediction accelerated by bayesian optimization. Physical Review Materials , 2 (1):013803, 2018.

Artem R Oganov, Chris J Pickard, Qiang Zhu, and Richard J Needs. Structure prediction drives materials discovery. Nature Reviews Materials , 4(5):331-348, 2019.

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-4535, 2020.

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):455902, 2021.

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.

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.

Brian L. Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi S. Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motifscaffolding problem. In The Eleventh International Conference on Learning Representations , 2023.

Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi S. Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. In The Eleventh International Conference on Learning Representations , 2023.

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.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, 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 W Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with AlphaFold. Nature , 596(7873):583-589, 2021. doi: 10.1038/s41586-021-03819-2.

Tian Xie and Jeffrey C. Grossman. Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties. Phys. Rev. Lett. , 120:145301, Apr 2018. doi: 10.1103/PhysRevLett.120.145301.

Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. arXiv preprint arXiv:2206.01729 , 2022.

Yanchao Wang, Jian Lv, Li Zhu, and Yanming Ma. Crystal structure prediction via particleswarm optimization. Physical Review B , 82(9):094116, 2010.

Yunwei Zhang, Hui Wang, Yanchao Wang, Lijun Zhang, and Yanming Ma. Computer-assisted inverse design of inorganic electrides. Physical Review X , 7(1):011017, 2017.

TL Jacobsen, MS Jørgensen, and B Hammer. On-the-fly machine learning of atomic potential in density functional theory structure optimization. Physical review letters , 120(2):026102, 2018.

Evgeny V Podryabinkin, Evgeny V Tikhonov, Alexander V Shapeev, and Artem R Oganov. Accelerating crystal structure prediction by machine-learning interatomic potentials with active learning. Physical Review B , 99(6):064114, 2019.

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.

Jordan Hoffmann, 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:1909.00949 , 2019.

Juhwan Noh, Jaehoon Kim, Helge S Stein, Benjamin Sanchez-Lengeling, John M Gregoire, Alan Aspuru-Guzik, and Yousung Jung. Inverse design of solid-state materials via a continuous representation. Matter , 1(5):1370-1384, 2019.

Jianjun Hu, Wenhui Yang, and Edirisuriya M Dilanga Siriwardane. Distance matrix-based crystal structure prediction using evolutionary algorithms. The Journal of Physical Chemistry A , 124(51):10909-10919, 2020.

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. CrystEngComm , 23(8):1765-1776, 2021.

Asma Nouira, Nataliya Sokolovska, and Jean-Claude Crivello. Crystalgan: learning to discover crystallographic structures with generative adversarial networks. arXiv preprint arXiv:1810.11203 , 2018.

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-1420, 2020.

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/10.1016/j.matt.2021. 11.032.

Kristof T Schütt, 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.

Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219 , 2018.

Fabian Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. Se(3)-transformers: 3d rototranslation equivariant attention networks. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International Conference on Machine Learning , pages 9323-9332. PMLR, 2021.

Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations , 2021.

Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane 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.1021/acscatal.0c04525.

Richard Tran, Janice Lan, Muhammed Shuaibi, Siddharth Goyal, Brandon M Wood, Abhishek 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.08917 , 2022.

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.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models, 2021.

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

Emiel Hoogeboom, Vıctor 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.

Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigenspecific 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.

Valentin De Bortoli, Emile Mathieu, Michael John Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modelling. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Chin-Wei Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron Courville. Riemannian diffusion models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

Detlef WM Hofmann and Joannis Apostolakis. Crystal structure prediction by data mining. Journal of Molecular Structure , 647(1-3):17-39, 2003.

Ralf W 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.

Thomas N Kipf and Max Welling. Variational graph auto-encoders. arXiv preprint arXiv:1611.07308 , 2016.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems , 33:12438-12448, 2020.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

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-3572, 2019.

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.

Ivano E Castelli, David D Landis, Kristian S Thygesen, Søren Dahl, Ib Chorkendorff, Thomas F 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): 9034-9043, 2012.

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):5814-5819, 2012.

Chris J. Pickard. Airss data for carbon at 10gpa and the c+n+h+o system at 1gpa, 2020.

Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, et al. Commentary: The materials project: A materials genome approach to accelerating materials innovation. APL materials , 1(1):011002, 2013.

Shyue Ping Ong, William Davidson Richards, Anubhav Jain, Geoffroy Hautier, Michael Kocher, Shreyas Cholia, Dan Gunter, Vincent L Chevrier, Kristin A Persson, and Gerbrand Ceder. Python materials genomics (pymatgen): A robust, open-source python library for materials analysis. Computational Materials Science , 68:314-319, 2013.

Colin W Glass, Artem R Oganov, and Nikolaus Hansen. Uspex-evolutionary crystal structure prediction. Computer physics communications , 175(11-12):713-720, 2006.

Niklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32 , pages 7566-7578. Curran Associates, Inc., 2019.

Gerhard Kurz, Igor Gilitschenski, and Uwe D Hanebeck. 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.

James Bergstra, Daniel Yamins, and David Cox. Making a science of model search: Hyperparameter optimization in hundreds of dimensions for vision architectures. In International conference on machine learning , pages 115-123. PMLR, 2013.

Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. In Machine Learning for Molecules Workshop, NeurIPS , 2020.

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.

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):6063-6081, 2020.

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 8748-8763. PMLR, 2021.

Peter E Blöchl. Projector augmented-wave method. Physical review B , 50(24):17953, 1994.

Georg Kresse and Jürgen Furthmüller. Efficiency of ab-initio total energy calculations for metals and semiconductors using a plane-wave basis set. Computational materials science , 6(1):15-50, 1996.

John P Perdew, Kieron Burke, and Matthias Ernzerhof. Generalized gradient approximation made simple. Physical review letters , 77(18):3865, 1996.

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.

Clement Vignac, Igor Krawczuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Discrete denoising diffusion for graph generation. In The Eleventh International Conference on Learning Representations , 2023.

Daniel W Davies, Keith T Butler, Adam J Jackson, Jonathan M Skelton, Kazuki Morita, and Aron Walsh. Smact: Semiconducting materials by analogy and chemical theory. Journal of Open Source Software , 4(38):1361, 2019.

Logan Ward, Ankit Agrawal, Alok Choudhary, and Christopher Wolverton. A general-purpose machine learning framework for predicting properties of inorganic materials. npj Computational Materials , 2(1):1-7, 2016.

Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu. EGSDE: Unpaired image-to-image translation via energy-guided stochastic differential equations. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

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.

Contents

A Theoretical Analysis

A.1

Proof of Proposition 1

We first introduce the following definition to describe the equivariance and invariance from the perspective of distributions.

Definition 4. We call a distribution p ( x ) is G -invariant if for any transformation g in the group G , p ( g · x ) = p ( x ) , and a conditional distribution p ( x | c ) is G-equivariant if p ( g · x | g · c ) = p ( x | c ) , ∀ g ∈ G .

We then provide and prove the following lemma to capture the symmetry of the generation process.

Lemma 1 (Xu et al. [10]) . Consider the generation Markov process p ( x 0 ) = p ( x T ) ∫ p ( x 0: T -1 | x t ) dx 1: T . If the prior distribution p ( x T ) is G-invariant and the Markov transitions p ( x t -1 | x t ) , 0 < t ≤ T are G-equivariant, the marginal distribution p ( x 0 ) is also G-invariant.

Proof. For any g ∈ G , we have

Hence, the marginal distribution p ( x 0 ) is G-invariant.

The proposition Proposition 1 is rewritten and proved as follows.

Proposition 1. The marginal distribution p ( L 0 ) by Eq. (3) is O(3)-invariant if ˆ ϵ L ( M t , t ) is O(3)equivariant, namely ˆ ϵ L ( QL t , F t , A , t ) = Q ˆ ϵ L ( L t , F t , A , t ) , ∀ Q ⊤ Q = I .

Proof. Consider the transition probability in Eq. (3), we have

where a t = 1 √ α t , b t = β t √ 1 -¯ α t , σ 2 t = β t · 1 -¯ α t -1 1 -¯ α t for simplicity, and ˆ ϵ L ( M t , t ) is completed as ˆ ϵ L ( L t , F t , A , t ) .

As the denoising term ˆ ϵ L ( L t , F t , A , t ) is O(3)-equivariant, we have ˆ ϵ L ( QL t , F t , A , t ) = Q ˆ ϵ L ( L t , F t , A , t ) for any orthogonal transformation Q ∈ R 3 × 3 , Q ⊤ Q = I .

For the variable L ∼ N ( ¯ L , σ 2 I ) , we have QL ∼ N ( Q ¯ L , Q ( σ 2 I ) Q ⊤ ) = N ( Q ¯ L , σ 2 I ) . That is,

For the transition probability p ( L t -1 | L t , F t , A ) , we have

As the transition is O (3) -equivariant and the prior distribution N (0 , I ) is O (3) -invariant, we prove that the the marginal distribution p ( L 0 ) is O (3) -invariant based on lemma 1.

A.2 Proof of Proposition 2

Let N w ( µ, σ 2 I ) denote the wrapped normal distribution with mean µ , variance σ 2 and period 1. We first provide the following lemma.

Lemma 3. If the denoising term ˆ ϵ F ( L t , F t , A , t ) is periodic translation invariant, and the transition probabilty can be formulated as p ( F t -1 | L t , F t , A ) = N w ( F t -1 | F t + u t ˆ ϵ F ( L t , F t , A , t ) , v 2 t I ) , where u t , v t are functions of t , the transition is periodic translation equivariant.

Proof. As the denoising term ˆ ϵ F ( L t , F t , A , t ) is periodic translation invariant (for short PTI), we have ˆ ϵ F ( L t , w ( F t + t 1 ⊤ ) , A , t ) = ˆ ϵ F ( L t , F t , A , t ) , for any translation t ∈ R 3 .

For the wrapping function w ( · ) , we have

For wrapped normal distribution N w ( µ, σ 2 ) with mean µ , variance σ 2 and period 1, and for any k ′ , k ′′ ∈ Z , we have

Let k ′ = 0 , k ′′ = w ( µ ) -µ , we directly have

For any t ∈ R , we have

Let k ′ = w ( x + t ) -( x + t ) , k ′′ = w ( µ + t ) -( µ + t ) , we have

For the transition probability p ( F t -1 | L t , F t , A ) , we have

The transition probability of the fractional coordinates during the Predictor-Corrector sampling can be formulated as

where p P , p C are the transitions of the predictor and corrector. According to lemma 3, both of the transitions are periodic translation equivariant. Therefore, the transition p ( F t -1 | L t , F t , A ) is periodic translation equivariant. As the prior distribution U (0 , 1) is periodic translation invariant, we finally prove that the marginal distribution p ( F 0 ) is periodic translation invariant based on lemma 1.

A.3 Proof of Proposition 3

We rewrite proposition 3 as follows.

Proposition 3. The score ˆ ϵ L by Eq. (9) is O(3)-equivariant, and the score ˆ ϵ F from Eq. (10) 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 L ⊤ L . For any orthogonal transformation Q ∈ R 3 × 3 , Q ⊤ Q = I , we have

For the Fourier Transformation, consider k is even, we have

Similar results can be acquired as k is odd. Therefore, we have ψ FT ( w ( f j + t ) -w ( f i + t )) = ψ FT ( f j -f i ) , ∀ t ∈ R 3 , i.e. , the Fourier Transformation ψ FT 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 ˆ ϵ F is periodic translation invariant. Let

(a) Translation on infinite 3D space.

(b) Periodic translation on a unit cell.

Above all, ˆ ϵ L is O(3)-equivariant, and ˆ ϵ F is periodic translation invariant. According to proposition 1 and 2, the generated distribution by DiffCSP in Algorithm 2 is periodic E(3) invariant.

A.4 Discussion on Periodic Translation Invariance

In Definition 3, we define the periodic translation invariance as a combination of translation invariance and periodicity. To see this, we illustrate an additional example in Figure 4. 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. 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 to the periodic translation defined in Definition 3.

Consider two origin points p 1 , p 2 ∈ R 3 × 1 and the lattice matrix L , the constructed unit cells by p 1 , p 2 can be represented as M 1 = ( A 1 , F 1 , L ) and M 2 = ( A 2 , F 2 , L ) , where F 1 , F 2 ∈ R 3 × N are fractional coordinates and

which means that the unit cells formed by different origin points actually represent the same infinite crystal structures [36]. We further construct a bijection T : M 1 →M 2 mapping each atom in M 1 to the corresponding atom in unit cell M 2 . For the pair ( a 1 ,i , f 1 ,i ) ∈ M 1 , ( a 2 ,j , f 2 ,j ) ∈ M 2 , we have T ( a 1 ,i , f 1 ,i ) = ( a 2 ,j , f 2 ,j ) iff ∃ k i ∈ Z 3 × 1 , s.t.

After proper transformation, we have

As f 1 ,i , f 2 ,i ∈ [0 , 1) 3 × 1 , we have

which means shifting the periodic boundaries by changing the origin point p 1 into p 2 is equivalent to a periodic translation F 2 = w ( F 1 + L -1 ( p 1 -p 2 ) 1 ⊤ ) .

B Implementation Details

B.1 Approximation of the Wrapped Normal Distribution

The Probability Density Function (PDF) of the wrapped normal distribution N w (0 , σ 2 t ) is

where x ∈ [0 , 1) . Because the above series is convergent, it is reasonable to approximate the infinite summation to a finite truncated summation [61] as

And the logarithmic gradient of f can be formulated as

To estimate λ t = E -1 x ∼N w (0 ,σ 2 t ) [ ∥∇ x log N w ( x | 0 , σ 2 t ) ∥ 2 2 ] , we first sample m points from N w (0 , σ 2 t ) , and the expectation is approximated as

For implementation, we select n = 10 and m = 10000 .

B.2 Adaptation of CDVAE

As illustrated in Figure 5, the original CDV AE [9] mainly consists of three parts: (1) a 3D encoder to encode the structure into the latent variable z 3 D , (2) a property predictor to predict the lattice L , the number of nodes in the unit cell N , and the proportion of each element in the composition c , (3) a 3D decoder to generate the structure from z 3 D , L , N, c via the Score Matching with Langevin Dynamics (SMLD, Song and Ermon [49]) method. The training objective is composed of the loss functions on the three parts, i.e. the KL divergence between the encoded distribution and the standard normal distribution L KL , the aggregated prediction loss L AGG and the denoising loss on the decoder L DEC . Formally, we have

We formulate L KL = βD KL ( N ( µ 3 D , σ 2 3 D I ) ∥N (0 , I )) for better comparison with the adapted method. β is the hyper-parameter 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, to take the composition as the condition, we apply an additional 1D prior encoder to encode the composition set into a latent distribution N ( µ 1 D , σ 2 1 D I ) and minimize the KL divergence between the 3D and 1D distribution. The training objective is modified into

3D Encoder

1D Encoder

z

3D

z

1D

𝜇

3D

,

𝜎

3D

𝜇

1D

,

𝜎

1D

ℒ

!"

3D Decoder

1D Encoder

z

1D

𝜇

1D

,

𝜎

1D

predict

𝑳

𝑁

5

𝑐

3D Decoder

predict

𝑳

𝑨

#$

Initialize

3D Encoder

z

3D

𝜇

3D

,

𝜎

3D

ℒ

!"

3D Decoder

Training

z

Inference

predict

𝑳

𝑁

5

𝑐

3D Decoder

predict

Initialize

𝒩

(0,

𝑰

)

𝒩

(0,

𝑰

)

𝑳

𝑁

5

𝑐

ℒ

%&&

ℒ

'()

ℒ

%&&

ℒ

'()

Original Framework

Adapted Framework

(a)

(b)

↑

(c)

(d)

During the inference procedure, as the composition is given, the latent variable z 1 D is sampled from N ( µ 1 D , σ 2 1 D I ) . For implementation, we apply a Transformer [51] without positional encoding as the 1D encoder to ensure the permutation invariance. Secondly, for the generation procedure, we apply the ground truth composition for initialization and keep the atom types unchanged during the Langevin dynamics to ensure the generated structure conforms to the given composition.

B.3 Algorithms for Training and Sampling

Algorithm 1 summarizes the forward diffusion process as well as the training of the denoising model ϕ , while Algorithm 2 illustrates the backward sampling process. They can maintain the symmetries if ϕ is delicately constructed. Notably, We apply the predictor-corrector sampler [50] to sample F 0 . In Algorithm 2, Line 7 refers to the predictor while Lines 9-10 correspond to the corrector.

Algorithm 1 Training Procedure of DiffCSP

1: Input: lattice matrix L 0 , atom types A , fractional coordinates F 0 , denoising model ϕ , and the number of sampling steps T .

2: Sample ϵ L ∼ N ( 0 , I ) , ϵ F ∼ N ( 0 , I ) and t ∼ U (1 , T ) . √

3: L t ← √ ¯ α t L 0 + 1 -¯ α t ϵ L

4: F t ← w ( F 0 + σ t ϵ F )

5: ˆ ϵ L , ˆ ϵ F ← ϕ ( L t , F t , A , t )

7: L F ← λ t ∥∇ F t log q ( F t | F 0 ) -ˆ ϵ F ∥ 2 2

6: L L ←∥ ϵ L -ˆ ϵ L ∥ 2 2

8: Minimize L L + L F

B.4 Hyper-parameters and Training Details

We acquire the origin datasets from CDVAE [9] 4 and MPTS-52 [57] 5 . We utilize the codebases from GN-OA [21] 6 , cG-SchNet [53] 7 and CDVAE [9] 8 for baseline implementations.

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 an Adam optimizer with learning rate 1 × 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 [62] 9 . Specifically, we choose observation quantile γ as 0.25 and the number of initial random points as 200 for BO. For PSO, we used scikit-opt 10 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 × 10 -4 and a Plateau scheduler with a decaying factor 0.5 and a 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 CDV AE, 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 2 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, Carbon-24, MP-20 and MPTS-52 respectively with an Adam optimizer with initial learning rate 1 × 10 -3 and a Plateau scheduler with a 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 k = 256 . We apply the cosine scheduler with s = 0 . 008 to control the variance of the DDPM process on L t , and an exponential scheduler with σ 1 = 0 . 005 , σ T = 0 . 5 to control the noise scale of the score matching process on F t . 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 γ in Langevin dynamics for the structure prediction task, we apply γ = 5 × 10 -7 for Perov-5, 1 × 10 -5 for MP-20 and MPTS-52, and for Carbon-24, we apply γ = 5 × 10 -6 to predict one sample and γ = 5 × 10 -7 for multiple samples. For the ab initio generation and optimization task on Perov-5, Carbon-24 and MP-20, we apply γ = 1 × 10 -6 , 1 × 10 -5 , 5 × 10 -6 , respectively. All models are trained on GeForce RTX 3090 GPU.

C Exploring the Effects of Sampling and Candidate Ranking

C.1 Impact of 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 DiffCSP outperforms the baselines methods under the arbitrary number of samples.

�

�

�

�

�

��

��

��

��

��

��

�����������������

�

��

��

��

��

���

��������������

�������

�������

�����������

�����

�

�

�

�

�

��

��

��

��

��

��

�����������������

�

��

��

��

��

���

��������������

���������

�������

�����

�����������

�

�

�

�

�

��

��

��

��

��

��

�

��

��

��

��

���

��������������

�����

�������

�����

�����������

�

�

�

�

�

��

��

��

��

��

��

�

��

��

��

��

���

��������������

�������

�������

�����

�����������

�����������������

�����������������

C.2 Diversity

We further evaluate the diversity by yielding the CrystalNN [65] fingerprint of each generated structure, calculating the L2distances 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 CDV AE and DiffCSP on Perov-5 and MP-20 in Table 5.

C.3 Ranking among Multiple Candidates

In § 5.1, we match each candidate with the ground truth structure to 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 models to score each sample for ranking from different perspectives: Energy Predictor (EP). Since lower formation energy typically leads to more stable structures, we directly train a predictor using energy labels and apply the negative of the predicted

energy as the confidence score. Match Discriminator (MD). Inspired by Diffdock [12], 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 a contrastive model between a 1D and 3D model to align the corresponding compositions and structures. The inner product of the 1D 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 6. 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 7 and 8. 1. For lattices, we originally use the cosine scheduler with = 0 . 008 , and we change it into the linear and sigmoid schedulers with β 1 = 0 . 0001 , β T = 0 . 02 . We find that the linear scheduler yields comparable results, while the sigmoid scheduler hinders the performance. 2. For fractional coordinates, we use the exponential scheduler with σ min = 0 . 005 , σ max = 0 . 5 , and we change the value of σ max into 0.1 and 1.0. Results show that the smallσ max variant performs obviously worse, as only sufficiently large σ max could approximate the prior uniform distribution. We visualize the PDF curves in Figure 7 for better understanding. 3. For atom 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.

E Learning Curves of Different Variants

Weplot the curves of training and validation loss of different variants proposed in § 5.3 in Figure 8 with the following observations. 1. The multi-graph methods struggle with higher training and validation loss, as 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 8) encounters obvious overfitting as the periodic translation invariance is violated, highlighting the necessity to leverage the desired invariance into the model.

F Comparison with DFT-based Methods

F.1 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

���

���

���

���

���

���

�

���

���

���

���

���

���

���

���

���

������

����

����

����

����

����

0

200

400

600

800

1000

0.8

1.0

1.2

1.4

1.6

Train Loss

MG w/o FT

MG w/ FT

DiffCSP w/o FT

DiffCSP w/ FT

0

200

400

600

800

1000

0.8

1.0

1.2

1.4

1.6

Val Loss

MG w/o FT

MG w/ FT

DiffCSP w/o FT

DiffCSP w/ FT

Epoch

Epoch

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 calculation is performed using one node with 48 cores (Intel(R) Xeon(R) CPU E5-2692 v2 @ 2.20GHz), while the 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-Burke-Ernzerhof (PBE) function [69].

F.2 Results

We select 10 binary and 5 ternary compounds in MP-20 testing set and compare our model with USPEX. For our method, we sample 20 candidates for each compound following the setting in Table 1. 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 Table 9, and additionally 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 10. The results show that DiffCSP correctly predicts more structures with higher match rate and significantly lower time cost.

G Extension to More General Tasks

G.1 Overview

Our method mainly focuses on CSP, aiming at predicting the stable structures from the fixed composition A . We first enable the generation on atom types and extend to ab initio generation task in § G.2, and then adopt the energy guidance for property optimization in § G.3. Figure 9 illustrated the differences and connections of CSP, ab initio generation, and property optimization. Besides, CDVAE [9] proposes a 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 a 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 ∈ R h × N in 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, 71],

(a) Crystal Structure Prediction (CSP)

(b) Ab Initio Generation

(c) Property Optimization

(d) Reconstruction

2#

2

3+3

O

where 1 h ∈ R h × 1 , 1 N ∈ R N × 1 are vectors with all elements setting to one, A 0 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 t .

The corresponding backward generation process is defined as

where

and ˆ ϵ A ∈ R h × N is predicted by the denoising model. We further find that specific for t = 1 , A 0 = argmax (ˆ ϵ A ( M 1 , 1)) works better.

The training objective for multinomial diffusion is

One-hot Diffusion for Continuous Generation Another approach is to simply consider the composition A as a continuous variable in real space R h × N , which enables the application of standard DDPM-based method [41]. Similar to Eq. (2)-(4), the forward diffusion process is defined as

And the backward generation process is defined as

where µ A ( M t ) = 1 √ α t ( A t -β t √ 1 -¯ α t ˆ ϵ A ( M t , t ) ) , σ 2 A ( M t ) = β t 1 -¯ α t -1 1 -¯ α t . The denoising term ˆ ϵ A ( M t , t ) ∈ R h × N is predicted by the model.

The training objective for one-hot diffusion is

The entire objective for training the joint diffusion model on L , F , A is combined as

We select λ L = λ F = 1 , λ A = 0 for the CSP task as A is fixed during the generation process, and λ L = λ F = 1 , λ A = 20 for the ab initio generation task to balance the scale of each loss component. Specifically, we do not optimize L A on the Carbon-24 dataset, as all atoms in this dataset are carbon.

Sample Structures with Arbitrary Numbers of Atoms As the number of atoms in a 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 is similar to Hoogeboom et al. [41]. The sampled distribution p ( M ) can be more concisely described as p ( M , N ) = p ( N ) p ( M| N ) . The former term p ( N ) is sampled from pre-computed data distribution, and the latter conditional distribution p ( M| N ) is 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.5Å, and the generated composition is valid if the entire charge is neutral as determined by SMACT [72]. Coverage : It measures the structural and compositional similarity between the testing set S t and the generated samples S g . Specifically, letting d S ( M 1 , M 2 ) , d C ( M 1 , M 2 ) denote the L 2 distances of the CrystalNN structural fingerprints [65] and the normalized Magpie compositional fingerprints [73], the COVerage Recall (COV-R) is determined as COV-R = 1 |S t | |{M i |M i ∈ S t , ∃M j ∈ S g , d S ( M i , M j ) < δ S , d C ( M i , M j ) < δ C }| where δ S , δ C are pre-defined thresholds. The COVerage Precision (COV-P) is defined similarly by swapping S g , S t . Property statistics : 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 d ρ , d E and d elem, 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 CDV AE on MP-20 are provided in Table 11. We observe that DiffCSP-D yields slightly lower validity and coverage rates than DiffCSP-C. Moreover, DiffCSP-D tends to generate structures with more types of elements, which is far from the data distribution. Hence, we select DiffCSP-C for the other experiments in Table 4 (abbreviated as 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.

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 t , F t , A t , t ) with the same message passing blocks as Eq. (6)-(8). And the final prediction is acquired by the final layer as

And the gradients w.r.t. L , F , A additionally guide the generation process. As the inner product term L ⊤ L is O(3) invariant, and the Fourier transformation term ϕ FT ( f j -f i ) is periodic translation invariant, the predicted energy E is periodic E(3) invariant. That is,

Algorithm 3 Energy-Guided Sampling Procedure of DiffCSP(-C)

1: Input: denoising model ϕ , energy predictor E , step size of Langevin dynamics γ , guidance magnitude s , input structure before optimization M = ( L , F , A ) , number of sampling steps T ′ , maximum number of sampling steps T .

2: if T'=T then 3: Sample L T ′ ∼ N ( 0 , I ) , F T ′ ∼ U (0 , 1) , A T ′ ∼ N ( 0 , I ) . 4: else 5: Sample L T ′ ∼ q ( L T ′ | L ) , F T ′ ∼ q ( F T ′ | F ) , A T ′ ∼ q ( A T ′ | A ) according to Eq. (2),(5) and (22). 6: end if 7: for t ← T ′ , · · · , 1 do 8: Sample ϵ L , ϵ F , ϵ A , ϵ ′ F ∼ N ( 0 , I ) 9: ˆ ϵ L , ˆ ϵ F , ˆ ϵ A ← ϕ ( L t , F t , A t , t ) . 10: Acquire ∇ L E, ∇ F E, ∇ A E from E ( L t , F t , A t , t ) . 11: L t -1 ← 1 √ α t ( L t -β t √ 1 -¯ α t ˆ ϵ L ) -sβ t · 1 -¯ α t -1 1 -¯ α t ∇ L E + √ β t · 1 -¯ α t -1 1 -¯ α t ϵ L . 12: A t -1 ← 1 √ α t ( A t -β t √ 1 -¯ α t ˆ ϵ A ) -sβ t · 1 -¯ α t -1 1 -¯ α t ∇ A E + √ β t · 1 -¯ α t -1 1 -¯ α t ϵ A . 13: F t -1 2 ← w ( F t +( σ 2 t -σ 2 t -1 )ˆ ϵ F -s σ 2 t -1 ( σ 2 t -σ 2 t -1 ) σ 2 t ∇ F E + σ t -1 √ σ 2 t -σ 2 t -1 σ t ϵ F ) 14: _ , ˆ ϵ F , _ ← ϕ ( L t -1 , F t -1 2 , A t -1 , t -1) . 15: d t ← γσ t -1 /σ 1 16: F t -1 ← w ( F t -1 2 + d t ˆ ϵ F + √ 2 d t ϵ ′ F ) . 17: end for 18: Return L 0 , F 0 , argmax ( A 0 ) .

Taking gradient to both sides w.r.t. L , F , A , respectively, we have

which implies that the gradient to L t is O(3) equivariant, and the gradient to F t and A t 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-guided 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. [9] 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 of the property distribution. We select the formation energy per atom as 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 10 and calculate the energies of the structures before and after optimization by V ASP [68]. Results show that our method is capable to search novel structures with lower energies compared to existing ones.

������������������

�������������������

������������������

������������������

������������������

�������������������

������������������

������������������

������������������

������������������

��������������������

������������������

������������������

������������������

�������������������

������������������

������������������

������������������

H Computational Cost for Inference

We provide the GPU hours (GeForce RTX 3090) for different generative methods to predict 20 candidates on the 4 datasets. Table 13 demonstrates the diffusion-based models (CDVAE and our DiffCSP) are slower than P-cG-SchNet. Yet, the computation overhead of our method is still 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 in Table 14. Our findings indicate that both models exhibit improved performance with an increased number of steps. Notably, DiffCSP with 1,000 steps outperforms CDVAE with 5,000 steps.

I Error Bars

We provide a single run to generate 20 candidates in Table 1. We apply two more inferences of each generative method on Perov-5 and MP-20. Table 15 shows similar results as Table 1.

J More Visualizations

In this section, we first present additional visualizations of the predicted structures from DiffCSP and other generative methods in Figure 11. In line with Figure 3, our DiffCSP provides more accurate predictions compared with the baseline methods. Figure 12 illustrates 16 generated structures on Perov-5, Carbon-24 and MP-20. The visualization shows the capability of DiffCSP to generate diverse

structures. We further visualize the generation process of 5 structures from MP-20 in Figure 13. We find that the generated structure M 0 is periodically translated from the ground truth structure, indicating that the marginal distribution p ( M 0 ) follows the desired periodic translation invariance. We provide the detailed generation process in the Supplementary Materials.

Ground Truth

DiffCSP

CDVAE

P-cG-SchNet

Perov-5

MP-20

MPTS-52

Carbon-24

?

8：

一

Q

Perov-5

MP-20

Carbon-24

?

𝑡

= 1000

𝑡

= 750

𝑡

= 500

𝑡

= 250

𝑡

= 0

Translated

Ground Truth

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nDiffCSP 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThey listed the limitations in the main paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nDiffCSP 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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/\n\nQUESTIONS:\nConcrete 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\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

comments: SUMMARY:
This paper addresses the crystal structure prediction problem, one of the important problems in this field.
The authors propose a diffusion model generating both lattices and fractional coordinates given chemical composition.
The proposed denoising model satisfies O(3) equivariance and periodic translation invariance.
The experimental results show that the proposed method outperforms the conventional methods in generating existing crystals.

STRENGTHS:
* The paper is well-organized and well-written.
* The representation of crystals is reasonable.
* The empirical result is strong, and the baseline methods are appropriate for the CSP task.
* The results compared with USPEX, a strong GA-based method,  attract the CSP community's interest.
* 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.

A minor comment

"$\psi_L$ in Eq. (10)" in the main text of section 4.3 seems to be a typo in Eq. (9).

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nSee Weaknesses. An additional one is-- how's the efficiency compared with CDVAE in empirical studies?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}