CrystalGRW: generative modeling of crystal structures with targeted crystallographic properties via geodesic random walks

Abstract

Determining whether a candidate crystalline material is thermodynamically stable depends on identifying its true ground-state structure, a central challenge in computational materials science. We introduce CrystalGRW, a diffusion-based generative model on Riemannian manifolds that proposes candidate crystal configurations in stable phases, validated through density functional theory calculations. Our model is designed for de novo generation, which creates crystal structures together with their compositions. The crystal properties, such as fractional coordinates, atomic types, and lattice matrices, are represented on suitable Riemannian manifolds, ensuring that new predictions generated through the diffusion process preserve the periodicity of crystal structures. We also incorporate an equivariant graph neural network to account for rotational and translational symmetries within the model. CrystalGRW generates crystal structures that are stable and closely resemble their density-functional-theory ground states. The model also supports conditional control, such as enforcing a specified crystallographic point group, thereby accelerating materials discovery and inverse design by providing symmetry-consistent, energetically stable candidate crystals for experimental validation.

Introduction

Crystal structures, defined as the periodic arrangement of atoms in a lattice, directly influence material properties such as stability, band gap, and mechanical strength. Understanding and predicting the behavior of materials from the atomistic level relies on knowing accurate structures. The search for novel materials involves identifying crystal structures and their compositions that are thermodynamically metastable. However, this task is complicated by the vast number of degrees of freedom that must be explored to identify local minima on the free energy surface^1,2.

Many advanced algorithms for sampling configurations on the potential energy surface have been developed for ground state structure prediction; these approaches typically combine stochastic sampling techniques such as evolutionary algorithms, ab-initio random structure searching, simulated annealing, metadynamics, basin hopping, and molecular dynamics simulations^2–8, with quantum-mechanics-based method, such as density functional theory (DFT) to accurately predict the energy of the sampled configurations. However, the cubic scaling with respect to the number of basis functions in DFT leads to a substantial computational time for searching for new candidate materials. Machine-learned interatomic potentials (MLIPs) are an accurate and efficient alternative to DFT^9–17, and therefore can speed up the prediction process. Nonetheless, improving the structure generation task remains an important challenge for materials discovery.

Instead of the traditional approach of starting from an initial structure and performing sampling to find new (possibly lower energy) structures, generative models enable a direct and much more efficient way to produce novel molecular and crystal structures because they can sample new materials given knowledge from the learned data distribution without being constrained by energy barriers^18–31. Diffusion-based generative models draw principles from non-equilibrium thermodynamics to transform random noise into structured data^32–36. These models generate samples through stochastic differential equations (SDEs) that consist of data-dependent drift and diffusion terms^37–42. The Brownian motion component introduces stochastic noise, enabling the generation of new crystal structures. Recent work has shown that properly designed diffusion models can generate crystal structures with realistic characteristics, such as nonoverlapping atoms, reasonable lattice parameters, and imposed crystallographic symmetries^22,23.

Moreover, materials discovery focuses not only on generating new structures but also on their properties. For subsequent property prediction, the candidate structures produced from generative approaches tend to closely resemble the DFT-optimized structures (similar to performing a structure relaxation with MLIPs), thereby reducing the time required for subsequent DFT relaxations. Additionally, generative models can be applied for inverse design, where the generated materials are tailored to possess specific, controlled properties^31,43–47.

To specify crystal structures, three properties are essential: fractional coordinates, atomic types, and lattice matrices. The first two properties reside in distinct non-Euclidean spaces, while the third property lies in Euclidean space. In particular, fractional coordinates lie on the 3D torus (Supplementary Appendix A 3), and atomic-type compositions lie on the d-dimensional simplex , with d ∼ 100 in many inorganic materials (Supplementary Appendix A 4). Introducing naive stochastic noise can corrupt these properties by moving them outside their natural domains. This issue can be circumvented by introducing noise distributions that align with their domains, such as wrapped normal distributions for fractional coordinates and categorical distributions for atomic types^48,49. Models like DiffCSP²⁴ and MatterGen³¹ employ these domain-aligned approaches to generate crystal properties, thereby improving crystal structure generation.

