Data-Driven Approach to Encoding and Decoding 3-D Crystal Structures
Authors: Jordan Hoffmann
Louis Maestrati
Yoshihide Sawada
Jian Tang
Jean Michel Sellier
Yoshua Bengio
Venue:  (2020)

Abstract: Generative models have achieved impressive results in many domains including image and text generation. In the natural sciences, generative models have lead to rapid progress in automated drug discovery.  Many of the current methods focus on either 1-D or 2-D representations of typically small, drug-like molecules. However, many molecules require 3-D descriptors and exceed the chemical complexity of commonly used dataset. We present a method to encode and decode the position of atoms in 3-D molecules along with a dataset of nearly 50,000 stable crystal unit cells that vary from containing 1 to over 100 atoms. We construct a smooth and continuous 3-D density representation of each crystal based on the positions of different atoms. Two different neural networks were trained on a dataset of over 120,000 three-dimensional samples of single and repeating crystal structures. The first, an Encoder-Decoder pair, constructs a  compressed latent space representation of each molecule and then decodes this description into an accurate reconstruction of the input. The second network segments the resulting output into atoms and assigns each atom an atomic number. By generating compressed, continuous latent spaces representations of molecules we are able to decode random samples, interpolate between two molecules, and alter known molecules.

Under review as a conference paper at ICLR 2020
DATA-DRIVEN APPROACH TO ENCODING AND DECOD- ING 3-D CRYSTAL STRUCTURES
Anonymous moonymous authors Paper under double-blind review
ABSTRACT
Generative models have achieved impressive results in many domains including image and text generation. In the natural sciences, generative models have lead rapid progress in automated drug discovery. Many of the current methods focus on either 1-D or 2-D representations of typically small, drug-like molecules. However, many molecules require 3-D descriptors and exceed the chemical complexity of commonly used datasets. We present a method encode and decode the position of atoms in 3-D molecules along with dataset of nearly 50,000 stable crystal unit cells that vary from containing 1 to over 100 atoms. We construct a smooth and continuous 3-D density representation of each crystal based on the positions of different atoms. Two different neural networks were trained on a dataset of over 120,000 tiree-dimensional samples of single and repeating crystal structures. The first, an Encoder-Decoder pair, constructs a compressed latent space representation of each molecule and then decodes this description into accurate reconstruction the input. The second network segments the resulting output into atoms and assigns each atom an atomic number. By generating compressed, continuous latent spaces representations of molecules we are able to decode random samples, interpolate between two molecules, and alter known molecules.
INTRODUCTION
Generative models have recently seen tremendous success in generating 2-D images of every day
objects (Kingma & Welling 2013 Goodfellow et al. 2014 Brock et al. 2018 Razavi et al. (019). The size and accuracy of generated results has greatly improved to the point where samples from the latent space decode to photo-realistic samples (Brock et al. 2018 Razavi et al. 2019). A very exciting and important future avenue for generative models is the generation of 3-D structures, like in the world around us. Adversarial networks and autoencoders have been extended into 3-D and have shown they able to encode useful representations of everyday objects (Wu et al.) 2016 Zhuu et al. 2018, Brock et al. 2016)] Achlioptas et al.) 2017] Valsesia et al. 20188] Meng et al. 2019] Maturana & Scherer 2015). Representations using point clouds have gained popularity as way capture the distribution of different types of objects (Achlioptas et al. 2017 Yang et al. 2019 Wu et al. 20018).
As machine learning approaches are able to understand and recreate the underlying distributions for many different types of objects in our world, the application of these tools in the physical sciences is very exciting direction. A clear field where generative models can have a tremendous impact is in material discovery, which is very important for designing new batteries or carbon capture devices for fighting climate change.
One very successful area of applying generative models to the sciences has been the field of drug discovery (Jin et al., 2018, Homeaz-Bombarelii et al., 2018 Assouel al. 2018 (see (Schwalbe- Koda & Gomez-Bombarell (2019) for an excellent summary of many methods). The success of machine learning in this domain has been enormously helpful as a process that was often painstakingly slow has been rapidly accelerated in searching unimaginably large space of possible drug compounds, estimate to be up to 10 (Polishchuk et al. 2013). Additionally, excellent datasets such as QM9 (Ramakrishnan et al. 2014) and ZINC (Irwin et al. 2012) which contain molecules of interest along with their precomputed properties have allowed for comparing different methods and developing state-of-the-art aftt tools.

---

Under review as a conference paper at ICLR 2020
Figure 1: Examples of crystal unit cells. Each examplc shows the unit cell of crystal. Different colors represent different atomic species. red, green, and blue line represent the of the crystal. Note that they vary in length and angle. Additionally, some unit cells have just one or two atoms while others have nearly 100. Visualizations were made with Mercury (Macrae et al. 2008
By learning a compressed, useful, latent space representation, this enormous space of molecules can be embedded in simpler latent space that is easier to search Kusner et al. (2017). Using generative models, compounds hypothesized to have specific properties can be rapidly generated and then a targeted subset of these can be experimentally tested. For example, in G.ómez-Bombalne an auxiliary property prediction task is introduced for a network separate from the encoder/decoder. This allows for optimization of properties in the latent space and then the resulting latent space vector decoded into a candidate molecule that can undergo more rigorous computational testing before experimentally synthesized. Using graph representation and reinforcement learning. You et al propose a novel method for molecule generation with specific properties You et al. (2018).
et al. (2018)
Most of the work combining generative models with chemical discovery has focused molecules that can be represented in either dimension (such as a SMILES string (Weininger 1988) or by leveraging a dimensional representation of a molecule, such graph. While I-D and 2-D representations have been very successful for many drug compounds, these representations are not sufficient describe all molecules. For example, with very different properties, to have identical graphs (Gebauer et al., 2019). Additionally, there are more complex classes of compounds where the 3-D structure integral to the molecule's properties. Therefore, in this work we focus on generating representations of molecules where the 3-D information essential for reconstruction. Generating 3-D structures still a relatively nascent field, when compared to image and text generation. The data requirements for modeling 3-D structures are larger than their 2-D counterparts and there are fewer standard datasets (Nguyen- Phuoc et al. 2019). Additionally, measure the properties of inorganic materials (like those considered here), costly DFT calculations are needed which require precise 3-D locations. There is an extra dimension in 3-D problems, providing additional symmetries that often need be learned (Weiler et al. 2018). New work from generates adjacency matrices which are, by construction, invariant to rotation Hoffmann & Noe (2019), In either case, for the automatic generation of viable complex 3-D structures of inorganic compounds one needs to accurately generate the locations of many atoms in 3-D space.
possible for different molecules.
In this work, we focus specifically on proposing an effective 3-D representation for this class of molecules that can not be represented in the same ways many smaller molecules. We hope to promote work in this still nascent field which has yet see the same level of success the generation of smaller molecules. In the Appendix, we have provided a clear comparison with 1-D and 2-D representations. Currently, no direct comparisons possible.
In addition to drug discovery, a promising avenue for molecular design is for minimizing environ- mental impact through the design of more efficient or more environmentally friendly compounds for a variety of applications. Designing more efficient materials for photo-vole voltaic panels and more environmentally friendly materials for batteries both very exciting research directions for using machine learning as a tool towards mitigating climate change (Niu et al. 2015) Tabor et al, 2018 Gebauer al. 2019 Rolnick et al. 2019). Many of the compounds of interest crystal structures
2

---

Under review as a conference paper at ICLR 2020
made of single small sub-units that repeat in all directions. This unit is referred to "unit cell" and can vary in size and shape as well as internal arrangement and chemical composition.
In this work, present method for encoding and decoding very important set of 3-D molccules that can not represented using standard methods used for drug-like molecules. We directly encode and decode 3-D volumes of density from two different data representations. In the first, use single unit cells, as shown Fig. Each unit cell is centered in the cube, but randomly rotated. In the second representation, we repeat the unit cell along all three axes such that the resulting sampled cube contains repeated unit cells of crystal structures, like those shown in Fig. 2 (and Fig. (S6C). Many very interesting and important molecules, such as materials for solar panels batteries, are composed of crystalline units. We train a variational autoencoder (VAE) and network to segment the decoded output based on the true locations of atoms in tandem. By coupling these two tasks, we are able to accurately encode and decode 3-D atomic positions and species. As far as we are aware, this is not an explored direction means to represent molecules for encoding and decoding their 3-D arrangement. We summarize our contribution as follows:
We propose method for encoding and decoding 3-D molecules that can not be represented by traditional string or graph based approaches. By jointly training VAE and a segmentation network on the output of the decoder, we train the entire network in end-to-end fashion.
We consider two different problems: (1) encoding and decoding single unit cells and (2) encoding and decoding repeated unit cells. In both cases accurately reconstruct the locations of atoms. We also achieve good results in atomic species identification and future work will improve this front. By sampling from z (0, 1), we generate complex structures that have physically realistic spacing between atoms.
file
S
M
M
$
VAE
Figure 2: Network Architecture, We encode and decode a 30 x 30 x 30 voxel grid representing 10 Á on each side. Each voxel contains the value of the density. The output of the decoder is passed into a 3-D U-Net Çiçek al. 2016). We train the two models in parallel. In the schematic, unit cell rather than single unit cell. The black arrows indicate deterministic transformations. From the file, the species matrix is constructed. From this, the density matrix is computed.
show the crystal represented as repeated
2 PRELIMINARIES AND RELATED WORK
There are many classes of compounds that are of interest for data-driven discovery. There has been significant work exploring the generation of molecules with potential medicinal properties. These molecules tend be organic molecules and can be represented efficiently using string or a graph. more complex class of molecules are inorganic crystal structures which vary in complexity across many different axes. Crystals are materials that are made up of a repeating pattern of a simpler "unit cell." Crystal structures are of key interest for many environmental problems, such materials for solar panels and batteries (Niu et al. 2015).
In contrast to problems in drug discovery, many crystals are not composed solely of organic molecules. Crystals commonly contain many heavy Hydrogydogg)n) atoms which make various quantum me- chanical calculations of their energetic properties very difficult and time consuming. For example, the common benchmark dataset QM9 for predicting quantum mechanical properties includes over 1 130,000 molecules contains fewer than 9 "heavy" atoms (Ramakrishnan et al. 2014).
3

---

Under review as a conference paper at ICLR 2020
A
B
C
D
predicted number
B-+A
Angstoom
predicted
MSE
epoch
target
target 2
percentile
target
prediction
high
slice
slice
slice
17
slice 25
target
tarpet
prodiction
prediction
Langet
prediction
10
target
20
203
prediction
Figure Single unit cell accuracy. (A) We show voxel wise reconstruction error (plotted with mean square error) during training. (B) For random molecules in the test set, we plot the number of true atoms and the number of recovered atoms after segmentation. (C) For the reconstructed atoms, we plot the predicted atomic number
the atomic number of the nearest true atom. (D) We compute the distance from each true the nearest predicted
(E) For three different crystals plot the predicted versus reconstructed density at each voxel. We also show 2-D slices through the target and prediction, along with the 3-D reconstructions. For plotting, the density on each figure is normalized between O and I though is not decoded as such.
and vice-versa (orange and blue, respectively).
There have been two parallel lines of work applying machine learning to material sciences. The first deals with the prediction of physical properties from compound without performing computation- ally expensive density functional theory (DFT) calculation.
Recently, Xie and Grossman introduced crystal graph convolutional networks (CGCNN) for accurate prediction of 8 different DFT-calculated properties of crystal structures (Xie & Grossman 2018). MatErials Graph Network (MEGNet) uses graph neural networks (Chen et al. 2019) predict variety of energetic properties on both the QM9 dataset (Ramakrishnan al. 2014) as well 69 000 crystal structures from the Materials Project (Jain et al. 2013). Another network, SchNet, uses 3-D spatial information to directly leverage the interactions between separate atoms. They then continuous- filter convolution for accurate property prediction (Schütt et al. 2018). These methods, and many have all achieved impressive results many standard benchmark datasets. Cubuk, et al. use transfer learning to search an enormous space of molecules for promising Lithium ion conductors (Cubuk et al. 2019).
Another line of work is concerned with generating molecules, often molecules that have specific properties. In the past decade, generative models for drug discovery have achieved impressive
4

---

Under review as a conference paper at ICLR 2020
results. Typically representation of a compound, such as SMILES string Jamnez-oombarelll et al. 2018; Segler et al. 2017] Alperstein et al. 2019) or a graph (Jin et al. 2018 Assouel et al.
2018: Mansimov et al. 2019) is encoded and decoded. Then, by sampling the resulting latent space, novel molecules can be generated. Fither by performing an auxiliary task with the latent space (as done in Gommez-Bombarelli et al. (2018) and then performing optimization or by conditioning the latent space on desirable properties, novel chemical structures are obtained.
One possible representation of drug-like molecules is a SMILES string. However, difficulty faced using the SMILES representation is ensuring that the decoder decoded a valid SMILES string Glomez-Bombarele et al. 2018 Kusner al., 2017). Additionally, there is not a strong notion of distance between molecules and their SMI LES representation (Jin et al.) 2018). One approach, Chem'TS, uses SMILES strings along with recurrent neural networks and Monte Carlo Tree Search to better ensure generated strings decode to real molecules (Yang et al. (017). By explicitly penalizing invalid SMILES strings in their reward function they are able to bypass the difficulty in their decoded molecules being non-physical.
More recently, graph based methods have proven very successful in generating synthetically attainable molecules. By building a database of molecule fragments (like LEGO bricks), a graph is formed based on the connectivity of different structures (Jin et al. 2018). These methods have far fewer difficulties ensuring that the decoder produces physically valid molecules. For generation of physical molecules, this is a very important consideration, as there are many hard constrains that most be obeyed. Other graph based approaches include flow based models (Madhawa et al.) 2019). Also using graph neural networks, E. Manismov and collaborators developed method to, given a graph, recreate conformations of 3-D molecule and predict its energetic properties (Mansimov et al. 2019). In Mansimov et al. (2019), the authors propose a graph-based generative method that is able to generate molecules with desired target properties. Impressively, they recover 3-D positional information for the atoms and show they achieve the required accuracy for relaxation.
An additional challenge for generative models in the physical sciences is ensuring that samples from the latent space, z. decode into physically plausible objects. The decoded objects need to obey physical constraints and an object that may appear physical is not certain to actually experimentally realizable. In many domains this is not of specific concern, for example a generated animal is scored on some likelihood based on its visual appearance, not whether or not the generated creature is genetically possible. Recently, there have been generative models leveraging the 3-D nature of the problem. G-SchNet, a generative model for 3-D molecules (Gebauer et al.] 2019)) leverages the SchNet architecture- start-of-the-art property prediction network (Schütt et al. 20188. They show they are able to generate molecules, placing atoms in 3-D space in a rotationally invariant manner. By appropriately conditioning the placement of new molecules based the location and identity of previous molecules, they ensure the symmetries that are required for the molecules to relax in a DFT calculation. They propose expanding G-SchNet to generating crystal structures as a future direction.
While there is a rich literature in the generation of organic compounds, recently there has been work in the generation of more complex crystal structures, such as CrystalGAN (Nouira et al., 2018). The authors of CrystalGAN introduce geometric constraints and show that their method is able to produce stable structures compared to methods that do not include such domain knowledge such as DiscoGAN (Kim et al. 2017). CrystalGAN uses a dataset with much larger selection of elements than much previous work, but they limit themselves to molecules of particular chemical structure, nor do they generate a 3-D arrangement. In this work, we do not limit the generated crystals in this way.
3 METHODS
We use a dataset containing 46,744 cif files which contains the information describing the unit cells of properly relaxed crystal structures from the Materials Project (Jain et al. 2013 Xie & Grossman 2018). We use 80% of the data for testing and the other 20% for training. Using the Python library pymatgen we preprocess all of the data into the density 3-D matrices (Ong et al. 2013) described below. The boundary box of a crystal structure is controlled by six different degrees of freedom (see Fig. Each side can have a different length and the angles between the three sides is also variable. Additionally, the internal complexity of each crystal can also vary widely- some unit cells in dataset contain a single atom while other unit cells can have over one hundred different atoms. This tremendous variability in structure makes universal representation difficult.
5

---

Under review as a conference paper at ICLR 2020
A
30
19
30
Figure 4: Species reconstruction and latent space interpolation. (A) For five different randomly selected crystal unit cells show the target and our reconstruction. Atoms are colored using default atomic colors, atoms with similar atomic number do not necessarily have similar colors. (B) We show the reconstruction of latent space interpolation between two molecules for 3 random sets of targets. For a video, see the Appendix.
3.1 DATA REPRESENTATION
As a first step for the problem of generating novel, physical crystal structures we begin with the simpler problem of encoding and decoding physical locations of atoms in space. We begin by considering cube with side length 10 Á, which we represent by M. We divide this cube into 30 equally spaced bins, resulting in 30 x 30 x 30 cube with each voxel representing 0.33 À on each side (see Fig. 2 Fig. 3). the appendix, we discuss this choice in more detail. This data representation is similar to many computer vision tasks, so we use similar convolution based network architectures.
We consider the crystal structures in our dataset where the maximum side length is less than 10 Å. We consider two different data representations of different complexity. the first, we shift each single unit cell to the center of our grid, then randomly rotate it. In this representation, and decode a single unit cell. We randomly sample different rotations for each crystal. In this case, the encoded structures typically have between 1 and 30 atoms, with few reaching between 50 and 100. In the second representation, we repeat these unit cells in each direction infinitely and randomly sample eube in space. This means that each cube contains at least unit cell, but some will contain more than one cell, since in all cases cube with 10 Ásides is chosen. We choose one representation where a unit cell begins at (0, O, 0) and randomly sample other cubes in this space resulting in dataset of over 100,000 examples. In this representation, typically encode the
encode

---

Under review as a conference paper at ICLR 2020
locations of between 20 to 100 atoms, with some having 200. An advantage of representation that we are agnostic to the number of atoms that we are encoding and decoding, something that is difficult for other approaches such graphs.
Using both these representations, we compute a density field, where each pixel (1,j, k) is defined to have value
d(Zm, (i) 20
M 0 (2 33/2
Σ Zm exp
(1)
Where -( represents the Euclidean distance between the two arguments. Zm represents the 3-D coordinates of atom Zm- We set a to 1.0 Å. When plotting, we multiply the output by の³ (2+)5/22 We do this because when = 1 this results in the value of Z (the atomic number) being present its location in space. Because of the structure of this representation, when the density of atoms increases, there are more interactions making it harder to reconstruct the true identity of atom. Ultimately, our aim is to be able to reconstruct the location and species of the different atoms in this grid. To that end, we also construct species matrix, S, where each voxel is either 0 equal to the atomic number of an atom that is within 0.5 Å. We construct neural networks that attempt to learn M and S in parallel. These two networks are trained together and are detailed in the following section.
3.2 NETWORK ARCHITECTURE
We variational autoencoder (Kingma & Welling 2013 encode and decode our 3-D density maps, M. We use a _3 multiplier on the Kullback- Leibler (KL) loss term to encourage no correlations between different elements of the latent space (Higgins et al. 2017). b. Our encoder, E(.),
convolutional neural network. The decoder, D--), uses upsampling and convolutions in favor of transposed convolutions (Odena et al. 2016). We use batch normalization (Toffe & Szegedy 2015). LeakyReLU activations, and the code is written in Pytorch (Paszke et al. (017). The optimization done using the Adam optimizer (Kingma Ba 2014). We use latent space size of 300 for all experiments.
Simultaneously with training the VAE, we train a 3-D U-Net segmentation model, Unet (-), with an attention mechanism to segment the output of the decoder (Ronneberger al. [2015] Oktay et al. 2018). For accurate segmentation, especially in the case of a repeating lattice, it is important to be able to capture the dependencies between different atoms. We found that U-Net model worked well, though future work will explore sequentially classifying atoms. Ideally, the primary work is done by the VAE and the role of U-Net is to perform the tedious optimization problem of recovering the species identity and locations that is most consistent with the reconstructed density field.
We include weighted (by 7) loss from the segmentation in the loss of the encoder/decoder. In experiments, we set to 0.1. We experimented with 7 = and 0.33 and found that 0.1 proved an acceptable intermediate. With 7 0.33 we found slightly improved segmentation results but worse density reconstruction results. When we completely removed the effect of the species loss from the density reconstruction, species reconstruction results degraded. This allows us train the entire network in an end-to-end fashion. The entire model is shown in Fig. 2
We use three terms in the loss function for the VAE with the most weight given to the reconstruction of the density matrices. The segmentation network is only concerned with the final segmentation of the reconstruction from the decoder. Currently, we treat each atom type as different class, with an additional (most common) class for "no atom." The loss function used for the VAE is
LVAE = LRE M, M) + ((DLL (2(()(((()))))) (z) (lll()))) + L
(2)
The first term is the reconstruction error and the third term is the binary cross entropy loss from the segmentation. The second term is the Kullback-Leibl divergence between the prior (set to p(z) .O, 1)) and q(2lM) the approximate posterior of the latent vector given input M. We use loss given by Lu-Net LBCE(S S. S) for the segmentation. In the above expression, M is the input density field. z = E(M) and M = (zz) is the reconstructed density field. The variable S represents a 1-hot species matrix. Lastly, $ Unei
MM represents the probability matrix from the segmentation routine, which will be of shape: [batch x classes x-dim y-dim x z-dim].
7

---

Under review as a conference paper at ICLR 2020
4 RESULTS
We find that we are able to accurately encode and decode 3-D representations of density fields (see Figs. and SS). Using the segmentation network, we segment the output of the decoder into distinet atoms. Using trained network, we show that ran- dom samples from the latent space decode to sam- ples that obey many of the same statistics as the training distribution (see Fig. 5 and Fig. S7). Addi- tionally, in the case of repeating unit cell, alter our training routine to condition the generation of erystals the largest atomic number present.
Table I: Location and species reconstruction enror. P denotes prediction and T' denotes target. Dis- tances in angstrom. For Z reconstruction, we show the percentile at cach absolute difference.
Locations Percentile
P P-)
T
0.
4.31 10
5.47 x 10
10.
5.03 10
5.02 x 10-2
20.
7.04 10
7.05 10
30.
9.12 x 10-2
9.13 x 10
4.1 PLACEMENT OF ATOMS
40.
1.14 x 10-1
1.14 10
50.
1.36 10
1.37 X 10
Using the output of the segmentation network, we
60.
1.60 10-1
1.61 X 10
take the $ argmax (S). From this representation,
70.
1.94 x 10-1
1.94 X 10
we find the connected components and use majority
80.
2.48 10
2.50 10
voting to assign each cluster atom identity (see
90.
3.94
10
4.03 x 10
Appendix for details). In Fig. 3 (and Fig S8) we
100.
2.59
1.89
show the aceuracy of all top-1 predictions. We
Species
compare the results of our segmented matrix, S with the true values, S. From the center of mass of each predicted atom we compute the distance the nearest true atom as follows
Abs. Diff. in Z
Percent
66.8 %
25.1 %
5.8
miny d(S,, Sj) A j and min, V.). (3)
> 3
2.2 %
Similarly, we can compute the distance from each true atom to the nearest predicted atom. By com-
puting this metric in both directions, we are able to verify that model is placing atoms in the correct locations but only in those locations.
4.2 ACCURACY ON UNIT CELL
In Fig. 3 we plot the results of applying our model to single unit cells of crystals (see also Table I. We find out model is able very accurately segment the locations of atoms with nearly 99% of atoms placed within 0.5 Ä of their number of atoms and 97% of unit cells are reconstructed with less than 2 atoms extra or missing, 66% species are correctly classified in the test set, with nearly 100 possible classes. When an atom is incorrectly classified, it is usually by one or two atomic numbers (see Table 1). Often adjacent elements (with atomic numbers that differ by 1) have very different chemical properties. We retrained the network to predict atomic group (where elements in the same group tend to have more similar properties) and found 93% accuracy. We report accuracy on as this a more difficult problem and in both cases we still unable to reconstruct crystals with the accuracy to relax in a DFT calculation. Missed atoms typically occur near the boundary.
location. Nearly 90% of unit cells are reconstructed with the correct
Next, will discuss the same network applied to repeating unit cells. In all cases, the results for single unit cells are improvements over the results of multiple unit cells. For most figures, there are corresponding panels between the single unit cell and repeated cell results.
4.3 ACCURACY ON REPEATING UNIT CELLS
In Fig. S8 we show the percentiles of all pairwise minimum distances in our reconstructed samples. We find that 50% of all reconstructed atoms are in 0.2 A. For both directions, the 75th percentile error was under 0.25 Á and the 90th percentile of reconstructed atoms were within 0.5 Å. For predicted atoms that are within 0.5 of true atom's location, we find that 65. 65.4% of the time our prediction matches the atomic number exactly Fig. S8 Again, nearly all errors occur near the boundary.
8

---

Under review as a conference paper at ICLR 2020
A 0.20
Demsty Beg
Real Crystals
Random
0.35
probability density
More
0.10
0.0b
horn VEE
distance
Figure Decode random latent space vectors. (A) look the spacing between nearest atoms from random draws from the latent space compared to those from real crystal structures. In blue we show the distribution of random reconstructions. In red we show the distribution of true inter-atomic spacing. (B) For 5 different random latent space vectors, show the reconstructed density field and the resulting segmentation. Notice that in some cases small amounts of density are predicted but are not segmented into an atom (for example, in columns and 4).
For 120 randomly selected reconstructed crystals from the test set, we compare the number of atoms after segmentation with the true number (Fig. (8D (Fig and 3B. In Fig. S&E we compare the predicted species compared with the species of the nearest true atom, independent of the distance to the nearest atom. We find that based on the results of Fig. S8 our current model is able to more accurately reconstruct the locations of atoms than obtain their true atomic number. This makes sense, given the enormous number of possible classes. However for repeating unit cells, when an atom is correctly predicted to be within 0.33 Á of a true atom location nearly 70% of the time the atom is assigned exactly the correct Z value. The correlation between the true and predicted value is 0.98 suggesting that when the network makes an incorrect assessment, it chooses a nearby atomic number.
Further improvements in this placement and identification of atoms necessary in achieving crystall that will appropriately relax in a DFT simulation for the calculation of quantum mechanical properties. However, a very exciting immediate application of this work in random structure searching Pickard & Needs 2011)
4.4 LATENT SPACE INTERPOLATION
We encode two different true density maps into vectors Z1 and zz, respectively. We then linearly interpolate between these two values in latent space, constructing intermediate vectors 2, We decode the resulting latent space vectors which result in predicted density maps, D(z). Passing these into the trained segmentation routine, we find that intermediate results segment into atoms (see results on a single unit cell in Fig. and Fig. 4). See results on a repeating lattice in Figs. $9 and S11
4.5 RANDOM DRAWS
We can randomly sample latent space vector ~ N 1) and we then decode the output. We do this for both single unit cells as well as for repeating lattices (see Appendix). Random structure searching is an important area of material science research an exciting direction for this work Pickard & Needs 1201 Using the network trained to encode and decode single unit cells, we decode a Z ~ N 1).
See anonymized link: http://tttteegggg....cooooeeeeee
home

---

Under review as a conference paper at ICLR 2020
In Fig. we show the decoded density fields along with the resulting segmented output. There are few desirable features: single cells are centered, atoms tend to be of similar atomic number, even in different parts of the reconstructed output, and atoms not ever too close together.
CONCLUSIONS AND FUTURE WORK
The ability encode and decode 3-D structures very interesting direction of recent work and increasingly important in helping create models that can understand the world that we live in. In current material design, lot of focus has been on molecules where the 3-D structure can be safely ignored. However, in many cases this is not the case. We do not know of an effective data representation for these 3-D molecules and find that by directly utilizing a proxy for density we are able to encode and decode the geometry of these molecules.
An important future direction of research is to come up with successful representations for encoding crystal structures such that they can be easily encoded and decoded. Coming up with representation of crystal unit cells (and 3-D structures in general) so that decoded molecules are physically plausible is an important aspect that needs to be carefully considered.
Our approach is currently unable to generate molecules that are physically stable, but there are promising directions in this direction. Using ideas similar to G-SchNet, using the decoded density field, atoms can be placed sequentially, conditioned on the placement of previous atoms (Gebauer al. 2019). Working towards decoding molecules that able relax is very exciting direction.
Another improvement would be to alter the structure of our encoder/decoder. Currently, we use standard 3-D convolutions. However, using SE(3) equivariant kernels from M. Weiler al. is promising direction for improved performance (Weiler et al. 2018). We feel that our current approach could be extended beyond the domain of material science, as are able encode 3-D distance map, this approach could be broadly applied to many objects. In essence, we are able to learn distance transform from an object. This is more general than for encoding and decoding atomic structures. To try to facilitate further research in the creation of 3-D representation of more complex atomic structures, we will release the code shortly.
REFERENCES
Panos Achlioptas, Olga Diamanti, Ioannis Mitliagkas, and Leonidas Guibas. Learning representations and generative models for 3d point clouds. arXiv preprint arXiv: 1707.v7707.000 707.9999....... 2017.
Zaccary Alperstein, Artem Cherkasov, and Jason Tyler Rolfe. All smiles variational autoencoder, 2019.
Rim Assouel, Mohamed Ahmed, Marwin H Segler, Amir Saffari, and Yoshua Bengio. Defac- tor: Differentiable edge factorization-ban basedd probabilistic graph generation. arXiv preprint a8Xiv: 1.09666, 2018.
Andrew Brock, Theodore Lim, James M Ritchie, and Nick Weston. Generative and discriminative voxel modeling with convolutional neural networks. arXiv preprint arXiv: 2016.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint wrXiv:100911199: 2018.
Chi Chen, Weike Ye, Yunxing Zuo, Chen Zheng, and Shyue Ping Ong. Graph networks as universal machine learning framework for molecules and crystals. Chemistry of Materials, 31(9):35643355 3572, 2019.
Ricky TQ Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disentangle- ment in
Xi Chen, Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. InfoGAN: Interpretable representation learning by information maximizing generative adversarial nets. In D. D. Lee, U. V. Luxburg, I. Guyon, and R. Garnett (eds.), Advances In Neural Information Processing Systems 29, pp. 2172-2180. Curran Associates, Inc., 2016.
10

---

Under review as a conference paper at ICLR 2020
Özgün Çişek, Ahmed Abdulkadir, Soeren S Lienkamp. Thomas Brox, and Olaf Ronneberger. 3d u-net: learning dense volumetric segmentation from sparse annotation. In Interaational conference medical image computing and computer-assisted intervention, pp. 424 432. Springer, 2016.
Ekin D Cubuk, Austin D Sendek, and Evan Reed. Screening billions candidates for solid lithium-ion conductors: A transfer learning approach for small data. The Journal of chemical physics, 150(21):214701, 2019.
Niklas WA Gebauer, Michael Gastegger, and T Schütt. Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules. arXiv preprint arXiv: Stciv:/9006.0000 1900999000077, 2019.
Rafael Gómezz-Bobareelli, Jennifer N Wei, David Duvenaud, José Miguel Herráánde-Lobaal Benjamín Bennchez-Lengeling Dennis Sheberla, Jorge gquilera-Iparraginng Timothy D Hirzel, Ryan P Adams, and Alán Aspuru- Guzik. Automatic chemical design using data-driven continuous representation of molecules. ACS central science, 4(2):268-21) 276, 2018.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural informa- tion processing systems, pp. 2672-2680, 2014.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. ICLR, 2(5):6, 2017.
Moritz Hoffmann and Frank Noć. Generating valid euclidean distance matrices, 2019.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03111 2015.
John J Irwin, Teague Sterling, Michael M Mysinger, Erin S Bolstad, and Ryan G Coleman. Zinc: free tool discover chemistry for biology. Journal of chemical information and modeling. 52(7): 1757-1768, 2012.
Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, et al. Commentary: The materials project: materials genome approach accelerating materials innovation. Apl Materials, 1(1):01 1002, 2013.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. arXiv preprint 100.04180,.04364, 2018.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. arXiv preprint 2018.
Tacksoo Kim, Moonsu Cha, Hyunsoo Kim, Jung Kwon Lee, and Jiwon Kim. Learning to dis- cover cross-domain relations with generative adversarial networks. In Proceedings of the 34th International Conference on Machine Krarning-Volune 70, pp. 1857-1865. JMLR. org, 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 64800 2014.
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. arXiv:1312.6114 2013.
arXiv preprint
Matt J Kusner, Brooks Paige, and José Miguel Hemáández-Lobato. Grammar variational autoencoder. In Proceedings of the 34th International Conference on Machine Learning-Vor Volume 70, pp. 1945- 1954. JMLR.org, 2017.
Clare F Macrae, Ian Bruno, James A Chisholm, Paul R Edgington, Patrick McCabe, Elna Pidcock, Lucia Rodriguez-Monge Robin Taylor, van de Streek, and Peter A Wood. Mereury esd 2.0- new features for the visualization and investigation of crystal structures. Journal of Applied Crystallography, 41(2):666-470 2008.
11

---

Under review as a conference paper at ICLR 2020
Kaushalya Madhawa, Katushiko Ishiguro, Kosuke Nakago, and Motoki Abe. Graphnvp: An invertible flow model for generating molecular graphs. arXiv preprint arXiv:1905. 2019.
Elman Mansimov, Omar Mahmood, Seokho Kang, and Kyunghyun Cho. Molecular geometry prediction using a deep generative graph neural network. arXiv preprint 19Xivv1904.00.) 2019.
Daniel Maturana and Sebastian Scherer. Voxnet: 3d convolutional neural network for real-time object recognition. In 2015 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 922-928. IEEE, 2015.
Hsien- Yu Meng. Lin Gao, Yu-Kun Lai, and Dinesh Manocha. Vv-net: Voxel vae net with group convolutions for point cloud segmentation. In Proceedings of the IEEE International Conference on Computer Vision, pp. 8500-8508, 2019.
Thu Nguyen-Phuoe Chuan Li, Lucas Theis, Christian Richardt, and Yong-Liang Yang. Hologan: Unsupervised learning of 3d representations from natural images, arXiv preprint arXiv:1904.0133: 2019.
Guangda Niu, Xudong Guo, and Liduo Wang. Review of recent progress in chemical stability of perovskite solar cells. Journal of Materials Chemistry A, 3(17):9970-899) 2015.
Asma Nouira, Jean-Claude Crivello, and Nataliya Sokolovska. Crystalgan: Learning to discover crystallographic structures with generative adversarial networks. arXiv preprint arXiv: 118./0011 1203, 2018.
Augustus Odena, Vincent Dumoulin, and Chris Olah. Deconvolution and checkerboard arti- facts. Distill, 2016. doi: 10.23915/distill000000 URL http::/ wttpssiii.ppp///111111 pub/2016/
deconv chw.cerboocrrrdd
Ozan Oktay, Jo Schlemper, Loic Le Folgoc, Matthew Lee, Mattias Heinrich, Kazumari Misawa, Kensaku Mori, Steven McDonagh, Nils Hammerla, Bernhard Kainz, et al. Attention u-net: Learning where to look for the pancreas. arXiv preprint arXiv: 180.0398040 2018.
Shyue Ping Ong. William Davidson Richards, Anubhav Jain, Geoffroy Hautier, Michael Kocher, Shreyas Cholia, Dan Gunter, Vincent L. Chevrier, Kristin A Persson, and Gerbrand Ceder. Python materials genomics (pymatgen): robust, open-source python library for materials analysis. Computational Materials Science, 18.314-319,2013
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary De Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Chris Pickard and RJ Nceds. Ab initio random structure searching. Journal of Physics: Condensed Matter, 2355)0053 2011.
Pavel G Polishchuk, Timur Madzhidov, and Alexandre Varnek. Estimation of the size of drug-like chemical space based on gdb- 17 data. Journal of computer-aided molecular design, 2776)6675-679, 2013.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1:140022, 2014.
Ali Razavi, Aaron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. arXiv preprint arXiv: 1906600446 2019.
David Rolnick, Priya Donti, Lynn H Kaack, Kelly Kochanski, Alexandre Lacoste, Kris Sankaran, Andrew Slavin Ross, Nikola Milojevic-Dupont, Natasha Jaques, Anna Valdman-Brown, et al. Tackling climate change with machine learning. arXiv preprint 1rXiv:1906.054433 2019.
O. Ronneberger, P.Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmen- tation. In Medical Image Computing and Computer-Assisted Intervention (MICCAI). volume 9351 of LNCS, pp. 234-241. Springer, 2015. URL :attp:/a h/tp:///mmmm/// httpm/minffoo informatik ik. wni-uneeiuuuee de/Publ Comublicatoo io / 5a (available arXiv:1505.04597 [csCCV).
12

