CrystalDiT: A Diffusion Transformer for Crystal Generation

CrystalDiT is a simplified diffusion transformer architecture for crystal structure generation that achieves state-of-the-art performance by treating lattice and atomic properties as a single, interdependent system.

CrystalDiT achieves 9.62% SUN rate on MP-20, significantly outperforming existing methods like FlowMM (4.38%) and MatterGen (3.42%).

Key features:

• Unified Architecture: Joint attention processing of lattice and atomic features
• Chemical Representation: Two-dimensional atomic encoding using periodic table positions
• Balance Score: Novel model selection metric optimizing discovery potential vs. generation quality

This checkpoint represents the best-performing model selected via Balance Score methodology after training for 50,000 epochs on the MP-20 dataset.

Files in this Repository

• best_model.pt: Pre-trained CrystalDiT model checkpoint (best model selected via Balance Score)
• generate_crystals.tar: Generated crystal structures from all compared methods, containing:
  • CrystalDiT_crystals/: 10,000 structures from our method (9.62% SUN rate)
  • flowmm_crystals/: 10,000 structures from FlowMM baseline (4.38% SUN rate)
  • mattergen_crystals/: 10,000 structures from MatterGen baseline (3.42% SUN rate)
  • ADiT_crystals_mp20/: 10,000 structures from ADiT baseline (2.74% SUN rate)
  • diffcep_crystals/: 10,000 structures from DiffCSP baseline
  • diffcsp-pp_crystals/: 10,000 structures from DiffCSP++ baseline

All crystal structures are in CIF format and were used for the comparative evaluation in our paper. These are provided to facilitate reproducible research and fair comparison with future methods.

Usage

Extract the generated structures:

tar -xf generate_crystals.tar

Load the model checkpoint as described in the README on GitHub.

Performance

Method	SUN (%)	MSUN (%)	UN Rate (%)
FlowMM	4.38	20.16	87.66
MatterGen	3.42	23.91	89.89
ADiT	2.74	13.50	37.08
CrystalDiT	9.62	25.94	63.28

Citation

@article{yi2024crystaldit,
  title={CrystalDiT: A Diffusion Transformer for Crystal Generation},
  author={Yi, Xiaohan and Xu, Guikun and Xiao, Xi and Zhang, Zhong and Liu, Liu and Bian, Yatao and Zhao, Peilin},
  journal={arXiv preprint arXiv:2508.16614},
  year={2024},
  url={https://arxiv.org/abs/2508.16614}
}