Crystal Diffusion Variational Autoencoder for Periodic Material Generation

Generating the periodic structure of stable materials is a long-standing challenge for the material design community . This task is difficult because stable materials only exist in a low-dimensional subspace of all possible periodic arrangements of atoms : 1 ) the coordinates must lie in the local energy minimum defined by quantum mechanics , and 2 ) global stability also requires the structure to follow the complex , yet specific bonding preferences between different atom types . Existing methods fail to incorporate these factors and often lack proper invariances . We propose a Crystal Diffusion Variational Autoencoder ( CDVAE ) that captures the physical inductive bias of material stability . By learning from the data distribution of stable materials , the decoder generates materials in a diffusion process that moves atomic coordinates towards a lower energy state and updates atom types to satisfy bonding preferences between neighbors . Our model also explicitly encodes interactions across periodic boundaries and respects permutation , translation , rotation , and periodic invariances . We significantly outperform past methods in three tasks : 1 ) reconstructing the input structure , 2 ) generating valid , diverse , and realistic materials , and 3 ) generating materials that optimize a specific property . We also provide several standard datasets and evaluation metrics for the broader machine learning community . 1 INTRODUCTION . Solid state materials , represented by the periodic arrangement of atoms in the 3D space , are the foundation of many key technologies including solar cells , batteries , and catalysis ( Butler et al. , 2018 ) . Despite the rapid progress of molecular generative models and their significant impact on drug discovery , the problem of material generation has many unique challenges . Compared with small molecules , materials have more complex periodic 3D structures and can not be adequately represented by a simple graph like small molecules ( Figure 1 ) . In addition , materials can be made up of more than 100 elements in the periodic table , while molecules are generally only made up of a small subset of atoms such as carbon , oxygen , and hydrogen . Finally , the data for training ML models for material design is limited . There are only ∼200k experimentally known inorganic materials , collected by the ICSD ( Belsky et al. , 2002 ) , in contrast to close to a billion purchasable molecules in ZINC ( Irwin & Shoichet , 2005 ) . The key challenge of this task is in generating stable materials . Such materials only exist in a lowdimensional subspace of all possible periodic arrangements of atoms : 1 ) the atom coordinates must lie in the local energy minimum defined by quantum mechanics ( QM ) ; 2 ) global stability also requires the structure to follow the complex , yet specific bonding preferences between differ- ent atom types ( section 3.2 ) . The issue of stability is unique to material generation because valency checkers assessing molecular stability are not applicable to materials . Moreover , we also have to encode the interactions crossing periodic boundaries ( Figure 1 , middle ) , and satisfy permutation , translation , rotation , and periodic invariances ( section 3.1 ) . Our goal is to learn representations that can learn features of stable materials from data , while adhering to the above invariance properties . We address these challenges by learning a variational autoencoder ( VAE ) ( Kingma & Welling , 2014 ) to generate stable 3D materials directly from a latent representation without intermediates like graphs . The key insight is to exploit the fact that all materials in the data distribution are stable , therefore if noise is added to the ground truth structure , denoising it back to its original structure will likely increase stability . We capture this insight by designing a noise conditional score network ( NCSN ) ( Song & Ermon , 2019 ) as our decoder : 1 ) the decoder outputs gradients that drive the atom coordinates to the energy local minimum ; 2 ) it also updates atom types based on the neighbors to capture the specific local bonding preferences ( e.g. , Si-O is preferred over Si-Si and O-O in SiO2 ) . During generation , materials are generated using Langevin dynamics that gradually deforms an initial random structure to a stable structure . To capture the necessary invariances and encode the interactions crossing periodic boundaries , we use SE ( 3 ) equivariant graph neural networks adapted with periodicity ( PGNNs ) for both the encoder and decoder of our VAE . Our theoretical analysis further reveals an intriguing connection between the gradient field learned by our decoder and an harmonic force field . De facto , the decoder utilizes the latter to estimate the forces on atoms when their coordinates deviate from the equilibrium positions . Consequently , this formulation provides an important physical inductive bias for generating stable materials . In this work , we propose Crystal Diffusion Variational AutoEncoder ( CDVAE ) to generate stable materials by learning from the data distribution of known materials . Our main contributions include : • We curate 3 standard datasets from QM simulations and create a set of physically meaningful tasks and metrics for the problem of material generation . • We incorporate stability as an inductive bias by designing a noise conditional score network as the decoder of our VAE , which allows us to generate significantly more realistic materials . • We encode permutation , translation , rotation , and periodic invariances , as well as interactions crossing periodic boundaries with SE ( 3 ) equivariant GNNs adapted with periodicity . • Empirically , our model significantly outperforms past methods in tasks including reconstructing an input structure , generating valid , diverse , and realistic materials , and generating materials that optimize specific properties . 2 RELATED WORK . Material graph representation learning . Graph neural networks have made major impacts in material property prediction . They were first applied to the representation learning of periodic materials by Xie & Grossman ( 2018 ) and later enhanced by many studies including Schütt et al . ( 2018 ) ; Chen et al . ( 2019 ) . The Open Catalyst Project ( OCP ) provides a platform for comparing different architectures by predicting energies and forces from the periodic structure of catalytic surfaces ( Chanussot et al. , 2021 ) . Our encoder and decoder PGNNs directly use GNN architectures developed for the OCP ( Klicpera et al. , 2020b ; 2021 ; Shuaibi et al. , 2021 ; Godwin et al. , 2021 ) , which are also closely related to SE ( 3 ) equivariant networks ( Thomas et al. , 2018 ; Fuchs et al. , 2020 ) . Quantum mechanical search of stable materials . GNN-based property prediction models are fundamentally limited by the space of ∼200k known materials . Predicting the structure of unknown materials requires very expensive random search and QM simulations , and is considered a grand challenge in materials discovery ( Oganov et al. , 2019 ) . State-of-the-art methods include random sampling ( Pickard & Needs , 2011 ) , evolutionary algorithms ( Wang et al. , 2012 ) , substituting elements in known materials ( Sun et al. , 2019 ) , etc. , but they generally have success rates less than 1 % and require extensive computation even on relatively small problems . Material generative models . Past material generative models mainly focus on two different approaches , and neither incorporate stability as an inductive bias . The first approach treats materials as 3D voxel images , but the process of decoding images back to atom types and coordinates often results in low validity , and the models are not rotationally invariant ( Hoffmann et al. , 2019 ; Noh et al. , 2019 ; Court et al. , 2020 ; Long et al. , 2021 ) . The second directly encodes atom coordinates , types , and lattices as vectors ( Ren et al. , 2020 ; Kim et al. , 2020 ; Zhao et al. , 2021 ) , but the models are generally not invariant to any Euclidean transformations . Another related method is to train a force field from QM forces and then apply the learned force field to generate stable materials by minimizing energy ( Deringer et al. , 2018 ) . This method is conceptually similar to our decoder , but it requires additional force data and can only be applied to 1-2 elements due to the exponentially increased data need for more elements . Remotely related works include generating contact maps from chemical compositions ( Hu et al. , 2021 ; Yang et al. , 2021 ) and building generative models only for compositions ( Sawada et al. , 2019 ; Pathak et al. , 2020 ; Dan et al. , 2020 ) . Molecular conformer generation and protein folding . Our decoder that generates the 3D atomic structures via a diffusion process is closely related to the diffusion models used for molecular conformer generation ( Shi et al. , 2021 ) . The key difference is that our model does not rely on intermediate representations like molecular graphs . G-SchNet ( Gebauer et al. , 2019 ) is more closely related to our method because it directly generates 3D molecules atom-by-atom without relying on a graph . Another closely related work is E-NFs ( Satorras et al. , 2021 ) that use a flow model to generate 3D molecules . In addition , score-based and energy-based models have also been used for molecular graph generation ( Liu et al. , 2021 ) and protein folding ( Wu et al. , 2021 ) . Flow models have also been used for molecular graph generation ( Shi et al. , 2020 ; Luo et al. , 2021 ) . However , these generative models do not incorporate periodicity , which makes them unsuitable for materials . 3 PRELIMINARIES . 3.1 PERIODIC STRUCTURE OF MATERIALS . Any material structure can be represented as the periodic arrangement of atoms in the 3D space . As illustrated in Figure 1 , we can always find a repeating unit , i.e . a unit cell , to describe the infinite periodic structure of a material . A unit cell that includes N atoms can be fully described by 3 sets : 1 ) atom types A = ( a0 , ... , aN ) ∈ AN , where A denotes the set of all chemical elements ; 2 ) atom coordinates X = ( x0 , ... , xN ) ∈ RN×3 ; and 3 ) periodic lattice L = ( l1 , l2 , l3 ) ∈ R3×3 . The periodic lattice defines the periodic translation symmetry of the material . Given M = ( A , X , L ) , the infinite periodic structure can be represented as , { ( a′i , x′i ) |a′i = ai , x′i = xi + k1l1 + k2l2 + k3l3 , k1 , k2 , k3 ∈ Z } , ( 1 ) where k1 , k2 , k3 are any integers that translate the unit cell using L to tile the entire 3D space . The composition of a material denotes the ratio of different elements that the material is composed of . Given the atom types of a material with N atoms A ∈ AN , the composition can be represented as c ∈ R|A| , where ci > 0 denotes the percentage of atom type i and ∑ i ci = 1 . For example , the composition of diamond in Figure 1 has c6 = 1 and ci = 0 for i 6= 6 because 6 is the atomic number of carbon . Invariances for materials . The structure of a material does not change under several invariances . 1 ) Permutation invariance . Exchanging the indices of any pair of atoms will not change the material . 2 ) Translation invariance . Translating the atom coordinates X by an arbitrary vector will not change the material . 3 ) Rotation invariance . Rotating X and L together by an arbitrary rotation matrix will not change the material . 4 ) Periodic invariance . There are infinite different ways of choosing unit cells with different shapes and sizes , e.g. , obtaining a bigger unit cell as an integer multiplier of a smaller unit cell using integer translations . The material will again not change given different choices of unit cells . Multi-graph representation for materials . Materials can be represented as a directed multi-graph G = { V , E } to encode the periodic structures following ( Xie & Grossman , 2018 ; Wells et al. , 1977 ; O ’ Keeffe & Hyde , 1980 ) , where V = { v1 , ... , vN } is the set of nodes representing atoms and E = { eij , ( k1 , k2 , k3 ) |i , j ∈ { 1 , ... , N } , k1 , k2 , k3 ∈ Z } is the set of edges representing bonds . eij , ( k1 , k2 , k3 ) denotes a directed edge from node i at the original unit cell to node j at the cell translated by k1l1 + k2l2 + k3l3 ( in Figure 1 right , ( k1 , k2 , k3 ) are labeled on top of edges ) . For materials , there is no unique way to define edges ( bonds ) and the edges are often computed using k-nearest neighbor ( KNN ) approaches under periodicity or more advanced methods such as CrystalNN ( Pan et al. , 2021 ) . Given this directed multi-graph , message-passing neural networks and SE ( 3 ) -equivariant networks can be used for the representation learning of materials .

Summary: This paper aims to address challenging stable crystal materials generation problems via diffusion variational autoencoder with graph representation learning. Several recent advances in generative models and GNN are combined together to develop the entire workflow from data distribution learning, prediction to sample generation and property optimization. The authors demonstrate their method by using three datasets and compare with three baseline methods.