---

Under review as a conference paper at ICLR 2020
Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722,, 2018.
Daniel Schwalbe-Koda and Rafael Gömez-Bombarelli. Generative models for automatic chemical design. arXiv preprint arXiv:1907.01632. 2019.
Marwin HS Segler, Thierry Kogej, Christian Tyrchan, and Mark P Waller. Generating focused molecule libraries for drug discovery with recurrent neural networks. ACS central science, 4(1): 120-131, 2017.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In Advances in neural information processing systems, pp. 3483 3491, 2015.
Daniel P Tabor, Loic M Roch, Semion K Saikin, Christoph Kreisbeck, Dennis Sheberla, Joseph H Montoya, Shyam Dwaraknath, Muratahan Aykol, Carlos Ortiz, Hermann Tribukait, et al. Acceler- ating the discovery of materials for clean energy in the era of smart automation. Nature Reviews Materials, 3(5):5, 2018.
Diego Valsesia, Giulia Fracastoro, and Enrico Magli. Learning localized generative models for 3d point clouds via graph convolution. 2018.
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3d steerable cnns: Learning rotationally equivariant features in volumetric data. In Advances in Neural Information Processing Systems, pp. 10381-10392, 2018.
David Weininger. Smiles, a chemical language and information system. 1. introduction to methodol- ogy and encoding rules. Journal of chemical information and computer sciences, 28(1):31-36, 1988.
Jiajun Wu, Chengkai Zhang, Tianfan Xue, William T Freeman, and Joshua B Tenenbaum. Learning a probabilistic latent space of object shapes via 3d generative-adversariil modeling. In Advances in Neural Information Processing Systemx, pp. 82-90, 2016.
Wenxuan Wu, Zhongang Qi, and Li Fuxin. Pointcony: Deep convolutional networks on 3d point clouds, 2018
Tian Xie and Jeffrey C Grossman. Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties. Physical review letters, 120(14):144 145301, 2018.
Guandao Yang, Xun Huang, Zekun Hao, Ming-X Yu Liu, Serge Belongie, and Bharath Hariha- ran. Pointflow: 3d point cloud generation with continuous normalizing flows. arXiv preprint arXiv: /2320, 2019.
Xiufeng Yang. Jinzhe Zhang, Kazuki Yoshizoe, Kei Terayama, and Koji Tsuda. Chemts: an efficient python library for de novo molecular generation. Science and technology of advanced materials, 18(1):972-976, 2017.
Jiaxuan You, Bowen Liu, Zhitao Ying, Vijay Pande, and Jure Leskovee. Graph convolutional policy network for goal-directed molecular graph generation. In Advances in Neural Information Processing Systems, pp. 6410-6421, 2018.
Jun-Yan Zhu, Zhoutong Zhang, Chengkai Zhang, Jiajun Wu, Antonio Torralba, Josh Tenenbaum, and Bill Freeman. Visual object networks: image generation with disentangled 3d representations. In Advances in Neural Information Processing Systems, pp. 118-129, 2018.
13

