VAE for the prediction of the relaxed state of materials

Inspiration

The accelerating pace of technological innovation is revolutionizing every corner of society, including scientific research. With the increasing compute power, computational scientists are becoming evermore important to complement more traditional fields of science. Over half of a modern high-performance cluster (HPC) is used in the ab-initio calculation of material properties, and most of the calculation is spent on finding the relaxed configuration of materials. Currently, the typical process to do so is to start from an initial configuration and do some variation of gradient descent on the potential energy landscape using atomic forces until convergence is reached. This process is often very time-consuming and expensive, and a newer, faster method is warranted in order to further accelerate the pace of materials discovery.

What it does

The set of relaxed configurations occupy only a very slim fraction of the high-dimensional configurational space, and this fact inspires us to train a machine-learning model that learns the distribution of relaxed structures in order to generate a relaxed structure given a user input. We use a variational autoencoder (VAE) for this purpose. After it is trained, the user would be able to input an arbitrary configuration, and it would generate a relaxed configuration that is "near" the initial structure.

How we built it

As training data, we downloaded a 11 GB dataset of relaxation trajectories from a materials database called OpenMaterials. We extracted the initial (unrelaxed) structures and the final (fully relaxed) structures and formatted them in a dataframe. Meanwhile, we programmed the basic outline of the VAE on PyTorch.

Challenges we ran into

The construction and adjustment from the vanilla VAE architecture that needed to be made for our project was very challenging. Specifically, there are certain inductive biases that must be incorporated to the ML model for inputs of the physical nature, such as invariance or equivariance to translations and rotations.

Accomplishments that we're proud of

We learned a lot about new ML architectures, including VAEs. The time limit of the hackathon gave us the urgency to learn a lot of new things, and it was very effective.

What we learned

Same as above.

What's next for VAE for the prediction of the relaxed state of materials

We would like to incorporate these inductive biases to complete the construction of our model, and then we would like to submit a proposal to get the necessary compute resources, such as time on GPU clusters, in order to train our model on this large dataset.

Built With

• python