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: The paper proposed a new generative model for 3D periodic material molecular structure. They designed a special variational autoencoder, where the decoder is parameterized by the denoising score-matching framework. Experiments demonstrate the model can successfully generate valid and realistic material, and optimize desired properties.