Instead of selecting specific noise distributions that align with the natural domains of the properties, Bortoli et al.⁵⁰ introduced the Riemannian Score-Based Generative Model (RSGM), which extends diffusion models to operate on Riemannian manifolds. In RSGM, data undergo random walks along geodesic paths on manifolds that appropriately represent the data’s natural domain, rather than diffusing strictly within Euclidean space. This framework enables the use of manifolds that accommodate non-Euclidean crystal properties, such as a 3D torus for fractional coordinates and d-simplex for atomic types.

Recent advances have also extended flow-based generative models to Riemannian manifolds. Riemannian flow matching⁵¹ unifies diffusion and flow models within Riemannian manifolds, a framework employed by FlowMM²⁵ to generate crystals via deterministic ODE integration. In contrast, CrystalGRW employs stochastic geodesic random walks (SDEs), where the natural stochastic dynamics facilitate the exploration of complex likelihood (analogous to Langevin dynamics), distinct from FlowMM’s deterministic paths. Furthermore, we develop an exact hypercube bijection for atomic types rather than approximate mappings (see Supplementary Table A1 for a detailed methodological comparison).

In this work, we introduce CrystalGRW, a generative model designed for crystal structure generation with specifiable target crystallographic properties. CrystalGRW combines the Riemannian score-based generative model (RSGM) with EquiformerV2¹⁴, an equivariant graph neural network with an attention mechanism. The model generates three key properties of crystals: atomic coordinates, atomic types, and lattice parameters. We demonstrate that CrystalGRW can produce realistic crystal structures that are close to their ground states with accuracy comparable to existing models. Moreover, our model incorporates conditional control to specify properties of the generated structures using classifier-free guidance⁵². To illustrate this conditional control capability, we train CrystalGRW to generate structures with specific point groups as dictated by the input conditions, highlighting its potential for targeted materials discovery.

Results and discussion

CrystalGRW

The crystal generative model aims to generate crystal properties of N atoms in a unit cell where are fractional coordinates of each atom, are their atomic types, and is a lattice matrix. Similar to diffusion probabilistic models, we generate new data by reversing a diffusion process, wherein trained neural networks map noise distributions to structured data distributions. In this work, new crystals emerge through geodesic random walks on manifolds that capture each property: a 3D torus for fractional coordinates , a d-dimensional simplex for atomic-type probabilities , and Euclidean space for the lattice matrix . See “Riemannian score-based generative model” section and Supplementary Appendix A for details.

Although atomic-type vectors naturally reside in , directly performing random walks on the simplex can be cumbersome due to boundary conditions of its geometry. To circumvent this, we adopt a uniform-spacing bijection⁵³ that maps to , d-dimensional hypercube. We then carry out the geodesic random walk on , where reflections at the boundary are straightforward (see Supplementary Eq. (A11) and Supplementary Appendix A 4), and finally map back to . This approach preserves the simplicity of a hypercube-based walk while respecting the simplex geometry of atomic types. Similar treatments for confined systems have also been studied by considering diffusion with reflecting and other constrained boundaries^54–56.

For denoising crystal properties in reverse diffusion, we employ EquiformerV2¹⁴, an equivariant graph neural network (GNN) that accounts for rotational and translational equivariance to effectively capture tensorial and vectorial quantities, e.g. atomic forces^13,57–59. Specifically, EquiformerV2 uses l-degree irreducible representations (irreps) of the group⁶⁰. The fractional coordinate denoiser uses an l=1 output head that produces a 3D equivariant vector in Euclidean space, which is then mapped onto the 3D torus, while atomic type and lattice matrix denoisers come from l=0 output heads providing 100- and 9-dimensional invariant features, respectively. Further details are given in Supplementary Appendix B. The hyperparameters for EquiformerV2 and CrystalGRW are provided in Supplementary Table H1 in Supplementary Appendix H.

Fig. 1 illustrates the training and sampling schemes of CrystalGRW. During training, crystal properties are corrupted using the geodesic random walk (GRW), as described in Algorithm 1. A distorted structure is then constructed from the corrupted properties and passed through the graph neural network to learn the score function. For crystal structure generation, the three crystal properties are initially sampled from a random distribution and gradually denoised through a reverse GRW process, computed using the learned scores. When using the condition-guided scheme, specific conditions are embedded alongside the node features during both the training and generation processes.

Fig. 1.

