Paper: Fine-Tuned Language Models Generate Stable Inorganic Materials as Text (ICLR-2024)

Reviewer comment (Novelty):
Can the authors provide a more detailed comparison with existing generative models in materials discovery, discussing the advantages and limitations of the proposed LLM-based approach compared to other state-of-the-art methods?

Author response:
Thank you for your review and your supportive comments. We respond to your questions pointwise. Note we have a separate [general post](https://openreview.net/forum?id=vN9fpfqoP1&noteId=wbNkvYNin2) outlining the key contributions in the paper.

**Existing generative models**

State-of-the-art generative models of materials from recent work can be grouped in two categories (1) diffusion models trained to denoise atomic coordinates (2) language models trained from scratch on crystal data. 

**(1) Overview**

- *Diffusion models*: Methods like CDVAE [1] and DiffCSP [2] learn a distribution over crystal structures by learning a neural network that takes noise-corrupted coordinates and/or lattice parameters and predicts the original input without noise. During sampling, examples are generated from noise through repeated applications of the denoising model. Because these methods operate on continuous representations of the coordinates and lattice (as opposed to discrete string representation), they can naturally incorporate symmetries into the architecture of the denoising model, for example translational equivariance. 

- *Language models*: Prior work on language models for crystals, e.g. [3,4], focuses on training models from scratch on crystal data. In these papers, GPT-2 style models with 10-100 million parameters are trained on custom vocabularies to represent element identities and atomic coordinates. [3] uses a relatively small dataset of around 50,000 materials while [4] uses millions of CIF files. [3] applies augmentations to encourage invariance in the model while [4] does not. 

- *Our method*: We finetune a pretrained large language model using a relatively small number of additional trainable parameters (between 3 and 35 million). We use LLaMA-2’s default tokenizer, which allows us to train on both crystal data and English text prompts. We apply augmentations to encourage invariance and measure learned invariance with a new metric useful to many applications of LLMs on atomistic data.

**(2) Capabilities**

- *Diffusion models*: Because diffusion models are trained to correct noisy inputs, they are often directly applicable to infilling tasks, e.g. for editing a small part of an existing structure. Additionally, because diffusion models act directly on continuous representations of coordinates they often exhibit higher rates of structural validity because notions of distance between atoms are more easily accessible to the model. On the other hand, compared to autoregressive language models, diffusion models often display lower rates of compositional validity, likely because compositions are naturally discrete sequences, and autoregressive language models remain state of the art generative models for discrete sequence. 

- *Language models*: Because language models are commonly used for text processing, they naturally incorporate text-conditional generation, which is explored to a limited degree in [4] but not in [3]. By default, language models are slightly less natural for editing tasks, unless trained with carefully chosen prompts, which is not explored in prior work. Because language models operate on string representations of atomic coordinates, they often have lower rates of structural validity because notions of distance/overlapping atoms must be learned over discrete sequences.

- *Our method*: By using pretrained language models and carefully engineered prompts, we can achieve the best of both worlds. Our models are capable of editing existing structures and conditioning on complex English prompts. Text pretraining also leads to models that more easily learn programmatic abstractions over text, e.g. for analyzing positions in 3D space. Thus we see that finetuned LLMs, especially larger models like LLaMA-2 70B, achieve high rates of both structural and compositional validity. 

**(3) Generality**
		
- *Diffusion models*: Models like CDVAE [1] and DiffCSP [2] are tailored specifically to generating crystals. The same methodology can not be applied to other atomistic data without major changes to the architecture and training procedure. 

- *Language models*: Unlike the diffusion models above, the training procedures used for language modeling are easier to generalize across different forms of atomistic data, as shown in [3]. When using custom vocabularies, however, the model weights themselves cannot be shared, despite shared training and sampling procedures. 

- *Our method*: By using a general-purpose tokenization method and a pretrained language model, we create a procedure that applies equally to natural language and string-formatted atomistic data. Though we focused primarily on crystals, a single LLM could also be trained simultaneously on other data types, such as small organic molecules or proteins.

Reviewer comment (Evaluation):
Are there any potential drawbacks or limitations of using LLMs for materials discovery, such as computational complexity or interpretability of the generated structures?

Author response:
**Limitations of LLMs**

Currently LLMs are relatively compute expensive, although by a smaller margin than one might expect. As we outline in the general rebuttal note, smaller LLMs (13B) have throughput on par with many baseline models, while larger models (70B) are only 2-3 times slower in throughput. Recently released models, e.g. Mistral 7B [5], have also matched the performance of large models with significantly less parameters. We expect this general trend to continue and lead to faster sampling times, as the demand for tractable base models has only intensified.

Current conditional generation from LLMs is not perfect. There are cases in which the composition is specified but the resulting crystal is inconsistent with this composition. Strict instruction following can be an issue for LLMs in general, and authors have proposed approaches that we could also apply in our setting, such as classifier-free guidance [6]. If we wanted to strictly enforce composition constraints, we could adapt the infilling framework to generate positions while keeping atom identities fixed, but we chose to use the more general framework for simplicity and because vanilla composition conditioning was successful frequently enough to be practical.

Hallucinations of non-existent elements can occur when using vanilla sampling from the model. If these mistakes are left unhandled, they can waste compute on samples that have no chance of being valid. Luckily, it is easy to prevent these mistakes with simple constrained sampling techniques. In the simplest case, we can simply lower the likelihood of tokens that do not occur in atom identities to zero. For ambiguous tokens that can occur within element identities or other non-element strings, we can apply sampling utilities available in huggingface [7].

Reviewer comment (Experiments):
In the paper, the authors mention that "We chose not to augment the ordering of atoms because these variables often contained valuable information". How to ensure that the atomic order will not change the results?

Author response:
**Role of atomic ordering** 

Although the global ordering of the atoms is arbitrary, the ordering in the data typically contains meaningful structure. For example, consider the following crystal for Pd3(PbSe)2, taken directly from Materials Project [8],

```
Lattice
    abc: 6.29 6.29 6.34
    angles: 60.3 60.3 60.0
Coordinates (fractional)
	Pd (0.5, 0.5, 0.0)
	Pd (0.5, 0.0, 0.0)
	Pd (0.0, 0.5, 0.0)
	Pb (0.5, 0.5, 0.5)
	Pb (0.0, 0.0, 0.0)
	Se (0.21, 0.21, 0.36)
	Se (0.79, 0.79, 0.64)
```

The elements are grouped together, and there are patterns evident in the atom placements within each element type. For example, we can imagine placing the Palladium atoms first, then the lead atoms at opposite corners in a cubic structure, and lastly the Selenium atoms can be placed between the existing atoms. In earlier versions of our method, we experimented with permuting the order of the atoms, and the model was still able to learn effectively, but the sample metrics (e.g. validity) were worse by small margin, so we ultimately decided this augmentation wasn’t optimal. Our hypothesis is that structures like the ones described above help the model construct plausible crystals.

Reviewer comment (Writing):
In the paper, the authors mention that "we only run VASP calculations on materials that have already been predicted as metastable by M3GNet", while in the caption of figure 3, "we only run VASP on structures predicted to be stable by M3GNet." What are the differences between the "stable" and "metastable"?

Author response:
**Definition of stable and metastable**

We defined a bulk material as stable if $\hat{E}_{\text{hull}} < 0.0$ and metastable if $< 0.1$. Stability corresponds to being below the energy hull (a very restrictive condition), while metastability just corresponds with being near the energy hull (a less restrictive condition).

Reviewer comment (Experiments):
While one anticipates good performance from LLMs on standard evaluation metrics, especially with the likes of LLaMA-70B, the critical matter lies in the practical application in experiments.

Author response:
>“one anticipates good performance from LLMs on standard evaluation metrics, especially with the likes of LLaMA-70B” 

Although LLMs and parameter-efficient finetuning have become popular for many tasks in vision and natural language processing, they are not commonly applied to problems in materials science, a field that typically does not typically utilize foundation models [1]. Our application is very nonstandard for language models and many researchers would not expect language models to perform well on this task because they are not designed with modeling 3D space in mind. Our paper is surprising and exciting because it shows that text pretraining allows LLMs to model physical constraints very effectively, with a small number of trainable parameters on par with domain-specific models trained from scratch (4.5 million trainable parameters for CDVAE [2] vs. 3.5 million trainable parameters for LLaMA-2 7B with LoRA). 

>"the critical matter lies in the practical application in experiments"

We strongly agree that practical performance is critical, and this belief was the motivation behind our extensive evaluations with DFT. These simulations verify that the model can generate plausible materials at a much higher rate than competing approaches and therefore requires less filtering before deployment in real physical experiments. The fact that LLMs consistently generate stable materials is very strong evidence that LLMs are learning fundamental properties of materials and not simply superficial correlations. The computational burden of our methods is also not prohibitive, as we detail in our general rebuttal note. Generation from our smallest LLaMA model requires only about 24 mins for 10,000 samples, while achieving better rates of metastability under M3GNet and comparable rates of stability under DFT. There are countless teams worldwide deploying real applications on top of LLaMA models of all sizes, and LLaMA-2 70B was downloaded 177,000 times from Hugging Face last month alone. Use of these models, therefore, doesn’t preclude practical applications and can in fact make it easier to leverage general advances in training and deploying LLMs (e.g. the GPT-4 fine-tuning API).

Reviewer comment (Novelty):
The technical contribution is relatively limited, where the tokenization, the prompt design and the objectives for pertaining are all well studied in the past.

Author response:
Having clarified the significance of using pretraining, we now expand upon the unique facets of our method and their significance. While prompt design and tokenization are established research areas, they have not been extensively studied within generative models of materials, or generative models of molecules more broadly (e.g. protein structures). We cite two papers that use language models to sample novel materials, both of which were published within the last six months [2,4]. Neither of these papers explores prompt design, and only one ablates tokenization methods [2]. Both papers use models trained from scratch on domain-specific data. 

**Tokenization**

Notably, our approach to tokenization is distinctly different from prior work on modeling atomic structures with language models. Instead of using a special vocabulary and training models from scratch, we use LLaMA-2’s existing tokenizer. This choice allows us to easily process both encoded crystals and text data. In early experiments, we tried out many other approaches, including fine-tuning LLaMA-2 models with additional tokens specific to crystal data. These methods were more challenging to train and didn’t lead to any improvements over using a shared tokenizer. We include a set of example training losses below:

| | Epoch 1 | Epoch 2 | Epoch 3 | Epoch 4 | Epoch 5 |
| -------- | -------- | -------- | -------- | -------- | -------- |
| Special Crystal Tokens | 0.783 | 0.693 | 0.623 | 0.611 | 0.588 |
| Shared Tokenization | 0.457 | 0.432 | 0.424 | 0.401 | 0.385 |

**Prompt design**

There are many important decisions involved both in text formatting (e.g the choice of fractional or absolute coordinates) and augmentation of the input data (e.g. translation or permutation augmentations on coordinates). We refined our choices over weeks of training and evaluation. As a simple example, we provide average validity numbers (using low temperature sampling) from earlier experiments on LLaMA-2 7B models trained different formatting styles

| Setting | Structural Validity | Compositional Validity |
| -------- | -------- | -------- |
| Fractional coords | 0.914 | 0.832 |
| Absolute coords  | 0.908  | 0.805 |
| No permutations | 0.925 | 0.829 |
| With permutations | 0.892  | 0.817 |

Even the phrasing of the instruction prompt can have a significant impact on the ultimate rates of validity, for example by providing proper context for element symbols:

| Setting | Structural Validity | Compositional Validity |
| -------- | -------- | -------- |
| Prompt 1 | 0.877 | 0.795 |
| Prompt 2 | 0.929 | 0.839 | 
| Prompt 3 | 0.957 | 0.875 | 

Prompt 1: “Bulk material description:\n”

Prompt 2: “Below is a description of a bulk material, starting with a list of atom types for the elemental composition followed by the lengths and angles of the lattice vectors and finally the atom type and coordinates for each atom within the lattice:\n\n"

Prompt 3: "Please provide the description of a stable bulk material, starting with a list of atom types for the elemental composition followed by the lengths and angles of the lattice vectors and finally the atom type and coordinates for each atom within the lattice.\n\nAn example of a bulk material description is:\n\n{example}\n\nInclude the new material description below:\n\n"

The length of Prompt 3 was prohibitive, so we ultimately abandoned it, despite promising results. 

**In closing**

Ultimately, we think that simplicity is a virtue. Using our method, it is easy to adapt popular LLMs methods directly to improving generative models of materials. Because of the generality of our approach, it would also be straightforward to extend to other atomic structures, such as drug-like compounds, protein structures, or nucleic acids.

We hope that our rebuttal addresses the concerns outlined in your review. If you have remaining questions, we would be happy to discuss any details further. We would appreciate it if you would consider raising your score in light of our response.

**References**

[1] J. Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. ArXiv, abs/2106.09685, 2021.

[2] Daniel Flam-Shepherd and Alán Aspuru-Guzik. Language models can generate molecules, materials, and protein binding sites directly in three dimensions as xyz, cif, and pdb files. arXiv preprint arXiv:2305.05708, 2023

[3] Tian Xie, Xiang Fu, Octavian-Eugen Ganea, Regina Barzilay, and Tommi Jaakkola. Crystal diffusion variational autoencoder for periodic material generation. arXiv preprint arXiv:2110.06197, 2021.

[4] Luis M Antunes, Keith T Butler, and Ricardo Grau-Crespo. Crystal structure generation with autoregressive large language modeling. arXiv preprint arXiv:2307.04340, 2023.

Reviewer comment (Experiments):
Since the nature of the material data is quite different compared to the natural language data, I’m wondering whether the pretrained LLAMA2 is offering any additional value. It would be helpful if one can show the performance with and without loading the pretrained checkpoint from LLAMA2 tasks.

Author response:
Thank you for your review. We would like to clarify that text pretraining is essential to our method for two reasons.
1. It would be impractically expensive or computationally infeasible to train models with up to 70B parameters from scratch on our data. Using a pretrained model with LoRA [1] offers the benefits of model scale while maintaining tractability and limiting overfitting, as the actual number of trainable parameters can be relatively small.

2. Pretraining on text data yields a model that can be conditioned on text for free, and text conditioning opens up a huge new realm of exciting possibilities, like conditioning samples on desired properties. It would be challenging to achieve a similar result from scratch without significantly expanding the size of the dataset (to improve general text understanding) and without essentially training a general-purpose language model in the process. 

To better understand the first point, let’s quickly review the exact details of the finetuning procedure. We are using low-rank adapters (LoRA), as opposed to end-to-end finetuning, and this means we are adding a small number of additional parameters to an existing, frozen model. The easiest way to see the difference between this approach and training a model from scratch (as in [2]) is to compare the training loss over the first few epochs of training.

| Model | Epoch 1 | Epoch 2 | Epoch 3 | Epoch 4 | Epoch 5 |
| -------- | -------- | -------- | -------- | -------- | -------- |
| GPT-2 (from scratch) | 0.946 | 0.878 | 0.807 | 0.757 | 0.740 |
| LLaMA-13B (LoRA) | 0.457 | 0.432 | 0.424 | 0.401 | 0.385 |
| LLaMA-70B (LoRA) | 0.402 | 0.344 | 0.325 | 0.305 | 0.296 |

If we attempt to run LoRA finetuning with randomly initialized parameters for the LLaMA-2 7B model we observe an immediate and significant difference in the training losses:

| Model | 1 Iter | 0.33 Epochs | 0.66 Epochs | 1 Epoch  |
| -------- | -------- | -------- | -------- | -------- |
| Random | 13.46 | 1.53 | 0.81 | 0.78 | 
| Pre-trained | 1.57 | 0.47 | 0.41 | 0.39 | 

While LoRA finetuning is tractable because 99.95% of the model is frozen, finetuning a LLaMA-2 model end-to-end in half-precision would require at least 4 times as many GPUs, making it infeasible for all but a handful of researchers. When using LoRA, even though the base models are large the number of trainable parameters is very small. In fact, the LLamA-2 7B model has less trainable parameters than one of the baseline methods we compared (CDVAE) [3]. The number of trainable parameters for each of our models and the baseline models is shown below:

*CDVAE [3]*: 4.5 million trainable parameters (100% of total)

*LM-CH/AC [2]*: 1-100 million trainable parameters (100% of total)

*LLaMA-2 7B*: 3.5 million trainable parameters (0.05% of total)

*LLaMA-2 13B*: 6.5 million trainable parameters (0.05% of total)

*LLaMA-2 70B*: 35 million trainable parameters (0.05% of total)