[Paper Review] Crystal Structure Generation with Autoregressive Large Language Modeling

Luis M. Antunes, Keith T. Butler+1 more|arXiv (Cornell University)|Jul 10, 2023

Machine Learning in Materials Science18 citations

TL;DR

CrystaLLM trains an autoregressive Transformer on CIF text to generate plausible inorganic crystal structures, with MCTS-guided refinement using an energy predictor to improve structure quality and realism. It generalizes to unseen formulas and compares favorably with diffusion/VAE-based baselines on CSP benchmarks.

ABSTRACT

The generation of plausible crystal structures is often the first step in predicting the structure and properties of a material from its chemical composition. Quickly generating and predicting inorganic crystal structures is important for the discovery of new materials, which can target applications such as energy or electronic devices. However, most current methods for crystal structure prediction are computationally expensive, slowing the pace of innovation. Seeding structure prediction algorithms with quality generated candidates can overcome a major bottleneck. Here, we introduce CrystaLLM, a methodology for the versatile generation of crystal structures, based on the autoregressive large language modeling (LLM) of the Crystallographic Information File (CIF) format. Trained on millions of CIF files, CrystaLLM focuses on modeling crystal structures through text. CrystaLLM can produce plausible crystal structures for a wide range of inorganic compounds unseen in training, as demonstrated by ab initio simulations. The integration with predictors of formation energy permits the use of a Monte Carlo Tree Search algorithm to improve the generation of meaningful structures. Our approach challenges conventional representations of crystals, and demonstrates the potential of LLMs for learning effective 'world models' of crystal chemistry, which will lead to accelerated discovery and innovation in materials science.

Motivation & Objective

• Motivate fast, flexible generation of plausible inorganic crystal structures to accelerate CSP workflows.
• Develop a model that treats crystal structures as token sequences in CIF format and learns a world-model of crystal chemistry.
• Demonstrate generation of unseen structures and assess generalization across structural classes and space groups.
• Integrate with energy-predictor-guided Monte Carlo Tree Search to improve generation of physically meaningful candidates.

Proposed method

• Train a decoder-only Transformer on millions of CIF files (25M-parameter small model and 200M-parameter large model).
• Represent CIF as a token sequence and predict the next token autoregressively to generate new CIF files conditioned on prompts (cell composition, space group).
• Evaluate syntactic validity and chemical plausibility of generated CIFs against held-out test structures.
• Assess generalization by prompting with unseen formulas and space groups, and by using a challenge set of literature-derived structures not in training.
• Use Monte Carlo Tree Search (MCTS) with an ALIGNN formation-energy-per-atom predictor to guide sampling toward lower-energy structures.

Figure 1 : a Core concepts in training a Large Language Model of CIF files: A CIF file (left) is converted into a sequence of symbols, through tokenization. The sequence is processed by the model, which produces a list of probability distributions over the vocabulary, for each corresponding symbol i

Experimental results

Research questions

• RQ1 Can CrystaLLM generate syntactically valid CIF files for unseen inorganic structures?
• RQ2 How well does CrystaLLM generalize to unseen compositions and space groups?
• RQ3 Does including space-group conditioning improve generation quality and match to known structures?
• RQ4 How does CrystaLLM compare to diffusion/ VAE-based CSP models on benchmark datasets?
• RQ5 Can search strategies (MCTS with energy prediction) improve the quality of generated structures?

Key findings

• A small model achieved 93.8% valid CIF generation without space group and 94.0% with space group on the held-out test set; longest valid generated CIF length was 1145 tokens (no space group) and 970 tokens (with space group).
• With space-group conditioning, structural consistency metrics (e.g., space group consistency 99.1%, atom site multiplicity consistency 99.4%) were high in generated CIFs for the test set.
• Across a challenge set of 70 structures (58 unseen from literature, 12 seen in training), the small model achieved 85.7% successful generation without space group and 88.6% with space group, while the large model achieved 87.1% and 91.4%, respectively; unseen-match rates reached up to 41.4% for the large model with space group.
• CrystaLLM matched true structures in the challenge set for unseen compounds at rates up to about 40% with the large model (and higher when space group is provided).
• CrystaLLM outperformed CDVAE and DiffCSP on several CSP benchmarks in RMSE when using 20 samples per test composition, and demonstrated a unique ability to condition on symmetry space groups.
• The approach enables structure generation by analogy (e.g., generating ZrMn6Sn6-like motifs with substitutions) and can produce plausible structures for complex classes like rutiles, spinels, elpasolites, and pyrochlores.