Overview of the CrystalGRW workflow. (a) Training scheme. Starting from pristine crystal properties , which include fractional coordinates , atomic types , and a lattice matrix , the model applies noise via forward geodesic random walks on their respective manifolds, generating corrupted crystal properties at a time step t. Here N is the number of atoms in a unit cell. Then, an equivariant graph neural network (GNN) processes these corrupted properties to predict manifold-specific scores (tangent vectors on each manifold), which are used to compute the training loss . (b) Sampling scheme. Beginning from noisy crystal properties , the trained GNN computes the trained scores associated with these inputs. Guided by these scores, a reverse geodesic random walk iteratively denoises until reaching , which is the final crystal structure sample. Specific crystal-structure conditions can be added to the network during training and used to guide attributes (e.g., crystal symmetry) of the generated structures through a classifier-free guidance scheme. The 2D torus, 2-simplex (triangle), and 3D Euclidean space manifolds schematically depict the respective manifolds for fractional coordinates, atomic types, and the lattice matrix, serving only as a conceptual visualization rather than an actual representation of these higher-dimensional domains.

Generating realistic structures

We train and evaluate the model on the ALEX-MP-20 dataset, which contains 675,204 crystal structures³¹. Each structure contains up to 20 atoms per unit cell, includes up to 10 different elements per compound, and involves only elements with atomic numbers no greater than 83. The dataset is a union of materials from the Alexandria materials database^61,62 and Materials Project database⁶³ where all structures have an energy above the convex hull () less than 100 meV/atom. It is worth noting that this is a relatively loose criterion, and stricter rules may be applied to avoid the possibility of false positives⁶⁴. We generated 10 structures where the number of atoms for each structure is sampled from the same distribution as in the ALEX-MP-20 training set (see Supplementary Fig. D1a in Supplementary Appendix D). CrystalGRW starts with random atomic positions, atomic types, and lattice matrices, and then iteratively denoises them through GRW for K steps until they approach candidate crystal structures using Supplementary Algorithm C2.

Two criteria are used to evaluate whether generated structures are considered to be valid: 1) all interatomic distances are at least 0.5 Å (i.e., structural validity); 2) the total charge of the atoms is neutral²² (i.e., compositional validity). Although these criteria are relatively loose, as noted in Ref. ²², they provide a useful sanity check for the model’s outcomes. For the structures generated by CrystalGRW, all are structurally valid, and 90.74% of generated structures are compositionally valid, surpassing the performance of other models (see Table 1).

Table 1.

Structural and compositional validity for 10 generated structures using the models trained on the MP-20 dataset (and the ALEX-MP-20 dataset for indicated models). RMSD is averaged over 10 structures that were sampled from the generation set and relaxed using DFT.

Dataset	Model	Structural (%) ↑	Composition (%) ↑	Avg. RMSD (Å) ↓
MP-20	FTCP	100.00	100.00	1.492
	CDVAE	100.00	86.70	0.359
	DP-CDVAE23	99.59	85.44
	KLDM (D)30			0.283
	DiffCSP	100.00	83.25	0.232
	FlowMM (500 steps)25	96.86	83.24
	MatterGen	100.00	84.50	0.110
	EquiCSP29			0.0510
	DiT-S-Transformer27	99.58	90.46	0.0502
	DiT-S-EqV227	99.26	86.09	0.0296
	OMatG67			0.0741
	GCFM67			0.0288
	CrystalGRW (ours)	100.00	85.40	0.0520
ALEX-MP-20	DiffCSP			0.104
	MatterGen	100.00	88.45	0.021
	OMatG67			0.1260
	GCFM67			0.0243
	CrystalGRW (ours)	100.00	90.74	0.0053

^a https://github.com/microsoft/mattergen

We further identify the symmetry of the generated structures using SpacegroupAnalyzer from pymatgen with symmetry criteria of and ^65,66. Fig. 2 compares the relative frequency of each point group in the ALEX-MP-20 training set and the generated structures. CrystalGRW is able to generate both high-symmetry groups, such as , 4/mmm, 6/mmm, , and , and lower-symmetry groups, like 2/m and mm2, in proportions similar to the ALEX-MP-20 dataset. However, it produces low-symmetry structures excessively, such as m, while marginally generating structures in some symmetries, like and groups with a 3-fold principal axis (trigonal lattice). Nonetheless, this result indicates that generated structures possess crystallographic symmetry upon creation.

Fig. 2.

