Wyckoff Transformer: Generation of Symmetric Crystals

Authors: Nikita Kazeev, Ruiming Zhu, Ignat Romanov, Wei Nong, Andrey E Ustyuzhanin, Shuya Yamazaki, Kedar Hippalgaonkar
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: material design, machine learning, crystal generation, space group symmetry, Transformer, Wyckoff position, generative model, autoregressive model, permutation invariance

Abstract:
We propose Wyckoff Transformer, a generative model for materials conditioned on space group symmetry. Most real--world inorganic materials have internal symmetry beyond lattice translation. Symmetry rules that atoms obey play a fundamental role in determining the physical, chemical, and electronic properties of crystals. These symmetries determine stability, and influence key material structural and functional properties such as electrical and thermal conductivity, optical and polarization behavior, and mechanical strength. And yet, despite the recent advancements, state--of--the--art diffusion models struggle to generate highly symmetric crystals. We use Wyckoff positions as the basis for an elegant, compressed, and discrete structure representation. To model the distribution we develop a permutation--invariant autoregressive model based on Transformer and absence of positional encoding. Our experiments demonstrate that Wyckoff Transformer has the best performance in generating novel diverse stable structures conditioned on the symmetry space group, while also having competitive metric values when compared to model not conditioned on symmetry. We also show that it is competitive in predicting formation energy, band gap, mechanical properties, and thermal conductivity.

Decision: reject
Reviewer scores: 5, 5, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1729912008750

thread_type:
official_review

<thread_object>
creator: reviewer_hhjg

<content>
summary:
The paper proposes the **Wyckoff Transformer**, a generative model designed for creating highly symmetric crystal structures by leveraging **space group symmetry**. Recognizing that most inorganic materials exhibit inherent symmetries, the authors develop a model that encodes these symmetries to influence key material properties such as stability, conductivity, and optical behavior. The Wyckoff Transformer uses **Wyckoff positions** as a discrete, permutation-invariant representation of atomic locations, eliminating the need for explicit positional encoding and improving the model’s efficiency and alignment with crystal symmetries.

Key contributions of the paper include:

1. **Tokenization of Crystals**: The authors represent crystals as an unordered set of tokens, merging information on chemical elements and their Wyckoff positions, enabling symmetry-based generation.
2. **Permutation-Invariant Encoding**: The model encodes Wyckoff positions based on symmetry-defined point groups, allowing for the generation of stable structures without positional encoding.
3. **Transformer Architecture**: The Wyckoff Transformer combines **autoregressive probability factorization** with permutation invariance, enhancing diversity and stability in generated structures.
4. **Empirical Outperformance**: The model outperforms existing methods, generating novel, stable structures that adhere to space group symmetry.
5. **Predictive Accuracy**: Despite limited input information, the model accurately predicts formation energy and band gap values comparable to DFT (Density Functional Theory) standards.

The model demonstrates superior performance in symmetry-conditioned generation, creating a diverse set of stable crystal structures that respect the underlying physical symmetries. This approach addresses limitations in prior methods, which struggled to produce symmetry-compliant structures, and shows promise for accelerating material discovery in fields requiring stable, symmetric crystals. However, the model inherits typical dataset limitations in generative models, as it learns distributions only within the scope of the training data, which may omit some stable but out-of-domain structures.

strengths:
Strengths

1. Originality: The Wyckoff Transformer introduces a novel approach to crystal generation by utilizing Wyckoff positions to encode symmetries explicitly, making it unique among generative models. Unlike traditional methods, it avoids positional encoding and uses permutation-invariant tokenization tailored to space group symmetries, a creative and effective innovation for materials science.

2. Quality: The paper includes experimental results (although the results are not yet complete), showing the model’s success in generating symmetric crystal structures while achieving competitive accuracy in formation energy and band gap predictions. The evaluation is thorough, comparing the model’s performance to state-of-the-art methods across multiple metrics, demonstrating its robustness and effectiveness in real-world scenarios.

