Kinetic Langevin Diffusion for Crystalline Materials Generation
Authors: François R J Cornet
Federico Bergamin
Arghya Bhowmik
Juan Maria Garcia-Lastra
Jes Frellsen
Mikkel N. Schmidt
Venue: ICML 2025 poster (2025)
Keywords: generative models
diffusion models
crystals
materials
Primary area: deep_learning->generative_models_and_autoencoders

Abstract: Generative modeling of crystalline materials using diffusion models presents a series of challenges: the data distribution is characterized by inherent symmetries and involves multiple modalities, with some defined on specific manifolds. Notably, the treatment of fractional coordinates representing atomic positions in the unit cell requires careful consideration, as they lie on a hypertorus. In this work, we introduce Kinetic Langevin Diffusion for Materials (KLDM), a novel diffusion model for crystalline materials generation, where the key innovation resides in the modeling of the coordinates. Instead of resorting to Riemannian diffusion on the hypertorus directly, we generalize Trivialized Diffusion Model (TDM) to account for the symmetries inherent to crystals. By coupling coordinates with auxiliary Euclidean variables representing velocities, the diffusion process is now offset to a flat space. This allows us to effectively perform diffusion on the hypertorus while providing a training objective that accounts for the periodic translation symmetry of the true data distribution. We evaluate KLDM on both Crystal Structure Prediction (CSP) and De-novo Generation (DNG) tasks, demonstrating its competitive performance with current state-of-the-art models.

