SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Youzhi Luo, Chengkai Liu, Shuiwang Ji  
Department of Computer Science & Engineering  
Texas A&M University  
College Station, TX 77843  
{yzluo, liuchengkai, sji}@tamu.edu
We consider the problem of generating periodic materials with deep models. While symmetry-aware molecule generation has been studied extensively, periodic materials possess different symmetries, which have not been completely captured by existing methods. In this work, we propose SyMat, a novel material generation approach that can capture physical symmetries of periodic material structures. SyMat generates atom types and lattices of materials through generating atom type sets, lattice lengths and lattice angles with a variational auto-encoder model. In addition, SyMat employs a score-based diffusion model to generate atom coordinates of materials, in which a novel symmetry-aware probabilistic model is used in the coordinate diffusion process. We show that SyMat is theoretically invariant to all symmetry transformations on materials and demonstrate that SyMat achieves promising performance on random generation and property optimization tasks. Our code is publicly available as part of the AIRS library (https://github.com/divelab/AIRS).
Designing or synthesizing novel periodic materials with target properties is a fundamental problem in many real-world applications, such as designing new periodic materials for solar cells and batteries [1]. For a long time, this challenging task heavily relies on either manually designing material structures based on the experience of chemical experts, or running very expensive and time-consuming density functional theory (DFT) based simulation. Recently, thanks to the progress of deep learning techniques, many studies have applied advanced deep generative models [22, 13] to generate or discover novel chemical compounds, such as molecules [10, 39, 50] and proteins [48, 16]. However, while deep learning has been widely used in periodic material representation learning [51, 5, 40, 4, 57, 25], generating novel periodic materials with deep generative models remains largely under-explored.
The major challenge of employing deep generative models to the generation of periodic materials is capturing physical symmetries of periodic material structures. Ideally, deep generative models for periodic materials should maintain invariance to symmetry transformations of periodic materials, including permutation, rotation, translation, and periodic transformations [59, 9, 60]. To achieve this target, we propose SyMat, a novel symmetry-aware periodic material generation method. SyMat transforms the atom types and lattices of materials to symmetry-aware generation targets, which are generated by a variational auto-encoder model [22]. Besides, SyMat adopts novel symmetry-aware probabilistic diffusion process to generate atom coordinates with a score-based diffusion model [44]. Experiments show that SyMat can achieve promising performance on random generation and property optimization tasks.
Relations with Prior Methods. We note that several existing studies are related to our proposed SyMat, such as CDVAE [52]. Generally, the major difference of SyMat between them lies in the capture of symmetries. We will discuss the detailed differences in Section 3.5.
2.1 Molecule Generation
Designing or discovering novel molecules with target properties has long been a challenging problem in many chemical applications. Recently, motivated from the significant progress in deep generative learning [22, 12, 44, 13], a lot of studies have applied advanced deep generative models to this problem. Some early studies [19, 58, 41, 32, 27] consider molecules as 2D topology graphs and design graph generation models to discover novel molecules. More recently, considering the importance of 3D structures in many molecular properties, many studies have also proposed advanced generative mod
Material
Figure 1: An illustration of molecule (left) and periodic material (right) structures. Different from molecules, atoms in materials periodically repeat themselves infinitely in 3D space (we use a 2D visualization here for simplicity).
els that can generate 3D molecular conformations either from scratch [10, 39, 31, 15, 50], or from input conditional information [33, 56, 43, 53, 54, 42, 55, 20, 28]. However, all these existing studies are designed for non-periodic molecules while materials have periodic structures, in which atoms periodically repeat themselves in 3D space (see Figure 1 for an illustration of differences between molecules and periodic materials). Hence, they cannot be directly applied to the periodic material generation problem that we study in this work.
2.2 Periodic Material Generation
While molecule generation has been extensively studied, currently the periodic material design or generation problem remains largely under-explored. Early periodic material design methods [34, 7] only generate compositions of chemical elements in periodic materials but do not generate 3D structures. More recent studies have proposed to use variational auto-encoders (VAEs) [22] or generative adversarial networks (GANs) [12] to generate 3D structures of periodic materials. However, these methods use coordinate matrices [36, 21, 61] or voxel images [14, 6] of materials as the direct generation targets, which results in the violation of physical symmetries of periodic material structures in their captured probabilistic distributions. Also, the VAE and GAN models used by them are not powerful enough to capture the distribution of complicated 3D periodic material structures. Different from them, our method employs a more powerful score-based diffusion model [44] to generate 3D periodic material structures, and the model is designed to capture physical symmetries in materials.
2.3 Score-Based Diffusion Models for 3D Structure Generation
Diffusion models are a family of deep generative models that have achieved outstanding generation performance in a variety of data modalities, such as images [38, 37], audios [24], and molecular conformations [55, 15, 20]. Currently, score-based diffusion models [44] and denoising diffusion probabilistic models [13] are two most commonly used diffusion models, and our method employs score-based diffusion models for 3D periodic material generation. For a given data distribution  $p(\pmb{x})$  in the high-dimensional data space, the score-based diffusion model trains a parameterized score model  $s_{\theta}(\cdot)$  with parameters  $\theta$  to accurately approximate the score function  $\nabla_{\pmb{x}}\log p(\pmb{x})$  from  $\pmb{x}$ . To improve the score approximation accuracy in the regions where the training data is sparse, denoising score matching [46, 44] is proposed to more effectively train  $s_{\theta}(\cdot)$ . Specifically, let  $\{\sigma_t\}_{t = 1}^T$  be a set of positive scalars satisfying  $\sigma_1 > \sigma_2 > \dots >\sigma_T$ , the denoising score matching method first perturbs the data point  $\pmb{x}$  sampled from  $p(\pmb{x})$  by a sequence of Gaussian noises with noise magnitude levels  $\{\sigma_t\}_{t = 1}^T$ , i.e., sampling  $\tilde{\pmb{x}}$  from  $p_{\sigma_t}(\tilde{\pmb{x}} |\pmb{x}) = \mathcal{N}(\tilde{\pmb{x}} |\pmb{x},\sigma_t^2\pmb {I})$  for every noise level  $\sigma_t$ . The score model  $s_{\theta}(\cdot)$  is optimized to predict the denoising score function  $\nabla_{\tilde{\pmb{x}}}\log p_{\sigma_t}(\tilde{\pmb{x}} |\pmb{x})$  from  $\tilde{\pmb{x}}$  at each noise level, where the ground truth of  $\nabla_{\tilde{\pmb{x}}}\log p_{\sigma_t}(\tilde{\pmb{x}} |\pmb{x})$  can be easily computed from the mathematical formula
of  $\mathcal{N}(\tilde{\pmb{x}} |\pmb {x},\sigma_t^2\pmb {I})$  . Formally,  $s_{\theta}(\cdot)$  is optimized to minimize the following training objective  $\mathcal{L}$  ..
$$
\mathcal {L} = \frac {1}{2 T} \sum_ {t = 1} ^ {T} \sigma_ {t} ^ {2} \mathbb {E} _ {\tilde {\boldsymbol {x}}, \boldsymbol {x}} \left[ \left| \left| \frac {\boldsymbol {s} _ {\theta} (\tilde {\boldsymbol {x}})}{\sigma_ {t}} + \frac {\tilde {\boldsymbol {x}} - \boldsymbol {x}}{\sigma_ {t} ^ {2}} \right| \right| _ {2} ^ {2} \right]. \tag {1}
$$
As shown in Vincent [46], when  $\sigma_{t}$  is small enough, after  $s_{\theta}(\cdot)$  is well trained with the objective in Equation (1), the output  $s_{\theta}(x)$  can accurately approximate the exact score function for any input  $x$ . New data samples can be generated by running multiple steps of the Langevin dynamics sampling algorithm [49] with the score functions approximated by the score model.
Several existing studies have applied score-based diffusion models to 3D molecular and material structure generation. ConfGF [42] is the first score-based diffusion model for 3D molecular conformation generation. It proposes a novel theoretic framework to achieve roto-translational equivariance in score functions. Based on the theoretic framework of ConfGF, DGSM [30] further improves the performance of ConfGF by dynamic graph score matching. Though they have achieved good performance in 3D molecules, they do not consider periodic invariance so they cannot be applied to 3D periodic materials. Recently, a novel approach CDVAE [52] is proposed for periodic material generation, and score-based diffusion models are used in CDVAE to generate atom coordinates in materials. However, as we discussed in Section 3.5, CDVAE fails to achieve translational invariance in calculating the denoising score function, so it does not capture all physical symmetries.
3 SyMat: Symmetry-Aware Generation of Periodic Materials
While a variety of deep generative models have been proposed to capture all physical symmetries in 3D molecular conformations, it remains challenging to do it for periodic materials as they have more complicated 3D periodic structures. In this section, we present SyMat, a novel periodic material generation method. In the following subsections, we will first introduce the problem of symmetry-aware material generation, then elaborate how our proposed SyMat approach achieves invariance to all symmetry transformations of periodic materials.
3.1 Symmetry-Aware Generation
For any periodic material structure, we can consider it as periodic repetitions of one unit cell in 3D space, where any unit cell is the smallest repeatable structure of the material. Hence, for any periodic material  $M$ , we can describe its complete structure information with one of its unit cells and the three lattice vectors describing its periodic repeating directions. Specifically, assuming there are  $n$  atoms in any unit cell of  $M$ , we represent  $M$  as  $M = (A, P, L)$ , where  $A \in \mathbb{Z}^n$ ,  $P \in \mathbb{R}^{3 \times n}$ , and  $L \in \mathbb{R}^{3 \times 3}$  are the atom type vector, coordinate matrix, and lattice matrix, respectively. Here, the  $i$ -th element  $a_i$  of  $A = [a_1, \dots, a_n]$  and the  $i$ -th column vector  $\pmb{p}_i$  of  $P = [p_1, \dots, p_n]$  denote the atom type, i.e., atomic number and the 3D Cartesian coordinate of the  $i$ -th atom in the unit cell, respectively. The lattice matrix  $L = [\ell_1, \ell_2, \ell_3]$  is formed by three lattice vectors  $\ell_1, \ell_2, \ell_3$  indicating how the atoms in a unit cell periodically repeat themselves in 3D space.
In this work, we consider the problem of generating periodic material structures with generative models. Formally, we are given a dataset  $\mathcal{M} = \{M_j\}_{j=1}^m$  where each periodic material data  $M_j$  is assumed to be sampled from the distribution  $p(\cdot)$ . Because of the physical symmetry properties in periodic materials [59, 9],  $p(\cdot)$  is restricted to be invariant to the following symmetry transformations.
- Permutation. For any  $M = (A, P, L)$ , if we permute  $A$ , i.e., exchange elements in  $A$  to obtain  $A'$ , and use the same permutation order to permute the column vectors of  $P$  to obtain  $P'$ , we actually change nothing but atom orders in unit cell representations. Hence, let  $M' = (A', P', L)$ , the probability densities of  $M$  and  $M'$  are actually the same, i.e.,  $p(M) = p(M')$ .
- Rotation and translation. For any  $M = (A, P, L)$ , we can rotate and translate its atom positions in 3D space to obtain a new material  $M' = (A, P', L')$ , where  $P' = QP + b1^T$  and  $L' = QL$ . Here  $Q \in \mathbb{R}^{3 \times 3}$  is the rotation matrix satisfying  $Q^T Q = I$ ,  $b \in \mathbb{R}^3$  is the translation vector, and  $\mathbf{1}$  is an  $n$ -dimensional vector whose elements are all 1. It can be intuitively understood that rotation and translation do not change any internal structures of periodic materials, so  $M$  and  $M'$  are actually different representations of the same material, so  $p(M) = p(M')$  should hold.
Figure 2: Illustrations of training and generation process in our proposed SyMat method.
- Periodic transformation. Because every atom from a unit cell periodically repeats itself along lattice vectors in 3D space, any coordinate matrix  $\mathbf{P} \in \mathbb{R}^{3 \times n}$  has infinite number of periodically equivalent coordinate matrices. Formally, a periodically equivalent coordinate matrix  $\mathbf{P}^{\mathbf{K}}$  of  $\mathbf{P}$  can be obtained by periodic transformation as  $\mathbf{P}^{\mathbf{K}} = \mathbf{P} + \mathbf{L}\mathbf{K}$ , where  $\mathbf{K} \in \mathbb{Z}^{3 \times n}$  is an arbitrary integer matrix. Since  $\mathbf{P}^{\mathbf{K}}$  and  $\mathbf{P}$  are equivalent and the periodic transformations between them are invertible, we consider  $M = (A, P, L)$  and  $M^{\mathbf{K}} = (A, P^{\mathbf{K}}, L)$  as the same material so  $p(M) = p(M^{\mathbf{K}})$  should hold.
In the periodic material generation problem, our target is to learn a generative model  $p_{\theta}(\cdot)$  with parameters  $\theta$  from the given dataset  $\mathcal{M}$  so that the model can capture the real data distribution  $p(\cdot)$ , and the learned generation model can generate a valid material structure  $M$  with a high probability  $p_{\theta}(M)$ . In addition, we require the distribution  $p_{\theta}(\cdot)$  captured by the generation model to have the same symmetry-related restrictions as the real data distribution  $p(\cdot)$ . The main motivation of modeling symmetries comes from property optimization. A significant target of developing generative models for periodic materials is to generate novel periodic materials with desirable chemical properties (e.g., low energies). Notably, chemical properties are invariant to physical symmetry transformations. For instance, rotating a material structure in 3D space does not change its energy. In other words, chemical properties only depend on the 3D geometric information that is invariant to symmetry transformations (e.g., interatomic distances). Using symmetry-aware modeling forces generative models to only capture the distribution of these 3D geometric information, which facilitates searching materials with desirable properties in the latent space.
3.2 An Overview of SyMat
Generally, to generate a new material  $M = (A, P, L)$ , our proposed SyMat method first generates the atom type vector  $A$  and lattice matrix  $L$ , then generates the coordinate matrix  $P$  conditioned on  $A$  and  $L$ . The probability of generating  $M = (A, P, L)$  can then be described as
$$
p _ {\theta} (\boldsymbol {M}) = p _ {\theta} (\boldsymbol {A}) p _ {\theta} (\boldsymbol {L}) p _ {\theta} (\boldsymbol {P} | \boldsymbol {A}, \boldsymbol {L}). \tag {2}
$$
Since the structures of atom types and lattices are relatively simple, we find that a VAE [22] model already has adequate capability to capture their distributions. Hence, we use a VAE model to generate  $A$  and  $L$ . However, the structures of coordinates are much more complicated, so we use a powerful score-based diffusion model [44] for their generation. Next, we will elaborate the details of generating atom types and lattices in Section 3.3 and the details of generating coordinates in Section 3.4. An overview of the proposed SyMat method is illustrated in Figure 2.
3.3 Atom Type and Lattice Generation
For the generation of atom types and lattices, a natural idea is using the VAE model to directly generate and capture the distribution of  $\mathbf{A}$  and  $\pmb{L}$ . However, this fails to capture the invariance to symmetry transformations described in Section 3.1, including the permutation invariance of  $\mathbf{A}$  and rotation invariance of  $\pmb{L}$ . To tackle this issue, we transform  $\mathbf{A}$  and  $\pmb{L}$  to items that are invariant to symmetry transformations, and make these items be the direct generation targets of the VAE model. Specifically, for  $\mathbf{A} = [a_{1},\dots,a_{n}]$ , we can count the number of atoms for every atom type existed in  $\mathbf{A}$ , and represent  $\mathbf{A}$  with an unordered  $k$ -element atom type set  $c = \{(e_1,n_1),\ldots ,(e_k,n_k)\}$ , which indicates that for  $i = 1,\dots,k$ , there exist  $n_i$  atoms with the atom type  $e_i$  in  $\mathbf{A}$ . We can easily find that  $c$  is invariant to any permutations on  $\mathbf{A}$  so we use  $c$  as the direct generation targets. Besides, for the lattice matrix  $L = [\ell_1,\ell_2,\ell_3]$ , we do not generate the three lattice vectors  $\ell_1,\ell_2,\ell_3$  directly. Instead, we use six rotation-invariant items as the generation targets, including three lattice lengths  $\ell = [\ell_1,\ell_2,\ell_2]$ , where  $\ell_i = ||\ell_i||_2$ , and three lattice angles  $\phi = [\phi_{12},\phi_{13},\phi_{23}]$ , where  $\phi_{ij}$  is the angle between  $\ell_i$  and  $\ell_j$ . Overall, we transform  $\mathbf{A}$  and  $L$  to generation targets  $c,\ell ,\phi$  so that the generation process and captured distributions  $p_{\theta}(\mathbf{A})$ ,  $p_{\theta}(L)$  are ensured to be symmetry-aware.
The VAE model used for generating  $c$ ,  $\ell$ , and  $\phi$  consists of an encoder model and a decoder model. The encoder model is based on a 3D graph neural network (GNN) that takes a material  $M = (A, P, L)$  as inputs and outputs a  $d$ -dimensional latent variable  $z^A \in \mathbb{R}^d$  and an  $f$ -dimensional latent variable  $z^L \in \mathbb{R}^f$ . This 3D GNN employs the symmetry-aware SphereNet [29] architecture so  $z^A, z^L$  will not change if applying any transformations described in Section 3.1 to the input material. We will introduce more details of this 3D GNN model in Appendix B. The decoder model is composed of four multi-layer perceptrons (MLPs) including  $\mathrm{MLP}^e$ ,  $\mathrm{MLP}^k$ ,  $\mathrm{MLP}^n$ ,  $\mathrm{MLP}^L$ , which are used to predict  $c = \{(e_1, n_1), \dots, (e_k, n_k)\}$ ,  $\ell = [\ell_1, \ell_2, \ell_3]$ , and  $\phi = [\phi_{12}, \phi_{13}, \phi_{23}]$  from the latent variables  $z^A$  and  $z^L$ . Specifically, let  $E$  be the number of all considered atom types and  $N$  be the largest possible number of atoms that can exist in a material,  $\mathrm{MLP}^e$  predicts a vector  $p_e \in [0, 1]^E$  which contains the existing probability for each of the  $E$  atom type from the input  $z^A$ , and  $\mathrm{MLP}^k$  predicts the size  $k$  of the set  $c$  from the input  $z^A$ . The  $k$  atom types  $e_1, \dots, e_k$  with top- $k$  probabilities in  $p_e$  are chosen as the atom types in  $c$ , and  $\mathrm{MLP}^n$  predicts the number of atoms  $n_1, \dots, n_k$  for them, where each  $n_i$  is predicted from  $e_i$  and  $z^A$ . Here, the prediction of  $p_e$  is considered as  $E$  binary classification tasks for each of  $E$  atom types, and the prediction of  $k$  and  $n_1, \dots, n_k$  are both considered as  $N$ -category classification tasks. Besides,  $\mathrm{MLP}^L$  is used to predict lattice items  $[\ell, \phi] = [\ell_1, \ell_2, \ell_3, \phi_{12}, \phi_{13}, \phi_{23}]$  from  $z^L$ , which is considered as a regression task.
We now describe the training and generation process of the VAE model. When training the VAE model on a material dataset  $\mathcal{M}$ , for any material  $M = (A, P, L)$  in  $\mathcal{M}$ , the encoder model first maps it to latent variable  $z^A$  and  $z^L$ . Afterwards, the decoder model uses  $z^A$  and  $z^L$  to reconstruct the exact  $c, \ell$ , and  $\phi$  obtained from  $A$  and  $L$ . The VAE model is optimized to minimize the reconstruction error of  $c, \ell, \phi$  together with the KL-distance between  $z^A, z^L$  and the standard Gaussian distribution  $\mathcal{N}(\mathbf{0}, \mathbf{I})$ . The reconstruction error losses are defined as cross entropy loss and minimum squared error loss for classification and regression tasks, respectively. After the model is trained, we can generate  $A$  and  $L$  by first sampling latent variables  $z^A, z^L$  from  $\mathcal{N}(\mathbf{0}, \mathbf{I})$ , then using the decoder to map  $z^A, z^L$  to  $c, \ell, \phi$ . Finally, we can produce  $A$  from  $c = \{(e_1, n_1), \dots, (e_k, n_k)\}$  as
$$
\boldsymbol {A} = \underbrace {\left[ \begin{array}{c} e _ {1} , \ldots , e _ {1} \\ n _ {1} \end{array} \right]} _ {n _ {1}}, \underbrace {e _ {2} , \ldots , e _ {2}} _ {n _ {2}}, \ldots , \underbrace {e _ {k} , \ldots , e _ {k}} _ {n _ {k}} ],
$$
and produce  $L$  from  $\ell$ ,  $\phi$  by constructing three lattice vectors with lengths in  $\ell$  and pairwise angles in  $\phi$ . The detailed procedure of producing three lattice vectors is described in Appendix B.
3.4 Coordinate Generation
After the atom type vector  $\mathbf{A}$  and lattice matrix  $\mathbf{L}$  are generated, our method will generate the coordinate matrix  $\mathbf{P}$  conditioned on  $\mathbf{A}$  and  $\mathbf{L}$ . We use a score-based diffusion model for coordinate generation. However, because the original score-based diffusion model proposed by Song and Ermon [44] does not consider any invariance to symmetry transformations in data distributions, we cannot directly apply its theoretic framework to coordinate generation. To overcome this limitation, we propose a novel probabilistic modeling process to calculate the score function of  $\mathbf{P}$ , i.e., the score matrix  $\nabla_{\mathbf{P}} \log p(\mathbf{P}|\mathbf{A},\mathbf{L})$ , and design the architecture and training framework of the score model to
achieve invariance to all symmetry transformations of crystal materials. We will elaborate the details of our proposed coordinate generation framework in the next paragraphs.
We first elaborate the restrictions of the score matrix  $\nabla_P\log p(P|A,L)$  that is supposed to be captured by the score model. By the following Proposition 3.1 (see its proof in Appendix A.1), we can show that if the material distribution  $p(M)$  is invariant to all symmetry transformations described in Section 3.1, the score matrix  $\nabla_P\log p(P|A,L)$  will be invariant to translation and periodic transformations, and equivariant to permutation and rotation transformations.
Proposition 3.1. If the material distribution  $p(M)$  is invariant to permutation, rotation, translation, and periodic transformations described in Section 3.1,  $\nabla_{P} \log p(P|A,L)$  will satisfy the following properties: (1)  $\nabla_{P} \log p(P|A,L)$  is equivariant to permutations on  $A$  and  $P$ ; (2)  $\nabla_{QP + b\mathbf{1}^T} \log p(QP + b\mathbf{1}^T|A,QL) = Q\nabla_{P} \log p(P|A,L)$  holds for any  $Q \in \mathbb{R}^{3\times 3}$ ,  $Q^T Q = I$  and  $b \in \mathbb{R}^3$ ; (3)  $\nabla_{P + LK} \log p(P + LK|A,L) = \nabla_{P} \log p(P|A,L)$  holds for any  $K \in \mathbb{Z}^{3\times n}$ .