C RYSTAL D IFFUSION V ARIATIONAL A UTOENCODER FOR P ERIODIC M ATERIAL G ENERATION
A BSTRACT
Generating the periodic structure of stable materials is a long-standing challenge for the material design community.
This task is difﬁcult 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 deﬁned by quantum mechanics, and 2) global stability also requires the structure to follow the complex, yet speciﬁc 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 signiﬁcantly 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 speciﬁc property.
We also provide several standard datasets and evaluation metrics for the broader machine learning community.
I NTRODUCTION
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 signiﬁcant 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 cannot be adequately represented by a simple graph like molecular graphs (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 molecules in ZINC (Irwin & Shoichet, 2005).
C
C 1 (1,0,0) (0,0,0) (0,0,1) (0,1,0) ... ...
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 deﬁned by quantum mechanics (QM) ) global stability also requires the structure to follow the complex, yet speciﬁc bonding preferences between different 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 speciﬁc local bonding preferences (e.g., Si-O is preferred over Si-Si and O-O in
SiO 2 ).
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 ﬁeld learned by our decoder and an harmonic force ﬁeld.
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 signiﬁcantly 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 signiﬁcantly outperforms past methods in tasks including reconstructing an input structure, generating valid, diverse, and realistic materials, and generating materials that optimize speciﬁc properties.
R ELATED W ORK
Material graph representation learning.
Graph neural networks have made major impacts in material property prediction.
They were ﬁrst applied to the representation learning of periodic materials by Xie & Grossman (2018) and later enhanced by many studies including Sch¨utt et al.
; Chen et al.
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.
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; Glass et al., 2006), substituting elements in known materials (Hautier et al., 2011), etc., but they generally have low success rates 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 ﬁrst 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 ﬁeld from QM forces and then apply the learned force ﬁeld to generate stable materials by minimizing energy (Deringer et al., 2018; Chen & Ong, 2022).
This method is conceptually similar to our decoder, but it requires additional force data which is expensive to obtain.
Remotely related works include generating contact maps from chemical compositions (Hu et al., 2021;
Yang et al., 2021) and building generative models only for chemical 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; Xu et al., 2021b).
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 ﬂow 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.
P RELIMINARIES 3. P ERIODIC 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 ﬁnd a repeating unit, i.e. a unit cell, to describe the inﬁnite periodic structure of a material.
A unit cell that includes N atoms can be fully described by 3 lists: 1) atom types A
= ( a 0 , ..., a N ) ∈
A N , where A denotes the set of all chemical elements; 2) atom coordinates X = ( x 0 , ..., x N ) ∈ R N × 3 ; and 3) periodic lattice L =
( l 1 , l 2 , l 3 ) ∈ R 3 × 3 .
The periodic lattice deﬁnes the periodic translation symmetry of the material.
Given M =
( A , X , L ) , the inﬁnite periodic structure can be represented as,
[Equation]
where k 1 , k 2 , k 3 are any integers that translate the unit cell using L to tile the entire 3D space.
The chemical 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 ∈ A N , the composition can be represented as c ∈ R | A | , where c i > 0 denotes the percentage of atom type i and (cid:80) i c i = 1 .
For example, the composition of diamond in Figure 1 has c 6 = 1 and c i =
for i (cid:54) = 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 inﬁnite 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 (Wells et al., 1977; O’Keeffe & Hyde, 1980; Xie & Grossman, 2018), where V = { v 1 , ..., v N } is the set of nodes representing atoms and E = { e ij, ( k 1 ,k 2 ,k 3 ) | i, j ∈ { 1 , ..., N } , k 1 , k 2 , k 3 ∈ Z } is the set of edges representing bonds.
e ij, ( k 1 ,k 2 ,k 3 ) denotes a directed edge from node i at the original unit cell to node j at the cell translated by k 1 l 1 + k 2 l 2 + k 3 l 3 (in Figure 1 right, ( k 1 , k 2 , k 3 ) are labeled on top of edges).
For materials, there is no unique way to deﬁne 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. Encode Add noises Predict 0.50.
Composition
Lattice + Denoise # of atoms + Rand. init.
Langevin dynamics
[Figure]
Generation
Conditional
Conditional
3. P ROBLEM DEFINITION AND ITS PHYSICAL ORIGIN
Our goal is to generate novel, stable materials M =
( A , X , L ) ∈
A N × R N × 3 × R 3 × 3 .
The space of stable materials is a subspace in A N × R N × 3 × R 3 × 3 that satisﬁes the following constraints.
The materials lie in the local minimum of the energy landscape deﬁned by quantum mechanics, with respect to the atom coordinates and lattice, i.e. ∂E/∂ X = 0 and ∂E/∂ L = 0 .
The material is globally stable and thus cannot decompose into nearby phases.
Global stability is strongly related to bonding preferences between neighboring atoms.
For example, in SiO 2 , each Si is surrounded by 4 O and each O is surrounded by 2 Si.
This conﬁguration is caused by the stronger bonding preferences between Si-O than Si-Si and O-O.
Generally, ﬁnding novel, stable materials requires very expensive random search and quantum mechanical simulations.
To bypass this challenge, we aim to learn a generative model p ( M ) from the empirical distribution of experimentally observed stable materials.
A successful generative model will be able to generate novel materials that satisfy the above constraints, which can then be veriﬁed using quantum mechanical simulations.
3. D IFFUSION MODELS
Diffusion models are a new class of generative models that have recently shown great success in generating high-quality images (Dhariwal & Nichol, 2021), point clouds (Cai et al., 2020; Luo & Hu, 2021), and molecular conformations (Shi et al., 2021).
There are several different types of diffusion models including diffusion probabilistic models (Sohl-Dickstein et al., 2015), noiseconditioned score networks (NCSN) (Song & Ermon, 2019), and denoising diffusion probabilistic models (DDPM)
(Ho et al., 2020).
We follow ideas from the NCSN (Song & Ermon, 2019) and learn a score network s θ ( x ) to approximate the gradient of a probability density ∇ x p ( x ) at different noise levels.
Let { σ i } Li =1 be a sequence of positive scalars that satisﬁes σ 1 /σ
= ...
=
σ L − 1 /σ L > 1 .
We deﬁne the data distribution perturbed by Gaussian noise σ as q σ ( x ) = (cid:82) p data ( t ) N ( x | t , σ 2 I ) d t .
The goal of NCSN is to learn a score network to jointly estimate the scores of all perturbed data distributions, i.e. ∀ σ ∈ { σ i } Li =1 : s θ ( x , σ )
≈ ∇ x q σ ( x ) .
During generation, NCSN uses an annealed Langevin dynamics algorithm to produce samples following the gradient estimated by the score network with a gradually reduced noise level.
P ROPOSED M ETHOD
Our approach generates new materials via a two-step process: 1) We sample a z from the latent space and use it to predict 3 aggregated properties of a material: composition ( c ), lattice ( L ), and number of atoms ( N ), which are then used to randomly initialize a material structure
˜ M
=
( ˜ A , ˜ X , L ) .
We perform Langevin dynamics to simultaneously denoise ˜ X and ˜ A conditioned on z to improve both the local and global stability of ˜ M and generate the ﬁnal structure of the new material.
To train our model, we optimize 3 networks concurrently using stable materials M =
( A , X , L ) sampled from the data distribution.
A periodic GNN encoder PGNN E NC ( M ) that encodes M into a latent representation z . 2)
A property predictor MLP A GG ( z ) that predicts the c , L , and N of
M from z . 3)
A periodic GNN decoder PGNN D EC ( ˜ M | z ) that denoises both ˜ X and ˜ A conditioned on z .
For 3), the noisy structure ˜ M
=
( ˜ A , ˜ X , L ) is obtained by adding different levels of noise to X and A .
The noise schedules are deﬁned by the predicted aggregated properties, with the motivation of simplifying the task for our decoder from denoising an arbitrary random structure from over ∼ 100 elements to a constrained random structure from predicted properties.
We train all three networks together by minimizing a combined loss including the aggregated property loss
L A GG , decoder denoising loss L D EC , and a KL divergence loss L KL for the VAE.
To capture the interactions across periodic boundaries, we employ a multi-graph representation (section 3.1) for both M and ˜ M .
We also use SE(3) equivariant GNNs adapted with periodicity as both the encoder and the decoder to ensure the permutation, translation, rotation, and periodic invariances of our model.
The CDVAE is summarized in Figure 2 and we explain the individual components of our method below.
The implementation details can be found in Appendix B.
Periodic material encoder.
PGNN E NC ( M ) encodes a material M as a latent representation z ∈ R D following the reparameterization trick in VAE (Kingma & Welling, 2014).
We use the multigraph representation (refer to section 3.1) to encode M , and PGNN E NC can be parameterized with an SE(3) invariant graph neural network.
Prediction of aggregated properties.
MLP A GG ( z ) predicts 3 aggregated properties of the encoded material from its latent representation z .
It is parameterized by 3 separate multilayer perceptrons (MLPs).
1) Composition c ∈ R | A | is predicted by minimizing the cross entropy between the ground truth composition and predicted composition, i.e. − (cid:80)
i p
i log c i . 2)
Lattice L ∈ R 3 × 3 is reduced to 6 unique, rotation invariant parameters with the Niggli algorithm (Grosse-Kunstleve et al., 2004), i.e., the lengths of the 3 lattice vectors, the angles between them, and the values are predicted with an MLP after being normalized to the same scale (Appendix B.1) with an L 2 loss.
3) Number of atoms N ∈ { 1 , 2 , ... } is predicted with a softmax classiﬁcation loss from the set of possible number of atoms.
L A GG is a weighted sum of the above 3 losses.
Conditional score matching decoder.
PGNN D EC ( ˜ M | z ) is a PGNN that inputs a noisy material ˜ M with type noises σ A , coordinate noises σ X , as well as a latent z , and outputs 1)
a score s X ( ˜ M | z ; σ A , σ X
) ∈ R N × 3 to denoise the coordinate for each atom towards its ground truth value, and 2) a probability distribution of the true atom types
p A ( ˜ M | z ; σ A , σ X ) ∈ R N ×| A | .
We use a
SE(3) graph network to ensure the equivariance of s X with respect to the rotation of ˜ M .
To obtain the noisy structures ˜ M , we sample σ A and σ X from two geometric sequences of the same length: { σ A ,j } Lj =1 , { σ X ,j } Lj =1 , and add the noises with the following methods.
For type noises, we use the type distribution deﬁned by the predicted composition c to linearly perturb true type distribution ˜
A ∼ ( 1 1+ σ A p
A
+ σ A 1+ σ A p c ) , where p A ,ij = 1 if atom i has the true atom type j and p A ,ij =
for all other j s, and p c is the predicted composition.
For coordinate noises, we add Gaussian noises to the true coordinates ˜
X ∼ N ( X , σ 2 X I ) .
PGNN D EC is parameterized by a SE(3) equivariant PGNN that inputs a multi-graph representation (section 3.1) of the noisy material structure and the latent representation.
The node embedding for node i is obtained by the concatenation of the element embedding of ˜ a i and the latent representation z , followed by a MLP, h 0 i =
MLP( e a (˜ a i ) (cid:107) z ) , where (cid:107) denotes concatenation of two vectors and e a is a learned embedding for elements.
After K message-passing layers, PGNN D EC outputs a vector per node that is equivariant to the rotation of ˜ M .
These vectors are used to predict the scores, and we follow Song & Ermon (2019); Shi et al.
(2021) to parameterize the score network with noise scaling: s X ( ˜ M | z ; σ A , σ X ) =
s X ( ˜ M | z ) /σ
X .
The node representations h Ki are used to predict the distribution of true atom types, and the type predictor is the same at all noise levels:
[Equation]
Periodicity inﬂuences denoising target.
Due to periodicity, a speciﬁc atom i may move out of the unit cell deﬁned by L when the noise is sufﬁciently large.
This leads to two different ways to deﬁne the scores for node i .
Ignore periodicity and deﬁne the target score as x i − ˜ x
i ; or 2) Deﬁne the target score as the shortest possible displacement between x i and ˜ x
i considering periodicity, i.e. d min ( x i , ˜ x i ) =
min k 1 ,k 2 ,k 3 ( x i − ˜ x
i + k 1 l 1 + k 2 l 2 + k 3 l 3 ) .
We choose 2) because the scores are the same given two different ˜ X that are periodically equivalent, which is mathematically grounded for periodic structures, and empirically results in much more stable training.
The training loss for the decoder L D EC can be written as, 2 L
[Equation]
where λ a denotes a coefﬁcient for balancing the coordinate and type losses, L a denotes the cross entropy loss over atom types, p A denotes the true atom type distribution.
Note that to simplify the equation, we follow the loss coefﬁcients in Song & Ermon (2019) for different σ X ,j and σ A ,j and factor them into Equation
[Equation]
[Table]
[Table]
[Table]
[Table]
[Table]
[Equation]
X t ← back to cell( X (cid:48) t , L ) 13:
A t = argmax p A ,t 14:
X 0 ← X T , A 0 ← A T
Material generation with Langevin dynamics.
After training the model, we can generate the periodic structure of material given a latent representation z .
First, we use z to predict the aggregated properties: 1) composition c , 2) lattice L , and 3) the number of atoms N .
Then, we randomly initialize an initial periodic structure ( A 0 , X 0 , L ) with the aggregated properties and perform an annealed Langevin dynamics (Song & Ermon, 2019) using the decoder, simultaneously updating the atom types and coordinates.
During the coordinate update, we map the coordinates back to the unit cell at each step if atoms move out of the cell.
The algorithm is summarized in Algorithm 1.
Connection between the gradient ﬁeld and a harmonic force ﬁeld.
The gradient ﬁeld s X ( ˜ M | z ) is used to update atom coordinates in Langevin dynamics via the force term, α j s X ,t .
In Appendix A, we show that α j s X ,t is mathematically equivalent to 2 a harmonic force ﬁeld F ( ˜ X ) = − k
( ˜ X −
X ) when the noises are small, where X is the equilibrium position of the atoms and k is a force constant.
Harmonic force ﬁeld, i.e. spring-like force ﬁeld, is a simple yet general physical model that approximates the forces on atoms when they are close to their equilibrium locations.
This indicates that our learned gradient ﬁeld utilizes the harmonic approximation to approximate QM forces without any explicit force data and generates stable materials with this physically motivated inductive bias.
E XPERIMENTS
We evaluate multiple aspects of material generation that are related to real-world material discovery process.
Past studies in this ﬁeld used very different tasks and metrics, making it difﬁcult to compare different methods.
Building upon past studies (Court et al., 2020; Ren et al., 2020), we create a set of standard tasks, datasets, and metrics to evaluate and compare models for material generation.
Experiment details can be found in Appendix D.
Tasks.
We focus on 3 tasks for material generation.
1) Reconstruction evaluates the ability of the model to reconstruct the original material from its latent representation z . 2)
Generation evaluates the validity, property statistics, and diversity of material structures generated by the model. 3)
Property optimization evaluates the model’s ability to generate materials that are optimized for a speciﬁc property.
Datasets.
We curated 3 datasets representing different types of material distributions.
1) Perov
(Castelli et al., 2012a;b) includes 18928 perovskite materials that share the same structure but differ in composition.
There are 56 elements and all materials have 5 atoms in the unit cell. 2)
Carbon-24 (Pickard, 2020) includes 10153 materials that are all made up of carbon atoms but differ in structures.
There is 1 element and the materials have 6 - 24 atoms in the unit cells.
3) MP- (Jain et al., 2013) includes 45231 materials that differ in both structure and composition.
There are
[Figure]
MP
[Figure]
F N V Rh O
O F N V Rb
O O Sn Zr O
F O Ru Ba
O
C C C C C CC
CCC C C CCC
CCC C C C C C C
C C C C C C CC C CC C
S Ti V S TiTiS Na S
O Cr O O EuOOOOOOSbSbO Eu
NO O TiNNNTiONOOTiOTiOONOOOTiOOTiN
Mg Al Si Si AlAl
F N V Rh O
O F N V Rb
O O Sn Zr O
F O Ba RuO
C CCCC CCC C CCC
C C C C C C
CC C C C C CCC C
CCC C C C C CC C CCC
V S Ti S S H Na
O Cr
O O EuOOOOOOSb
Eu SbO
ONNOOOOONNONNNNOOOOOOTi TiO OOOOTi O N OOOOOOOOTi OOOTi NONNNOOOOOTi
Al Mg Si Si AlAl
F N V Rh O
O F N OV Rb
O O Pd ZrO
F O Ru Ba
O
N Rh F O V O F Rb V N
OO Zr
O O Sn
F Ru O O F Ba
C C C C CC CC C C C C C CC
C C C C CC C CCC C C C C C C C C C CCC
V S S S Na SVS
O EuOOCrOOCrOMoEuOOOCrEu
N O N TiNNNNNO N NNTiOOOTiNONNTiOOTiNOTiOOON
Al Si AlAlAl
Mg Si
[Figure]
[Table]
89 elements and the materials have 1 - 20 atoms in the unit cells.
We use a 60-20-20 random split for all of our experiments.
Details regarding dataset curation can be found at Appendix C.
Stability of materials in datasets.
Structures in all 3 datasets are obtained from QM simulations and all structures are at local energy minima.
Most materials in Perov-5 and Carbon-24 are hypothetical, i.e. they may not have global stability (section 3.2) and likely cannot be synthesized.
MP-20 is a realistic dataset that includes most experimentally known inorganic materials with at most 20 atoms in the unit cell, most of which are globally stable.
A model achieving good performance in MPhas the potential to generate novel materials that can be experimentally synthesized.
Baselines.
We compare CDVAE with the following 4 baselines, which include the latest coordinatebased, voxel-based, and 3D molecule generation methods.
FTCP (Ren et al., 2020) is a crystal representation that concatenates real-space properties (atom positions, atom types, etc.) and Fouriertransformed momentum-space properties (diffraction pattern).
A 1D CNN-VAE is trained over this representation for crystal generation.
Cond-DFC-VAE (Court et al., 2020) encodes and generates crystals with 3D density maps, while employing several modiﬁcations over the previous Voxel-VAE (Hoffmann et al., 2019) method.
However, the effectiveness is only demonstrated for cubic systems, limiting its usage to the Perov-5 dataset.
G-SchNet (Gebauer et al., 2019) is an auto-regressive model that generates 3D molecules by performing atom-by-atom completion using SchNet (Sch ¨ utt et al., 2018).
Since G-SchNet is unaware of periodicity and cannot generate the lattice L .
We adapt G-SchNet to our material generation tasks by constructing the smallest oriented bounding box with PCA such that the introduced periodicity does not cause structural invalidity.
P-G-SchNet is our modiﬁed G-SchNet that incorporates periodicity.
During training, the SchNet encoder inputs the partial periodic structure to predict next atoms.
During generation, we ﬁrst randomly sample a lattice L from training data and autoregressively generate the periodic structure.
5. M ATERIAL RECONSTRUCTION
Setup.
The ﬁrst task is to reconstruct the material from its latent representation.
We evaluate reconstruction performance by matching the generated structure and the input structure for all materials in the test set.
We use StructureMatcher from pymatgen (Ong et al., 2013), which ﬁnds the best match between two structures considering all invariances of materials.
The match rate is the percentage of materials satisfying the criteria stol=0.5, angle tol=10, ltol=0.3 .
The
RMSE is averaged over all matched materials.
Because the inter-atomic distances can vary signiﬁcantly for different materials, the RMSE is normalized by 3 (cid:112) V/N , roughly the average atom radius per material.
Note G-SchNet is not a VAE so we do not evaluate its reconstruction performance.
Results.
The reconstructed structures are shown in Figure 3 and the metrics are in Table 1.
Since our model is SE(3) invariant, the generated structures may be a translated (or rotated) version of the ground truth structure.
Our model has a lower RMSE than all other models, indicating its stronger capability to reconstruct the original stable structures.
FTCP has a higher match rate than our model.
[Table]
Perov
[Table]
Carbon
[Table]
MPG-SchNet
O O Sc OOAlOAlAlAlSc
GeN NN GeGeGeGeN NNCa
Cu N NNN CuCu NNN N Ta
NNN MoMoHN NSc
H
H H
C
Ca
K Fe KMo FeFeFeFeFeFeFeFeMoFeFeFeFeFeMoK MoMoMoMoMoMoMoK
La
AuAuPbAuPbPbPbPbPbAuPbAuPbAu
Te N N
N K N Os
O N N
O O LiLi Cr Si
S O Fe
Ru O
N Hf N HfHfN NN
Ca
F N Ca Be CaO
O O Hf S ORh CCCC C C CCCCCC C CC CCCCCC CC CC CCC CCCCCCCC C CC C CC C C C CC C C
C CC CC C CC C CCCCCCC
Al Al Al Al Al Al Al Al Al AlAlAlAlCuAlAl
Tb I P Tl I
Tl
I I Tl PNa Cr TlTlTlI Tl PTl Sn TlTl I TlAg
Se Ag V Se Yb Se
Se SeSe SeVSeV YbSeYb
Se SeVSeSe
Er Er
Er Rh Rh ErErEr
Er Rh Er RhEr Er Rh ErRhErRhRhRhErErErRhErErErErErErRhEr
N Al F Bi
O
O Cr O F OOK
N N Pt N Ru
Os F N OOCu
C C C C C C CC C CCC C CCCCCCCCCCCCC
CC CC C C CC C C CCCCCCCC C CC CCCCCCCCC CCCC CCCCCC CC C CC C CCCCCCCCC C CCCCCC
CCCCC CC C C CCC C BBBOs
B MnMnMnMnOs B OsOs
Li O O GeGeGeGeGeGe LiO Au O LiLiLiLiGeGeCuO
V C V BC Si C B SiSiC CV BC CVCCCVVV
Nb Nb AlNb NbNb Nb Nb Al NbAl Nb AlNbNbAl AlAl
Nb NbNbAlAlAlAlAlAlNbNbNbNbNbNbNbAl AlAl NbNbNbNbNbNbAlAlAlAlAlCo
PdO S MoO
O Al O N Te
O N Ag
As O
N Co N N N Sc
C CC C C C
C CC C CCC C C C CC C CCC CCCCC CCCCCCCC CC C C CCCCC CCCCCCC CCCCCCCCC
C C C C CCCC C C
Pb Se Se
Te Li Li Ca PbTe
H O HHHO EuO H Cu HHOOO
Zr Mg Mg Mg Zr MgZrTb
Ni NiNi Au RuNiNiAuRuRuAu
G
[Figure]
[Table]
This can be explained by the fact that the same set of local structures can be assembled into different stable materials globally (e.g., two different crystal forms of ZnS).
Our model is SE(3) invariant and only encodes local structures, while FTCP directly encodes the absolute coordinates and types of each atom.
In Figure 5, we show that CDVAE can generate different plausible arrangements of atoms by sampling 3 Langevin dynamics with different random seeds from the same z .
We note that this capability could be an advantage since it generates more diverse structures than simply reconstructing the original ones.
5. M ATERIAL GENERATION
Setup.
The second task is to generate novel, stable materials that are distributionally similar to the test materials.
The only high-ﬁdelity evaluation of stability of generated materials is to perform
QM calculations, but it is computationally prohibitive to use QM for computing evaluation metrics.
We developed several physically meaningful metrics to evaluate the validity, property statistics, and diversity of generated materials.
1) Validity .
Following Court et al.
(2020), a structure is valid as long as the shortest distance between any pair of atoms is larger than 0 .
5 ˚A, which is a relative weak criterion.
The composition is valid if the overall charge is neutral as computed by SMACT (Davies et al., 2019).
2) Coverage (COV) .
Inspired by Xu et al.
(2021a)
; Ganea et al.
(2021), we deﬁne two coverage metrics, COV-R (Recall) and COV-P (Precision), to measure the similarity between ensembles of generated materials and ground truth materials in test set.
Intuitively, COV-R measures the percentage of ground truth materials being correctly predicted, and COV-P measures the percentage of predicted materials having high quality (details in Appendix G).
3) Property statistics .
We compute the earth mover’s distance (EMD) between the property distribution of generated materials and test materials.
We use density ( ρ , unit g / cm 3 ), energy predicted by an independent GNN ( E , unit eV / atom ), and number of unique elements (# elem.)
as our properties.
Validity and coverage are computed over 10,000 materials randomly sampled from N (0 , 1) .
Property statistics is computed over 1,000 valid materials randomly sampled from those that pass the validity test.
[Table]
Results.
The generated structures are shown in Figure 4 and the metrics are in Table 2.
Our model achieves a higher validity than FTCP, Cond-DFC-VAE, and P-G-SchNet, while G-SchNet achieves a similar validity as ours.
The lower structural validity in P-G-SchNet than G-SchNet is likely due to the difﬁculty of avoiding atom collisions during the autoregressive generation inside a ﬁnite periodic box.
On the contrary, our G-SchNet baseline constructs the lattice box after the 3D positions of all atoms are generated, and the construction explicitly avoids introducing invalidity.
Furthermore, our model also achieves higher COV-R and COV-P than all other models, except in MP-20 our COV-P is similar to G-SchNet and P-G-SchNet.
These results indicate that our model generates both diverse (COV-R) and high quality (COV-P) materials.
More detailed results on the choice of thresholds for
COV-R and COV-P, as well as additional metrics can be found in Appendix G. Finally, our model also signiﬁcantly outperforms all other models in the property statistics of density and energy, further conﬁrming the high quality of generated materials.
We observe that our method tends to generate more elements in a material than ground truth, which explains the lower performance in the statistics of # of elems.
than G-SchNet.
We hypothesize this is due to the non-Gaussian statistical structure of ground truth materials (details in Appendix D.3), and using a more complex prior, e.g., a ﬂowmodel-transformed Gaussian
(Yang et al., 2019), might resolve this issue.
5. P ROPERTY OPTIMIZATION
Setup.
The third task is to generate materials that optimize a speciﬁc property.
Following Jin et al.
(2018), we jointly train a property predictor F parameterized by an MLP to predict properties of training materials from latent z .
To optimize properties, we start with the latent representations of testing materials and apply gradient ascent in the latent space to improve the predicted property
F ( · ) .
After applying 5000 gradient steps with step sizes of 1 × 10 − 3 , 10 materials are decoded from the latent trajectories every 500 steps.
We use an independently trained property predictor to select the best one from the 10 decoded materials.
Cond-DFC-VAE is a conditional VAE
so we directly condition on the target property, sample 10 materials, and select the best one using the property predictor.
For all methods, we generate 100 materials following the protocol above.
We use the independent property predictor to predict the properties for evaluation.
We report the success rate (SR) as the percentage of materials achieving 5, 10, and 15 percentiles of the target property distribution.
Our task is to minimize formation energy per atom for all 3 datasets.
Results.
The performance is shown in Table 3.
We signiﬁcantly outperform FTCP, while having a similar performance as Cond-DFC-VAE in Perov-5 (Cond-DFC-VAE cannot work for Carbonand MP-20).
Both G-SchNet and P-G-SchNet are incapable of property optimization 6 .
We note that all models perform poorly on the Carbon-24 dataset, which might be explained by the complex and diverse 3D structures of carbon.
C ONCLUSIONS AND O UTLOOK
We have introduced a Crystal Diffusion Variational Autoencoder (CDVAE) to generate the periodic structure of stable materials and demonstrated that it signiﬁcantly outperforms past methods on the tasks of reconstruction, generation, and property optimization.
We note that the last two tasks are far more important for material design than reconstruction because they can be directly used to generate new materials whose properties can then be veriﬁed by QM simulations and experiments.
We believe CDVAE opens up exciting opportunities for the inverse design of materials for various important applications.
Meanwhile, our model is just a ﬁrst step towards the grand challenge of material design.
We provide our datasets and evaluation metrics to the broader machine learning community to collectively develop better methods for the task of material generation.
R EPRODUCIBILITY S TATEMENT
We have made the following efforts to ensure reproducibility: 1) We provide our code at https:// github.com/txie-93/cdvae ;
2)We provide our data and corresponding train/validation/test splits at https://github.com/txie-93/cdvae/tree/main/data ; 3) We provide details on experimental conﬁgurations in Appendix D. A CKNOWLEDGMENTS
We thank Peter Mikhael, Jason Yim, Rachel Wu, Bracha Laufer, Gabriele Corso, Felix Faltings,
Bowen Jing and the rest of the RB and TJ group members for their helpful comments and suggestions.
The authors gratefully thank DARPA (HR00111920025), the consortium Machine Learning for Pharmaceutical Discovery and Synthesis (mlpds.mit.edu) and MIT-GIST collaboration for support.
A P ROOF FOR THE CONNECTION TO A HARMONIC FORCE FIELD
We assume the loss in Equation 2 can be minimized to zero when the noises are small, meaning that
[Equation]
where σ X ,j ∈ { σ X ,j } Lj =1 and any noise smaller than σ X ,J is considered as small.
The force term in the Langevin dynamics α j s X ,t can then be written as
[Equation]
If we write (cid:15)/σ 2 X ,L = k , then,
[Equation] [Equation]
If the noises are small enough that atoms do not cross the periodic boundaries, then we have d min ( X , ˜ X ) =
X − ˜ X .
Therefore,
[Equation]
B I MPLEMENTATION DETAILS B. P REDICTION OF LATTICE PARAMETERS
There are inﬁnitely many different ways of choosing the lattice for the same material.
We compute the Niggli reduced lattice (Grosse-Kunstleve et al., 2004) with pymatgen (Ong et al., 2013), which is a unique lattice for any given material.
Since the lattice matrix L is not rotation invariant, we instead predict the 6 lattice parameters, i.e. the lengths of the 3 lattice vectors and the angles between them.
We normalize the lengths of lattice vectors with 3 √ N , where N is the number of atoms, to ensure that the lengths for materials of different sizes are at the same scale.
B. M ULTI - GRAPH CONSTRUCTION
For the encoder, we use CrystalNN (Pan et al., 2021) to determine edges between atoms and build a multi-graph representation.
For the decoder, since it inputs a noisy structure generated on the ﬂy, the multi-graph must also be built on the ﬂy for both training and generation, and CrystalNN is too slow for that purpose.
We use a KNN algorithm that considers periodicity to build the decoder graph where K = 20 in all of our experiments.
B. GNN ARCHITECTURE
We use DimeNet++ adapted for periodicity (Klicpera et al., 2020a;b) as the encoder, which is SE(3) invariant to the input structure.
The decoder needs to output an vector per node that is SE(3) equivariant to the input structure.
We use GemNet-dQ (Klicpera et al., 2021) as the decoder.
We used implementations from the Open Catalysis Project (OCP)
(Chanussot et al., 2021), but we reduced the size of hidden dimensions to 128 for faster training.
The encoder has 2.2 million parameters and the decoder has 2.3 million parameters.
C D ATASET CURATION C. P EROV -
Perovskite is a class of materials that share a similar structure and have the general chemical formula
ABX 3 .
The ideal perovskites have a cubic structure, where the site A atom sits at a corner position, the site B atom sits at a body centered position and site X atoms sit at face centered positions.
Perovskite materials are known for their wide applications.
We curate the Perov-5 dataset from an open database that was originally developed for water splitting (Castelli et al., 2012a;b).
All 18928 materials in the original database are included.
In the database, A, B can be any nonradioactive metal and X can be one or several elements from O, N, S, and F. Note that there can be multiple different X atoms in the same material.
All materials in Perov-5 are relaxed using density functional theory (DFT), and their relaxed structure can deviate signiﬁcantly from the ideal structures.
A signiﬁcant portion of the materials are not thermodynamically stable, i.e., they will decompose to nearby phases and cannot be synthesized.
C. C ARBON -
Carbon-24 includes various carbon structures obtained via ab initio random structure searching (AIRSS)
(Pickard & Needs, 2006; 2011) performed at 10 GPa.
The original dataset includes 101529 carbon structures, and we selected the 10% of the carbon structure with the lowest energy per atom to create Carbon-24.
All 10153 structures in Carbonare relaxed using DFT.
The most stable structure is diamond at 10 GPa.
All remaining structures are thermodynamically unstable but may be kinetically stable.
Most of the structures cannot be synthesized.
C. MP-
MP-20 includes almost all experimentally stable materials from the Materials Project (Jain et al., 2013) with unit cells including at most 20 atoms.
We only include materials that are originally from
ICSD (Belsky et al., 2002) to ensure the experimental stability, and these materials represent the majority of experimentally known materials with at most 20 atoms in unit cells.
To ensure stability, we only select materials with energy above the hull smaller than 0.08 eV/atom and formation energy smaller than 2 eV/atom, following Ren et al.
Differing from Ren et al.
(2020), we do not constrain the number of unique elements per material.
All materials in MP-20 are relaxed using DFT.
Most materials are thermodynamcially stable and have been synthesized.
D E XPERIMENT DETAILS D. R EASONS FOR THE UNSUITABILITY OF SOME METRICS FOR SPECIFIC DATASETS
In Table 2, property statistics are computed by comparing the earth mover’s distance between the property distribution of generated materials and ground truth materials.
So, they are not meaningful for ground truth data.
Materials in Perov-5 have the same structure, so it is not meaningful to require higher structure diversity.
Materials in Carbon-24 have the same composition (carbon), so it is not meaningful to require higher composition diversity.
In addition, all models have ∼ 100% composition validity, so it is not compared in the table.
D. C OMPOSITION VALIDITY CHECKER
We modiﬁed the charge neutrality checker from SMACT (Davies et al., 2019) because the original checker is not suitable for alloys.
The checker is based on a list of possible charges for each element and it checks if the material can be charge neutral by enumerating all possible charge combinations.
However, it does not consider that metal alloys can be mixed with almost any combination.
As a result, for materials composed of all metal elements, we always assume the composition is valid in our validity checker.
For the ground truth materials in MP-20, the original checker gives a composition validity of ∼ 50%, which signiﬁcantly underestimates the validity of MP-20 materials (because most of them are experimentally synthesizable and thus valid).
Our checker gives a composition validity of ∼ 90%, which is far more reasonable.
We note again that these checkers are all empirical and the only high-ﬁdelity evaluation of material stability requires QM simulations.
D. N ON -G AUSSIAN STATISTICAL STRUCTURE OF MATERIALS
The material datasets are usually biased towards certain material groups.
For example, there are lots of lithium-containing materials in MP-20 because it started with battery research.
We also ﬁnd that our decoder tends to underﬁt the data distribution with a larger β in Equation 9.
We believe these observations indicate that the statistical structure of the ground truth materials are far from Gaussian.
As a result, sampling from N (0 , 1) may lead to out-of-distribution materials, which explains why our method tends to generate more elements per material than the ground truth.
D. H YPERPARAMETERS AND TRAINING DETAILS
The total loss can be written as,
[Equation]
We aim to keep each loss term at a similar scale.
For all three datasets, we use λ c = 1 , λ L = 10 , λ N = 1 , λ X = 10 , L A = 1 .
We tune β between 0 . 01 , 0 .
1 for all three datasets and select the model with best validation loss.
For Perov-5, MP-20, we use β = 0 .
01 , and for Carbon-24, we use β = 0 .
For the noise levels in { σ A ,j } Lj =1 , { σ X ,j } Lj =1 , we follow Shi et al.
and set L = 50 .
For all three datasets, we use σ A , max = 5 , σ A , min = 0 .
01 , σ X , max = 10 , σ X , min = 0 .
During the training, we use an initial learning rate of 0.001 and reduce the learning rate by a factor of 0.6 if the validation loss does not improve after 30 epochs.
The minimum learning rate is 0.0001.
During the generation, we use (cid:15) = 0 . 0001 and run Langevin dynamics for 100 steps at each noise level.
E V ISUALIZATION OF MULTIPLE
RECONSTRUCTED
[Table]
Sample
Sample
Sample
PerovCarbon-
MPF N V
Rh O
O F N V Rb
O O Sn Zr O
F O Ru Ba
O
C C C C C CC
CCC C
C CCC CCC C C
C C C C
C C C C C C CC C CC C
S Ti V
S TiTiS Na S
O Cr
O O
Eu O OOOOOSbSbO Eu
N O O
TiNNNTiON O O TiOTiOON O OOTiOO TiN
Mg Al Si Si AlAl
V F O N VRh
N O Rb V F
O Zr O O Sn
Ba F O O BaBaRu
C C C C C C
CC C CCC C C
C C C C C C C C
C C
C C C C C C C CC
S Ti V TiTiS S SNa S
O Eu O OCr O O
Sb CrO Eu O OOCrEu
O Ti N TiO O N TiNO
O N OO TiTi OOOTiN TiTiTi
O TiNNTiTiN
Al Mg
Si MgMgSi AlAlAl
O F V Rh N
O V F N O Rb
O O Sn O Zr O
Ba O F O BaBaRu
C C C C C
C C C
C C C C C C
C C CCC C C C C C C C
C C C C C C C C C C
S S S V Na SS Ti
O Eu O OCr O O
Sb CrO Eu O OOCrEu
N TiNNNNNTiOTi N N OTiTiNNNTi Ti NO N TiTiTi
NN NNTiNNNNNO
Si Mg Al
Al MgSi
O F V Rh N
F V O N F Rb
OO Zr
O O Sn
Ba F O O BaBaRu
C C C
C CC C C C
C C CCC C CC
C C C C
C C C C C
C C C C C CC CC C C C CCC
Na S NaNaNaS S S V STi
OCr
O O O K
O OOOCrCrEuOCrCrOOCrOCrOSb
O OTiN O N OOOTiO TiON
O OO TiTiTiTi TiTiNNNTiNTiTiOOONNTi
Si Al SiSiSiAl AlSi Mg F S AMPLING SPEED FOR MATERIAL GENERATION We summarize the speed for generating 10,000 materials for all models in Table 4. FTCP is signiﬁcantly faster, but the quality of generated materials is very poor as shown in Table 2.
CondDFC-VAE is faster than our method in Perov-5, but has a lower quality than our method and only works for cubic systems.
It is also unclear how it will perform on larger materials in Carbonand MP-20, because the compute increases cubicily with the increased size of the density map.
GSchNet/P-G-SchNet have a comparable sampling time as our method, but have a lower quality.
We also note that we did not optimize sampling speed in current work.
It is possible to reduce sampling time by using fewer sampling steps without signiﬁcantly inﬂuencing generation quality.
There are also many recent works that aim to speed up the sampling process for diffusion models (Nichol &
Dhariwal, 2021; Kong & Ping, 2021; Salimans & Ho, 2022).
[Table]
G C OVERAGE METRICS FOR MATERIAL GENERATION
Inspired by Xu et al.
(2021a)
; Ganea et al.
(2021), we deﬁne six metrics to compare two ensembles of materials: materials generated by a method { M k } k ∈ [1 ..K ] , and ground truth materials in test data { M ∗ l } ∈ [1 ..L ] .
We use the Euclidean distance of the CrystalNN ﬁngerprint (Zimmermann & Jain, 2020) and normalized Magpie ﬁngerprint (Ward et al., 2016) to deﬁne the structure distance and composition distance between generated and ground truth materials, respectively.
They can be written as
D struc .
( M k , M ∗ l ) and D comp .
( M k , M ∗ l ) .
We further deﬁne the thresholds for the structure and composition distance as δ struc .
and δ comp . , respectively.
Following the established classiﬁcation metrics of Precision and Recall, we deﬁne the coverage metrics as:
[Equation]
[Equation]
where COV is ”Coverage”, AMSD is ”Average Minimum Structure Distance”, AMCD is ”Average
Minimum Composition Distance”, and COV-P (precision), AMSD-P (precision), AMCD-P (precision) are deﬁned as in above equations, but with the generated and ground truth material sets swapped.
The recall metrics measure how many ground truth materials are correctly predicted, while the precision metrics measure how many generated materials are of high quality (more discussions can be found in Ganea et al. (2021)).
We note several points on why we deﬁne the metrics in their current forms.
1) COV requires both structure and composition distances to be within the thresholds, because generating materials that are structurally close to one ground truth material and compositionally close to another is not meaningful.
As a result, AMSD and AMCD are less useful than COV.
We use ﬁngerprint distance, rather than RMSE from StructureMatcher (Ong et al., 2013), because the material space is too large for the models to generate enough materials to exactly match the ground truth materials.
StructureMatcher
ﬁrst requires the compositions of two materials to exactly match, which will cause all models to have close-to-zero coverage.
For Perov-5 and Carbon-24, we choose δ struc .
= 0 .
2 , δ comp .
= 4 .
For MP-20, we choose δ struc .
= 0 .
4 , δ comp .
= 10 .
In Figure 6, Figure 7, Figure 8, we show how both COV-R and COV-P change by varying δ struc .
and δ comp .
in all three datasets.
[Table]
1.0 Struc
[Figure]
[Figure]
[Figure]