Kinetic Langevin Diffusion for Crystalline Materials Generation
François Cornet"
Federico Bergamin 12 Jes Frellsen 12 Mikkel N. Schmidt 12
Arghya Bhowmik
Juan Maria Garcia Lastra
Abstract
search the space for feasible materials and to accurately estimate their properties. Conventional approaches usually combine random structure search methods with ab-initio Quantum Mechanics (QM) methods (Oganov et al., 2019), such as Density Functional Theory (DFT) (Kohn Sham, 1965); however, structural optimization and property evalu- ation with DFT can be computationally expensive. Recently, the field has witnessed paradigm shift where deep genera- tive models have been introduced to supplement traditional search methods (Anstine & Isayev, 2023), such as random search (Pickard & Needs, 201 1) or evolutionary algorithms (Glass et al., 2006; Wang et al., 2010). Deep generative mod- els learn to approximate underlying probability distributions from existing material data, and in turn, can be sampled from generate novel materials based on the learned pat- terns.
Generative modeling of erystalline materials us- ing diffusion models presents series of chal- lenges: the data distribution is characterized by in- herent symmetries and involves multiple modali- ties, with some defined on specific manifolds. No- tably, the treatment of fractional coordinates rep- resenting atomic positions in the unit cell requires careful consideration, as they lie on a hypertorus. In this work, we introduce Kinetic Langevin Dif- fusion for Materials (KLDM), a novel diffusion model for crystalline materials generation, where the key innovation resides in the modeling of the coordinates. Instead of resorting to Riemannian diffusion on the hypertorus directly, we generalize Trivialized Diffusion Model (TDM) account for the symmetries inherent to crystals. By coupling coordinates with auxiliary Euclidean variables representing velocities, the diffusion process is now offset to a flat space. This allows us to effec- tively perform diffusion on the hypertorus while providing a training objective that accounts for the periodic translation symmetry of the true data distribution. We evaluate KL.DM on both Crystal Structure Prediction (CSP) and De-novo Genera- tion (DNG) tasks, demonstrating its competitive performance with current state-of-the-an models.
Among deep generative models, diffusion models have been successful on a variety of data modalities relevant to the sciences, ranging from Partial Differential Equations (PDE) simulations (Lippe et al., 2023; Rozet & Louppe, 2023; Shysheya et al., 2024) to molecule generation (Hoogeboom et al., 2022; Xu et al., 2023; Cornet et al., 2024). Unlike molecules, crystalline materials consist of periodic ar- rangement of atoms, typically described by umit cell, a parallelepiped that serves as the fundamental building block repeated to tile the entire space. A unit cell is typically rep- resented by three vectors that define its edges and the angles between them, along with the coordinates and species of the atoms inside it. It can be considered data type, specifically geometric graph combining discrete and continuous features (Joshi et al., 2023). Notably, the atomic positions are described by coordinates that lie on a hypertorus. Dealing with non-Fuclidean data in the dif- fusion setting requires careful consideration: for example, restricting Brownian motion to manifold requires incorpo- rating its geometric strueture ensure trajectories remain the manifold, typically using projection specialized approximations (Lou et al., 2023). Current diffusion mod- els for crystals either handle this by working in real space, through multi-graph representations (Xie et al., 2022), or on fractional coordinates (Jiao et al., 2023). In addition to this, the data distribution is governed by inherent symmetries, including permutation invariance (swapping atom indices
multi-modal
1. Introduction
The discovery of novel compounds with desired properties is critical several scientific fields, such molecular discov- ery (Bilodeau et al., 2022) and materials design (Merchant et al., 2023; Zeni et al., 2025). In the case of crystalline materials, the search space is vast, but only a fraction of it is physically plausible. The main challenges are to efficiently
'Equal contribution 'Technical University of Denmark. Kon- gens Lyngby, Denmark 'Pioneer Center for Artificial Intelligence, Copenhagen, Denmark. Correspondence to: François Cornet <fr dk>, Federico Bergamin < <fedboldtu dk>.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

Kinctic Langevin Diffusson for Crystalline Materials Generation
or lattice bases), translation invariance (shifting atom coor- dinates), and rotation invariance (rotating atom positions).
where 1 is a vector of ones with size K, and & (N1,, Ag ka) translates the unit cell to tile the entire space. The lattice vec- tors I can also be compactly represented by a 6-dimensional vector consisting of the three lattice vector lengths and their interior angles, such that the representation is invariant to rotation.
When operating on fractional coordinates, the main chal- lenge is to ensure the periodic translation invariance of the learned distribution. This is usually enforced by parameter- izing the score network with a periodic translation invariant architecture. However, while existing models have success- fully demonstrated the potential of applying diffusion to crystalline materials generation, previous work has high- lighted the existence of mismatch between architecture and supervision signal, resulting in a suboptimal training ob- jective (Lin et al., 2024). While the issue has been acknowl- edged in the literature, in this paper we show an alternative way to mitigate this mismatch that effective at low levels of noise.
Tasks Let p(z) denote the true data distribution over the unit cells, and let ((x) = N Σi= 62, be the empirical data distribution defined by the N available samples. The goal of generative models to learn an approximate model po(2) of p(æ) using g(z). In the realm of crystalline materials generation, there two tasks of interest.
Crystal Structure Prediction (CSP) aims at finding low- energy atomic arrangements (f I), Lee stable structures, for a given atomic composition a. This is framed as a conditional generation task where a model po (ffff((a))))) is trained using samples {(f, L, a)) )/m approximate the true conditional distribution p(f: laa.. Given a specific atomic composition a, the model is used to generate possible coordinates f and lattice parameters L.
Contributions In this work, we introduce Kinetic Langevin Diffusion for Materials (KIDM), novel diffusion model for crystalline materials generation, where the key innova- tion resides in the modeling of the fractional coordinates. Instead of resorting to Riemannian diffusion on the hyper- torus directly as in previous work (Jiao et al., 2023), we generalize Trivialized Diffusion Model (TDM) (Zhu al., 2024) to model symmetric distributions over coordinates. Using the structure of the torus, the diffusion process is offset to auxiliary Euclidean variables representing veloci- ties. We propose a specific parameterization of the resulting diffusion process leading to faster training convergence and improved performance while mitigating a mismatch in the training objective of previous diffusion models for materials. Finally, we show that KLDM offers competitive performance on Crystal Structure Prediction (CSP) and De-novo Genera- tion (DNG).
De-novo Generation (DNG) aims at discovering novel and stable materials. The goal is to generate samples from p(f,l,a) We approximate the distribution by training generative model pe (f L, a) on samples {(f, L, a.)); and evaluate the samples generated by po-
Symmetries of crystalline materials Given a symmetry group G. distribution is G-invariant if for any group el- ement g G. p(g x) p(z), with denoting the group action. A conditional distribution is G-equivariant if for any G. p(g zgg y) = p
2. Background
A number of transformations leave a material x unchanged, meaning the true data distribution has inherent symmetries:
2.1. Crystalline materials and related symmetries
Permutation of atom indices p(f,1,a) a) = p(g f. a), Rotation of lattice vectors
In this section, we introduce the data modality we are inter- ested in, along with the relevant symmetries.
Ag E SK: (3)
Unit cell We are interested in learning the distribution of crystalline materials, described as the repetition of unit cell in 3D-space. We describe unit cell containing K atoms as a fully connected geometric graph x,
p(f a.) = =p(ff g L,a), Permutation of the lattice hasis
4g E SO(3); (4)
p(f, a, p(g f.g L,a), Periodic translation of fractional coordinates
Vg E Szj (5)
x == (f)
(1)
p f (. p(g f.
Vg T /
(6)
where f = =f......, fx) E [0. 1)3R denotes the frac- tional coordinates, 1 = (1, 12) E 123x3 refers to the lattice vectors of the unit cell, and a = (a.,..., ax) ) E zk encodes the chemical composition.
We want our model, po. to inherit these symmetries. Eq. (3) is naturally addressed by using graph neural networks, com- bined with a factorized prior distribution with no depen- dency on the index. The combination of fractional coordi- nates and rotation-invariant lattice representations addresses Eq. (4). Specific neural network architectures or additional
The infinite periodic structure can be represented as,
((a', I')la' = a; f f + +1 k 2²),
2

---

Kinctic Langevin Diffusion for Crystalline Materials Generation
loss terms can address Eq. (5)-see Remaining is to handle Eq. (6), we detail how it can be done in Section 2.3.
see e.g. (Lin al., 2024).
2.3. Existing diffusion models for fractional coordinates
As previously mentioned, the fractional coordinates define hypertorus f E [0, 1)3x = Taxk Most existing work has built upon DIFFCSP (Jiao et al., 2023), which leverages the score-based framework of Song et al. (2021) extended to Riemannian manifolds (De Bortoli et al., 2022; Jing et al., 2022).
2.2. Diffusion models
Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021) are generative models that Iearn distributions through a hierarchy of latent variables, cor- responding to corrupted versions of the data at increasing noise scales. Diffusion models consist of a forward and a reverse (generative) process. The forward process per- turbs samples from the data distribution over time through noise injection, resulting in a trajectory of increasingly noisy latent variables (1)))e[0,r]) Given an initial condi- tion, 2n ~ (a()) = Paata (z), the conditional distribution of (2))te(o,r] can be characterized by Stochastic Differ- ential Equation (SDE),
In what follows, we present the key ingredients of DIFFCSP for modeling fractional coordinates.
Transition kernel DIFFCSP implements specitic case of Eq. (7)-i.e. Sariance-Explodiie SDE (VE-SDE), with f(t) = 0 and g(t) = (do2(t)/dt, where at)) = min " with miin and mmax being hyperparameters.
σ 1-t
In practice, as sampling noisy fractional coordinates f: given fo using a normal distribution does not capture the bounded and cyclical nature of p(f), the solution con- sists in first sampling from the normal distribution f. ~ N (fo. o² (t)) and then wrapping the samples f. = w(f)) where w(+) = L-] is the wrapping function.
dæ;- = (t)xz dt +g(t) duc..
where Ha E Rd denotes the latent variable at time 6, f(t and g(t) are scalar function of time t, and w, is stan- dard Wiener process in Rdd Due to the linearity of the drift term, for any t 2 0, the corresponding transition kernel admits a closed-form expression (Särkkä & Solin, 2019). For instance, in the Variance- Presscvvnn SDE (vP- SDE) setting (Song et al., 2021), where f() = 2 (t) and g(t = vB(t) for a fixed schedule B(t), the ker- nel writes (2)0(2))200) = la I) with o₂ = exp(-0.5 f (s)dss) and of 1 - exp([o (ss) and the process defined by Eq. (7) converges geometrically from low-veriance Gaussian distribution centered around the data to the standard Gaussian distribution pr (axt) = (xx, 00, II), which can be therefore interpreted as an unin- formative prior distribution.
The transition kernel corresponding to this two-step proce- dure corresponds to a wrapped normal distribution,
I1 So - k0 M (10)
(fffff
Σ kEZ
exp
20
whose gradient can then be approximated by truncating the series, allowing for training using denoising score-matching, Due to the wrapping operation, Eq. (10) converges to a uniform distribution over the hypertorus as T.
Denoising score-matching Given the transition kernel Prjo (felfo)) the approximate score function can be opti- mized by minimizing a denoising score-matching objective, which at time writes
The rime-reversal of Eq. (7) is another diffusion process described by the following reverse-time SDE (Anderson, 1982).
Lg. =E= = Ju. [
56 (a). 0)1]] (1)) (11)
where fo Pu(f). ft ~ (el(f.).o) of (x,, t) denotes the output of the score network corre- sponding to fractional coordinates, and A(t) is a (pre-
dæ;- = [f(t)x; - loggpt)))) dt +t()) duž,, (8)
with pr() being the density of æ, and dwē, a time- reversed Wiener process. Sampling from the prior distri- bution pr (x:)). and simulating Eq. (8) results in a sample from p(æ). In practice, the score V. n, log (x) is not avail- able and it is approximated using a score network sp(x,, t), whose parameters 0 are trained via Denoising Score Match- ing (DSM),
1/Esrepip(fals)) [ log ( computable) scale factor ensuring that the loss magnitude is constant in expectation across time (Jing et al., 2022; Jiao et al., 2023).
Invariant approximate distribution As introduced in Eq. (6), the true data distribution p(z), with = (f, a)) as defined in Section 2.1, is periodic translation invariant i.e. unit cells equivalent up to periodic translation shall be equally likely. In DIFFCSP (Jiao et al., 2023), the learned dis- tribution Pa Pa(x) is made invariant by combining an invariant prior - i.e. uniform distribution on the hypertorus with an approximate reverse process that is equivariant to periodic
LUsM Ex(e) [|sp(z;;,)) t)- V log pe (zx (x) [x0)
(9)
where A(t) is a positive denendepnndnnt weighting function and the expectation is taken over the joint distribution t ~ 24(0 (0, 1). 2o ~ p(x),, and Ta ~ Pt (Z0))
3

---

Kinctic Langevin Diffusion for Crystalline Materials Generation
translation-iee _ = polg- Ft-119 2t), Vg E R$ R$/Z². In practice, this is realized by parameterizing the score using neural network that is invariant to periodic translations-iee i.e. (xx,,) t) = s (g- ze, t), Vg E R²/2². If the learned score s/ (xx,,) t) correctly approximates the guaran-
é are 2 x 2 skew-symmetric matrices, whose anti-diagonal element, v, is real-valued.
Forward process As the Lie algebra is isomorphic to Euclidean space, a standard diffusion process can be de- fined for 6. Given its coupling with f determined by left- trivialization (Zhu et al., 2024), the resulting forward pro- cess is defined as.
score of the true target distribution, then teed to model distribution pe (a:) that exhibits the same symmetries as p(æ:) from which new samples can be drawn.
Saf. = dt, dē. = --ú dl + 22 dw..
(12)
3. Kinetic Langevin Diffusion for Materials generation
where the Ordinary Differential Equation (ODE) describes the time evolution of the fractional coordinates f. (that lie on a manifold) through a coupling with the velocity variables et The latter evolve according to an SDE similar to that of Eq. (7), with constant drift f(t) -y and constant volatil- ity g()) /2y, and where weft denotes standard Wiener process the Lie algebra g. We note that, in principle, f and g could be generalized to time-dependent functions but that we consider constant here. In summary. Eq. (12) corrupts fi living on a hypertorus via a standard Euclidean diffusion process on the auxiliary variables, while guarantee- ing that the trajectory (firre (0,7] remains on the manifold at all times. Eq. (12) converges to a product of easy-to-sample distributions: a uniform distribution on G, and standard Normal distribution on g.
We now introduce Kinetic Langevin Diffusion for Materials (KLDM), and particularize our exposition to the fractional coordinates f E [0, 1)3* K T3xK, used to define the positions of atoms inside the unit cell.
Given the isomorphism between the torus T and SO(2)', the fractional coordinates can alternatively be deseribed as collection of 2 x 2 rotation matrices. We denote this al- ternative representation f. Since f is defined on a direct product of connected compact Lie groups', we propose to generalize Trivialized Diffusson Model (TDM) (Zhu et al., 2024) to operate on geometric graphs similar to those de- fined in Eq. (1). While TDM in principle allows for proper treatment of the fractional coordinates out-of-the-box, we find its direct application to crystalline materials generation to result in slow convergence and subpar performance. We consequently propose a set of modifications designed to facilitate faster convergence and enhance empirical results.
Reverse (generative) process The time-reversal of Eq. (12) is given by
df: = fddd
3.1. Trivialized Diffusion Model (TDM) for fractional coordinates
dó: = -yor - -2 log bt)] dt
(13)
+ +V2 deog.
Building upon previous work on momentum-bass based opti- mization (Tao & Ohsawa, 2020) and sampling (Kong & Tao, 2023), TDM is specifically tailored to data defined on Lie groups, and exploits the particular group structure, specifi- cally the left-trivialization operation, to effectively perform diffusion on the manifold via the Lie algebra. The main idea is to couple the variables of interest defined on the manifold with auxiliary variables representing velocities defined on the Lie algebra. More precisely, the fractional coordinates f. elements of the group G SO(2)3> K are coupled with the Lie algebra g. The latter
where t flows backwards and Ve, Pe úz) denotes the (true) score. The latter is unavailable, but can be approxi- mated with a neural network sp (fi, 6,, () trained using DSM as in Eq. (11). However, we stress that the gradient is now with respect to ê.-i.e. an Euclidean variable, in contrast to Eq. (11) where it was with respect to f: fi-see see Eq. (17). As in Eq. (12) there is no direct noise in the forward dynamics of ft. its time-reversal in Eq. (13) does not need score-based correction and corresponds to the reverse ODE.
velocities e E g defined corresponds to the tangent space T&G at the identity element of the group E G, and crucially can be thought of as Fu- clidean space, g S x K In the present setting, velocities
Practically, we split the vector field of the generative dynam- ics in Eq. (13) into two vector fields: one for fa where ú, is considered fixed-thi -ths is a simple linear ODE; and one for é, where f. is considered fixed-this is a semi-linear SDE, as usually encountered in diffusion models.
We provide an intuitive explanation of this correspondence in Appendix B.I.
Discretized update of the fractional coordinates The ODES in Eqs. (12) and (13) describe the dynamics associated with the fractional coordinates. Since they lie on a manifold, integrating the dynamics from time t to + dt involves solving Riemannian exponential map. Intuitively, this
A Lie group is smooth manifold equipped with group structure denoted by G and smooth group operations. We provide short informal introduction to Lie groups and Lie algebra in Appendix A.
4

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
operation generalizes the concept of moving on straight line the manifold case. By considering an initial velocity ot, initial position fr and the step-size dt, the update of the positions can be written as
variables, dag unlike the objective presented in Eq. (11) where the gradient was instead computed with respect to the coordinates fr. It writes,
Oper. sv.
Vv, Pogo V Ver, log WN (reldrr of.1)
(17)
fitd: = f. (frb,dl) = ft expm(6,d(), (14)
Vv, log.No (vx (Uil c I..
where the matrix multiplication fo returns the tangent vector in fr. and the matrix exponential is the exact solution of the exponential map (Zhu et al., 2024).
where we shortened (ejo ve fos vo) as Prjo- Eq. (17) is the target used to train our score network with the DSM objective defined in Eq. (9).
The formalism previously introduced considered f as a collection of rotation matrices. However, the usual neural networks used parameterize the score have been designed to process fractional coordinates (Jiao et al., 2023; Lin et al., 2024). We therefore want to work the hypertorus T directly, using the original representation of the fractional coordinates f. Concretely, if we consider the exponential map defined in Eq. (14) and that f is a single rotation matrix by an angle 0 E [-x,,), the matrix exponential corresponds to a periodic translation as follows.
3.2. Modeling choices
Initial zero velocities Defining p(vo) is a design choice, and we find that setting the initial velocities to zero, L.e. p(vo) = 8(vo), to be beneficial. This intuitively corre- sponds to considering coordinates f initially rest, at time = 0. with noise gradually propagating from the velocities to the coordinates over the course of the diffusion process. A similar benefit was also observed by Dockhorn et al. (2022), for small initial velocities.
f' eexpm(6' - wr(0 + + df),
(15)
Simplified score parameterization While predicting the full score in Eq. (17) directly is possible, alternative pa- rameterizations exist
where 6' is a skew-symmetric matrix and o² is its anti- diagonal element, while w=-) atan2( (tinn) cos(-)) is the wrapping function with alan2 denoting the signed alan function-ses Appendix B.2 for more details. Since 0² in Eq. (15) corresponds to scaled version of the correspond- ing coordinate ʃ', we continue the presentation with all operations expressed in terms of f-
what is done in standard diffusion models. By carefully inspecting the score expres- sion, we note that the target writes sum of two terms, where some of the quantities are known,
V Vv, log (fe velfu, vo)
(18)
1-e-t Ite Mr, WN (rallrr,,
Eu
Transition kernel (Zhu al., 2024) The transition kemnel corresponding to the forward dynamics of Eq. (12) writes
vnn
de
with e. denoting the reparameterization noise sampled to obtain vy. When considering zero initial velocities, i.e. ao O, we have that pr. = 0, Vl. Using the relationship, Hv. + Ee, the second term in Eq. (18) can be computed in closed-form-a aa we have e,, = u,/ Te.. and does not need to be learned. We are then left with a single unknown term, and the simplified parameterization writes,
Pt o(fi, mr|fo: va) = WN(relMe. I..I) No (or (0 I.
(16)
where T;- = logm(fo ff), WN( [Pm, oz, 1) denotes the density of the Wrapped Normal distribution with mean 8
Mr. i + vo) and variance oZ. = 2t + टमा 4,
while No is Gaussian distribution with 1 Ug and o?, =1 22
1 - e s t)
1 1 (x
(19)
o
Practically, noisy samples ft can be obtained from the tran- sition kernel in Eq. (16) as follows. First, we sample the auxiliary variables vt. We then sample r; = logm( f ft). Finally, fr is obtained as ft = foexpm(r:))
where superscript f in so (xx,,)) t) refers to the output of the score network corresponding to fractional coordinates. More details are provided in Appendix C.
The joint distribution in Eq. (16) evolves from the product p(f)- p( eu), tractable limiting distribution that is the product of a uniform distribution over T3*K in f and standard Normal in V.
Empirically, we find the combination of the zero initial velocities and the simplified parameterization to be ben- eficial for convergence and performance, as discussed in Appendix E.
Denoising score-matching target As the diffusion process only acts on the (Euclidean) velocity variables, we only need to compute the score of Eq. (16) with respect to the velocity
Architecture of the score network As imposed by the sym- metries inherent to the target distribution, we parameterize our score network, so (zt, t), using a graph neural network
5

---

Kinctic Langevin Diffussin for Crystalline Materials Generation
architecture, with backbone similar that of previous work (Jiao et al., 2023). The sole difference is that the score network now takes the auxiliary variables, Ve. as additional inputs. The network is made invariant to periodic transla- tion of ft, by featurizing pairwise differences of fractional coordinates with periodic functions of different frequencies, or similarly with shortest distances on the flat torus. More details about the architecture are provided in Appendix I.
and a more detailed analysis.
3.3. Other modalities
Lattice vectors We represent lattices as 6-dimensional vectors, collecting side lengths and angles. We follow pre- vious work (Jiao et al., 2023; Lin et al., 2024) and rely on standard Euclidean diffusion (VP-SDE) as defined in Eqs. (7) and (8), with drift and diffusion functions defined by a linear schedule.
Velocity fields with zero net translation Due to the transla- tional symmetry of the target distribution and consequently the translation invariance of the score parametrization we want to prevent velocities (0))<E(o.7)) from inducing a net overall translation of the system at every time step. Since velocities are defined on the Lie algebra and can therefore be viewed as Euclidean variables, we consider velocity fields living in the mean-free linear subspace. This is, in essence, similar to subtracting the center of gravity when working with molecules in Euclidean space (Xu et al., 2022; Hooge- boom et al., 2022), except that it now takes place on the velocities defined on the Lie algebra.
Compositions In the DNG setting, we employ three distinct representations for atomic compositions a, as per previous work: continuous diffusion using one-hot encoding (Jiao et al., 2023); continuous diffusion on analog-bits (Chen et al., 2023; Miller et al., 2024); and discrete (absorbing) diffusion (Austin et al., 2021; Jiao et al., 2023).
4. Related work
Early deep generative models for crystal generation lever- aged image representations (Hoffmann et al., 2019; Court et al., 2020). ad-hoc frequency space representations (Ren et al., 2022) or 3D coordinates without considering their ge- ometric nature (Nouira et al., 2018; Kim et al., 2020; Yang et al., 2024b).
In practice, this implies that the distribution of velocities, N₁ in Eq. (16), corresponds a projected Normal where all samples are mean-free- account for the constraint on the velocities. We also remove the mean from the sampled rh, and project the score in Eq. (17), to ensure that it does not introduce a net translation. Finally, the corresponding velocity score output s (Ix, t) in Eq. (19)- and hence s/f (Iz, t)- is constructed to be mean- free as well.
is, they satisfy k Uk 0, to
Since then, several works have leveraged diffusion models operating on geometric graphs. The seminal approach of Xie et al. (2022) initially worked in real space, and resorted to nulti-graphs to account for periodicity (Xie & Grossman, 2018), 8). The diffusion process was limited to coordinates, with fixed lattice parameters and composition predicted by a VAE. While such a model has shown practical usefulness, e.g. to find novel 2D materials (Lyngby & Thygesen, 2022), more recent diffusion models instead defined a (more flexi- ble) joint diffusion process for coordinates, lattice structure, and atom types (Jiao et al., 2023; Zeni et al., 2025), and accounted for periodicity by operating on fractional coor- dinates. Miller et al. (2024) generalized Riemannian flow matching (Lipman et al., 2023; Chen & Lipman, 2024) to the same setup. Sriram et al. (2024) further leveraged the flexibility of flow matching and used fine-tuned LLM as a base distribution. Others Shemh-Shepherd & Aspuru-Guzik 2023; Gruver et al., 2024; Antunes et al., 2024) also trained or fine-tuned LLMS on text representation of materials and demonstrated the ability of such models to generate valid de- scriptions, sometimes outperforming domain-specific meth- ods.
Ideally, we would like to preserve the group element g used to represent the clean sample fo. since the score network is insensitive to translation. While restricting the velocity fields to the mean-free subspace prevents global drift, it does not always ensure that the noisy coordinates, ft, maintain the mean' of the original sample, fo-
However, we empirically observe that the mean tends to be preserved at moderate noise levels (see Figs. 2 and 3). This helps mitigate the potential mismatch that exists between the translation-invariant score parameterization and the non- invariant training target.
This issue has been pointed out in previous. Lin et al. (2024), for example, proposed a so-called Periodic CoM-free Nois- ing scheme, in which the noise is carefully designed to training target. The resulting
ensure a transition kernel is formulated as a von Mises distribution and estimated numerically via Monte Carlo simulations.
Finally, a recent line of work (Jiao al., 2024; Levy et al., 2024) has sought to exploit space- information to re- strict generation to the smallest asymmetric part of the unit cell only, including for disordered materials (Petersen et al., 2025).
We refer the reader to Appendix D for a longer discussion
IIn this case, we are referring to the mean defined on manifold and not the Euclidean mean. See Appendix D for more precise definition.
6

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
Table 1: Crystal Structure Prediction (CSP) results. Baseline results are extracted from the respective papers. @ indicates the number of samples considered to evaluate the metrics, e.g. @20 indicates the best of 20. Error bars for @1 represent the standard deviation over the mean at sampling time across 20 different seeds. Most notably, KLDM compares favorably to the competing models, achieving performance comparable to better than state-of-ther afrt (SOTA) methods.
PEROV-5
MP-20
MPTS-52
MODEL
MR [%] t
RMSE 4
MR [%]] 1
RMSE +
MR [%] f RMSE
METRICS @
CDVAE
45.31
0.1138
33.90
0.1045
5.34
0.2106
DIFFCSP (PC)
52.02
0.0760
51.49
0.0631
12.19
0.1786
EQUICSP (PC)
52.02
0.0707
57.59
0.0510
14.85
0.1169
FLOWMM
53.15
0.0992
61.39
0.0566
17.54
0.1726
KLDM-E (EM)
53.14 +
0.0758
61.72 =
0.0686
17.71 +..
0.2023
KLDM-E (PC)
52.72 8
0.0678
2. 002
65.37 #11
0.0455
+. 001
21.46 4:2
0.1339
= 002
KLDM-20 (EM)
52.44 +.7
0.0698
5.002
62.92 2:
0.0833
+: 2.002
21.13 #:2
0.1800
.000
KLDM-2o (PC)
52.14
0.0647
65.83
0.0517
23.93
0.1276
METRICS @ 20
CDVAE
88.51
0.0464
66.95
0.1026
20.79
0.2085
DIFFCSP (PC)
98.60
0.0128
77.93
0.0492
34.02
0.1749
FLOWMM
98.60
0.0328
75.81
0.0479
34.05
0.1813
KLDM-E (EM)
99.97
0.0152
83.68
0.0532
39.04
0.1865
KLDM-E (PC)
99.94
0.0226
81.08
0.0440
39.81
0.1462
KLDM-Z0 (EM)
99.89
0.0186
82.94
0.0575
37.77
0.1673
KLDM-20 (PC)
99.92
0.0255
80.18
0.0453
37.10
0.1394
5. Experimental results
5.2. Crystal Structure Prediction (CSP) task
5.1. Settings
Model setup We present two different models that differ in terms of the score parameterization for the lattice parameters I. We consider a model that uses the e-parameterization (KLDM-E) and one using the ao parameterization (KLDM- Øg). For the latter, we also standardize the value of the lattice parameters following (Miller et al., 2024).
Tasks and Datasets We evaluate KLDM the tasks outlined in Section 2.1: Crystal Structure Prediction (CSP) and De-novo Generation (DNG).
We follow previous work (Jiao et al., 2023) and evaluate KLDM across 4 datasets: PEROV-5 (Castelli et al., 2012) including perovskite materials with atoms per unit cell (ABX,), all sharing the same structure but differing in com- position; MP-20 including almost all experimentally stable materials from the Materials Project (Jain et al., 2013), with unit cells containing at most 20 atoms; and MPTS-52 also extracted from the Materials Project (Jain et al., 2013), with unit cells containing up to 52 atoms.
Metrics We follow the evaluation procedure of Xie et al. (2022) to assess the quality of the structures gen- erated by KLDM. We report Match Rate (MR), measuring the proportion of reconstructions from (f l(a) that are satisfactorily close to the ground truth structures as per StructureMat
(Ong et al., 2013); and Root-Mean- Square-Error (RMSE), quantifying the RMSE between coor- dinates of matching reconstructions and ground truth struc- tures.
We describe the exact experimental setting Appendix and release public code repository with our implementa- tion of KLDM.
Baselines We compare KLDM to the following recent gener- ative models: CDVAE (Xie et al., 2022), DIFFCSP (Jiao et al., 2023), EQUICSP (Lin et al., 2024), and FLOWMM (Miller et al., 2024).
Sampling schemes As DIFFCSP (Jiao et al., 2023) and EQUICSP (Lin et al., 2024) both rely on a Predictor- Corrector (PC) integrator (Song et al., 2021), we consider two different integration schemes for KLDM: the first com- bines Euler Maruyaaa (EM) for the lattice parameters with an exponential integrator for the velocities, whereas the sec- ond applies PC scheme to the velocities and EM for the other modalities. We provide detailed pseudo-code for the training and sampling procedures in Appendix H.
Results We perform CSP all datasets and present the corresponding results in Table 1. We note that this task constitutes an ideal test bench for measuring the effect of the novel treatment of the fractional coordinates specific to KLDM. On the simpler PEROV-5 dataset, KLDM performs on par with the compared models for the experiment, and

---

Kinctic Langevin Diffusion for Crystalline Materials Generation
Table 2: De-novo Generation (DNG) results on MP-20. Stability metrics are evaluated using MATTERGEN's pipeline (Zeni et al., 2025) with benchmark, KLDM was trained on the original MP-20 dataset, whereas MATTERGEN MP** and DIFFCSP* were trained on a re-optimized version. We report average and standard deviation across 3 sampling seeds. Notably, KLDMwith analog -it or discrete diffusion outperforms DIFFCSP in terms of RMSD, energy above the hull, and stability, while being slightly subpar on RMSD.
(Yang et al., 2024a). Baseline results are taken from MATTERGEN 's own
RMSD [ +
AVG. ABOVE HULL [eV/atom] +
STABLE [%]1
S.U.N. [%] f
MATTERGEN-MP*
0.147
0.201
47.05
25.76
DIFFCSP*
0.413
0.189
41.25
20.13
KLDM-BU (C)
0.371 4.01
0.269 1.01
38.62 =.1
16.67 #.1
KLDM-Co (C-AB)
0.296 1.01
0.187 1.01
49.84 =.1
17.91 1.1
KLDM-2y (D)
0.283 4.01
0.155 4.01
59.21 =.1
18.52 #.1
yields improved results for @20. Notably, the PC sampler does not enhance sample quality on this dataset. In contrast, on the more realistic datasets MP-20 and MPTS-52, KLDM already yields competitive MR for the EM sampler. The PC sampler further improves performance, improving MR and reducing RMSE. Additional gains are achieved on these datasets by standardizing the lattice parameters and adopting the tto-parammtteizzation for the lattice parameters.
pipeline from MATTERGEN (Zeni et al., 2025) using MatterSim-vi-IM (Yang et al., 2024a). Sample qual- ity is assessed in terms of the following metrics: RMSD between the generated samples and their relaxed structure, where lower values mean generated structures are closer to equilibrium; the average energy above the hull, with lower values meaning that generated materials are closer to thermodynamic (meta-)stabiligy stability, measured as the proportion of samples with an energy above the convex hull below 0.1 e//ato;; and S.U.N. (stable, unique, novel) measuring the percentage of generated samples that satisfy all three criteria, identifying promising candidates. For each variant of KLDM, we generated 10000 samples, from which discarded samples that contained elements not supported by the validation pipeline.
For completeness, we report results on the CARBON-24 24 dataset in Appendix F, Table 4. The @1 task is inherently ill- defined, as the dataset consists exclusively of carbon atoms, allowing multiple distinct structures to satisfy the same composition. On the @20 task, KLDM compares favorably to baselines.
5.3. De-novo Generation (DNG) task
Baselines We compare KLDM to DIFFCSP and MATTERGEN- MP on MP-20. For the baselines, we report numbers bor- rowed from MATTERGEN's own benchmark (Zeni et al., 2025). Notably, the baseline models were trained on a re-optimized version of MP-20 in which certain chemical el- ements were removed, specifically noble gases, radioactive elements, and elements with atomic number greater than 84. Additionally, samples with energy above the hull bigger than 0. eVaatom were also filtered out. In contrast, our model was trained on the original, unfiltered MP-20.
Model setup As described in Section 3.3, the CSP model formulation can be readily extended to the DNG task, by introducing an additional diffusion process operating on the atom types, a. To sample from (x),,, we first sam- ple the number of atoms K in the unit cell from the em- pirical distribution observed in the training set, i.e. from Po(a (ax K) ())pp(K) (K) (Hoogeboom et al., 2022).
For the lattice parameters we adopt the KLDM-00 variant, as it demonstrated the best performance in the CSP task the larger datasets. Similarly, we use the PC sampler as it consistently yielded improved results.
Results The results are summarized in Table 2. Notably, when relying on analog-bits or discrete diffusion to model the atom types, KLDM outperforms DIFFCSP in terms of RMSD, energy above the hull and stability while being slightly subpar on S.U.N..
Regarding the atom types a, we compare three approaches for modeling diffiusion over these discrete attributes: (C) con- tinuous diffusion on one-hot encoded atom types (Jiao et al., 2023), (C-AB) continuous diffusion on analog bits (Chen et al., 2023), and (D) discrete diffusion with absorbing state (Austin et al., 2021; Shi et al., 2024).
Compared to MATTERGEN-MP. KLDM generates structures that are more thermodynamically stable and closer to the convex hull on average, but with slightly higher RMSD, indicating larger deviations from the relaxed geometries. It also produces slightly fewer unique and stable samples, as reflected by lower S.U.N. score. We attribute this
Metrics We evaluate generated samples using a machine- learning interatomic potential, based on the open-source
8

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
60
Match Rate [5b]
40
0,5
Direct
2k
0
20
0.1
Simplifedd
Σ 2
0
Ok
1k
2k
3k 0k
5k
10k
15k
Ok
1k
2k
3k
Epoch
Epoch
Epoch
Figure 1: Ablation study results. We report the mean and standard error across seeds of the validation Match Rate (MR). (Left) Impact of the initial velocity distribution variance. variance of 0 corresponds to the zero-initial velocity case (i.e.. delta distribution). (Center) Effect of the score parameterization on convergence. (Right) Impact of enforcing zero-net translation in velocity fields under zero-initial velocities. Notably, all design choices shown contribute to performance improvements.
remaining performance gap to several factors: (1) the more expressive denoiser architecture used by MATTERGEN-MP. which operates in real space, (2) its uses of PC sampler for
6. Conclusion
We introduced KLDM, a novel diffusion model for periodic crystal structure generation, whose key innovation lies in the modeling of the fractional coordinates through cou- pling with auxiliary variables representing velocities. We evaluated KLDM on both the CSP and DNG tasks, where performed on par with, or outperformed, several state-of- the-art (SOTA) models. Notably, the CSP task, KLDM demonstrated significant improvements on the real-scale datasets MPTS-20 and MPTS-52, at similar computational cost to the compared models. On the DNG task, it matched current SOTA performance as evaluated by machine-learning interatomic potentials. Further validation with DFT simula- tions is natural next step confirm these findings.
the lattice parameters, and (3) the impact of dataset pre- processing in the re-optimized MP-20 used for training.
Additional results using standard proxy metrics are provided in Appendix G.
5.4. Ablation study
In this section, we analyze the impact of three key design choices: (I) the distribution of the initial velocities, (2) the simplified score parameterization in Eq. (19) compared to direct parameterization of Eq. (18), and (3) the enforce- ment of a zero-net translation velocity field.
Despite its strong performance, there are opportunities for further improvement of KLDM. Exploring alternative pro- cesses to noise the velocity variables in Eq. (12) could lead to better results. Additionally, optimizing the score net- work architecture to better account for the periodic nature of crystal structures (Lin et al., 2023) and investigating modality-specifet noise schedules (Qiu et al., 2025) may result in even further gains. We also envision that incorpo- rating space-group information (Jiao et al., 2024) Wyck- off positions (Levy et al., 2024) could improve practical performance, particularly on the DNG task.
We find that using zero-initial velocities significantly im- proves both performance and convergence speed, shown in Fig. (Left). This improvement may also stem from the ability to use the simplified parameterization from Eq. (19), which only compatible with zero zero-initial velocities. As shown in Table 3, this simplified parameterization consis- tently yields better final performance across different sam- pling schemes and datasets, while also accelerating conver- gence (Fig. 1, Center).
Finally, enforcing zero-net translation velocity field pro- vides additional improvements in terms of MR on the valida- tion set, albeit to a smaller extent. These gains are observed both when using the direct parameterization (Fig. 4) and when combining zero-initial velocities with the simplified parametrization (Fig. 1, Right).
We believe that KLDM provides a strong foundation for diffusion- based crystal structure generation. Future work will focus targeted material generation with specific prop- erties, well extending KLDM to larger systems such as metal-organic frameworks (MOFs) by incorporating rota- tional frame modeling (Kim et al., 2025).
We refer the reader to Appendix E for more detailed anal- ysis.

---

Kinetic Langevin Diffusson for Crystalline Materials Generation
Impact statement
Bilodeau, C., Jin, W., Jaakkola, T., Barzilay, R., and Jensen, K. F. Generative models for molecular discovery: Recent ad adnceees and challenges. Wiley Interdisciplinary Reviews: Computational Molecular Science, (2(5)ee6608 2022.
Our paper presents a generative model for novel crystalline materials generation. As mentioned in the introduction, the ability to accurately generate novel compounds with tar- geted properties has significant implications across multiple scientific domains, particularly in molecular discovery and materials design. While advances in materials discovery might benefit society, from more efficient energy storage to improved catalysts, we acknowledge the potential for unin- tended applications. However, the considerable gap between the model prediction the successful material synthesis acts as natural safeguard. Therefore, in the immediate term, our work mostly impacts researchers working the same topic.
Castelli, I. E., Landis, D. D., Thygesen, K. S., Dahl, S., Chorkendorff, I., Jaramillo, T. F., and Jacobsen, K. W. New cubic perovskites for one-and two-photon water splitting using the computational materials repository. En- ergy & Environmental Science, 5(10):9034-904) 2012.
Chen, R. T. T. and Lipman, Y. Flow matching on general geometries. In The Twelfth International Conference on Learning Representations, 2024.
Chen, T., Zhang, R., and Hinton, G. Analog bits: Gen- erating discrete data using diffusion models with self- conditioning. In The Eleventh International Conference on Leurning Representations, 2023.
Acknowledgements
The authors would like to thank the anonymous reviewers for their feedback that helped improve the paper. They would also like to thank Johanna Marie Gegenfurtner for insightful discussions on Lie groups
Cornet, F. R., Bartosh, G., Schmidt, M. N., and Naesseth, C. A. Equivariant neural diffusion for molecule genera- tion. Advances in Neural Information Processing Systems, 37, December 2024.
FC, AB, and JMGL are supported by the Independent Re- search Fund Denmark (DELIGHT, Grant No. 0217-00326B; "TeraBatt: Data-driven quest for TWh scalable Na-ion bat- tery". Grant No. 2035-00232B). FB and JF are supported by the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in Life Science (MLLS, grant No. NNF20OC0062066), MNS acknowledges support from the Novo Nordisk Foundation ("BNNs for molecular dis- covery", Grant No. NNF220C0076658).
Court, C. J., Yildirim, B., Jain, A., and Cole, J. M. 3-d inorganic crystal structure generation and property pre- diction via representation learning. Journal of Chemical Information and Modeling, 60(10):4518-4 -43555151155, 2020.
Davies, D. W., Butler, K. T., Jackson, A. J., Skelton, J. M., Morita, K., and Walsh, A. Smact: Semiconducting ma- terials by analogy and chemical theory. Journal of Open Source Software, 4(38):1361, 2019.
References
De Bortoli, V., Mathieu, E., Hutchinson, M., Thornton, J., Teh, Y. W., and Doucet, A. Riemannian score- bare-based generative modelling. Advances in Neural Information Processing Systems, 35:2406-244 2422, 2022.
Anderson, B. D. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313- 326, 1982.
Anstine, D. M. and Isayev, O. Generative models as an emerging paradigm in the chemical sciences. Journal of the American Chemical Society, 145(16):776--8750 2023.
Dockhorn, T. Vahdat, A., and Kreis, K. Score-basee gener- ative modeling with critically-damped langevin diffusion. In The Tenth International Conference on Learning Rep- resentations, ICLR, 2022.
Antunes, L. M., Butler, K. T., and Grau- Grau-Crespo, R. Crystal structure generation with autoregressive large language modeling. Nature Communications, 15(1):1-16, 2024.
Flam-Shepherd, D. and Aspura-Guzit A. Language models can generate molecules, materials, and protein binding sites directly in three dimensions as xyz, cif, and pdb files. arXiv preprint arXiv:2305.0577 2023.
Arvanitogcörgos, A. An introduction to Lie groups and the geometry of homogeneous spaces, volume 22. American Mathematical Soc., 2003.
Fréchet, M. Les éléments aléatoires de nature quelconque dans espace distancié. In Annales de I' l'institut Henri Poincaré, volume 10, pp. 215-211311 1948.
Austin, J., Johnson, D. D., Ho. J., Tarlow, D., and Van Den Berg, R. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 144:17911-1799) 2021.
Glass, C. W., Oganov, A. R., and Hansen, N. Us- pex evolutiooon crystal structure prediction. Computer physics communications, 175(11-12):713-22 2006.
10

---

Kinctic Langevin Diffusion for Crystalline Materials Generation
Gruver, N., Sriram, A., Madotto, A., Wilson, A. G., Zitnick, C. L.., and Ulissi, 7. W. Fine-tuned language models generate stable inorganic materials text. In The Twelfth International Conference on Learning Representations, 2024.
Kong, L. and Tao, M. Convergence of kinetic langevin monte carlo lie groups. In Agrawal, S. and Roth, A. (eds.), The Thirty Seventh Annual Conference on Learn- ing Theory (COLT), volume 247, pp. 301 1-1633..3. PMLR, 2023.
Ho. J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840-6851, 2020.
Levy, D., Panigrahi, S. S., Kaba, S.-O., Zhu, Q. Galkin, M., Miret, S., and Ravanbakhsh, S. SymmCD: Symmetry- preserving crystal generation with diffusion models. In Al for Accelerated Materials Design NeurlPS 2024, 2024.
Hoffmann, J., Maestrati, L.., Sawada, Y., Tang, J., Sellier, M., and Bengio, Y. Data-driven approach to encod- ing and decoding 3-d crystal structures. arXiv preprint arXiv: 1909.00949 2019.
Lin, P., Chen, P., Jiao, R., Mo, Q., Jianhuan, C., Huang, W., Liu, Y., Huang, D., and Lu, Y. Equivariant diffusion for crystal structure prediction. In Forty-first International Conference on Machine Learning, 2024.
Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d. In International conference machine learning, pp. 8867- 8887. PMLR, 2022.
Lin, Y., Yan, K., Luo, Y., Liu, Y., Qian, X., and Ji, Effi- cient approximations of complete interatomic potentials for crystal property prediction. In International confer- ence on machine learning, pp. 11260-21227, PMLR, 2023.
Jain, A., Ong, S. P., Hautier, G., Chen, W. Richards, W. D., Dacek, S., Cholia, S., Gunter, D., Skinner, D., Ceder, G., et al. Commentary: The materials project: A materials genome approach to accelerating materials innovation.
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023.
APL materials, 1(1), 2013.
Jiao, R., Huang, W., Lin, P., Han, J., Chen, P., Lu, Y., and Liu, Y. Crystal structure prediction by joint equivariant diffusion. Advances in Neural Information Processing Systems, 36, 2023.
Lippe, P., Veeling, B. S., Perdikaris, P., Turner, R. E., and Brandstetter, J. PDE-refiner: Achieving accurate long rollouts with neural PDE solvers. In Thirty-seventh Con- ference on Neural Information Processing Systems, 2023.
Jiao, R., Huang, W., Liu, Y., Zhao, D., and Liu, Y. Space group constrained crystal generation. In The Twelfth International Conference on Learning Representations, 2024.
Lou, A., Xu, M., Farris, A., and Ermon, S. Scaling rieman- nian diffusion models. Advances in Neural Information Processing Systems, 16:80291-80308 2023.
Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T. S. Torsional diffusion for molecular conformer genera- tion. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022.
Lyngby, P. and Thygesen, K. S. Data-driven discovery of 2d materials by deep generative models. npj Computational Materials, 8(1): 2022.
Merchant, A., Batzner, S., Schoenholz, S., Aykol, M., Cheon, G., and Cubuk, E. D. Scaling deep learning for materials discovery. Nature, 624(7990):80 85, 2023.
Joshi, C. K., Bodnar, C., Mathis, S. V., Cohen, T., and Lio, On the expressive power of geometric graph neural networks. In International conference on machine learning, pp. 15330-15355, PMLR, 2023.
Miller, B. K., Chen, R. T., Sriram, A., and Wood, B. M.
Flowmm: Generating materials with riemannian flow matching. In Forty-first International Conference on Machine Learning, 2024.
Kim, N., Kim, S., Kim, M., Park, J., and Ahn, S. Mofflow: Flow matching for structure prediction of metal-organic frameworks. In The Thirteenth International Conference on Learning Representations, ICLR, 2025.
Nouira, A., Sokolovska, N., and Crivello, J.-C. Crys- talgan: learning discover crystallographic structures with generative adversarial networks. arXiv preprint arXiv: 1810. 12033, 2018
Kim, S., Noh, J., Gu, G. H., Aszuku-Guzik, A., and Jung, Y. Generative adversarial networks for crystal structure prediction. ACS central science, 608):1412-1420, 2020.
Kohn, W. and Sham, L. J. Self-consistent equations includ- ing exchange and correlation effects. Physical review, 140(4A): A)AA 1 1133, 1965.
Oganov, A. R., Pickard, C. J., Zhu, Q., and Needs, R. J. Structure prediction drives materials discovery. Nature Reviews Materials, 4(5):331-344 2019.