3. Significance: This work is impactful for both machine learning and materials science in general. By generating materials that are symmetric and physically plausible, the Wyckoff Transformer can accelerate material discovery for applications requiring stable, symmetric structures (e.g., in semiconductors and optoelectronics). The model’s potential for symmetry-conditioned generation highlights a promising direction for future research in material informatics and generative modeling.

weaknesses:
I observed the following weakness in the paper and my recommendations towards constructively improving this paper:

1. The write-up does not qualify for the levels of ICLR acceptance, it has too many inconsistent notation, typing errors, vague statements and hypothesises without proofs. Following is a non-exhaustive list: Line 174 WP 4a becomes (m-3m, 0), what does this line mean? later in Figure 4 the authors mention m-3m without any tuple, and later again in Figure 5 where the authors mention similar symbols as 1a = [m-3m] just to show and example of inconsistency in this paper, in Line 51 , ``atom position'' use correct quatotation marks, Figure 1 is vague no details on how it is constructed or taken from other source, line 227 coset (closest) representative (if you wish to mention grouping theory coset, then kindly use proper mathematical notations), Section 3.1 line 269, 270, 308, 315, 327, 329 no enumeration. Kindly go through the paper again, and correct all of them.

2. The paper revolves around Wyckoff Transformer, but no explanation of the model architecture either in writting or in schematic has been given in the paper, section 2.2 mentions the title Model architecture but did not mention anything other than the Transformer paper by Vaswani et el which was used for Neural machine translation, moreover the authors are wrong in stating it to be an encoder only architecture. The authors seem to be lacking the knowledge about Encoder only, Decoder only and Encoder-Decoder models. This confusion seem to prpogate through the training section (2.3) where they mention (a) De novo generation , (b) Property prediction , where they are unable to state the difference between these two tasks and how a single kind of model cannot handle these two separate taks without any further layers (some hints of this is mentioned on Line 222, but still vague while following the whole paragraph.). Please attend to all these points, this would make the paper more readable and understandable.

