Space Group Equivariant Crystal Diffusion

Thank you for the helpful feedback. Below we aim to address your concerns. Please let us know if we can clarify anything further.

Why not sample the training distribution of space groups?

We agree. For simplicity, we have changed our code to sample space groups directly from the training set, observing no changes to our results.

What difficulties would arise in learning a joint distribution?

We believe there are a few difficulties. (1) Joint modeling requires tuning per-component loss weights. (2) Jointly modeling number of atoms with atom attributes requires unnatural “mask states”. (3) Space groups impose hard constraints on lattices (lines 212-217), available Wyckoff positions (lines 120-130), and atom coordinates. Sampling these constrained and constraining variables jointly will yield constraint violations with high probability. Several works try learning space group constraints instead of enforcing them, but we explain in lines 21-45 and show empirically through the $JSD_G$ and $JSD_{d_{Wyckoff}}$ metrics why this might not be a good idea.

While SymmCD diffuses atomic coordinates jointly with Wyckoff positions, the model requires post-hoc projections to correct constraint violations. This complicates the model distribution (since marginalizing out projections is intractable) and possibly contributes to the lower structural validity metric (Table 1) as the model cannot recover from projections to invalid structures.

sampling $n$ using the transformer in Section 4.3 is not described

Number of atoms $n$ is sampled implicitly via a stop token (p(stop | …) in Eq 7). We will state this in line 245.

Would it be more appropriate to avoid imposed atom orderings?

We agree, although “all models are wrong, but some are useful.” Autoregressive models with stop tokens can naturally handle variable numbers of atoms, unlike pure diffusion models. We highlight the strong precedent of applying sequential and non-permutation invariant architectures to unordered data across AI, often achieving SOTA results [9,10,11]. While one can train sequential models on all orderings, marginalizing out orderings is intractable for even 20 atoms.

[9] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR 2021

[10] Scaling Autoregressive Models for Content-Rich Text-to-Image Generation. TMLR 2022

[11] Scalable Diffusion for Materials Generation. ICLR 2024

How were features for space groups and Wyckoff positions derived?

We will add the following to our appendix.

For the space groups, we used one-hot features indicating the 6 lattice centering types, 6 crystal families, 32 crystallographic point groups, whether the space group is chiral, and whether the space group is centrosymmetric. Since this representation ignores fractional translation symmetries and thus cannot differentiate certain space groups, we counted the max number of collisions between space group embeddings (16), arbitrarily indexed colliding embeddings, and created additional one-hot features to avoid collisions. This yielded 6+6+32+1+1+16=62 features. While the last 16 features are not physically motivated, they give the model the capacity to differentiate space groups and learn their relationships.

To featurize a Wyckoff position, we used our space group features (62); one-hot features for the site symmetry symbol (100), number of degrees of freedom in atom coordinates (4), and Wyckoff multiplicity (17); and 48 Fourier features of Wyckoff shapes’ midpoints and vertices (see Sec 4.1). This yielded 62+100+4+17+48=231 features. To compute Fourier features $\nu \in \mathbb{R}^{48}$, we collected the midpoint and vertices of each Wyckoff shape belonging to the Wyckoff position. We used these points $\\{p_i\in\mathbb{R}^{3\times1}\\}_{i=1}^{n_w}$ to compute the following:

$

\begin{aligned} f &=2\pi \cdot \operatorname{linspace}(\text{start}=1,\text{end}=32,\text{steps}=8) \in \mathbb{R}^{1 \times 8}\\ \nu_i &=\operatorname{flatten}(p_i f)\in \mathbb{R}^{24}\\ \nu &=\frac{1}{n_w}\sum_{i=1}^{n_w} [\sin(\nu_i),\cos(\nu_i)]\in\mathbb{R}^{48} \end{aligned}

$

is 1731 the total number of Wyckoff positions across space groups?

Yes. We will edit line 255 accordingly.

SymmCD symmetry representations could be an alternative. Was any ablation performed on token design?

We performed the ablation, replacing our space group and Wyckoff features with SymmCD’s. On MP20, our features yielded better validity, Wasserstein distance for density, structural divergence, and validation log-likelihoods of Wyckoff positions. SymmCD features yielded a better UN rate, Wasserstein distance for number of elements, and compositional diversity. Results with standard deviations over 3 training runs are below. We believe the result on Wyckoff log-likelihoods motivates our featurization.

Is LDiff equivalent to the representation of lattice matrices in DiffCSP++?

Yes. We will edit lines 330-331 for clarity.

