Crystal Structure Prediction by Joint Equivariant Diffusion

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.

Paper

References (75)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer pv4b5/10 · confidence 4/52023-06-28

Summary

DiffCSP is proposed in this paper for the task of crystal structure prediction task, which generates the stable crystal structures of given compositions. The diffusion process of DIffCSP considers both lattice matrix and fractional coordinates of atoms in the unit cell. Experimental results demonstrate the effectiveness and efficiency of this method.

Strengths

- The proposed method shows good performances compared to previous methods in the task of Crystal Structure Prediction. - The problem setting is new to ML, and interesting. - Writing is clear with informative figures, formulas, and proofs. - Experiments are comprehensive with additional adjustments to change the current framework to the generative model setting and compare with previous generative models. Ablation studies are provided.

Weaknesses

- Limited novelty: the proposed method is straightforward by further considering lattice matrix diffusion; the novelty is somehow limited. The way they diffuse lattice matrix is similar to previous works for molecule conformer generation. And the way they diffuse fractional atom coordinates is similar to CDVAE. - 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. - 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.

Questions

- Maybe a typo: In line 277, it is '20 structures of the lowest energy or' or of? Suggestions: - 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.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

They listed the limitations in the main paper.

Reviewer pLZt8/10 · confidence 4/52023-07-01

Summary

DiffCSP is a novel diffusion model for predicting crystal structures conditioned on the crystal's atom composition; it can also be extended to an initio crystal generation (sampling both atom composition and positions). Crystal modeling is challenging due to its periodic symmetry. DiffCSP proposes to solve periodicity constraints by diffusing on fractional coordinates and jointly generating the lattice and atom coordinates. Fractional coordinate diffusion is achieved with torsional diffusion that constraints each intermediate state to stay within a bounding box, i.e. a fixed unit cell. When compared with existing methods, DiffCSP is state-of-the-art across DFT and ML-based methods in predicting the crystal structures on three datasets -- Perov-5, Carbon-24, and MP-20 -- while requiring less computation. It is also state-of-the-art on ab initio generation for the same datasets.

Strengths

- Using fractional coordinates to handle periodicity is supported by several prior studies on crystal structures. It is encouraging to see DiffCSP explore this representation in diffusion models. The authors do a in-depth comparison with CDVAE [1] which alternatively explored a multi-graph construction for handling periodicity. - 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. - 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. - DiffCSP achieves state-of-the-art in crystal structure prediction and ab initio crystal generation. [1] https://arxiv.org/abs/2110.06197 [2] https://arxiv.org/pdf/2102.09844.pdf

Weaknesses

- The authors claim DiffCSP is the first generative model for crystal structure prediction (CSP). I am curious to know why a generative viewpoint of CSP is beneficial. Discussion of why diffusion is important for this task would help, i.e. a discussion about uncertainty in the problem and how it manifests. - 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]. - 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. - 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. [1] https://pubmed.ncbi.nlm.nih.gov/14691322/

Questions

Concrete questions pertaining to my weaknesses: - 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? - 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. - Why is DiffCSP chosen to be E(3) equivariant but not SE(3) equivariant like in CDVAE? Is this a limitation? - Was the noising schedule of the atom types, lattice, and coordinates important? Did it make a practical difference? Other questions: - 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. - How diverse are the crystals when multiple samples are drawn? Typoes/suggestions: - Line 33. Reference 11 does protein structure generation, not prediction. - 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. - 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]. [1] https://arxiv.org/abs/2202.02763

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

- A potential limitation is SE(3) vs. E(3) equivariance. The authors do not discuss this. I have mentioned this already in the questions. - 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.

Reviewer eNoa7/10 · confidence 4/52023-07-03

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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes, the limitation is discussed in Section 6.

Reviewer 1HBn5/10 · confidence 5/52023-07-08

Summary

The paper studies crystal structure prediction, an important topic in the field of AI for scientific discovery. Diffusion models are designed considering necessary symmetries for crystal study, with diffusing lattice vectors and fractional coordinates for generating lattices and atoms, respectively. Experiments show the effectiveness of the proposed method.

Strengths