Relative frequency of 32 crystallographic point groups appearing in the ALEX-MP-20 training (blue) and generated structures (pink). The point groups are sorted from low order (left) to high order (right) symmetry.

Generating stable, unique, and novel structures

To present the performance of the generative models for crystalline materials, the generated structures should be stable, unique, and novel (S.U.N.)³¹. Firstly, in thermodynamics, materials evolve to the global minimum on the free energy landscape; however, they could be metastable by falling into local minima instead. In terms of energy, materials under certain temperature and pressure are thermodynamically stable if they lie on the convex hull. Nevertheless, the stability of materials can be preliminarily evaluated by . In this work, materials are regarded as stable if their values are below 100 meV/atom after the DFT relaxations at 0 K and 0 GPa³¹. It is worth noting that this is a relatively loose criterion; in practice, the synthesizability of materials must also account for dynamical and mechanical stability⁶⁸. Secondly, uniqueness refers to a structure being distinct from other structures in the generation set. Lastly, the novelty is an evaluation of the number of generated structures that do not match the training set (see Supplementary Appendix G).

We evaluate novelty by comparing the structures in the generation set with the ALEX-MP dataset and uniqueness by comparing the structures within the generation set (details in Supplementary Appendix G). Fig. 3(a) shows that approximately 48% of the generated structures (from to samples) are novel. Additionally, around 90% of structures are unique, with this percentage decreasing to 76% for samples. Consequently, the percentage of generated structures that are both unique and novel is 44.6% for samples, which drops to 36.9% for samples. Although 117,652 disordered structures from ICSD database are not used for novelty evaluation, the unique and novel performance is lower than MatterGen which is 61% for 10 generated structures. This may be because CrystalGRW tends to generate highly symmetric structures, as illustrated in Fig. 2. Such symmetry can lead to degeneracy, potentially reducing the uniqueness of the generated structures.

Fig. 3.

(a) Percentage of generated structures that are unique (green), novel (blue), and novel and unique (red) with respect to the number of samples. (b) Distribution of for unique and novel structures computed using the GGA/GGA+U mixing scheme. is referenced with the convex hull of materials in ALEX-MP dataset.

Moreover, we compute the for 1436 structures using DFT calculation (see “DFT setting” section). Fig. 3b illustrates a distribution of the where 95.0% of the relaxed structures have lower than 100 meV/atom. We find that the S.U.N. rate is 37.72%. In contrast, MatterGen reported 38.57% for the S.U.N. rate, where novelty was determined by comparing the generated structures with the combined Alex-MP-ICSD dataset, which includes 107.7K disordered structures requiring a license for access³¹. Although our uniqueness and novelty are diminished by symmetrical degeneracy, most of our generated structures are stable, achieving a S.U.N. rate comparable to MatterGen’s. We present examples of S.U.N. structures–metal alloys, carbides, oxides, and hydrides–in Supplementary Table G1 and Supplementary Fig. G4 in Supplementary Appendix G.

Generating DFT ground state structures

One of the primary goals of generative models is to generate crystal structures that are close to DFT ground states. Ground state structures are in equilibrium under thermodynamic conditions where the total force acting on every atom should be 0 eV/Å. From the DFT calculation, we extracted the forces acting on atoms in the generated structures and averaged the magnitude of the force over all atoms. As shown in Fig. 4(a), 18.69% and 28.96% of the samples have forces and eV/Å, respectively. DFT relaxation is typically considered converged when the norm of all the forces acting on every atom in a unit cell is eV/Å (see “DFT setting” section), indicating that about one fifth of the structures were already close to their DFT ground states upon generation. Machine learning force fields^9–12 enable fast and accurate force computations, so we also employ them to calculate the forces of the generated structures, as illustrated in Supplementary Fig. G3, whose results exhibit similarly to the DFT-calculated forces.

Fig. 4.

(a) Distribution of DFT force magnitudes (in eV/Å), averaged over all atoms for each generated structure. (b) RMSD between generated and DFT-relaxed structures.

Another metric to determine the degree of structural similarity between generated and DFT-relaxed structures is the root mean square distance (RMSD), which measures the average atomic displacement between two structures. It can be evaluated using StructureMatcher from pymatgen. Table 1 presents the average RMSD over structures is 0.052 ± 0.094 and 0.0053 ± 0.0173 Å for the models trained on MP-20 and ALEX-MP-20 datasets, respectively. Fig. 4(b) illustrates the RMSD of structures generated by the model trained on ALEX-MP-20 dataset, grouped into two categories. For unique and novel structures, the average RMSD is 0.011 Å with a few outliers above 0.1 Å. Supplementary Table G2 lists structures with RMSD >0.1 Å, most of which are oxides with 14-20 atoms per unit cell and low-symmetry space group, i.e., P1, Cm, and C2/m. For non-novel structures, the average RMSD is 0.0014 Å. This demonstrates that CrystalGRW can generate both non-novel and novel structures that closely match their DFT relaxations, outperforming other models in this task.

Generating structures with targeted properties

One advantage of the generative models is their ability to perform inverse design, generating structures with specific targeted properties. To achieve this, we implement the classifier-free guidance method (CFG)⁵², which enhances our model to sample crystal structures on given conditions. Among the key structural properties, crystallographic symmetry plays a crucial role in determining the properties of crystal structures. For instance, piezoelectricity is absent when crystals have a center of inversion. Recent work has controlled symmetry by constraining the generations with Wyckoff positions, creating structures with a particular space group^67,69,70. In contrast, we demonstrate that our model can control the 32 crystallographic point groups of the output structures through label embedding using the CFG method. Since the number of point groups in the dataset is unbalanced, we create 7 labels, consisting of , , for three rotational axes, , , for horizontal, vertical, and diagonal mirror planes, respectively, and i for an inversion, to represent each point group⁷¹. Supplementary Figs. G1b-c show that the training set is more balanced when labeled by these 7 labels. (Details on the implementation of conditional control within the graph neural network are provided in Supplementary Appendix D.)

We generate 128 structures for each point group, specifying 2 and 4 atoms per unit cell and guidance strength w=0.5⁵². Each row of the heat maps in Fig. 5 presents the output point groups of the generated structures guided (conditioned) by the input point groups. Notably, entries below the diagonal indicate cases where the output symmetry is lower than the input symmetry, while entries above the diagonal indicate higher output symmetry relative to the input. For both 2 and 4 atoms per unit cell, the generated structures match the input point groups for frequently occurring symmetries in the dataset, such as , , 6/mmm, , mmm, and 2/m (as shown in Fig. 2). Additionally, mm2 and 3m also well match for 4-atom structures. However, for less common point groups such as , , , , and , the generated structures appear to have higher symmetry that share the same principal rotation axis such as , 6/mmm, , , and 2/m, respectively.

Fig. 5.

Heat maps illustrate the input conditions, represented as the guiding point groups of output structures (vertical axis), versus the actual point groups of the generated structures (horizontal axis). The generated structures are constrained to contain 2 atoms per unit cell (top) and 4 atoms per unit cell (bottom). Black diagonal blocks indicate cases where the input and output point groups match, while orange diagonal blocks highlight instances where the input and output point groups have the same symmetry order. The color bar denotes the number of structures, ranging from 0 to 128.

Supplementary Table G3 reports success rates (recalls) of the output structures whose point group symmetry matches the input point group, alongside the percentage of occurrence in the training set for 2 and 4 atoms per unit cell. The results show that point groups with low success rates, such as 4, 23, and , rarely appear or are absent from the training set, whereas point groups with high success rates tend to be those that appear more frequently. Micro F1 scores for each point group are also computed, where true positives are defined as matches between generated and input point groups; false negatives occur when the point group does not match the input; and false positives occur when structures controlled for other point groups appear as the considered point group. Therefore, the point group yields a low F1 score despite its high success rate, as structures tend to be even when they are controlled to be other point groups (especially from absent point groups). On average, the macro F1 scores are 0.130 and 0.216 for 2 and 4 atoms per unit cell, respectively. The weighted F1 scores are 0.343 and 0.376 for 2 and 4 atoms per unit cell, respectively, where the weights correspond to the ratios of each point group’s occurrence in the training set. Supplementary Tables G4 and G5 present examples of structures with 2 and 4 atoms per unit cell, respectively, where the output space group reflects the symmetry imposed by the input point group.

