Kinetic Langevin Diffusion for Crystalline Materials Generation

We thank the reviewer for their positive consideration and suggestions to improve the paper. We address questions and comments below.

De-Novo generation task results Due to the limited character, we have to refer to the answer provided to reviewer MrGy about this topic.

Zero-net translation intuition We agree with the reviewer that the intuition was lacking from the submitted manuscript. Here, we provide a simple example to build intuition. We will include it in the updated version.

Consider a datapoint with a single atom in 1D, i.e. $\boldsymbol{f}_0$ consists of just a single coordinate. In this simple setup, every point $\boldsymbol{f}_t$ can be seen as a periodic translation of any another $\boldsymbol{f}^{'}_t$, hence also of the clean sample $\boldsymbol{f}_0$ itself.

With no constraint on the velocity field, the forward dynamics results in noisy samples $\boldsymbol{f}_t$ corresponding to periodic translations of $\boldsymbol{f}_0$ (i.e. almost surely represented by different group elements) with non-zero target scores pointing back to $\boldsymbol{f}_0$. Since all $\boldsymbol{f}_t$ effectively represent the same datapoint, modelling this degree of freedom is unecessary.

By constraining the velocity field to be zero-mean, the single velocity has to be zero for the constraint to be satisfied. By simulating the forward dynamics, all noisy samples, $\boldsymbol{f}_t$, are exactly $\boldsymbol{f}_0$ (i.e. they share the same group element) with an associated zero target.

Ablation of design choices
While the proposed constraint of the velocity field is an important part of the paper, we do not see it as being the main contribution. The core of the paper is instead the extension of the TDM framework to crystalline materials generation. To obtain fast convergence and competitive results, we find that zero initial velocities and the resulting simplified parameterization are key elements (see Figure). We show that non-zero initial velocities (see Figure) systematically lead to subpar performance.

As suggested by the reviewer, we provide an ablation of the effect of the zero net translation for zero initial velocities (see Figure) and non-zero ones (see Figure). By removing this unecesseray degree of freedom, we observe a benefit in all cases, in particular with non-zero initial velocities.

We thank the reviewer for their positive consideration and suggestions to improve the paper. Thanks for pointing out some typos, we will correct them in the updates version. We address questions and comments below.

Invariant network and equivariant target inconsistency We agree with the reviewer that in settings where no symmetries are involved, there is no issue with the denoising score matching loss. In the present case (i.e. target distribution with translational symmetry), the problem stems from the use of a periodic translation invariant score network to match an equivariant target. Considering for example a noisy point-cloud and a periodic translated version thereof, these two datapoints are equivalent from the network's perspective while the target scores are going to be different. Although this is averaged out over the course of the training and does not prevent models from learning a useful score approximation (e.g. DiffCSP and MatterGen), this is undesirable. For an intuition on this, refer to the reply "Zero-net translation intuition" given to Reviewer WGaL and or an alternative discussion, see also [1].

To further support this, we ablate the effect of the zero net translation in the next paragraph.

Ablation on zero net translation and initial zero velocity We investigate the effect of the zero net translation in terms of the match rate on the validation set of MP-20 (see Figure), where we obtain (slightly) better results by enforcing zero net translation.

We also present an analysis about the impact of non-zero initial velocities for different variances (see Figure and Figure). We observe that the zero net translation get better results no matter the initial distribution, and that by forcing the initial velocity to be zero the model converges faster and get better results in terms of match rate on the validation set.

Error bars We agree with the reviewer that we are not consistent as we present errors bars only for some of the experiments. For the baselines, results are taken from the previous papers. We will add error bars also for the DNG task in the updated version. For CSP@20, this was due to the computational cost, but we can add them in the updated version.

New metrics de novo generation task Due to the limited character, we have to refer to the answer provided to Reviewer MrGy about this topic.

Why does this work? We hypothesize that the added momentum on the fractional coordinates dynamics is the main driver behind the improved performance over DiffCSP. We find that zero initial velocities and velocity fields zero net translation are critical for better results and faster convergence. Exploring different noise schedules for the velocities is an interesting direction for further improving KLDM.

Possible future applications Our model can also be applied to other tasks that involve the generation of periodic systems. A natural application can be surfaces or other lower dimensional periodic systems, e.g. 2D or 1D materials. The generation of metal-organic frameworks (MOF) is another interesting future application, with the main challenge being the additional modelling of rotational frames.

References

[1] Lin, Peijia, et al. "Equivariant diffusion for crystal structure prediction." ICML 2024.

We thank the reviewer for their positive consideration and suggestions for improving the paper. We address their questions and comments below.

RMSE computation Similar to previous work, we compute the RMSE of the generated samples wrt. ground truth using StructureMatcher from pymatgen, after filtering for structural and compositional validity. The algorithm internally accounts for the symmetries in the data.

Simplified parameterization To support this, we provide a plot showing the evolution of the match rate on the validation set of MP-20 (see Figure), where the simplified parameterization is shown to converge significantly faster and to higher values than the direct one.

Other design choices ablation We note that this simplified parameterization is only possible when $\boldsymbol{v}_0=0$. To further support this design choice, we evaluate the effect of the initial velocity standard deviation on the convergence / performance of the model (see Figure). When $\boldsymbol{v}_0\neq0$, the models do not reach convergence within the allocated budget of $3$k epochs -- as for the direct parameterization in the case $\boldsymbol{v}_0=0$.

Architecture compared to previous models KLDM and DiffCSP are comparable in terms of the NN architecture. We use the same backbone as that of DiffCSP (and EquiCSP), with the minor difference being that now our score network receives an additional input representing the velocity $\boldsymbol{v}_t$, resulting in a limited increase in learnable parameters.

Matrix exponential computation and difference with other Riemannian score based models (RSBMs) The main difference with DiffCSP is that our diffusion process is defined on the velocity variables and not directly on the fractional coordinates. Our transition kernel has an additional distribution (wrapped normal + normal), resulting in the modelling of $(\boldsymbol{v}_t, \boldsymbol{f}_t)$ instead of $\boldsymbol{f}_t$ only. Compared to RSBM (Algorithm 1 in [1]), which DiffCSP builds upon, our process (Eq 12., and Algorithm 3 in the submitted manuscript) has an additional momentum term, resulting in velocities displaying some inertia.

Intuitively, this can be thought of as the difference between gradient descent ($\sim$DiffCSP) and gradient descent with momentum ($\sim$KLDM).

Regarding the expontial map, our implementation follows what we presented in Eq. 15 in the paper and Appendix C.2. In the case of a torus, this is simply equivalent to a translation and wrapping operation.

De-Novo generation task results We acknowledge the limitations of the presented metrics, and therefore provide more meaningful discovery-related metrics in this new table. Given the timeline and the available resources, the evaluation is performed using a machine-learning interatomic potential, based on the open-source MatterGen pipeline.

For completeness, we compare $3$ ways of performing diffusion on the discrete atom types: continuous diffusion on one-hot encoded atom types (C), continuous diffusion on analog bits (C-AB), and discrete diffusion with absorbing state (D). Notably, when relying on analog-bits or discrete diffusion to model the atom types, KLDM performs better than DiffCSP in terms of RMSD (lower values means generated structures closer to relaxed ones), energy above the hull (lower values means generated materials closer to stability) and stability, while being slightly subpar on S.U.N..

We however note that that the compared DiffCSP and MatterGen-MP were trained on a re-optimized version of MP-20 where some chemical elements have been removed, specifically noble gases, radioactive elements and elements with atomic number greater than 84. Samples with energy above the hull bigger than 0.1 eV / atom have also been filtered out. Our model was trained on the original MP-20.

Regarding Mattergen-MP, we believe that the gap can be explained by different elements: (1) a more expressive denoiser operating in real space, (2) a PC sampler on the lattice parameters, and (3) effect of the pre-processing of MP-20.

References

[1] De Bortoli, Valentin, et al. "Riemannian score-based generative modelling." NeurIPS 2022