---

Kinctic Langevin Diffusson for Crystalline Materials Generation
Ong, S. P., Richards, W. D., Jain, A., Hautier, G., Kocher, M., Cholia, S., Gunter, D., Chevrier, V. L., Persson, K. A., and Ceder, G. Python materials genomies (pymatgen); A robust, open-source python library for materials analysis. Computational Materials Science, 68:314-319, 2013.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In Interaaioonal Conference on Learning Representations, 2021.
Sriram, A., Miller, B. K., Chen, R. T. Q., and Wood, B. M. FlowLLM: Flow matching for material generation with large language models as base distributions. In The Thirty- eighth Annual Conference on Neural Information Pro- cessing Systems, 2024.
Petersen, M. H., Zhu, R., Dai, H., Aggarwal, S., Wei, N., Chen, A. P., Bhowmik, A., Lastra, J. M. G., and Hip- palgaonkar, K. Dis-CSP: Disordered crystal structure predictions. In Al for Accelerated Materials Design- ICLR 2025, 2025.
Tao, M. and Ohsawa, T. Variational optimization lie groups, with examples of leading (generalized) eigen- value problems. In International Conference on Artificial Intelligence and Statistics, pp. 4269-4280. PMLR, 2020.
Pickard, C. J. Airss data for carbon at 10gpa and the c+ n+ h+ system at ggpa. Materials Cloud Archive, doi: 10. 00.24433/attrriasscoodd.020
0026/v1, 2020.
Pickard, C. J. and Needs, R. Ab initio random structure searching. Journal of Physics: Condensed Matter, 23(5): 053201, 2011.
Wang. Y., Lv, J., Zhu, L., and Ma, Y. Crystal structure prediction via particle-swarm optimization. Physical Re- view B B-Condensed Matter and Materials Physics, 82(9): 194116,200 2010.
Qiu, K., Song. Y., Fan, Z., Liu, P., Zhang, Z., Zheng, M., Zhou, H., and Ma, W.-Y Piloting structure-based drug optimal schedule. arXiv
Ward, L.,, Agrawal, A., Choudhary, A., and Wolverton, C. general- purposs machine learning framework for predict- ing properties of inorganic materials. npj Computational Materials, 2(1):1-7, 2016.
design via modality-specine preprint ar 25t.:::20500228800 2025.
Ren, Z., Tian, I. P., Noh, J., Oviedo, F., Xing, G., Li, J., Liang. Q.. Zhu, R., Aberle, A. G., Sun, S., al. An invertible crystallographic representation for general in- verse design of inorganic crystals with targeted properties. Matter, 5(1 1::14-33333333 2022.
Xie, T. and Grossman, J. C. Crystal graph convolutional neural networks for accurate and interpretable predic- tion of material properties. Physical review letters, 120 (14): 1000:15501,, 2018.
Rozet, F. and Louppe, G. Score-based data assimilation. In Thirty-seventh Conference on Neural Information Pro- cessing Systems, 2023.
Xie, T., Fu, X., Ganea, O.-E., Barzilay, R., and Jaakkola, T. S. Crystal diffusion variational autoencoder for peri- odic material generation. In International Conference on Learning Representations, 2022.
Särkkä, S. and Solin, A. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Xu, M., Yu, L., Song, Y., Shi, C., Ermon, S., and Tang, J. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
Satorras, V. G., Hoogeboom, E., and Welling, M. E (n) equivariant graph neural networks. In International con- ference on machine learning, pp. 9323-9332. PMLR, 2021.
Xu, M., Powers, A. S., Dror, R. O., Ermon, S., and Leskovee, J. Geometric latent diffusion models for 3d molecule gen- eration. In International Conference on Machine Learn- ing, pp. 38592-38610. PMLR, 2023.
Shi, J., Han, K., Wang, Z., Doucet, A., and Titsias, M. Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 17:103131-10314 2024.
Yang, H., Hu, C., Zhou, Y., Liu, X., Shi, Li, Li, G., Chen, Z., Chen, S., Zeni, C., al. Mattersim: A deep learning atomistic model across elements, temperatures and pressures. arXiv preprint arXiv:2405.0499 044400 2024a.
Shysheya, A., Diaconu, C., Bergamin, F., Perdikaris, P., Herández--Lobato J. M., Turner, R. E., and Mathieu, E. On conditional diffusion models for PDE simulations. In The Thirty-eighth Annual Conference on Neural Informa- tion Processing Systems, 2024.
Yang, S., Cho, K., Merchant, A., Abbeel, P., Schuurmans, D., Mordatch, I., and Cubuk, E. D. Scalable diffusion for materials generation. In The Twelfth International Conference on Learning Representatiinss, 2024b.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International conference on machine learning, pp. 0256-2265, PMLR, 2015.
Zeni, C., Pinsler, R., Zügner, D.,, Fowler, A., Horton, M., Fu, X., Wang, Z.., Shysheya, A., Crabbé, J., Ueda, S., Sordillo,
12

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
R., Sun, L., Smith, J., Nguyen, B., Schulz, H., Lewis, S., Huang, C.-W., Lu, Z.,, Zhou, Y., Yang. H., Hao, H., Li, Yang, C., Li, W., Tomioka, R., and Xie, T. A generative model for inorganic materials design. Nature, Jan 2025. ISSN 1476-4687. doi: 10. 10.1088644166--025-88866633
Zhu, Y., Chen, T., Kong, L.,, Theodorou, E. A., and Tao, M. Trivialized momentum facilitates diffusion generative modeling on lie groups. arXiv preprint arXiv:2405.16381, 2024.
Zimmermann, N. E. and Jain, A. Local structure order parameters and site fingerprints for quantification of ordination environment and crystal structure similarity. RSC advances, 10(10);6033-601,, 2020.
13

---

Kinctic Langevin Diffussin for Crystalline Materials Generation
Organization of the Supplementary Material
Appendix A We provide short informal introduction to Lie groups and manifolds. We provide intuition behind key equations Eqs. (12) and (13), and the TDM and KLDM models more broadly.
Appendix B We discuss why data torus be represented either as angles rotation matrices. We also show that, in this case, the matrix exponential simplifies to rotation matrix, enabling a coordinate update that corresponds to a translation following by a wrapping operation.
Appendix C We derive, step-by-step, the target used to train the score network in KLDM.
Appendix D We analyze how enforcing zero-net translation constraint affects the mean of the system.
Appendix E We provide an additional discussion and analysis of the ablation study.
Appendix F and Appendix G We present additional results for the CSP and DNG tasks, respectively.
Appendix H We present pseudo-code for all the core algorithms in KLDM.
Appendix I Experimental settings required for reproducibility, along with brief discussion of the differences between KLDM and the baseline methods.
A. A primer on Lie groups and manifolds
In this section, we provide an intuitive overview of key concepts used in the main text. For crystalline materials, the fractional coordinates f. which define atomic positions within the unit cell, lie on a hypertorus: i.e. f E [0, 1)5* = T The hypertorus a direct product of Lie groups. In the main paper, we introduced Lie groups as smooth manifolds equipped with group structure G and smooth group operations. Here, we revisit some of the building concepts underlying KLDM in more detail, from Riemannian geometry perspective rather than from the Lie group one.
Lie group In mathematics, group is denoted by (G, -,, with G being the non-empty set of elements that belong to the group and being the group operation that combines any two elements a, b in the group G' and results in an element of G. The group operation has to satisfy three different properties:
(a-b)-c=aa a c) Va, b, c E (associativity property)
For every element a E G, there should be an unique element e E G, such that e a = a and a e = a (existence of an unique identity element)
For every element a E G there is an element b E G such that a b = e and b a = e. (existence of an inverse element)
If the group operation is also commutative, i.c. a b bb b.a. Va, b E G, then the group is called an Abelian group. This property holds in the case of the hypertorus T3xK and allows us to define the transition kernel in Eq. (16) in closed-form (and hence a simulation-fia training procedure). For complete derivation of this, we refer Zhu et al. (2024).
We can formalize two different operations that can be done with elements of the group G that will be useful to give some intuition about why data on manifold can be modeled through quantities defined in Euclidean space. Let us consider two elements a, E G of the group. The left translation of by a is defined as the operation L2- G + G which takes g as input and returns a g, or equivalently + a In the same way, we can define the right translation as R. : G - G. + a.
Tangent space Before describing other key elements of a Lie group, we introduce the concept of tangent space, which stems from the fact that Lie group is also manifold. Possible ways understand the notion of manifold involve thinking about it as a hypersurface embedded in higher-dimensional space or as a collection of points that are somewhat connected to generate a surface. A tangent space is a vector space containing all the vectors that are tangential to a specific point in the manifold, i.e. any element y E G in the case of a Lie group. It is usually denoted as TyM for manifolds and G,G for Lie Groups. The tangent space offers linearized view of the manifold in the neighborhood of the point g and therefore it is usually considered as Euclidean space-i.e Lee. T = Rd, with d being the dimension of g. If we consider all the elements in the group G. the set of all the tangent spaces associated with these elements form the tangent bundle, denoted with T M or TG, depending we are working with a manifold or Lie group.
14

---

Kinctic Langevin Diffusion for Crystalline Materials Generation
Lie algebra, left-invariant vector fields and exponential maps The Lie algebra g is defined the tangent space the identity element of the group, G: TeG. Since tangent spaces Euclidean well.
approximately Euclidean, treat the Lie algebra g as
To understand how to define the dynamics in Eq. (12) and introduced by Zhu et al. (2024), we first introduce left- infarinvar vector fields.
vector field, X tangent space at g.
: G: - TG, is a function that associates every element of the group g E G to a vector X, E T of the
Recall the left-translation operation Lu - G, which maps group element g E G to g G for any group element E G. The differential of this operation, denoted as (dLa)g TgG + TagGG, acts on tangent vectors instead of group elements. A vector field is left- left-invariant if:
(dLa), =Xag: for any two group elements a, E G.
(20)
Intuitively, this means that the vector field behaves predictably under left-translations Specifically, if g is translated to g, the vector Xg E GgG will be mapped to Xa.g. Thus, left-invariant vector looks the same at all points in the group.
In particular, it can be shown that there exists a linear isomorphism between the Lie algebra g and the tangent space of the identity element e E G, i.e. g = T& (Arvanitogeôrgos, 2003). This means that an invariant vector field is determined entirely by its value on the tangent space of the identity element e E G. Using this, we can rewrite Eq. (20) in terms of the identity element e G and its tangent space TeG, which is g:
(dLa)e: : ToG + + Ta. TGG (dLo))(X..) Xave Xa
(21)
where a = a. Eq. (21) shows that any vector X. E T&G can be translated a tangent space X. T&G using (dLa)e(Xe). Now, with both the group element a E G and the tangent vector X, T.G, we can use the Riemannian exponential map' expa T.G + G' to compute the group element a' G reached by starting from with velocity Xa in one unit time. This map extends the Euclidean notion of motion along a straight line to the manifold.
How does this relate to modeling data living on the torus? We can take advantage of the torus being a Lie group and use the operations from the previous section. Specifically, we use the Lie algebra g (a Euclidean vector space) and the operation mapping vectors in to the tangent space of any point of the manifold (via Eq. (21)). In the forward dynamics of KLDM (Eq. (12)), we couple fractional coordinates (elements of our Lie group) with velocity vectors in g. Since the torus is a Lie group, we can map any vector in g the tangent space of any group element. This allows us to define a standard Euclidean diffusion model over these velocities in g and use Eq. (20) to get the corresponding tangential velocities. By applying the exponential map, we can update the fractional coordinates, explaining the coupled process in Eq. (12).
The computation of Eq. (21) for the torus is straightforward, as both the vectors in the Lie algebra g and group elements can be represented as d x d matrices. The differential operation in Eq. (21) simplifies to matrix multiplication (Zhu et al., 2024). Given a vector X, E g and an element a E G, the tangent vector in a is simply X. = aX, E T.G.
Additionally, the Riemannian exponential map, involving constant velocity, has closed-form solution using the matrix exponential. The group element a' reached from a with velocity Xa is given by a' expa(aXe) expx aexpm(X.) where expm(X.) (X.) denotes the matrix exponential.
B. Sampling from the noising and denoising processes
B. Intuition about the correspondence between T and SO(2)
In the following, we provide an intuitive explanation of the isomorphism between T and SO(2). Let us consider a variable x E [a, b), we are going to work on an alternative representation thereof, that we will call a representation in SO(2).
*Note that the exponential map operation can be defined from a Riemannian manifold perspective and from the Lie group one. In the Riemannian case, we have that it is defined as exp, : T.G - G, i.e. it takes velocity living in the tangent space of an element of the manifold (which in this case is also a group) and map it another element of the manifold. The Lie group exponential map instead function exp. : g - that takes a vector in the Lie algebra g T&G and maps it to the group.
15

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
First, we going map x to angle follows,
such that [[ m)
Then, construct the representation, as,
cos 0
sin 0 cos 0
9
sin 0
From we can readily recover 0,
0 sign(go,1) arccos 9o,0.
and in turn x,
1
(b a)-
2T
B.2. Update of positions is equivalent to periodic translation
As have mentioned in the previous section, the update of the coordinate in the forward process be expressed as fi fo expm(r:) where then r. WN( o Therefore, the update involves matrix exponential operation. In the case of the torus, we have that the structure of rr can be written as a 2 x 2 skew-symmetric matrix of the form:
0
r:
Ti E
rt 0
In our case, we are interested computing the matrix exponential of rs. For simplicity, we assume that rt 1, we are interested in computing matrix exponential with the following structure:
([:
1
expm (At) expm
This corresponds to the following infinite sum
expm (At) = expm (L- 1 1
1 0
0 1
0
1
12
0
-13
[2
+
t
0
+
2
0
-2
+
31
0
+
4!
0
....
which can be rewritten as follows:
expm (
1 1)
2
1
13
5
2
4
37
1
31
5!
2
4!
where we can recognize the Maclaurin series for sin(t) and cos(t). Therefore, we can conclude that:
expm ([)
cos(t) sin(t)
sin cos(t)
0 1)
In forward and backward integration the dynamics described by Eq. (7) and Eq. (13) the update of the coordinates involves the computation a matrix exponential. In the following, we are going to derive step-by-stey the closed-form solution for the matrix exponential, highlighting how this will be equivalent to a translation and a wrap of the current position.
16

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
Forward process following update:
In the forward process, assuming discretization step denoted by dl, we interested in computing the
ft Si- fr-iexp((dloo)
Ut 0
([
vrdt]] !ol))
fi-1 expm
dt
= f. expm
o,dt
cos 1zdt
sin utdt] cos v,dl
cos A
sin f
cos n,dt
sin v,dt cos o,dtt
f(-1
sin v,dt
sin 0
cos
sin o,dt
cos 0 cos ipdt - sin f sin o, sin oydl cos - cos v.dl sin
cos O sin 4 + sin A cos vtdt 1
sin 0 sin o,dl + cos 0 cos vzdl
cos(O + vdd() sin(0 + oydt)
sin(O + vydl) cos + v,dt)]
where we used the rotation matrix representation for the fractional coordinates on the torus, as we explained the beginning of this section.
Reverse process
In the backward process, instead, we are interested in computing the following update:
ft ft 11-iexpm(-ddaa
-dtt
0
U/
(1 md
ut )
fi expm
= fi expm
et
0
coos v,dit
sin of cos v,dtt
cos
sin
cos v,dtt sin v,dl
sin o, cos v,dt
fi-1
sin ve
sin
cos
cos 9 cos vidt
+ sin sin d,dt sin oydl cos sin - v,d cos dtd
cos sin 1tdt + sin & cos nf sin 0 sin vill + cos cos vzdl
cos sin
cos(0- v,dt) s0 - v,d
C. Derivation of the target of KLDM
We recall that the transition kernel of our KLDM forward process defined in Eq. (12) can be obtained in closed form and it is given by (Zhu et al., 2024):
Prjo (f. fo. vo) WN(logm(fa"' o2.1) No (w:lMe 02 I),
where the mean and the variance of the two distributions are the following:
1
Mr. 1 - (v: + vo)
Nv. = -
(22)
8
a = 2t +
0 - 2h
(23)
e +1
Intuitively, the normal distribution describes how should noise the initial velocity vo to get sample U2, while the wrapped-normn normal distribution implicitly defines how to get a noisy sample for the fractional coordinates fa starting from ft). (f- and by taking the matrix-exponential on both sides, we can get the following update rule ft = fo expm(r:) where then r; ~ WN(r:lMrr o2,
fo. The trick is to define r, = logm(fo
The target score used in the DSM loss (Eq. (9)) for KLDM is V., log log fc, velfo, vo). In the following, we derive step-by-step. highlighting at the end the parameterization used to train our score network. The target score can therefore be
17

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
computed follows:
Vv, (ff vo) = Vv, v, log [WN( logm(fu fi) pr., 12.1) No (Dille (ool of Vv, log WN (ogm fc) Uro o Vv, log log
(24)
coordinates term
By velocity torm
We start by deriving the score for the second term SR, which can be expressed in multiple way:
pe Ut
8e
(25)
o
dv
We can now focus on the 8c term, where for simplicity we define r, = logm(soo fi) as above:
Se Ve vi log WN logm(ff f() Mrr of Vv log WN I2.T)
+00
exp (- (r' pr. + 2TTk)22)
vv, log
Σ
V2TTrr k=-00
20
to0
(rr Mr, + +2xk)²
vv log
Σ
exp
20
A=-00
to0
(- (rr Mr, kTx
Ve,
Σ exp A=-
(r,- Mr. 2x
20
Dk=-oc exp
20%
t
Vo, exp (- (rr- Mr, + 2TT
C.
20?,
toc
(r: Mr. + 2TT
(rrs Mr. + 2mx))²
CC.
exp
20
20
k=-00
400
(T'- Mr. + 2TT/) 2
Vee (-
(r= 1-e
-
(Ue vo) + kxk)² )
exp
20
20
x
(ri Mr. + 22Tk)²
1- 1
1 - e
C.
exp
20
1 lte of
P/:
et
(ve vv) + 27Tr)
A=
400
c
1- -t
(we + vo) + 2 exp (- (rr (re- pr. + 2
T'
Eter +00
(Ti (re-Mr, 2xk
1 1 -
1 e-
20
k= exp
207
A
ct 1
r pr. + exp (
(r: Mr, + 2xrk2²
Dtseo exp
(ri-Mr 2xk)2
-2nk)2)
lte 02
20
201
h=-00
--
t00
(rr - Mr, + kTT
r; - pr. + 2 exp
Tte-1 Ek=-00 exp
(Mi-14m +2xkj²
o
20
20%
1=-00
--
1+ Vun, log WN(rrlMug o2
Òpr Vp., log WN(Ti(Mr.; I)
(26)
dv
Therefore, the target score is given by summing Eq. (25) and Eq. (26) together. If we now assume that the velocities are
18

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
zeros at time = O, we can notice that the score 8, can be rewritten as follows:
U + pv + e e
(27)
0
d
d
where we used the definition of j4u e and the fact that we assumed that vo 0, resulting in Muu 0.
D. Effect of the zero-net translation velocity field
Mismatch between translation-invariae score parameterization and non-invariant training target As noted by Lin et al. (2024), there is a potential mismatch between the translation-invarian parameterization of the score and the non-invariant training target. For instance, a noisy point cloud and its periodic translation are considered equivalent by the network, but the target scores, as computed in Eq. (18) and used by DIFFCSP (Jiao al., 2023), will differ. While this mismatch is averaged out during training and does not prevent the model from learning useful score (as DIFFCSP performs well on the tasks), it still affects gradient variance during training.
Constraining the velocity field In Section 3.2, we enforce a velocity field with zero net translation by imposing zero-mean velocities ( Uk- = 0). Since our score network is periodic conslation-invariin the goal is to remove any translation component from the dynamics, as the network cannot distinguish between translations.
To illustrate this, consider the simple case of a one-dimensional system with a single atom (i.e. fo is just one coordinate). In this scenario, any noisy ft can be seen as a periodic translation of any other fú and hence even of the clean sample fo itself. Without constraining the velocity field, the dynamics defined in Eq. (12) produce noisy samples f. corresponding to periodic translations of fo. which the network cannot distinguish. By imposing zero-mean velocity field, the velocity coupled with fo must, in this case, be zero, preventing the forward dynamics from introducing an overall translation of the crystal.
Effect of the constraint on forward dynamics To understand the impact of this constraint on the forward dynamics in the general case, we track the mean of the noisy samples generated by using the transition kernel defined in Eq. (16). Since our data resides on a torus, the appropriate generalization of the concept of arithmetic mean to non- non-Euulidean spaces is the Fréchet mean (Fréchet, 1948). The Fréchet mean ensures that the resulting point lies on the manifold but requires solving an optimization problem. Given N datapoints {So, fNN E M on manifold, we define the Karcher means as the set of solutions to:
N d (p. fi),
finen arg min PEM
(28)
where -(;,-) is the geodesic distance on the manifold. When the optimization yields a unique solution, it is the Fréchet mean.
We consider one-dimensional crystal composed of 10 atoms and two different noise levels, namely 0² = 0.1 and x² = 0.7. We sample 1000 points from the transition kernel defined in Eq. (16). For each sample, we check whether the original Fréchet mean is preserved when enforcing zero-net velocity field. As shown in Fig. 2, at low noise levels, the Fréchet mean is generally preserved, indicating that explicitly ensuring zero-net velocity field prevents the translation of the whole system. However, at higher levels of noise, we note that the Fréchet mean jumps between n different values, each of them separated by 2x, Despite this, the noisy samples typically share the same mean as the noiseless reference. In contrast, when no constraint is imposed on the velocity field (Fig. 3), the Fréchet mean is never preserved for both noise levels, and it varies continuously. Our approach, which constrains the velocity field, limits the Fréchet mean to a discrete set of possible values, as opposed to the continuous variation observed without the constraint.
n
The above analysis shows that sampling from the transition kernel during training can introduce translations, meaning the mismatch may not be fully mitigated across all noise levels. However, we argue that this mismatch is significantly reduced at low noise levels when the vector field constraint is enforced, which benefits training (as shown in Appendix E). Enforcing the constraint also turns the issue of shifting the mean of the system into a discrete problem, rather than a continuous one.
19

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
Mcans of transition kemel samples
Means of transition kernel samples for high noise level
for low noise level
Objective system
Menn after noising
50
50
Utjective noiseless system
X
Original mein
Possible
Histogram of means after noising step
Histogram of means after noising step
1000
1000
Possible moans
Distibbtionn
-3
Figure 2: Effect sampling from the transition kernel Eq. (16) with a zero-net translation velocity field.. We consider a ID crystal made of 10 atoms and two different noise levels: low noise (sampled from (0, 0.1,, left) and high noise (sampled from (0, 0.7), right). For each noise level, we generate 1000 noisy samples. low noise, enforcing zero-net translation velocity field preserves the original mean, preventing any system translation. At higher noise levels, while the mean is typically conserved, it occasionally shifts between 10 distinct values, each separated by 2
Means of transition kemel samples
Means of transition kernel samples for high noise level
low noise level
Objective noisad system
Mean noising step
Ohjective moiscless system
Original mean
-3
Histogram means after noising stop
Histogram means after noising step
100
100
Noisy distabution
-3
-2
Figure 3: Effect of sampling from the transition kernel Eq. (16) without a zero-net translation velocity field. We consider 1D crystal made of 10 atoms and two different noise levels: low noise (sampled from (o, left) and high noise (sampled from (0, 0.7), right). For each noise level, generate 1000 noisy samples. We observe that the mean is preserved at either low or high noise levels, changes continuously.
20

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
E. Ablation study
In this section, we analyze the effect of the different design choices our method. We focus on the initial velocity distributions, the score parameterization, and the translation-free velocity fields.
Initial-velocity distributions KLDM uses zero-initial velocities, which intuitively corresponds to considering the coordinates at time ( = 0 being at rest, allowing the noise gradually propagate from the velocity to these variables. This also allows us to use the simplified parameterization for the score formulated in Eq. (19) ablated above. In Fig. 1 (Left), we compare the choice of using zero initial velocities against having a zero-mean Gaussian distribution with three different variance values. We note there is a strong benefit in using zero initial velocities, potentially explained by the simplified parameterization.
Velocity field The last design choice we consider is the enforcement velocity field with a zero-net translation. We analyze the effect of these choices both for initial velocities and therefore simplified score parameterization, and in the case of non-zero initial velocities and direct parameterization. We report results in Fig. (Right). By removing the degree of freedom of modeling overall translations of the system, we observe gain, albeit marginal, in terms of validation set match rate during training in all cases, in particular with non-zero initial velocities, as displayed in Fig. 4.
Match Rate [8b]
0.4
0.3
U 7
Ug
0
U/
0
0.2
k
Uk
0
k 2/ 0
0k
1k
2k
3k
Ok
1k
2k
3k
Ok
1k
2k
3k
Epoch
Epoch
Epoch
Figure 4: Ablation of velocity fields with zero net translation for non-zero initial velocities. (Left) Initial velocities are sampled from (0 0.1 I), i.c. o2 0.1. (Center) Initial velocities are sampled from (0, 0.5 1), i.c. o 0.5 (Right) Initial velocities sampled from W/O0, 1 I), i.e. of 1.0. We report the mean and standard error from the mean of the validation Match Rate over 6 seeds. We note that regarding the initial distribution variance. there is benefit given by enforcing a zero net translation velocity field.
0.5
Score parameterization We compare the simplified parameterization of the score formulated in Eq. (19) compared to the direct parameterization of Eq. (18). We present in Fig. 1, the validation match-rate during training on the MP-20 dataset. From Fig. 1 (Center), we note that the simplified parameterization leads to faster convergence and better performance. The direct parameterization slowly makes the gap smaller if trained for long enough. We also compare fully trained models on MP-20 and MPTS-52 terms of MR and RMSE on the test set. We note Table that no matter the sampling scheme, the simplified parameterization always outperforms the direct one, in addition to being much faster to converge.
Table 3: Impact of the score parameterization the CSP task performance, on the realistic MP-20 and MPTS-52. The simplified parameterization introduced in Eq. (19) improves significantly upon the direct parameterization from Eq. (18).
MP-20
MPTS-52
sampler
MR [%] t
RMSE +
MR [%] f
RMSE
DIRECT PARAMETERIZATION
EM
56.25
0.1071
11.55 42
0.2535
#.4 63.29
2.001
t.004
PC
1.2
0.0591
1.002
15.85
1.2
0.1666
1.004
SIMPLIFIED PARAMETERIZATION
EM
61.72
0.0686
17.71 43
0.2023
1.2
2.001
+ 4.005
PC
65.37 1.1
0.0455
21.46 1.2
0.1339
1.001
21

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
F. Additional CSP results
Previous work also evaluated their approach on CARBON-24 (Pickard, 2020) dataset, which consists carbon-only materials with 24 atoms per unit cell. Since all structures are composed solely of carbon atoms, computing the MR on the is inherently ill-defined-.ce multiple distinct structures may satisfy specific formula.
For completeness, we report results in Table 4, where KLDM performs comparably to existing methods on the @ @1 task experiment, but outperforms them on @20. Similar to observations on the PEROV-5 dataset, we find that the PC sampler does not improve sample quality over the standard EM sampler in this setting.
Table 4: Crystal Structure Prediction (CSP) task results on the CARBON-24 dataset. Baseline results are extracted from the respective papers. @ indicates the number of samples considered to evaluate the metrics, e.g. @20 indicates the best of 20. Error bars for @ 1 represent the standard deviation over the mean at sampling time across 20 different seeds. For this dataset, the CSP@1 task is ill-defined due to its one-to one-to-many ta-y nature, and we only report the obtained values for completeness.
CARBON-24
MODEL
MR [%] f RMSE
METRICS
CDVAE
17.09
0.2969
DIFECSP (PC)
17.54
0.2759
EQUICSP (PC)
FLOWMM
23.47
0.4122
KLDM-E (EM)
18.04 +.8
0.3188
#.008
KLDM-E (PC)
17.26
0.2827
+.006
KLDM-20 (EM)
18.58 1.8
0.3278
1:0 000
KLDM-Bo (PC)
17.69 1.8
0.2915
0000
METRICS 20
COVAR
88.37
0.2286
DIFFCSP (PC)
88.47
0.2192
FLOWMM
84.15
0.3301
KLDM-E (EM)
90.19
0.2154
KLDM-E (PC)
85.86
0.1988
KLDM-Zo (EM)
89.66
0.2198
KLDM-Eo (PC)
84.58
0.2011
G. Additional DNG results
In this section, we present additional results regarding the DNG task. These results were computed using the KLDM-E model trained on the non-standardized values of the lattice parameters and implemented using the e-parameterization for their score. We evaluate samples in terms of validity, coverage, and property statistics. A sample is deemed structurally valid if all pairwise distances are above 0.5A, while SMACT (Davies et al., 2019) is used to determine compositional validity, by checking the overall charge neutrality. The coverage metrics are computed using fingerprints: CrystalNM structural fingerprints (Zimmermann Jain, 2020) and Magpie compositional fingerprints (Ward et al., 2016). COV-R (recall) and COV-P (precision) are obtained by comparing the distances between generated and test fingerprints. Finally, the property statistics are obtained by comparing distributions of properties, computed on the generated samples and respectively: atomic densities d, and number of unique elements datem-
We provide results in Table 5, where it can be observed that, unlike for the CSP task, the difference in performance between the compared methods is less pronounced. Nevertheless, KLDM remains competitive with the other models.
22

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
Table Results for the De-novo Generation (DNG) task. Baseline results are extracted from the respective papers.
VALIDITY [%]1
COVERAGE [%]
PROPERTY 1
STRUC. COMP.
COV-R COV-P
dp delem
PEROV-5
CDVAE
100.0
98.59
99.45
98.46
0.1258
0.0628
DIFFCSP
100.0
98.85
99.74
98.27
0.1110
0.0128
EQUICSP
100.0
98.60
99.60
98.76
0.1110
0.0503
KLDM-S
99.97
98.83
99.60
98.61
0.2970
0.0478
CARBON-22
CDVAE
100.0
99.80
83.08
0.1407
DIFFCSP
100.0
99.90
97.27
0,0805
EQUICSP
100.0
99.75
97.12
0.0734
KLDM-6
100.0
99.90
98.86
0.0658
MP-20
CDVAE
100.0
86.70
99.15
99.49
0.6875
1,432
DIFFCSP
100.0
83.25
99.71
99.76
0.3502
0.3398
EQUICSP
99.97
82.20
99.65
99.68
0.1300
0.3978
FLOWMM
96.85
83.19
99.49
99.58
0.2390
0.0830
KLDM-S
99.88
84.86
98.94
99.50
0.4658
0.1280
23

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
H. Algorithms
This section presents all the algorithms at the core of our method. We show how one can sample the noisy inputs and the training targets Algorithm 1, we then present the training loop used to train the parameters of the score network se (fr, ut, l,, az) in Algorithm 2. We then focus on sampling, presenting both sampling scheme that uses an exponential integrator for simulating the reverse SDE of the dynamics of the velocities Vy and fractional coordinates fi proposed by Zhu et al. (2024) while using classic Euler-Maruyana step for lattice parameters I, and atom type in Algorithm 3. In algorithm Algorithm 4, instead, we present how we sample from our model using the predictor-corrector steps as proposed in (Song et al., 2021).
Algorithm training target:s(f, v, a, l): Routine for sampling fi, Ut, 4, 0z from the transition kernels and the corresponding target scores
Require: task (either CSP or DNG), timestep t, scheduler a(t) and (t) for I and a, scheduler a.,(t) and (t()) for u, scheduler Mr. (t, vo, vy) and or. (t, vo, vt) for r. Initial sample z0 (fo, aa) and initial velocities vo. our experiments considered to 0 (i.e. initial velocities are 0). ## Sampling and ft.
Sample (0, I), 6r, ~ (0.
>. No is a normal distribution such that Σ, 0.
vr = 0z (1) + (4) 6v if initial velocities are zero then target, (1)
else
target, (t)
See Eq. (26).
end if
w(rr (1, vo, vz) + or. (1, (o, er.)
indicates the wrap function.
ft = wfo + +rr)
ft = center (fr)
b center(-) keeps the center of gravity fixed. D Equivalent computation of Eq. (26).
target, (1 exp(-l)))((C exp(-1)) Vr(oiN
target, target,, + target,
# Sampling 4
Sample e -N(0.)) I)
4 = ( (l)
target, =
if task DNG then
## Sampling d;
Sample <a ~ ~N(011) I) a,
- =a)aa + (f) ta
target, ca
return (Ut, fcr Le, az), (target,, target, target,)
Return both noisy samples and training targets.
else
return (ve, fr, Lt), (target,, target,)
end
24

---

Kinctic Langevin Diffusion for Crystalline Materials Generation
Algorithm 2 Training algorithm
Require: score network (t, fi, ve. 4, a,), fractional coordinates transition kernel Pejo(fe. 8; fo. vo). lattice parameters transition kernel (r)o(4e(o)) empirical distribution g(x)) & Σi oze and distribution over initial velocities po(v). In our experiments we considered p(vv) = 5(vo) (i.e. initial velocities are O). Losses weights X dt, 7a We always used X, and \ 1. For DNG task we require also an atom type transition kernel for training iterations do
2ig = [[foo,4 lo, am)) ~ q(x)) 1~ U(I), vn ~ po(v)
indicates the batch size.
if task is DNG then
(ve, fú, L,, az), (target,,, target,, t traiining u, a),
outs, out, = (t, fú, Vz, Le, ax) & The network takes I, fú, Ut. l input and output all the scores.
La = louut target,
else
(v,, fc. 4)), (target,, target,) =traininggtargggessfff v, N,Noe, t)
out, out s0(1, fr, (1,
end if
L |lout; target, 113 out (1 exp(-t))/(L + exp(-t)) outy v1/02
Construct the score, Eq. (19) D X(t) computed as Jiao et al. (2023)
f. lout, target, 11
if task is DNG then
Lior = XuCo + NeCi XaCa
Depending on how we model the atom types, da has different values
else
Lit >ooo + NeL
end if
Compute gradients of Liot with respect to 0 and perform gradient step. end for
25

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
Algorithm 3 Sampling algorithm
Require: trained score network st, ft, ur, Lt, az), discretization steps, di = 1/N step-size, prior distributions over velocities (r( (0, I). over fractional coordinates pr = ((0 1), over lattice parameters pr No I). and over atom types (r(a) (), I). We require also the knowledge of the forward drift f()) and the diffusion coefficient g(4) of the SDEs describing the evolution of L and
## Note: in the paper we use 0 as index for samples at / 0. However, here it will be slightly different notation.
Sample from the prior ~ No. (0, I), fo 4(0, 1), to ~ NN I) Set fu = w(fu) if task DNG then Sample ag (O,
b w indicates the wrap function.
end if
for = 1, NN do
if task is DNG then
out (n-1) out (u (= outa (n-1) so((1 - (n 1) dt), fn-1, Un 1h--1,n-11 1n-1,,,
an-1)
else
out !n out out (n-1)
= so((1 (n - 1) * df), fn- 1n in
end if
## Update step for and f
out, (1 exp(-(1 - (n 1)dtt))) exp(e exp((111 (n- ()dtt)) out, v./o Follow Eq. (19)
Sample Eu ~ (0,1)
> No is normal distribution such that Σ, 0.
n, = -exp(d)) 2(exp(2dl) ((exp//(211((
Dour () + exxp(2dl) 16p
b Update on Update on f
fn. -w(f--1 v,,df)
## Update step for I
Sample 4 I)
In = 4--1- (J(1) g² (()s(om((=1))))) + vdlas
b EM step for t
task is DNG then
## Update step for a
Sample Ea I)
O, an-1 (f(t) g ("))(oou"""")))) + Vdtea
> EM step for
end end for
if task DNG then
return A crystalline material sample (fN, aN)
else
return A crystalline material sample (fN,Ix)
end if'
26

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
Algorithm Sampling with a single Predictor-Correctoor step (PC) similar to Rozet & Louppe (2023, Algorithm 4)
Require:
trained score network st,t, ft. Vt, lt, az), N discretization steps, dt = 1/N step-size, prior distributions over velocities pr(v) = NO, (0, I). over fractional coordinates pr = (0, 1), over lattice parameters Pr(1) = (o I), and over atom types pr pr(a) (a) = ON .). We require the scheduler a (t) and o (t) for v and also the knowledge of the forward drift (t) and the diffusion coefficient g(t) of the SDEs describing the evolution of I and a. We also require hyperparameter
0.
## Note: in the paper use index for samples t 0. However, here it will be slightly different notation. Sample from the prior vo (0, I), fo U(0, 1), I0,
First steps similar Algorithm 3 b w indicates the wrap function.
Set fo = w(fu) if task DNG then Sample ag (0,
end
for
I.,,,..,N do if task is DNG then
out (n (n-1) out (n-1) outt (m-1) so((1 (n 1) dt), fn-1, Un 1n-1, In- an
else
out,
(n-1) out (n-1) out(n-)) = s0((1 - (n - 1) * d(), fn- Un 1n-,, an-1)
end
# # Update step for v and f
# Prediction step on velocities v and coordinates f
Compute a. (n)) (n 1)
Compute (ro (n
o (n))o.(n
On r + cout 1n
pred
f = w(fn + padd _ di)
# ## Correction step on velocities v and coordinates f
task is DNG then out,, out,, outa = s0((1 * df), fhred ufd ,n- an-1)
else
out out, s0((1 ct), f w.fc
pred 4n-1,an-1)
end if
out, (1 - exp(-(p111 n
n dt))) out, ve/o
>> Follow Eq. (19)
Compute 8 T dim (out..
Tout
Sample Go (0, I)
No is normal distribution such that Σ, v2 0.
U = pred + Sout, + V2cc. fn w( pprod dt)
>> Update on Update on
# # Update step for I
Sample & (),
in = 4--1 (f(t) 9" (f)s(out)))dt + ddter
P. EM step for /
if task is DNG then
## Update step for a
Sample Ea ~ NO, I)
an an-1 (f(() - 9 (1))(((o)))) + Vdlea
D EM step for
end if end for
if task is DNG then
return A crystalline material sample (fx aN)
else
return A crystalline material sample (fN,IN)
end
27

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
I. Experimental details
I.1. Hardware
All experiments presented in this paper can be performed on a single GPU. We relied on a GPU cluster with a mix of RTX 3090 and RTX A5000, with 24GB of memory.
1.2. Architecture
As mentioned in Section 3.2, our score network parameterized using an architecture whose backbone is similar to that of previous work (Jiao et al., 2023, DIFFCSP) which enforces the periodic translation invariant by featurizing pairwise fractional coordinate differences with periodic functions of different frequencies. There are two main differences compared the architecture of DIFFCSP: as represents the velocity, we need the network to also take this velocity variable vn additional inputs. The network is then outputting the score for all the diffusion processes involved in the model: the score related to the velocities, the one for the lattice parameters, and the one for the atom types. The additional modification we do is to consider a two-layer network instead of a single layer to predict the score related to the velocities. Regarding the network parameters, we considered 4 peasaag--assingg layers for PEROV-5, while we increased them to 6 for the remaining three datasets. In all the experiments, we considered the hidden dimension to be 512, the time embedding to be a 256-dimensional vector and we used s activation with layer norm. While the presentation in the paper is done in terms of continuous time diffusion models, the implementation is done in discrete time to guarantee an apples-to-apples comparison with basclines such DIFFCSP and EQUICSP.
coupling the fractional coordinates with additional auxiliary variable that
DIFFCSP employs a graph-neural network as a score network that adapts EGNN from Satorras et al. (2021) to fractional coordinates. the following, we are going to present all the components that form this architecture.
Lattice parameters pre-processing We follow the same pre-processing steps for the lattice parameters (both lengths and angles) used by Lin et al. (2024, EQUICSP). Lengths are usually defined from [0, +oo) while angles are defined in the (0, T) interval. However. The diffusion process defined in Eq. (7) operates in the (-o0, +o0) domain, and therefore can result in unreasonable lattice parameters. Therefore, we use a logarithmic transformation for the lengths, mapping them from (0, +00) to to (-00, +oo). For angles, we map them using the following operation tan( - π x/2) from (0, x) to (-00, +00).
Components of EGNN Let consider ft, Ur L, a being the input of our network. The input features are computed by
NN(Satom: fpon(t))
where fatom; Spos are the atomic embedding and sinusoidal positional embedding and NN is an MLP.
Then the input features are processed by series of 8 message-passing layers that compute
m
(8)
4m
(n) (5-1) (s-1) ,,,,, SimusoilaFEmbeedtingg. f:))
.
(s)
Σ
(s)
m
mi
ij
j=1
h(t) = h (1-1
+ pn (o-1) (*))
(8) where miy and h (a-1) represent the messages layer between nodes and j. Pm and yo two MLPs. Compared to the DIFFCSP implementation, we want to highlight that we are also passing the velocity as input.
The SimusoidalFimbeeddingg is sinusoidal embedding layer defined as
SinusoidalEEmbeddinggg xx)
(2in(2nr::)) cos (2 k
with being Mirey being an hyper-parameter.
28

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
After S steps of message passing, we compute all the different scores by doing the following:
Pe Po(h)
(h!
N
81 = Pi
h
(8)
N
im1
8 = pa (a(h (5))
where we want to stress that Po is a 2-layer neural network while fo and Pa are single-layer MLPs.
For training, we used the same loss weights that were used by DIFFSCP. We consider X, = 1 and A = 1 for the CSP task. For the DNG task, instead, as we consider three different ways for modelling the discrete atom type features, we used different weights depending on the modelling choice. we use one-hot encoding for the atom types, we still rely on the DIFFSCP weights given by A, = L, D L, and A 20. In the case of analog- bits, we used d = while when using discrete diffusion for the atom types, scale the losses using Xa 0.33. The scaling is needed make the different loss terms have similar magnitudes.
I.3. Parameters for the KLDM forward process
We kept the drift coefficient (l) constant at 1 in all the experiments presented in the paper following Zhu al. (2024). In their experiments, they also tuned the time horizon T of the process in Eq. (12) depending experiments for material generation, we kept the time horizon constant at T = 2. In terms of the implementation, as we implemented everything in discrete time, we discretize the time in the interval [0, 2) for the diffusion process. For lattice parameters and atom types, we rely on the standard Euclidean diffusion model where the drift and diffusion coefficients are defined by a linear schedule on the interval [0, 1). We trained all the networks using Acamii with the default PyTorch parameters, without gradient clipping and by performing early stopping based on metries computed on a subset of the validation set: match-rate for the CSP task and valid structures for the DNG task.
the considered task. In our
Table 6: Dataset hyperparameters
INFO
PEROV-5
CARRON-24
MP-20
MPTS-52
Max Atoms
5
24
20
52
Total Number of Samples
18928
10153
45231
40476
Batch Size
1024
256
256
256
L.4. Major difference between KLDM and competitors
In this section, we compare the key differences between our proposed KLDM and the bascline methods discussed in Section 5, namely DIFFCSP, EQUICSP, and FLOWMM. Among these, we think that DIFFCSP is the closest to our approach. The main differences are that we model fractional coordinates differently and we use linear noise schedule, as opposed to their cosine noise schedule. In addition that, DIFFCSP employs matrix representation for the lattice parameters, while treat them as a vector of six scalars.
EQUICSP builds upon DIFFCSP with two main modifications: it introduces additional losses to ensure the lattice permutation invariance of the learned distributions, and it defines a different noising mechanism called Periodic CoM-free Noising scheme. This scheme ensures that the sampled noise does not induce a translation of the center of mass of fo. thereby preserving the periodic translation invariance in the target score. contrast, we define the forward process on the coupling given by fractional coordinates and associated velocity variables, and we do not account for lattice permutation invariance. leaving that direction for future work.
FL.WWMM generalizes Riemannian Flow matching for material generation, which, although closely related to diffusion, is a different modeling approach. In addition to that, it considers an informed prior distribution over the lattice parameters, and in the DNG task, it represents atom types using analog bits (Chen et al., 2023), in contrast to DIFFCSP and EQUICSP, which use a one-hot encoding. We present results by considering the same informed prior and also by using three different representations for the discrete atom type features. Additionally, for DNG, FLOWMM provides extra inputs to the score
29

---

Kinetic Langevin Diffusion for Crystalline Materials Generation
network that neither KLDM nor the other baselines use. The additional input represents the cosine of the angles between the Cartesian edges between atoms and three lattice vectors
30


Reviews:
[{"review_id": "Up376zTzUX", "reviewer": "Reviewer_MrGy", "summary": "This paper proposes a diffusion model tailored for crystalline material generation. It utilizes the specific manifold structure of the data, and applies the framework of Trivialized Diffusion model, which is a diffusion model that works on Lie groups. This framework avoids doing Riemannian diffusion by taking the tangent space and defining the noising process on the velocity, which lies in an Euclidean space, largely simplifies the computation. It demonstrates empirical performance on structure prediction and de novo generation tasks, with comparable performance with existing methods.\n## Update after rebuttal\nThank you for adding these empirical results, comparison and explanations. I have raised my score accordingly.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive consideration and suggestions for improving the paper. We address their questions and comments below.\n\n**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.\n\n**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](https://anonymous.4open.science/r/rebuttal_icml_kldm-36FF/simplified_vs_direct_parametrization.pdf)), where the *simplified* parameterization is shown to converge significantly faster and to higher values than the *direct* one.  \n\n**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](https://anonymous.4open.science/r/rebuttal_icml_kldm-36FF/init_velocity_when_zero_cog.pdf)). 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$. \n\n\n**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. \n\n\n**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. \n\nIntuitively, this can be thought of as the difference between gradient descent ($\\sim$DiffCSP) and gradient descent with momentum ($\\sim$KLDM).\n\nRegarding 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.\n\n**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](https://anonymous.4open.science/r/rebuttal_icml_kldm-36FF/dng_additional_results.pdf). Given the timeline and the available resources, the evaluation is performed using a machine-learning interatomic potential, based on the open-source [MatterGen pipeline](https://github.com/microsoft/mattergen). \n\n\nFor 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..\n\nWe however note that that the [compared DiffCSP and MatterGen-MP](https://github.com/microsoft/mattergen/tree/main/benchmark/metrics) 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.\n\nRegarding 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.\n\n\n#### **References**\n[1] De Bortoli, Valentin, et al. \"Riemannian score-based generative modelling.\" NeurIPS 2022", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The methods that specifically design diffusion process for the coordinate parametrization of the crystalline data structures makes sense.\n\nFor the structure prediction task, it compares RMSE with ground truth and Match Rate. For RMSE computation, I wonder if it considered the symmetry of the coordinates as described in section 2.1.\n\nThe Metric for de novo generation makes sense and aligns with literature.", "theoretical_claims": "I checked the main ideas, the transition kernels and objectives, and they make sense. I did not look into the details of the derivation in the appendix.", "experimental_designs_or_analyses": "The experimental designs are sound. The structure prediction and de novo generation make sense. The ablation study shows the simplified parameterization improves the accuracy of the prediction. The paper also mentions “the simplified parameterization” leads to faster convergence, but I did not see quantitative results supporting this.", "relation_to_broader_scientific_literature": "How are the key contributions of the paper related to the broader scientific literature? Be specific in terms of prior related findings/results/ideas/etc.\n\nThis paper mainly uses the Trivialized Diffusion Model, which enables simpler training of diffusion model for data with a Lie group structure. It provides an interesting direction of designing the diffusion process specific to the algebraic and geometric structure of crystals. It has application in structure prediction and crystal generation.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths: The idea of designing the diffusion process and score-matching objective specific to the crystal problem is novel, and the application of the trivialized diffusion model for this data with the group structure is interesting.\n\nWeaknesses: Analysis and result of complexity, convergence, are missing, which would support the benefit of this approach over existing methods. Especially given the fact that it does not outperform them for de novo generation tasks.", "other_comments_or_suggestions": "For de novo generation, the performance is not as good as existing methods. Maybe a future direction would be adding some guidance of those desirable structures.", "questions_for_authors": "Can you provide a complexity analysis and comparison with the existing methods, especially the ones using Riemannian Diffusion models on manifolds? Is the matrix exponential step slow to compute, or are they simplified with the trigonometric functions? Compared to existing methods (DIFFCSP), are there less parameters?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RRd2tmU23s", "reviewer": "Reviewer_xaaD", "summary": "This paper proposes a new diffusion model for modeling crystalline materials.\n\nThe model is built upon a Kinetic Langevin Diffusion on the fractional coordinates, and standard Euclidean diffusion for the lattice vector and atom types (one-hot embedded).\n\nThe core contributions of the paper are:\n- proposing to use a velocity noising process in the fractional coordinate diffusion to make the noising process itself invariant to fractional translations.\n- proposing a simplified score parameterisation for the combined model.\n- application of the model to standard benchmark tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n\n\n**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]. \n\nTo further support this, we ablate the effect of the zero net translation in the next paragraph.\n\n\n**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](https://anonymous.4open.science/r/rebuttal_icml_kldm-36FF/zero_net_translation.pdf)), where we obtain (slightly) better results by enforcing zero net translation. \n\nWe also present an analysis about the impact of non-zero initial velocities for different variances (see [Figure](https://anonymous.4open.science/r/rebuttal_icml_kldm-36FF/init_velocity_when_zero_cog.pdf) and [Figure](https://anonymous.4open.science/r/rebuttal_icml_kldm-36FF/v0_not_zero_ablation.pdf)). 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.\n\n\n\n**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. \n\n\n**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.\n\n\n**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. \nWe 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.\n\n\n**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. \n\n**References**\n\n[1] Lin, Peijia, et al. \"Equivariant diffusion for crystal structure prediction.\" ICML 2024.", "claims_and_evidence": "For the most part the claims are well supported.\n\nThe claim I have most issue with is the discussion around issue in the subsection *Score parametrization and targets*.\n\nIt is not clear to me that the issue is that the conditional target scores can be different for different translations of the same $f_t$. Is this not an expected result of score matching? The point of the denoising score matching loss is to minimise the average square error over the conditional scores to give you the score function. Perhaps I have misunderstood the issue?\n\nThe solution proposed still seems interesting to me - but in that it reduces the complexity of the learnt function by quotienting out an additional symmetry of the model, namely by the noising process.\n\nThis feature is not ablated in the experimental results, although the simplification in the parameterisation of the score function is, and I think it would be quite important to show that this procedure does indeed help with better model performance.", "methods_and_evaluation_criteria": "The benchmarks are in line with prior work in the area, and appear to be sufficient.", "theoretical_claims": "The pieces of analysis in the paper, such as the loss derivation, are correct. There are no other claims made.", "experimental_designs_or_analyses": "Overall the design is sound. \n\nI have one small nitpick and that is that for some of the experiments there are error bars computed, and for others there are not. Could the authors explain why?\n\nAdditionally, in the De Novo generation task the majority of the methods appear to be very close together in performance for the majority of metrics. Could the authors comment on which of these metrics is most important, and why there is little gain on this task compared to the tasks presented in table 1.", "relation_to_broader_scientific_literature": "The paper builds upon other work in the crystal structure generation literature, and is compared well to other baseline methods such as CDVAE, DIFFCSP, EQUICSP, FLOWMM. The paper is most related to DIFFSCP, where it replaces the diffusion on the fractional coordinates with the Kinetic Langevin Diffusion.", "essential_references_not_discussed": "None to my knowladge.", "other_strengths_and_weaknesses": "I appreciate the value in the combination of previous ideas presented here, and the innovations in the modelling process regarding the score function parameterisation and noising process. The results in Table 1 tasks suggest that the new model does perform better than competitors in some settings.\n\nThere are quite a few typos in the paper:\n- 248R differ -> defer.\n- 407R does not make sense.\nFor example.", "other_comments_or_suggestions": "None", "questions_for_authors": "Other than the questions posed in other boxes, could the Authors discuss why they think this method appears to be working better than previous methods? Could they discuss why this does not appear to be the case for the denovo task?\n\nCould the Authors discuss if they see impact for this work and the modeling developments outside the application area of crystal structure generation?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cz5Sd3zdpE", "reviewer": "Reviewer_WGaL", "summary": "The paper presents Kinetic Langevin Diffusion for Materials (KLDM), a groundbreaking diffusion model designed for generating crystalline materials. KLDM tackles the challenge of modeling fractional coordinates on a hypertorus by introducing auxiliary Euclidean velocity variables,  eliminating the need for approximations inherent in Riemannian diffusion and ensures consistent training objectives. The model is tested on two key tasks—Crystal Structure Prediction (CSP) and De-novo Generation (DNG)—and achieves competitive results compared to state-of-the-art models, especially on large datasets such as MP-20 and MPTS-52.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive consideration and suggestions to improve the paper. We address questions and comments below.\n\n\n**De-Novo generation task results** Due to the limited character, we have to refer to the answer provided to reviewer **MrGy** about this topic.\n\n\n**Zero-net translation intuition**\nWe 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.\n\nConsider 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. \n\nWith 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. \n\n\nBy 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. \n\n\n\n**Ablation of design choices**  \nWhile 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](https://anonymous.4open.science/r/rebuttal_icml_kldm-36FF/simplified_vs_direct_parametrization.pdf)). We show that non-zero initial velocities (see [Figure](https://anonymous.4open.science/r/rebuttal_icml_kldm-36FF/init_velocity_when_zero_cog.pdf)) systematically lead to subpar performance. \n\nAs suggested by the reviewer, we provide an ablation of the effect of the zero net translation for zero initial velocities (see [Figure](https://anonymous.4open.science/r/rebuttal_icml_kldm-36FF/zero_net_translation.pdf)) and non-zero ones (see [Figure](https://anonymous.4open.science/r/rebuttal_icml_kldm-36FF/v0_not_zero_ablation.pdf)). By removing this unecesseray degree of freedom, we observe a benefit in all cases, in particular with non-zero initial velocities.", "claims_and_evidence": "The claims in this paper are supported by clear evidence.", "methods_and_evaluation_criteria": "This article provides a fairly detailed explanation of the methodology, and the evaluation is also quite reasonable.", "theoretical_claims": "The theoretical claims presented in this paper are well-founded.", "experimental_designs_or_analyses": "This paper conducts experiments on the CSP and DNG benchmarks, comparing with the mainstream models. The results serve as evidence supporting the effectiveness of the method.", "relation_to_broader_scientific_literature": "This paper focuses on addressing the issue of inconsistent training objectives in crystal generation tasks, a problem that had not been adequately resolved in previous works such as DiffCSP [A] and EquiCSP [B]. Compared to these models, this paper introduces the Kinetic Langevin Diffusion process, inspired by TDM [C]. By incorporating an  auxiliary velocity $v$, the modeling of fractional coordinates is simplified, eliminating the need to focus on Riemannian manifolds. Compared to TDM [3], this paper extends the method to the crystal generation task, which requires considering additional symmetries.\n\n[A] Jiao, Rui, et al. \"Crystal structure prediction by joint equivariant diffusion.\" NeurIPS 2023.\n\n[B] Lin, Peijia, et al. \"Equivariant diffusion for crystal structure prediction.\" ICML 2024.\n\n[C] Zhu, Yuchen, et al. \"Trivialized Momentum Facilitates Diffusion Generative Modeling on Lie Groups.\" ICLR 2025/", "essential_references_not_discussed": "The references in this paper are sufficient, requiring no further supplementation.", "other_strengths_and_weaknesses": "Strengths: \n1. This paper tackles the challenge of inconsistent training objectives in crystal generation tasks. \n2.  The paper aligns both datasets and metrics with prior models, achieving state-of-the-art (SOTA) results in CSP tasks and comparable results in DNG tasks.\n\nWeaknesses:\n1. Technically, the diffusion framework they employ is derived from TDM [C], while the backbone model is based on DiffCSP [A]. Although the paper tackles a key issue and achieves good results, the level of technical innovation appears to be somewhat limited.", "other_comments_or_suggestions": "See \"Questions for Authors\".", "questions_for_authors": "1. Why can a zero-mean $v$ ensure that $f_0$ and $f_t$ share the same group element $g$? Please provide more intuition and explanation.\n2. The core operation of this paper is the introduction of the zero-mean $v$, which ensures the consistency of the target score function $s_\\theta$. Why does the paper not provide an ablation study for this operation? I believe such an experiment could highlight the key contribution of the work.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper addresses the interesting problem of crystalline material generation using the state of the art tools of diffusion modelling. While the reviewers had some initial doubts, the rebuttal addressed all their concerns. The properly conducted evaluations indeed show the advantage of the presented approach. All reviewers recommend acceptance and the AC agrees. AC kindly asks that the paper reflects all discussions provided in the rebuttal.
Decision: Accept (poster)