[Paper Review] Fine-Tuned Language Models Generate Stable Inorganic Materials as Text

Nate Gruver, Anuroop Sriram+4 more|arXiv (Cornell University)|Feb 6, 2024

Modular Robots and Swarm Intelligence34 citations

TL;DR

Fine-tuned large language models can generate 3D crystal structures as text that are valid and metastable by ML-potential or DFT standards, outperforming CDVAE in metastable/stable yield at scale.

ABSTRACT

We propose fine-tuning large language models for generation of stable materials. While unorthodox, fine-tuning large language models on text-encoded atomistic data is simple to implement yet reliable, with around 90% of sampled structures obeying physical constraints on atom positions and charges. Using energy above hull calculations from both learned ML potentials and gold-standard DFT calculations, we show that our strongest model (fine-tuned LLaMA-2 70B) can generate materials predicted to be metastable at about twice the rate (49% vs 28%) of CDVAE, a competing diffusion model. Because of text prompting's inherent flexibility, our models can simultaneously be used for unconditional generation of stable material, infilling of partial structures and text-conditional generation. Finally, we show that language models' ability to capture key symmetries of crystal structures improves with model scale, suggesting that the biases of pretrained LLMs are surprisingly well-suited for atomistic data.

Motivation & Objective

• Demonstrate that fine-tuned large language models can generate valid crystal structures encoded as text.
• Evaluate stability of generated materials using ML potentials and DFT relaxations.
• Compare against diffusion-model baselines and analyze symmetry learning and model scale effects.
• Showcase text-conditioned generation and infilling as practical design tools for materials discovery.

Proposed method

• Encode crystals as newline-delimited strings and fine-tune base LLMs (LLaMA-2) with a multitask curriculum and translation augmentations.
• Use low-rank adapters (LoRA) and 4-bit quantization to train on Materials Project data.
• Prompt design enabling unconditional generation, text-conditioned generation, and infilling.
• Augment training with random translations to encourage translational invariance.
• Evaluate samples with structural/compositional validity, diversity, and energy above hull estimates from ML potentials (M3GNet) and DFT (VASP).
• Propose Increase in Perplexity under Transformation (IPT) as a metric for translation invariance.

Experimental results

Research questions

• RQ1 Can fine-tuned LLMs generate valid and metastable inorganic crystal structures when encoded as text?
• RQ2 How does model scale affect the ability to learn crystal symmetries and generate stable materials?
• RQ3 How do text-conditioned generation and infilling perform for targeted compositions, space groups, and stability?
• RQ4 How do LLMs compare to diffusion-based crystal generators (CDVAE) in stability, diversity, and sampling speed?

Key findings

• Fine-tuned LLaMA-2 models generate high validity rates for both structural and compositional aspects.
• LLaMA-2 70B with τ=0.7 achieves 49.8% metastable samples and 10.6% stable samples based on Table 1 metrics, outperforming CDVAE in metastable/stable yield.
• Larger models better learn translational invariance (IPT metric) and diversity/novelty tradeoffs.
• Text-conditioning reliably yields materials with the desired composition and, to a lesser extent, space group and stability, given oracle labels.
• Infilling with the 70B model guides mutations toward higher metastable yields without sacrificing diversity.
• Compared to CDVAE, LLaMA-2 models deliver faster generation and often greater sample diversity.