3. The paper does not include any structural schematic of the generated structure, if any. The paper does inlcude a few figures on Wyckoff positions and Wyckoff representation of a toy 2D Crystal and SrTiO3 in Figure 2 and 5 respectively. But none for their actually generated structures (if generated any). Claims on generative capability stand weak when figures are not included (kindly refer to the papers which they've cited, they have shown a wide set of generated structures with their corrosponding inductive biases.) Since the authors mention about their novel representation, they should support the validity of those representations and inductive biases (representational consistency) in their outputs. If possible I would like this point to be addressed in the revision.

4. The section covering related work is weak. The authors must do a good survey of the past work in cystal generations particularly in the field of crystal genration in the representaion space. Some of which I was able to find by searching for representation based genrative model in citations to CDVAE (Xie et. al.) paper are: 1. https://arxiv.org/abs/2306.04510, 2. https://arxiv.org/abs/2403.10846,  3. https://arxiv.org/abs/2408.07213 (kindly read and search for more). I request the authors to kindly include papers which are in the same field to address the concerns in this paper and how your research aligns or complements with these papers, so that this work becomes complete. Kindly include a paragraph discussing how their approach compares to or builds upon these specific papers mentioned in the comments.

5. Lack of clarity: Reading through the paper multiple times, I have found out that the paper is written very poorly, and fails to convey the message of the authors. Upon my earlier assesment I had mention this as a strength (note: I had confidence = 3 and later 4), but going into the review process and reading the paper multiple times I am confident that the paper lacks clarity (updated confidence = 5).

questions:
The following are my questions and suggestions for this paper:

1. The structure of the network, inputs, outputs and loss function (including tokeization, loss function computation needs to be defined properly with proper mathamtical notation and schematics). Kindly include a schematic diagram of their Wyckoff Transformer architecture, clearly showing how it differs from standard Transformer models. Additionally, kindly clarify how their model handles both de novo generation and property prediction tasks, possibly by explaining any additional layers or modifications to the base architecture.

2. How did the authors plot Figure 1? Kindly include other generated structures in Figure 2, it will be best to show how the generated structures also follow these symmetris and where do ther lie in terms of space group number. Kindly provide a clear caption explaining the source and construction of the figure. Kindly include a figure showing examples of structures generated by their model, possibly comparing them to real structures or those generated by baseline models. This would help demonstrate the model's capabilities and the effectiveness of their novel representation.

3. In section 2.2 clearly mention the assumtpions for your model, as of now the reviewer was not able to find the assumptions which the authors have taken. These need to be mentioned in a list.

4. What was the training objective, were they two different models for task of De novo genration and Property prediction? If so, then how were they both trained? (Objective funvtion, optimiser hyper-parameters, Input data, valdiation metrics, regularisers, hardware specifications etc.)

5. The algorithm explained for (i) Tokenization, (ii) De-nove generation, (iii) Structure genration and (iv) Metric computation is vaguely defined. These need to be defined propely in algorithm sections. As of now they don't make any sense and are without mathematical notation.

6. Which DFT calculation has been used in the paper? Did the authors perform DFT of their own or are using previously reported results, if yes, then it will still need to be described and also how are they reporting.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730578603965

thread_type:
official_review

<thread_object>
creator: reviewer_rfbw

<content>
summary:
The paper presents a transformer-based architecture to generate symmetric crystals conditioned on space groups in a two-stage process. Initially, it generates tokens representing elements and their site symmetries, followed by lattice and coordinate predictions for these tokens with existing methods. The results include comparisons with multiple baselines, showing competitive performance across established proxy metrics. Finally, the paper also proposes metrics to assess the symmetry of the generated crystals and highlights further gains over baseline approaches.

strengths:
- The paper emphasizes the importance of generating symmetric crystals and highlights challenges with existing methods.
- The evaluation methods, including the newly proposed methods for evaluating symmetry, form a compelling discussion section.
- The method demonstrates effective gains for the symmetry metrics and is competitive for widely-used proxy metrics.
- The paper proposes a novel representation of crystal symmetry that could facilitate learning of crystal symmetry with deep learning approaches.

weaknesses:
- **Presentation and writing**: Essential concepts (site symmetry, wyckoff positions, space groups) are not appropriately introduced, which would create difficulty for readers unfamiliar with the field. Several works are cited in the related works section, but neither described nor highlighted the difference from their approach. Figures for architecture and pseudo-code to describe the training and generation pipeline would greatly benefit the understanding of the work.

- **Generalization of the approach**: 
  - The paper heavily focuses on the MP-20 dataset and does not provide any experiments with other datasets. For instance, it states permutation invariance was achieved because the number of WPs in the MP-20 dataset is small. How can this method be extended (or how does it fare in terms of performance) for crystals with a very high number of WPs?
   - There are no precise details on how many tokens were formed from the MP-20 dataset after tokenization. It would be interesting to discuss this number and other statistics about the tokens, e.g., which tokens are present more often (for some of the high symmetry space groups) and how the distribution of tokens affects training. 
   - It is also important to add how many new tokens the method generates or if it just predicts the fixed set of tokens in different combinations (and these combinations result in more template novelty than just sampling existing templates from training data). For instance, naively thinking about it, how will your model generate tokens that are not present in its dictionary? 
  - Finally, in Table 1a, please also provide the number of novel templates as absolute numbers instead of percentages.

- **Architecture**:
  - Please provide at least a pseudo-code of the training/generation algorithm and a figure explaining the training/generation process with a sample crystal example. The central algorithm is not clear from the text. 
  - Please mention the size of the model and computational and memory consumption (training time, memory required during training and generation). The paper lists that it is trained for $150k$ epochs, which seems to be a very long training process compared to existing methods (~$1k$ epochs). Can you explain this behaviour along with the set of hyperparameters used?

- **Evaluation**:
  - Can WyCryst be considered a fair baseline for comparison since it supports a limited number of unique elements per structure? For instance, it wouldn't compare with other methods that generate an arbitrary number of methods because it would result in poorer metrics (as seen in Table 1).
  - Is CHGNet used both to relax the generated structures and determine the energy in Table 1? 
  - Please mention the percentage of novel but structurally invalid generations from your method.

questions:
- **Two-stage approach**: Can you explain the benefits of a two-stage approach instead of a one-shot (such as DiffCSP) prediction of the site symmetries, elements, their positions and lattice parameters? If problems exist with a one-shot prediction, please explain and motivate the need for a two-step approach. For instance, can we (as an example) predict all the tokens (with elements, site symmetry, enumeration) followed by the lattice parameters and the fractional coordinates of the tokens or are there inherent issues with this approach? This question becomes more necessary since the generation of crystals would be slow for the proposed "sequential two-stage" approach. 
- **Crystal Structure Prediction (CSP)**: The paper focuses on generating crystals conditioned on space group. How could this method be extended to the CSP task, which is also crucial and could potentially benefit from using crystal symmetry?
- **Dataset fragmentation**: Although the tokens can be shared across different space groups, there will still be dataset fragmentation when the approach is conditioned on the space group. Is the training (and then generation) not affected by how many samples are present within each space group?

Some of the other questions are listed in the Weaknesses section. I will be happy to improve the score if the authors address the questions and weaknesses with supportive evidence during the discussion phase.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730581684851

thread_type:
official_review

<thread_object>
creator: reviewer_otui

<content>
summary:
This paper highlights the problem of generative models for crystals not generating symmetric crystals, which is an important property of these materials. This results in less realistic materials as well as inability to model some properties of crystals correctly. The authors propose to address this limitation by generating materials in a two-stage process. First, they train a Transformer model to output occupied Wyckoff positions in the crystal. Then either a method based on DiffSCP++ or PyXtal is used for atomic coordinates. The authors verify experimentally that this allows to generate more symmetric and diverse crystals.

strengths:
- The work tackles an important limitation of generative models for crystals
- The proposed solution is simple and sound
- The experimental evaluation shows that the method addressed the limitation. The evaluation metrics for symmetry and novelty of structures based on Wyckoff templates is also valuable.

weaknesses:
- Just using the Wyckoff positions is not a complete representation, especially for atoms in the general position. The sentence "reducing the number of parameters by an order of magnitude without information loss" is false. I also don't think that statement that desired properties can be obtained from the discrete values alone is accurate or substantiated by enough evidence. I therefore encourage the authors to substantially nuance that section. The experiments on property prediction indicate a degradation of performance in property when discarding coordinates.
- The model is claimed to be invariant with respect to the choice of coset representative and to permutations. This formulation is too strong, since this is achieved through data augmentation. A correct statement would be that the model is encouraged to be invariant.
- The proposed representation for Wyckoff positions is universal across space groups but might not allow proper generalization since the "enumeration" variable is not grounded on physical information but on an arbitrary convention. Therefore, if a group is rare in the training data (this is indeed the case for datasets like MP20), there is no reason that the model will learn to capture that variable correctly. The authors should discuss this limitation appropriately. 
- I did not find the discussion of the related works to be sufficient. The authors should expand that discussion so that the readers understand the differences and similarities with the proposed method better.
- I find that the explanation of Wyckoff position in the third paragraph of the introduction is not easy to understand. It may be too early in the paper to go into such an explanation.
- There are some typos and mistakes that the authors should look into correcting. For example, "lattice transition -> lattice translation" or "   Cordiality -> Cardinality".

questions:
- I don't see what the footnote 1 adds to the discussion, I find it more confusing than anything. Could the authors clarify it, or consider simply removing it?
- I don't understand the expression in the abstract "These symmetries form energy configurations". What is meant there?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730585988134

thread_type:
official_review

<thread_object>
creator: reviewer_7mtl

<content>
summary:
This paper proposes a transformer-based approach that leverages Wyckoff positions to encode material symmetries efficiently. This is done by primarily encoding the discrete symmetries of space groups without using atomic coordinates. The discussion on WyFormer, including tokenization and (extensive) metrics, is detailed. Their main contribution is to represent a crystal as an unordered set of tokens and make de novo material and property predictions. Furthermore, four new metrics are proposed (P1, Template Novelty, Space Group, and S.S.U.N.) to judge the ability to reproduce symmetry properties accurately. Results indicate that WyFormer outperforms other methods in terms of template novelty, space group distribution, and fraction of asymmetric structures.

Overall I think the paper could be a good step in the direction of using symmetries for property prediction, provided certain clarifications on experimental details and broader evaluations are addressed.

strengths:
- Crystal representation for tokenization.
- Material property prediction results are surprisingly good when compared against neural nets trained for energy prediction.
- Four new metrics provide a new way of looking at models' ability to generate symmetry properties.
- Justification in the appendix for the selected two structure generation methods.

weaknesses:
- The scope of material property prediction- authors focus on just two (energy and band gap). If feasible, can the authors provide some insight on which other properties could be predicted, purely from a correlation with crystal structure perspective?
- The proposed method is evaluated on a single dataset, MP-20, and makes it hard to judge the generalizable nature of WyFormer from it. Are there other datasets on which performance can be evaluated?

questions:
1. In section 1.3, line 138, "...our main differences are listed in the discussion of our contributions 1.2."  where are the main differences listed in section 1.2? Or am I missing something?

2. Can the authors explain why the Space Group value for WyForDiffCSP++ is high while the S.S.U.N. value is similar to WyCryst in Table 1a? 

3. A discussion on computational cost would be good to have, given that the authors mention that the entire dataset fits into GPU memory (training time and memory requirements)

4. Are there methods apart from CHGNet that improve crystal structure generation?

5. Have the authors tried other property prediction experiments besides energy and band gap?

Additional Feedback:
1. line 280: "they to be" -> "they are"?
2. line 282: percetage -> percentage?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730712314581

thread_type:
official_review

<thread_object>
creator: reviewer_deyq

<content>
summary:
The paper focuses on the tasks of de novo materials generation and materials property prediction. The main contribution is a Wyckoff representation tokenization and model training strategy. For de novo generation, once the transformer generates a Wyckoff position then PyXtal and CHGNet are used to generate/relax the structure. In particular, the model is good at generating materials with the proper space group.

strengths:
- The application of ML to materials discovery is interesting and timely. 
- The Wyckoff representation builds in crystal symmetries in a natural way. 
- Good empirical results for template novelty, P1, and Space Group metrics.
- The paper is written well

weaknesses:
- Little improvement in standard de novo generation metrics. SUN actually goes down compared to DiffCSP.   
- Additionally, de novo generation metrics were computed with a ML potential instead of DFT.  
- The property prediction benchmark is not particularly compelling because there are better benchmarks out there with other more recent models as baselines (e.g. CHGNet), such as Matbench discovery. 
- If not there, it would be good to include this citation (https://arxiv.org/abs/2106.11132).

questions:
- You found a nice way to tokenize a Wyckoff representation, would it be better to fine-tune a LLM with this representation than train a transformer from scratch? The CrystalLLM paper (https://arxiv.org/abs/2402.04379v1) had some nice results that could potentially be improved with your representation. 
- Another important axis is the inference speed or cost of de novo generation, how does WyFormer or WyForDiffCSP++ compare to DiffSCP/FlowMM? 
- Is there a way to more concretely show the benefit template novelty?
- Can you plot the distribution of space groups generated from WyFormer compared to MP-20 distribution?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733302837888

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Closing argument

comment:
We are deeply thankful to the reviewers for their time and effort that made WyFormer a better model with an even larger number of experiments to prove it. Here is our closing argument that summarizes the discussion and makes the final case for our paper.

**Strengths**
1. Application of ML to materials discovery is an important goal [DeYQ, hhjg]
2. Good representation:
    1. The Wyckoff representation builds in crystal symmetries in a natural way. [DeYQ]
    2. Strengths: Crystal representation for tokenization. [7MTL]
    3. Originality: The Wyckoff Transformer introduces a novel approach to crystal generation by utilizing Wyckoff positions to encode symmetries explicitly, making it unique among generative models … a creative and effective innovation for materials science. [hhjg]
3. Introduction and motivation of symmetry-based metrics and good performance according to them:
    1. Good empirical results for template novelty, P1, and Space Group metrics. [DeYQ]
    2. Four new metrics provide a new way of looking at models' ability to generate symmetry properties. [7MTL]
    3. The evaluation metrics for symmetry and novelty of structures based on Wyckoff templates is also valuable. [otui]
    4. The evaluation methods, including the newly proposed methods for evaluating symmetry, form a compelling discussion section. The method demonstrates effective gains for the symmetry metrics and is competitive for widely-used proxy metrics. [rFbw]
    5. The paper includes rigorous experimental results, showing the model’s success in generating symmetric crystal structures … The evaluation is thorough, comparing the model’s performance to state-of-the-art methods across multiple metrics, demonstrating its robustness and effectiveness in real-world scenarios. [hhjg]
4. Good property prediction performance:
    1. Material property prediction results are surprisingly good when compared against neural nets trained for energy prediction. [7MTL]
5. Low inference computational resources requirements:
    1. Another important axis is the inference speed or cost of de novo generation … it seems there is a case to be made for WyFormer in terms of the computational cost for generating SUN/SSUN structures. [[DeYQ](forum?id=ursX3k1rTO&noteId=Kl8xL7aUbA)]
6. The proposed solution is simple and sound [otui]

**Neutral**

1. Performance according to the standard metrics. DeYQ considers the lack of gain to be a weakness, while rFbw cosiders competitive performance to be a strength
2. Training resources requirements (7MTL, rFbw). WyFormer takes similar time to train compared to diffusion-based baselines, but an order of magnitude less GPU memory [[anonymized url](https://wyckoff.notion.site/c)].

**Weaknesses**

1. Evaluation on a single dataset (7MTL, rFbw). We address it in three ways:
    1. Conceptual. We have only one physical world, and the goal of the model is to generate novel stable materials in this world; MP-20 is a diverse dataset and contains a sizable fraction of all known materials.
    2. We show WyFormer generative performance on MPTS-52, a more challenging dataset containing materials with up to 52 atoms per unit cell; it is the only generative model that demonstrated the capability to handle it that we are aware of. [[anonymized url](https://wyckoff.notion.site/mpts)]
    3. We show WyFormer’s competitive performance in property prediction on AFLOW dataset (Table 5)
2. Data fragmentation because of conditioning on symmetry space group (rFbw, otui). We address this by using transferable space group and Wyckoff position representations. In the initial version of the paper one part of the representation was not transferable between space groups, during the review period we have developed a transferable representation for this part as well (Appendix B).
3. Use of CHGNet for stability evaluation (DeYQ). We have run DFT evaluations that, despite the limited sample size, show competitive performance of WyFormer according to stability metrics, and positive correlation between CHGNet and DFT estimations of stability (Table 1)
4. Only predicting energy and band gap (7MTL, DeYQ). We added experiments showing competitive performance of WyFormer in predicting thermal conductivity, Debye temperature, bulk modulus, and shear modulus on AFLOW dataset (Table 5)
5. Related work discussion (otui, hhjg). We have expanded, clarified and systematized the relevant work section.
6. Presentation (otui, rFbw, hhjg); DeYQ, however, commended the writing. We have extensively edited, clarified, and expanded the writing, added and improved illustrations.

In summary, WyFormer has a unique and powerful synergy of attributes, proven by extensive experimentation: best-in-class symmetry-conditioned generation, physics-motivated inductive bias, competitive stability of the generated structures, competitive property prediction quality, and unparalleled inference speed. Altogether, they make it a great starting point in the material design pipeline.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors propose Wyckoff Transformer, a generative model for material crystals that obey space group symmetry. The key algorithmic idea is that transformers can provide a permutation-invariant, autoregressively-generated, discrete representation of inorganic material structures. Using this representation the authors show performance of the model on de novo material generation along with several property prediction tasks.

The method proposed in the paper is intuitive, and addresses an important problem in de novo crystal generation with current AI models (ie, the lack of symmetry). However, issues were raised regarding the clarity of the exposition; the strength of the results (both the de novo generation and the property prediction tasks); potential limitations with the representation that may not have been properly highlighted in the narrative; and conceptual overlap/attribution of previous work. 

This was a close call, but I hope that the authors will take into account the extensive feedback provided by the reviews while preparing revisions.