---

Under review as a conference paper at ICLR 2020
A
SUPPLEMENTAL INFORMATION
A.1
VIDEOS
See
anonymoized project site at:
sttss://stees.googee..olll/
encode-decode -eooood--dccddd---------a-aaaaa- wd-c-rdatsrttsseeeeeeeeeeeee
B
CHOICE OF DATA REPRESENTATION
We choose to use a Angstrom cutoff as kept over 90% of the data we had available without becoming too computationally demanding, as we needed to maintain relatively small grid spacing between voxels (0.33 Angstrom). Because the problem is 3-D, a 303 input equivalent to 164 164 image, in terms of pixels.
C COMPARISON OF 1-D, 2-D, AND 3-D REPRESENTATIONS
For small molecules, 1-D SMI LES and 2-D graph based representations have been very effectively used for property prediction and for molecule generation tasks. In SMILES representation, specific grammar is learned but from a valid string, specific molecule is generated based on the bond arrangement Weininger (1988), GGómez-Bombrede et al. (2018). Similarly, graph based representation is able to place atoms in space without specifically assigning each atom a specific coordinate Jin et al. (2018) Additionally, graph representations have been shown to be successful for property prediction of crystals Xie & Grossman (2018), but not for generating crystals. For small molecules, these representations have proved very effective, especially since both these representations effectively reduce the dimension the problem.
Specifically, SMILES and graph representations both do not encode the specific spatial arrange- ment of atoms. Instead, these representations are cleverly created in such a way that the spatial arrangements reconstructed. Unfortunately, for many larger molecules these representations not sufficient. Therefore, we focus on developing a representation that is able to directly reconstruct both species and locations in a 3-D space.
D SPECIES PREDICTION
Each voxel is assigned an identity based on the argmax of the U-Net. The background has label First, find all connected components of non-background voxels. For each connected component, we tabulate the species predicted at each voxel and each cluster is assigned an identity baased on the majority vote. The location is based the centroid of the cluster. We experimented with more complicated procedures (such as weighting by the probability vector) and did not find meaningful improvement.
E TRAINING DETAILS
We used a Tesla V100 for training. We used learning rate 1e-5 and a batch size of 24.
E. NETWORK DETAILS
1.
ENCODER
3D convolution, kernel size of 5 and stride of 2. 16 channels. Batch Norm and LeakyReLU activation.
3D convolution, kernel size of 3 and stride of 1. 32 channels. Batch Norm and LeakyReLU activation.
3D convolution, kernel size of 3 and stride of 1. 64 channels.
Batch Norm and LeakyReLU activation.
3D convolution, kernel size of 3 and stride of 2. 128 channels.
14