1. The paper presents novel and shining ideas, e.g., considers fractional coordinates rather than Cartesian coordinates. This eliminates the need of constructing multi-edge graphs and naturally guarantees periodic invariance. The theory for processing fractional coordinates in the diffusion model looks rigid and solid. 2. The joint learning fashion for lattices and coordinates should be the right way to go in the field. 3. Experiments are solid and performance is impressive.

Weaknesses

1. One major concern is-- the technical contribution is not clear or strong enough. Diffusion models for 3D scientific data ate not new. The authors may need to clearly state the significant and unique design of the proposed diffusion model to distinguish it from other diffusion models. 2. Another concern is-- necessary symmetries are not completely considered. For example, similar to CDVAE, translation invariance (not periodic translation invariance) may not be considered. 3. 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. 4. 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. Overall, this is a neat work. I would raise my scores if my concerns are fully addressed.

Questions

See Weaknesses. An additional one is-- how's the efficiency compared with CDVAE in empirical studies?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Reviewer pLZt2023-08-10

Response

Thank you for the very detailed response. I have read the global rebuttal and enjoyed all the additional analysis. I have raised my score as a result. I believe this problem formulation will be a useful and interesting contribution to NeurIPS.

Authorsrebuttal2023-08-12

Thank you

Dear Reviewer pLZt, Thank you very much! Your valuable comments and suggestions have helped improve the paper. Thanks again! Best, Authors

Reviewer 1HBn2023-08-13

Given the clarifications the authors provided as well as the promise that certain content will be added in the revision, I raise my rating to 5. Considering the work is not significant enough in the context of existing studies, 5 is the best score I can give.

Authorsrebuttal2023-08-15

Thank you

Thanks for raising the score. We will highlight the contributions in the revised paper.

Reviewer eNoa2023-08-14

Thanks for the comments

Thank you for your reply. > Nice suggestion! In fact, we have equipped our method with a guidance model to generate structures with lower energies. Please refer to Appendix F.3 for more details. Here according to the reviewer's suggestion, we fix the composition of the initial crystal and only optimize its structures for lower energy. Results are visualized in Figure R2 (General Response), where the energies of the structures before and after optimization by VASP are shown. It reads that our method does search novel structures with lower energies compared to existing ones. It sounds promising, but the impact of this paper could be enhanced more by listing novel crystal structures that possess lower energy than those currently registered in crystal databases, such as the Crystallography Open Database or Material Project. > Moreover, the training time for our method is approximately 5.2 hours on MP-20, and we do acknowledge that this is an important consideration for users when implementing our model. Even taking the training time into consideration, our method still showcases more efficiency against USPEX (the total time is 5.2h + 10s vs 12.5h). The total time, i.e., 5.2h + 10s, is still impressive, and I recommend including the training time in the main text since many readers would like to know it.

Authorsrebuttal2023-08-15

Thank you very much for your feedback. > It sounds promising, but the impact of this paper could be enhanced more by listing novel crystal structures that possess lower energy than those currently registered in crystal databases, such as the Crystallography Open Database or Material Project. We really agree with your suggestion of enhancing our paper by listing novel crystal structures with lower energy compared with the Crystallography Open Database [1] or Material Project [2], which is actually our ultimate goal of developing generative models for the CSP task. Given that our current framework can only optimize the structure of the input crystal composition one by one, it would be more challenging to discover a totally novel and more stable structure in constrast to all structures in the Crystallography Open Database and Material Project that contain a vast number of structures (504,942 and 144,595 entries respectively). Optimizing each individual structure with DFT verification can be a time-consuming process, and it may also require larger datasets to train a more generalized diffusion model. While we do understand and appreciate the kind suggestion by the reviewer, we are more willing to leave it for future studies. [1] http://www.crystallography.net/cod/ [2] https://legacy.materialsproject.org/ > The total time, i.e., 5.2h + 10s, is still impressive, and I recommend including the training time in the main text since many readers would like to know it. Thank you once again. We will certainly integrate the training time into the revised version of the main paper.

Reviewer eNoa2023-08-15

Thank you

Thank you, and I am looking forward to seeing follow-on research in the future.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC