GenMS: Generative Hierarchical Materials Search

Language-to-Structure as Multi-Objective Search

A wealth of high-level semantic text data related to materials (e.g., Wikipedia articles, research papers, textbooks) are used to train LLMs, which can generate formulas from natural language (e.g., "give me the formula for a stable, chalcogenide with atom ratio 1:1:2 that's not in the ICSD database"). Meanwhile, previous work has used crystal databases (e.g., Materials Project, ICSD) with low-level information of crystal structures (e.g., locations of each atoms) to train diffusion models to generate crystal structures from chemical formulas. Meanwhile, materials databases (e.g., Materials Project, ICSD) contain low-level crystal structures, which can be used to train diffusion models. We can formulate the language-to-structure problem as a multi-objective search problem:

where $\pi_\text{hi}$ and $\pi_\text{lo}$ are LLM and diffusion models, and $R_\text{hi}$ and $R_\text{lo}$ are heuristic functions that can be specified by users (e.g., formation energy). $z$ and $x$ represents chemical formulas and crystal structures, respectively. We can optimize this objective using hierarchical search:

Diffusion with Compact Crystal Representation

In order to support efficient tree search at inference time, we represent each crystal structure as a point cloud using an Ax4 tensor, where A represents the number of atoms in a crystal, and the inner 4 dimensions representing the location of an atom along with its atomic number. We then train a diffusion model to predict predict atom locations and atomic number as a continuous value, and round the continuous value to the closest atomic number. The diffusion model architecture with this compact crystal representation is shown below. Each atom undergoes blocks consisting of multi-layer perceptrons followed by orderinvariant self-attention. The MLP and self-attention blocks are repeated times where each repetition increases the dimension of the hidden units. The concatenation of skip connections are employed as in other U-Net architectures.

Evaluation

We first conduct end-to-end evaluation of generating crystal structure from natural language for three families of crystals. Generative Hierarchical Materials Search (GenMS) significantly outperforms LLM prompting baselines in producing unique and low-energy (predicted by GNN) structures that satisfy user request. We further conduct DFT calculation to compute formation energy in eV/atom averaged across structures generated by GenMS. Values before "/" in the row "(GNN/DFT)" represent GNN predicted formation energy, and after ``/'' represent DFT computed formation energy. We report formation energy from GNN prior to relaxation, and formation energy from DFT post relaxation. DFT calculations for baselines are eliminated as many structures from the baselines do not follow user instruction. N/A represents $E_f$ predicted by GNN falling outside of the reasonable range.

Furthermore, when we increase the amount of compute during inference (i.e., sample more formula and structure candidates), GenMS can find more stable crystal structures with lower formation energy:

Generated Structures

We test GenMS on a set of ad hoc language inputs to generate plausible examples from well-known crystal families. GenMS is able to search for the corresponding structures that satisfy user requests and have plausible initial geometries. Visualization provided by VESTA.