---

Under review as a conference paper at ICLR 2020
Batch Norm and LeakyReLU activation. Fully connected layer with bottleneck size 300.
E.1.2 DECODER
Fully connected layer.
Reshape 128 5,5,5
Trlinear upsample by factor of two.
Conv3D with 64 channels, kernel size 5, leakyReLU. Trlinear upsample by factor of two.
Conv3D with 32 channels, kernel size 5, leakyReI.U. Trlinear upsample by factor of two.
Conv3D with 16 channels, kernel size 4, IeakyReLU. Conv3D with channels, kernel size of 4, ReLU.
F RESULTS FROM THE ENCODER-DECODER
Our encoder-decoder architecture is able to accurately reconstruct the voxel-wise density (Fig. S8 and Fig. (S6). find a strong correlation between the predicted and target voxel value (Fig. S6A) This is essential for achieving accurate segmentation of the correct atoms and their corresponding positions. Due to the strong penalty on the KL term we find that there is an increased spread in the resulting predictions when attempting to encode and decode repeating unit cells. This is less of an issue when encoding and decoding single unit cells.
G ACCURACY OF RANDOM DRAWS IN A REPEATING LATTICE
To test whether random samples from the latent space, ž, decode to physically realistic molecules, we trained discriminator. We trained an auto-encoder that was able to very accurately reconstruct molecules by greatly reducing the penalty on 3 (see Fig. S122. Then, we compute real latent space representations of our different molecules, z. We then draw a U (0, 1) and construct a new latent space vector
2 XŽ + 101-X)ze -)Zeeal
(4)
where ž is random sample from unit normal and Zreal is random "true" latent space. We pass this through the decoder and provide this label (1 - X, A). then randomly draw z ~~ NN(, 1) from our trained network and pass this reconstruction D into the discriminator network, which outputs prediction of the distance from true crystal reconstruction versus random draw from latent space of a previously trained network. Applying this network to random samples from z ~ WNO0 1) we get mean value of 0.84 with standard deviation 0.05. For true decoded samples, a mean value of 0.90 with a standard deviation of 0.01. In Fig. S13 we show results of the trained network and predictions on three different datasets. We find that the target latent space decodes to high-scoring samples (see Fig. ()13B) and that out of distribution samples typically decode to much lower scores.
find
G.1 CONDITIONAL VAE
Moving forward, an important step will be creating structures that have specific chemical properties. Moving towards this direction, we attempt to condition the generation of molecules on the largest species present (Sohn et al. 2015). To do this, we multiply the input and output of the bottleneck by the largest present density. By doing this, are able to control the magnitude of the resulting density field without perturbing the geometry (see Fig. 3). We attempted to condition the encoder and decoder by concatenation though we found this was not sufficient to generate density fields with the desired property.
still able to get the network to train but find that
The ability to condition specific properties has obvious applications in the targeted generation of molecules. In this case, we choose to condition on the largest present density as a way to generate molecules that do not have atom with more than specified maximum atomic number present. We
15

---

Under review as a conference paper at ICLR 2020
A
Sarget
prediction
slice
slice 9
slice 17
slice 25
low
high
low
high
target
preditctio
prediction
prediction
prediction
prodiction
Figure S6: Accuracy of Model. (A) For each position in 3-D space, plot the predicted and target density for 4 different random erystals from the test set. The red dashed line an identity line. (B) For each of panels (A), we show 4 different z-slices through the true and predicted density fields. (C) We show the full 3-D reconstruction of the prediction and the true density field.
find that by varying different parameters in the latent space, some vary the field in such way that appears to correlate with the location of atoms (see Fig. (SDOC,D). However, unlike as in work on 2-D problems, the changes in the 3-D scalar field are much harder to interpret (Chen et al. 2016 Higgins et al. [007p. Future work will attempt to use improved factorizing techniques (Kim & Mnih 2018; Chen et al.). Additionally, one could seek to condition the generation of single unit cells quantum mechanical properties or by adding auxiliary loss using the latent space to predict quantum mechanical properties.
16

---

Under review as a conference paper at ICLR 2020
Top
Front
Atomic Number
Small
Large
Figure S7: Random Decoded Samples. We randomlyl žample the latent space N(O, 1). We show series of examples (from three different views) of different complexity. Note that while symmetry is not preserved, various geometric motifs are.

---

Under review as a conference paper at ICLR 2020
A
B
80
70e
0.19
0.05
60
MSE error
50
6C+
top-1 emor
0:00
X
0.06
0.02
006
epoch
epoch
C
D
E
100
A-B
B-A
predicted number
Angstrom
0.500
predicted
0.100
0.050
0.00
0.05
40
80
100
100
150
60
percentile
target number
target
Figure S8: Repeating unit cell accuracy. (A) each position 3-D space, we compute the difference between the truth and the reconstruction. We show different percentile bands the reconstruction error between the target and predicted density, plotted using the mean square error (MSE). (B) For the species matrix from U-Net ask what the error of top-1 predictions is between predictions and the ground truth. (C) For segmented matrices, ask the distance from the distance errors by percentile for both the nearest true atom each predicted atom and vice-versa (orange and blue, respectively). (D) After segmenting the reconstructed density maps we show the predicted and true number of atoms. (E) We plot the predicted nearest atom species versus the true species of the closest corresponding Á. We find 65. 4% are correctly predicted.
segmented
the ground truth. We show the
atom, as long as distance less
18

---

Under review as a conference paper at ICLR 2020
50
75 68
24
Figure S9: Interpolation between two molecules. In (A) we show the interpolation between two crystal density maps. We show three equally spaced intermediates the corresponding segmentation. In (B), we show the same interpolation but highlight different dimensional slices (y-axis). Along the x-axis label the fraction of the way between the two latent space vectors.
19

---

Under review as a conference paper at ICLR 2020
max:: 17
max:: 31
max::
maxc 54
max: 6G
max: 73
max: 78
A
B
generatod max
100
normalized
fixed
target max
D
+0.5
20
-0.5
10
20
Figure S10: Conditional generation molecules. We multiply the input and output of the bottleneck by the maximum density during training. Then, randomly decoding samples of oz where z ~ NNO0, 1) and a is random target, we find we able to generate density maps that decode appropriately. If we take true decoded value and re-scale the output find target and generated max value. (B) Using an encoded z, we multiply by different target values and decode them. We find that we are able to change seale without affecting the geometry. (C) We show decoded sample. We cut away the right most corner facing the viewer. (D) For the molecule shown in (C), we show the results of varying value in the latent space from -3 to 3. We plot the difference the slice when compared the unaltered, decoded, for slice that is in-plane with many atoms.
able control scale without affecting geometry. (A) We show the
20

---

Under review as a conference paper at ICLR 2020
re0 c
iv
Figure S11: Interpolation between two molecules. From i show two views of the output from the segmentation routine decoding the latent space between two different molecules. the primed corresponding figures we show the density fields that output decoder. We show the output of each grid normalized between 0-1 in bluc-green and on a fixed color map in magma. We also vary the opacity to show the locations of predicted atoms.
21

---

Under review as a conference paper at ICLR 2020
A
Slice
Density Fields
Torget
Prediction
Prediction
True
Torget
Prediction
Prediction
Target
Prediction
Prediction
12
30
80
True
Figure $12: Results with smaller 3. We decrease 3 by factor of 10. reducing the penalty the Kullback- Leibler (KL) able generate accurate reconstructions, However, sampling N(,0 1) does not fully sample the enended space of molecules.
1.0
Decaded Samples 1|
0.8
districution draws (0))
Prediction
Precicted
Roadom from Unit Motmml
0.4
Out
0.2
02
0.0
0.2
0.4
0.8
1.0
True À
Figure S13: prediction for different samples. Using the network in Fig, S12 we train a discriminator network to predict the "distance" from a real crystal. We apply this discriminator to real decoded samples, random latent space draws, and out of distribution draws (right).
22

---

Under review as a conference paper at ICLR 2020
Figure S14: Decoded samples of repeating unit cells. Results from the encoder-decoder.
Figure S15: Segmented Output from the decoded samples. The segmented output from Fig. S14
23


Reviews:
[{"review_id": "Syl-yGZzcH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "PAPER SUMMARY: This paper addresses the problem of encoding and decoding 3D chemical structures, with the ultimate goal of generating 3D crystal structures. The authors propose an auto-encoder framework for encoding the 3D locations of atoms in the crystal to a latent representation and then decoding that representation back into 3D structure. The paper's contributions can be summarized as follows:\n1) A data representation that converts the 3D atom locations to a 3D voxel density map, so that they can be encoded by a standard 3D convolutional network.\n2) A decoder network that first estimates a 3D density map from the latent vector using upsampling and convolutions and then classifies the atom type (atomic number) per voxel using a 3D segmentation.\n3) The network is applied on unit cells of crystals or repeated unit cells from a dataset of crystal structures and the network is shown to be able to accurately reconstruct the 3D density maps, predict the number of atoms in the cell and perform fairly well in their classification into atomic types.\n\nI appreciate that the paper addresses an interesting problem that is not sufficiently explored and can motivate the development of novel methods that generate 3D molecules with particular structure and multiple types of atoms, however the current work combines existing methods, without any architectural modifications that exploit the new domain.  If the presentation of results and experiments is improved, this could be a good application paper in material science with an interesting combination of techniques from machine learning and computer vision. It would be more appropriate to submit this paper to a domain-specific venue rather than to ICLR. Therefore, I cannot justify its acceptance to ICLR in its current format.\n\nStrengths\n-----------------------\n1.\tNew domain: The paper addresses an interesting problem in a new domain. Previous work on generative models for molecules have concentrated on molecules with 1D structure (which can be represented as strings) or 2D structure (which can be represented as planar graphs). Generating 3D molecule structures is a challenging and interesting problem.\n2.\tEfficient data representation: The authors bypass the difficulties associated with modeling sets of 3D points with arbitrary structure by proposing a canonical, voxel-based density representation.\n3.\tJoint VAE + UNet training: The joint training of the encoder-decoder VAE and the 3D segmentation network results in a decoder that can reconstruct atom locations and atom types, being robust to mistakes in the density map reconstruction.\n4.\tQualitative results: There are nice visualizations of the reconstructed molecules and of the effect of the latent variable z. Figure 3 in particular does a great job at elucidating the outputs of the network and the reconstruction quality.\n\nWeaknesses\n------------------------\n1.\tLimited methodological novelty: The methods used in the paper, i.e. the data representation (see detailed comments), the encoder network, the decoder network and the segmentation network are all existing methods without any (or with only minor) modifications.  Their combination also seems straightforward, except for the joint training of the VAE and the segmentation network, which has not been tried before to the best of my knowledge. Networks have not been modified to exploit the intricacies of the new domain/task, such as symmetries, repetition of unit cells, large number of atoms.\n2.\tLack of comparison with other methods/baselines: There is no quantitative comparison with alternative baselines or methods or even discussion of such alternatives. It is understandable that the paper addresses a relatively new problem, however the method could be compared to CrystalGAN. Also, there is no justification of the advantages of using a voxel-based density map representation along with 3D convolutions vs, for instance, a graph representation along with graph convolutions (Xie & Grossman, [3]) or a point cloud GAN (Achlioptas). \n3.\tNo quantitative evaluation of the generative capabilities of the network: In the case of drug molecule generation, logP and QED scores  [4,5] are used to evaluate their drug-likeness. Without such scores in the crystal generation domain, it is not easy to judge how good a generated crystal is.  The need for some type of quantitative evaluation is especially important, since it is not as easy to qualitatively judge the quality of crystals obtained by sampling random latent vectors (Fig 5B) as in the case of generated images or text. The only such evaluation presented in the paper is related to the distribution of distances between atoms which seems to be close to their actual distribution in nature. \n4.\tPresentation of quantitative results: The quantitative results are scattered in text throughout the results section without being summarized in a table. Results about the distance of generated atoms w.r.t to their true location, the predicted atom counts, predicted atom types etc. for the cases of single unit cells and repeated unit cells should be added to a table to complete a rigorous experimental evaluation. \n\nAdditional Comments\n-------------------------\n1.\tLack of novelty: The data representation, which is presented as one of the core contributions of the paper, is not new. Radial basis functions have been used to generate 3D density maps (e.g. [2]) and 3D Convolutional VAEs (Brock2016, [1]) have been employed on density maps (e.g. occupancy maps) on computer vision tasks. \n2.\tConfusing description of the repeated unit cells case: It is not clear, especially for a general audience, how the repeated unit cells data are generated (in which directions are the unit cells repeated, how many times). It is also not clear how the training routine (conditioning) and the output post-processing (connected components + majority voting) is modified.\n3.\tAfter Eq. (2), q is the --approximate-- posterior …\n\nSuggested References\n--------------------------\n\n1.\tVoxNet: A 3D Convolutional Neural Network for Real-Time Object Recognition\n2.\tVV-NET: Voxel VAE Net with Group Convolutions for Point Cloud Segmentation\n3.\tDynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs\n4.\tGraph Convolutional Policy Network forGoal-Directed Molecular Graph Generation\n5.\tGrammar Variational Autoencoder\n6.\t3D U-Net: Learning Dense VolumetricSegmentation from Sparse Annotation\n7.\tPointConv: Deep Convolutional Networks on 3D Point Clouds\n", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "N/A", "review_assessment:_checking_correctness_of_experiments": "N/A", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "Hkl1FAmAFr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper deals with accurately encoding and decoding 3D atomic positions and the crystal’s species using 2 sets of neural networks a) a VAE that builds a compressed latent space representation of a crystal and b) a UNET for segmenting the latent space into atoms and assigns each atom to its atomic number. Experiments were conducted on over 120K 3D samples of crystals and the results seem to be promising.\n\nThe paper is neatly written and well organized.\n\nComments:\n\nA) Figure 2: For completion, consider marking M and S as outputs of the VAE and U-Net respectively. \n\nB) In Section 3.1, why do the authors use a cube with side of ’10’ Angstrom? And why divide the cube into ’30’ bins?\n\nC) The paper revolves on vanilla 3D convolutions of the crystal structures. Have the authors considered how the results would change if SO(3) rotation invariant convolutions were used instead. The SO(3) convolutions would empower capturing all possible rotations of the crystal other than only its canonical form.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "ByxjbyqV5H", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors describe a method to encode and decode the position of atoms in 3-D molecules. An encoder-decoder architecture is used to create a representation of a molecule and to reconstruct the molecule from its representation. Then a second Neural Network segments the output and assigns an atomic number. Prior work on this task has used 1D (SMILES) and 2D (Graph) representations. The authors argues that exploiting 3D structure can create better representations.\n\nAs the paper's related work section shows, this is not the first attempt to use 3D structure to create molecular representations. Unfortunately, the paper does not compare their work to prior work on 3D structure representations (e.g Gebaur et al 2019). Also, it is not clear whether the 3D representation is better than 1D or 2D representations especially since there have been many new 1D models that perform very well for tasks like molecular property prediction (For example All SMILES VAE https://arxiv.org/abs/1905.13343 ). I think the community will benefit if the authors perform a comparison with state of the art 1D and 2D models. I think this is a main drawback of this work.\n\nAnother concern is that the authors claim that the errors in atomic numbers differ only by 1 or 2. But doesn't this show that the network has not learnt a good representation? Because atoms that differ in atomic number by 1 or 2 will have different valencies and hence exhibit different properties? On the other hand, if the authors can show that the errors in atomic numbers suggest they correspond to similar atoms (may be along the same column in the periodic table), then one can have better confidence that the network has learned a meaningful representation.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents an encoder-decoder based approach to construct a compressed latent space representation of each molecule. Then a second neural network segments the output and assigns an atomic number. Unlike previous works using 1D or 2D representations, the proposed method focuses on the 3D representations.

The reviewers have several major concerns. Firstly, the novelty of the paper seems to be limited as the proposed method mainly use the existing techniques. Secondly, there is no clear baseline to compare with. Finally, there is no clear quantitative results to measure the proposed method. The rebuttal did not well address these problems.

Overall, this paper did not meet the standard of ICLR and I choose to reject the paper.

Decision: Reject