"next best" model is misleading. Given acknowledged variability in SUN rate, lines 334–335 is not well-supported

We will rephrase to “achieved the highest evaluated SUN rate”. We are currently running ~800 more DFT calculations per model to address potential variability.

Beyond efficiency, are there other advantages to the asymmetric unit? The conventional cell may capture more structural information

Computational efficiency is the main advantage to using the ASU. We used the conventional cell and its periodic boundaries for graph construction (lines 1183-1185) to avoid losing structural information and to symmetrize the model with space group equivariance (Eq 3). Our efficiency gains from using the ASU stem from only computing score matching losses on ASU atoms, only storing ASU atoms in the dataset, and only generating atom types/Wyckoff positions of ASU atoms. We tried excluding fractional coordinates from the representation such that atoms in the same orbit have the same embeddings, removing the need to compute embeddings for atoms outside the ASU. However, performance was better with fractional coordinates.

A possible approach to remove the need to compute “non-ASU” atom embeddings while maintaining performance is to replace periodic sinusoidal embedding functions with space group invariant functions [12]. We defer this to future work.

[12] arXiv:2306.05261 (2023)

Table 1

What does T represent? What is T in the SGEquiDiff row?

T is the inference-time temperature for all discrete distributions, and T=1.0. We will add these to the table.

There appears to be an inverse relationship between composition validity and UN rate

We observed the uniqueness rate was nearly 100% for all models. Thus, the relationship occurs because exotic compositions are more likely to be novel but also invalid.

It is incorrect that SymmCD uses fixed templates

Thank you, we will fix that.

Please report variance across multiple seeds for SGEquiDiff

We trained SGEquiDiff with 3 random seeds per dataset and will update the tables with average metrics and standard deviations, as below.

Why is Wasserstein distance for density significantly higher for SGEquiDiff

Good question - we aren’t sure. We checked 1D histograms of unit cell volumes and numbers of atoms per unit cell for generated vs ground truth crystals. We did not see qualitative differences.

Some insight can be drawn from the LDiff ablation, which has a comparable density metric with the best baselines. It's possible that (1) iterative refinement from lattice diffusion improves the density metric (aligned with the result of CDVAE, which doesn't diffuse lattices), (2) the higher density metric is an artifact of discretizing the lattice, and/or (3) choosing the lattice before atom types/Wyckoff positions is harder to learn than the reverse.

Regardless, we don't think the higher density metric is significant in practice since both CDVAE and SGEquiDiff attain better SUN rates (Table 1) and RMSDs (Table 3) than DiffCSP++ despite having higher Wasserstein distances for density.

The table caption would benefit from more description

We will add the following to table 1 (similar for table 2): “We report metrics of validity for 10,000 sampled crystals; uniqueness and novelty for 1,000 valid crystals; diversity and goodness-of-fit for valid, unique, and novel (UN) crystals; and stability for 100 UN crystals. Stability is defined as having a DFT-predicted $E_\mathrm{hull}$ < 0.1 eV/atom compared to Materials Project v2024.11.14. Error bars are the standard deviation of three training runs. $T$ represents inference-time temperature of discrete distributions in SGEquiDiff.”

---

We hope this addresses your questions. If our responses are helpful, we would be grateful if you would consider raising your score.

Thank you for your helpful feedback and questions. Below we respond to your questions and concerns.

W1 : The experiments are conducted on only two datasets, whereas related work such as [1] and [2] typically report results on at least three datasets (e.g., Perov-5, MP20, and Carbon-24).

We note that Perov-5 and Carbon-24 are not used in more recent works [4,5,6,7,8] (including a followup work to [3]) because they are toy datasets containing unstable crystals; all crystals in Perov-5 have the same fractional coordinates and all crystals in Carbon-24 have only 1 element. In contrast, crystals in MP20 and MPTS52 have all been experimentally observed and span 89 elements of the periodic table with a large diversity of structures. These datasets are thus much more challenging and relevant for materials discovery, becoming the de facto benchmarks for crystal generative models.

[4] A generative model for inorganic materials design. Nature 639, 624–632 (2025)

[5] SymmCD: Symmetry-Preserving Crystal Generation with Diffusion Models (ICLR 2025)

[6] All-atom Diffusion Transformers: Unified generative modelling of molecules and materials (ICML 2025)

[7] Wyckoff Transformer: Generation of Symmetric Crystals (ICML 2025)

[8] FlowLLM: Flow Matching for Material Generation with Large Language Models as Base Distributions (NeurIPS 2024)

One of the selected datasets—though challenging—lacks thorough baseline comparisons. While the authors note that previous models have not reported results on this dataset, I would expect them to run those baselines using publicly available code to enable a fairer and more comprehensive comparison. Relying on results from just one baseline weakens the experimental support for the proposed method.

We agree. To strengthen our results, we modified the codes of DiffCSP and DiffCSP++ and benchmarked them on MPTS-52, using the same diffusion corrector step size as for MP20 and dropping the default batch size of DiffCSP++ to 96 to avoid OOMs. We additionally ran SymmCD on our own hardware. We failed to train CDVAE on MPTS-52 due to errors during graph construction arising from noised atoms becoming isolated within the model’s cutoff radius. As an additional diversity metric, we followed SymmCD and report the fraction of unique and novel templates with respect to the training data (where a template is defined as a space group and a multiset of occupied Wyckoff positions) out of 1000 structurally and compositionally valid sampled crystals. We found that SGEquiDiff had a significantly higher unique and novel template rate than baselines while maintaining high structural validity. Although DiffCSP++ yielded slightly higher structural validity, the model uses ground truth templates from training data which explains its extremely low rate of novel templates. The updated table 2 is below:

Q1 : Why does SGEquiDiff achieve better performance on the S.U.N. metric compared to baselines, yet underperform on most other evaluation metrics? How should the importance of the S.U.N. metric be interpreted relative to the others?

We believe that our inclusion of a comprehensive set of metrics (including new ones) is a strength of our work. We also respectfully disagree that SGEquiDiff underperforms on most other evaluation metrics - we believe its performance is on par with other SOTA models that we benchmarked. None of the baseline models achieved SOTA performance on more than 4 of the 12 metrics (although we don’t believe that disqualifies them from being useful and interesting). SGEquiDiff achieved SOTA (or comparable) performance on 4: Wasserstein distance for number of elements, Jensen-Shannon divergence to space groups and Wyckoff dimensions, and the S.U.N. metric. On several other metrics, SGEquiDiff’s performance is competitive in any practical sense (e.g., structure validity, composition validity, composition diversity). We highlight that SGEquiDiff (5.5M parameters) is the second smallest model (CDVAE has 4.9M, DiffCSP(++) has 12.3M, SymmCD has 60.4M, and Mattergen has 44.6M), which we will add to the paper.

Ultimately, the most important metric measures the ability to produce novel and synthesizable crystals. For this reason, the S.U.N. rate [4] is the most important metric as it leverages quantum chemistry calculations to assess stability.

W3 : It is unclear how well the proposed model performs on the Crystal Structure Prediction (CSP) task—a central objective in crystal generation—unlike other models such as [1], [2], and [3], which explicitly evaluate this aspect. Q2 : Can the proposed method be applied to Crystal Structure Prediction? If so, how does it perform compared to existing baselines on this task?

CSP was not a focus of our work because we believe the current formulation of the task in ML papers is slightly misaligned with real world application. DiffCSP [2], FlowMM [3], and other models assume that the number of atoms per unit cell will be known beforehand, which is not true in practice. For example, iron exists in BCC (2 atoms per conventional cell), FCC (4 atoms per conventional cell), and at high pressures, HCP (6 atoms per conventional cell). DiffCSP++ [1] relies on training examples of templates which specify number of atoms, Wyckoff positions, and atom types; thus, performance will worsen in composition spaces with sparse training data. However, these sparse spaces are exactly where experimental materials discovery campaigns are aimed.

Nevertheless, we agree that CSP is an important problem for materials discovery and have done a preliminary evaluation of SGEquiDiff in the same way as prior crystal generative model papers, calculating match rate and RMSE based on a single generated crystal per composition. Due to the rebuttal time constraints, we only evaluated the CSP task on MP20 using the pre-trained SGEquiDiff+LDiff model from Table 1 with DiffCSP++’s CSPML method for template selection. Without any CSP-specific training or hyperparameter tuning (unlike the baselines), we observed competitive performance (see table below). We would be happy to revisit this experiment, performing full tuning to ensure representative and meaningful results for the revised paper.

We note that SGEquiDiff could be alternatively adapted for CSP by reordering the factorization in Eq 5 or by supplying composition as conditioning during training. Learning composition (rather than strictly enforcing it) may be desirable since an arbitrary composition is not guaranteed to host a stable crystal (accordingly, experimentally observed compositions can be averages of multiple phases).

---

We hope this addresses your concerns. If our responses are helpful, we would be grateful if you would consider raising your score.

Thank you for the helpful feedback and questions. Below we aim to answer questions and address your concerns. Please let us know if we can clarify anything further.

The notation used for functions and group actions is sometimes overlapping, which can be misleading at times

Thank you for pointing this out. If we understand correctly, you are referring to the usage of $g(t)$ in equations 1 and 2 overlapping with the usage of $g$ as a group element. To avoid confusion, we will replace $g(t)$ in equations 1 and 2 with $\omega (t)$.

Line 185 to 188: it's not completely clear what the message is there.

We agree this is a confusing point and will update the text for clarity. The intended message is that, for example, if a 0D Wyckoff position (a point) lies in a 1D Wyckoff position (a line), then a space group equivariant vector field cannot move an atom along the line through the point. This is problematic because we observed examples of space groups where traversing between symmetrically unique regions of a 1D Wyckoff position requires moving through a 0D Wyckoff position. In contrast to, e.g., ODE-based sampling, the use of Gaussian noise (projected onto the tangent space of the Wyckoff position) in SDE-based sampling allows models to stochastically overcome these barriers while maintaining that the marginal distribution over atom coordinates is still space group invariant.

Line 237 to 240: it's not completely clear how the model used to sample the lattice parameters is trained. Once one lattice parameter is generated, are you using a noisy version of it as conditioning to generate the others? Clarifying that process would be helpful

During training, we maximize the log probability of the next ground truth lattice parameter conditioned on noisy versions of the previous ground truth lattice parameters (where noise is masked such that space group constraints are respected). At sampling time, noising is not applied. We will add this to the revised paper.

The table captions could be expanded to be more self-contained, for example by being be more specific about what is presented in them.

We agree and will append the following text to the caption of table 1: “We report metrics of validity for 10,000 sampled crystals; uniqueness and novelty for 1,000 valid crystals; diversity and goodness-of-fit distributional distances for valid, unique, and novel crystals; and stability for 100 valid, unique, and novel crystals. Stability is defined as having a DFT-predicted $E_\mathrm{hull}$ < 0.1 eV/atom with respect to the Materials Project v2024.11.14. Error bars are the standard deviation of three training runs. $T$ represents the inference-time temperature of discrete distributions in SGEquiDiff.” We will append similar text to the caption of table 2.

1. I am bit confused about eq. 10 and eq. 11 in appendix. The two equations means that the function $f(x)$ is at the same time invariant and equivariant to $R$. But then this should hold just for trivial $R$, like if $R$ is the identity, right? Am I missing something here? Or is it just because we are considering group actions from the group stabilizer (which leaves the input unchanged and therefore the output is fixed) that this holds? Clarifying this would help avoid confusion.

Yes, as you mentioned, equations 10 and 11 hold because as in line 1133, we are only considering group actions from the stabilizer group. To avoid confusion, we will update the notation to read $R_{g_x}$ and $v_{g_x}$ instead of $R$ and $v$.

2. I'm curious about the decision to learn $p(G)$ (the distribution over the space groups), rather than setting it to the empirical distribution of the training set. Doesn't the optimization of the logits lead to the empirical distribution?

We agree (although we applied early stopping based on validation performance as a soft regularizer). For simplicity, we have changed our code to sample space groups directly from the training distribution. We observed no changes to our experimental results.

3. I think it will be interesting and useful for the community to understand where the most gains in performance comes from. Although there are some differences in the input used across all the different generative models, many of them uses the same EGNN backbone architecture (from Satorras et al.). Is the FaeNet specifically needed to enforce some of the constraints you want to impose? Have you try to compare the performance using the same backbone architecture as the other baselines, to isolate the benefit of the different architecture?

Thank you for the question and suggestion. FAENet is not required to enforce space group equivariance. From lines 174-175, any periodic translation invariant architecture suffices, including the EGNN backbone used in DiffCSP and others. We have conducted the ablation on MP20 as suggested using the aforementioned EGNN backbone (”CSPNet”). For fair comparison, we held constant the number of message passing steps (5), the number of Gaussian edge frequencies (96), and the number of trainable parameters (2.2M) by reducing CSPNet’s hidden dimension from 512 to 224. We observed our architecture performed slightly better than CSPNet on the structural validity and U.N. metrics (possibly due to our use of Cartesian distances as edge features) while performing the same on other structural metrics. We will include the results (below) in our appendix, where standard deviations over 3 training runs are reported and best result is bold.

4. Space groups define specific constraints to the lattice angles. How do you enforce those constraints during sampling? Are they hard-coded?

Constraints are enforced in the same way during training and sampling with hard-coded masking and addition of any pre-determined angles. We will update lines 219-220 to reflect this.

5. Is the featurization of the space group different from the one used in SymmCD? If so, I think it would be helpful to explicitly state that and highlight differences.

Yes, our space group featurization is different. SymmCD used binary features indicating the 26 space group symmetry operations, the 15 symmetry axes, and the 7 lattice systems. This yielded 7 + 26 x 15 = 397 features. We used more abstracted one-hot features indicating the 6 lattice centering types, 6 crystal families, 32 crystallographic point groups, whether the space group is chiral, and whether the space group is centrosymmetric. Since this representation ignores fractional translations from glide and screw symmetries, we counted the maximum number of collisions between space group representations (16), arbitrarily indexed colliding space groups, and created additional one-hot features to avoid collisions. This yielded 6 + 6 + 32 + 1 + 1 + 16 = 62 features. While the last 16 features are not physically motivated, they give the model the capacity to differentiate space groups and learn their relationships.

As an ablation, we replaced our space group and Wyckoff position features with those used in SymmCD. On MP20, we found our featurization yielded better structure validity, composition validity, Wasserstein distance for atomic density, structural divergence, and validation log-likelihoods of ground truth Wyckoff positions. The SymmCD features yielded a slightly better U.N. rate, Wasserstein distance for number of elements, and compositional diversity. The results (with standard deviations over 3 training runs) are summarized below where best result is bold. We believe the particularly strong result on Wyckoff log-likelihoods motivates our featurization.

6. The EquiCSP paper highlights that the score used by DiffCSP is not periodic translation invariant, meaning that the score for a position and a periodic translation of it are different. Does your score formulation address this issue or does it suffer still from this mismatch?

Yes, our score is periodic translation invariant. Following our paper’s notation, the space group wrapped normal (SGWN) distribution and its score can be written as

$

\begin{aligned} p_\mathrm{SGWN}(x_t|x_0) &\propto \sum_{t_L \in \mathbb{Z}^3} \sum_{\{R_j, t_j\} \in G/T_L} \exp\left(\frac{-\lVert x_t - (R_j x_0 + t_j + t_L) \rVert_F^2}{2\sigma_t^2} \right) \\ \nabla_{x_t}\log p_\mathrm{SGWN}(x_t|x_0) &= \frac{\sum_{t_L \in \mathbb{Z}^3} \sum_{\{R_j, t_j\} \in G/T_L} \exp\left(\frac{-\lVert x_t - (R_j x_0 + t_j + t_L) \rVert_F^2}{2\sigma_t^2} \right) (R_jx_0 + t_j + t_L - x_t)}{\sigma_t^2 p_\mathrm{SGWN}(x_t|x_0)}. \end{aligned}

$

Applying a lattice translation $t^\prime_L \in \mathbb{Z}^3$ as $x_t \rightarrow x_t + t^\prime_L$ will leave the SGWN and its score invariant since $t^\prime_L$ can be absorbed by the infinite sum $\sum_{t_L \in \mathbb{Z}^3}$.

Missing the log in the score matching loss and also in line 280

Thank you for catching this!

---

We hope this addresses your questions. If our responses are helpful, we would be grateful if you would consider raising your score.

Thank you for the thoughtful review and encouraging comments. Below we try to address your concerns.

As noted by the authors, one can relax the outputs of non-"space group equivariant" generative models to get high-symmetry structures. My interpretation of this is that the proposed modelling approach is not strictly necessary to achieve their stated outcomes.

If we understand correctly, this comment refers to our statement in lines 375-376. To clarify, we believe there are interesting avenues for future work that explore the boundary between space group-constrained models and non-constrained ones to potentially combine the strengths of both. For example, one might train a non-equivariant model with the equivariant scores of our Space Group Wrapped Normal. Furthermore, one might use a non-equivariant model during training, but then symmetrize it with equation 3 to yield space group equivariant scores at inference time. Such approaches could potentially unlock a broader diversity of model architectures able to obtain high symmetry structures. We believe that this capacity to inspire new modeling approaches is a core strength of our work rather than a weakness.

The empirical advantage is not so clear. Only N=100 DFT relaxations were used in assessing their SUN metric. Therefore I have some doubts about the strength of the findings.

We agree and are currently running several hundred (~800) more DFT calculations per model. Unfortunately we could not obtain the results by the end of the rebuttal period.