For comparison, Supplementary Fig. G2 shows 1056 generated structures for each of the 2 and 4 atoms per unit cell without point-group conditioning. For the 2-atom structures, most point groups are , with a small number of and 6/mmm. For the 4-atom structures, the dominant point groups are 4/mmm, , and , with only a few instances of 4mm, mmm, and . These results show that several point groups do not appear or rarely appear without conditioning. In contrast, Supplementary Tables G4 and G5 show high success rates for 2/m and , for example, which cannot be freely generated with an unconditional model, indicating that the conditioning scheme is effective. Nevertheless, the performance in controlling point groups drops when the number of atoms per unit cell is large (see Supplementary Fig. G1) where the generated structures tend to have a lower symmetry than the specified point group. An exception is the 7 atoms per unit cell, which achieves an F1 score of 0.408 and can generate structures with the 422 point group at a 57% match rate. (422 is a non-centrosymmetric point group that is difficult to generate with other numbers of atoms due to its scarcity in the training set.)

Conclusions

In summary, we introduced CrystalGRW, a generative model leveraging diffusion on Riemannian manifolds and EquiformerV2 to produce candidate thermodynamically stable crystal structures with targeted properties, illustrated by specified crystallographic point groups. A key methodological distinction of our framework is the use of SDEs through geodesic random walks in specific manifolds (constrained domains), allowing atomic positions, atomic types, and lattices to be treated under a unified theoretical framework. This contrasts with approaches that rely on diffusion using specific noise distributions (e.g., wrapped Gaussian or categorical distributions) or on deterministic flow matching. By employing a rigorous framework together with the highly equivariant GNN, CrystalGRW generates physically valid structures with high symmetry that are close to DFT ground states, with below 100 meV/atom in the majority of structures.

Moreover, the crystallographic point groups of the generated structures can be steered through conditional controls, enabling targeted inverse design. However, the efficiency of point-group control is limited to a small number of atoms per unit cell, since lower symmetry tends to emerge when the number of atoms is larger than 10 atoms per unit cell. Another limitation of our work is that it is currently more effective for categorical conditions (e.g. point groups), while improving the efficiency of continuous property control remains a direction for future work.

Methods

Riemannian score-based generative model

We consider data , where is a Riemannian manifold that may be compact, such as a torus (for fractional coordinates) or a finite-volume hypercube (for atomic types), or non-compact Euclidean space (for lattice matrices). Our goal is to learn a generative model for through a score-based diffusion process defined on . In the continuous-time limit, the forward noising mechanism can be viewed as an SDE on whose reverse-time counterpart depends on the score . In particular, one performs a forward (noising) diffusion phase that starts with real data and evolves it over [η ,T] until it resembles a tractable noise distribution at t=T where . Then, one applies a reverse (sampling) diffusion phase, beginning from that noise distribution at t=T and integrating backward to yield new samples at t=0.

To implement these phases in discrete time, we employ a geodesic random walk (GRW) that updates along geodesics in , adding drift and noise within the tangent space.

Forward geodesic random walk (noising) Let γ = T/K be the step size, with so that t = kγ. The forward update is

1

where is an i.i.d. Gaussian random vector in a basis of the tangent space orthonormal with respect to the Riemannian metric, and the noise schedule specifies how the noise amplitude evolves with time (details in Supplementary Appendix E). The exponential map displaces by the tangent vector . A crucial choice is the manifold-dependent drift .

On compact manifolds, such as tori or closed hypercubes, we typically set , which is sufficient to make the random walk become ergodic yielding the late-time noise distribution that approaches the uniform distribution.

On non-compact manifolds, the late-time density will converge to a Riemannian normal distribution⁷² by choosing with the potential , where is the geodesic distance and is a mean location. The forward noising process then follows overdamped Langevin dynamics that eventually converges to the stationary density that is proportional to . This Riemannian normal distribution reduces to the standard Gaussian distribution on when the geodesic distance becomes the usual Euclidean norm. In practice, we further simplify by ignoring and setting directly (as in a variance-exploding approach³⁵) to ensure a wide final distribution. For this no-drift approximation, see Supplementary Appendix F b.

After K iterations of Eq. (1), the state approximates a tractable noise distribution: uniform on compact manifolds, or Gaussian-like on non-compact ones.

Reverse geodesic random walk (sampling) According to Föllmer’s approach^38,39, reversing the forward process yields

2

where is the score of the intermediate distribution . Note that the sign of switches in reverse time. We next discuss how to approximate with a learnable score function of each manifold .

Training via Riemannian score matching To train , we adopt a Riemannian score-matching approach, such as Varadhan’s asymptotic^50,73, which gives a manifold-specific loss function

3

Here, one uniformly samples a random time , applies Eq. (1) up to that time t, and compares to

4

the latter being the manifold-specific inverse exponential map from to , whose analytical expressions for relevant manifolds are given in Supplementary Appendix A.

We set rather than 0 for three key reasons. First, it prevents division by zero in the term and helps avoid pathologically large gradients arising from Varadhan’s asymptotic near t=0. Second, it ensures a nontrivial denoising task at the earliest time t=η (instead of an identity mapping at t=0). Lastly, it promotes stable training signals and faster convergence, as the model never sees fully uncorrupted data in its noised phase.

Variation on adaptive timesteps The RSGM algorithm⁵⁰ suggests that each diffusion step can be split into several smaller random-walk increments to enhance numerical stability. However, in our experiments, we find that training the model by performing only a single (final) random-walk update per diffusion step is both more computationally efficient and converges faster. Moreover, in reverse diffusion, generated crystal structures become more realistic by using an adaptive timestep () for sampling like in previous studies such as in⁷⁴ and⁷⁵. See the details in Algorithm 3 and Supplementary Appendix F.

Implementation in CrystalGRW

In our CrystalGRW implementation, we let reside on three distinct manifolds , reflecting the natural domains of fractional coordinates, atomic types, and lattice parameters, respectively. We define a manifold-specific drift that vanishes on the compact manifolds but equals for the Euclidean one, and train an equivariant graph neural network (EquiformerV2¹⁴) which provides separate outputs for each manifold slice. To learn the score function jointly across all manifolds, we minimize the total loss

where each term follows the Riemannian score-matching objective in Eq. (3) and each is a hyperparameter that weights the contribution of the loss on each manifold. However, for the finite hypercube manifold describing atomic types, we find it both more efficient and more accurate to train an estimator network in place of a direct score predictor; details of this atomic-type denoising approach appear in Supplementary Appendix C. Supplementary Appendix B details the equivariant network architecture, and Algorithm 2 provides the step-by-step training procedure.

After completing the training procedure, we sample from the trained generative model by first drawing a final noisy state from either a uniform prior (if the manifold is compact) or a Gaussian prior (if it is Euclidean). We then run the reverse geodesic random walk in Eq. (2), using the learned score , to generate a new sample . The complete procedure is detailed in Algorithm 3.

Algorithm modifications A key challenge lies in generating atomic types. As discussed in Supplementary Appendix C, our experiments revealed that using Algorithms 2 and 3 to generate atomic types, particularly for perovskite compositions, resulted in suboptimal performance. To address this, we adapted the process by training atomic types against their ground truths, as outlined in Supplementary Algorithm C1, and sampling them using Supplementary Algorithm C2. Detailed discussions and methodologies are provided in Supplementary Appendix C. All reported results in Figs. 2-4 are based on the modified Supplementary Algorithms C1 and C2.

Condition guidance When conditional sampling is desired (e.g., for controlling point groups or compositions), we apply classifier-free guidance⁵² in our manifold-aware pipeline. Let be the score for a condition , and let be the unconditional score (i.e., with a null condition). We construct a condition-guided score by combining a conditional prediction with an unconditional one as

5

where w is the guidance strength. During training, we randomly replace a fraction (e.g., 10%) of the conditional samples with a null condition ∅, obtained by zeroing out the condition vector. Thus, if the batch size is 10, about 1 sample per batch is trained with ∅ while the others include the true condition . In sampling process, we interpolate between and according to w using Eq. (5), thereby controlling how strongly the generation is influenced by . This procedure naturally extends to our manifold setting by applying in the reverse diffusion step Eq. (2), and ensures consistent conditional guidance across all manifold components.

Algorithm 1.

Forward Geodesic Random Walk (Noising) for Multi-Manifold States

Algorithm 2.

Training Multi-Manifold Crystal Properties

Algorithm 3.

Reverse Geodesic Random Walk (Sampling) for Multi-Manifold States

Crystal graphs

Here we describe how and are integrated into inputs of graph neural networks. Since GNNs incorporate inductive biases via graph structures⁷⁶, we represent each crystal as a multigraph . The node set includes both atom-centric features and positions, while the edge set captures interactions between pairs of atoms. In this multigraph setting, the same pair of nodes can be connected by multiple edges if different crystallographic directions link those atoms. Concretely, we define

