CrystalFormer-CSP: Thinking Fast and Slow for Crystal Structure Prediction

TL;DR

CrystalFormer-CSP tackles the CSP problem by unifying space-group-aware generative modeling with energy-based relaxation via a universal MLFF, augmented by reinforcement learning. The method generates candidate crystal structures conditioned on a chemical formula, relaxes them with MLFFs, and ranks by energy above hull to assess stability, achieving competitive benchmarks. The approach demonstrates strong space-group prediction and improved CSP match rates on two datasets, with RL further boosting performance and enabling discovery of lower-energy polymorphs. It provides practical usage via Colab and LLM integration, and offers a promising framework for scalable, end-to-end crystal structure prediction.

Abstract

Crystal structure prediction is a fundamental problem in materials science. We present CrystalFormer-CSP, an efficient framework that unifies data-driven heuristic and physics-driven optimization approaches to predict stable crystal structures for given chemical compositions. The approach combines pretrained generative models for space-group-informed structure generation and a universal machine learning force field for energy minimization. Reinforcement fine-tuning can be employed to further boost the accuracy of the framework. We demonstrate the effectiveness of CrystalFormer-CSP on benchmark problems and showcase its usage via web interface and language model integration.

Figures (4)

• Figure 1: (a) The CrystalFormer-CSP framework. The process starts from a user-specified chemical formula and proceeds as follows: ➀ Generation: the pretrained autoregressive transformer samples a batch of crystal structure candidates; ➁ Relaxation: the machine learning force field relaxes the structures and provides energy estimates of the final structures; ➂ Ranking: the relaxed candidates are ranked according to their energy above hull ($E_{\mathrm{hull}}$), providing an assessment of their thermodynamic stability. Optionally, one can further fine-tune the generative model with reinforcement learning using the relaxed energy as the reward signal. (b) The generation step samples from the conditional probability distribution of crystal structure given a chemical formula. The chemical formula is represented as a weighted summation of learned embedding vectors for the chemical elements. The crystal structure is represented as a sequence that consists of space group number, Wyckoff letters, atom species, fractional coordinates, and lattice parameters. In practice, only these underlined bold tokens need to be sampled; the remaining ones are fixed by the preconditions.
• Figure 2: The correlation plot of the number of atoms in the conventional cell and the number of Wyckoff sites. The histograms are shown in the side panels.
• Figure 3: Schematic illustrations of (a) Colab notebook interface, (b) Natural language model integration.
• Figure 4: (a) The average and minimal energy above hull of relaxed LiP(HO$_2$)$_2$ versus reinforcement fine-tuning steps. The horizontal line indicates the energy above hull of the ground truth structure. (b) The ground truth structure and discovered structures along their Wyckoff sequences and energy above hull predicted via DFT calculations.