6

where is a node attribute (e.g. atomic number derived from the atomic type ), is the Cartesian position of atom i, is a translation vector, and N is the number of atoms in a unit cell. The lattice matrix maps fractional coordinates to . Although some previous works represented in ^25,31, we find that diffusing all nine elements of yields similar performance. For a detailed discussion of crystal graphs defined by Eq. (6), see Ref.⁷⁷.

For each node, we concatenate its node attributes (e.g. atomic numbers) with a time step t to reflect the diffusion stage. Empirically, appending the corrupted lattice variables to the node attributes further improves model performance. For the conditional control, we similarly append condition vector embeddings (e.g. point groups) to these node features (see Fig. 1), thereby allowing the GNN to handle a condition alongside the atomic attributes.

DFT setting

We evaluated energies and performed structural relaxations using density functional theory (DFT)^78,79, adopting the projector augmented wave (PAW) method^80,81 as implemented in the Vienna Ab Initio Simulation Package (VASP)^82,83. Our calculations followed the MPRelaxSet, a Materials Project workflow⁸⁴, which includes exchange-correlation functionals and sets a plane-wave energy cutoff of 520 eV. The energy convergence threshold is set to  eV/atom, so EDIFF is  eV, and the force convergence threshold is ten times EDIFF, in the unit of eV/Å. Brillouin zone sampling was achieved via Monkhorst-Pack k-point grids⁸⁵, maintaining a k-point density of 1000 k-points reciprocal per atom for different unit-cell sizes. We used the Perdew–Burke–Ernzerhof (PBE) exchange-correlation functional for the generalized gradient approximation (GGA)⁸⁶ and the GGA with Hubbard U corrections (GGA+U)^87,88, where the spin polarization was enabled for the latter method. To evaluate the thermodynamic stability⁸⁹, we computed the using the GGA/GGA+U mixing scheme and constructed phase diagrams with PhaseDiagram tool from pymatgen^90,91. We found that unary compounds such as Mo, Pd, Ir, and Pt have significantly lower than those in the Materials Project database, despite following the MPRelaxSet. Therefore, we excluded structures containing these elements to avoid artificially low in the stability evaluation.

Supplementary Information

Supplementary Information. ^{(4.3MB, pdf)}

Author contributions

K.T., T.P., and T.C. designed the methodology. K.T., T.P., and N.C. developed the software. K.T., T.P., and C.A. performed the validation and formal analysis. A.E., B.A., C.S., T.B., and T.C. provided resources. T.B. supervised the project. K.T., T.P., and T.C. wrote the original draft. All authors contributed to the review and editing of the manuscript.

Funding

K.T. and T.B. acknowledge funding support from Thailand Science Research and Innovation Fund Chulalongkorn University (IND_FF_69_220_2300_049). T.C. acknowledges the funding support from the NSRF via the Program Management Unit for Human Resources & Institutional Development, Research and Innovation (Grant No. B39G690073) and from Thailand Science research and Innovation Fund Chulalongkorn University (IND_FF_69_ 258_2300_062). B.A. acknowledges the financial support from the Swedish Government Strategic Research Area in Materials Science on Functional Materials at Linköping University (Faculty Grant SFOMatLiU No. 2009-00971). T.P. and C.S. acknowledge financial support by the U.S. Army Research Office (ARO) and the U.S. Air Force Office of Scientific Research (AFOSR) under Grant Number FA9550-23-1-0505 "Unraveling the Role of Cation Solvation in Aqueous Zn-Ion Batteries: A Combined Theoretical and Experimental Approach". The views and conclusions in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of ARO or AFOSR or the U.S. Government.

Data availability

The ALEX-MP dataset for novelty evaluation and the ALEX-MP-20 dataset for training are available at https://github.com/microsoft/mattergen. The MP-20 dataset for training, including point group labels for conditional control, is available at https://github.com/trachote/crystalgrw. The datasets generated and/or analysed during the current study are available in the GitHub repository, https://github.com/trachote/crystalgrw.

Declarations

Competing interests

The authors declare no competing interests.

Supplementary Information

The online version contains supplementary material available at https://doi.org/10.1038/s41598-026-62470-x.