Atomic structure generation from reconstructing structural fingerprints

To reconstruct the atomic structure from a given representation, we treat it as a global optimization problem where the Cartesian positions of the atoms are the variables to be optimized and the objective is to minimize the difference or the loss function between the representation of the target structure and the ground truth solution (figure 2). This is often a challenging problem where the global minimum cannot be easily found, especially for cases with many local minima. To significantly reduce the number of evaluations, we accelerate the optimization using gradients obtained from automatic differentiation, where the derivative of the loss function is taken with respect to the atomic positions. This assumes the representation is a function of the positions and is differentiable, which is almost universally true for these classes of structure representations. This yields a structure reconstruction workflow as follows. Given a representation R (to be reconstructed into Cartesian positions), the number of atoms N and the cell lattice parameters C, a separate set of atomic positions x are first randomly initialized with the same N and C. From the initialized positions, its representation R*(x) is computed, and the loss L between R and R* can be obtained. Here, both R*(x) and L(R*, R) are auto-differentiable functions, which allows us to obtain the derivative of L(R*, R) with respect to x. The positions x are then updated as x' with an optimizer using the gradients s_x calculated via auto-differentiation. The procedure is repeated with R*(x) until a convergence criterion is met. The overall process is outlined in algorithm 1 and figure 2.

Zoom In Zoom Out Reset image size

. 

Algorithm 1. Structure reconstruction from its representation via gradient-based optimization.
1:  Input: representation R, the number of atoms N and (optionally) the dimensions of the periodic cell lattice parameters C
2.  while initializations < max_ initializations do
3:   Randomly initialize a set of Cartesian positions $\left\{ {{x_i}} \right\}_{i = 1}^N$ within C
4.   while hops < max_hops do
5:    while iterations < max_iterations do
6:     calculate R*(x)
7:     calculate L (R*, R)
8:     obtain gradients sx from L (R*, R) using automatic differentiation
9:     use sx to update x to x' with optimizer
10.   Add noise to x' to push positions out of local minima
11. Select x' with the lowest L (R*, R)
12: Construct crystal structure with x' and C

For a demonstration, we choose the classic ACSF [36] to compute the representation R. In this work, R is a concatenation of the outputs of the radial ACSF function G [2] and the angular ACSF function G [5]:

$\begin{equation}\begin{array}{*{20}{c}} {G_i^2 = \mathop {{{\mathop \sum \nolimits }}}\limits_j {{\text{e}}^{ - \eta {{\left( {{D_{ij}} - {D_{\text{s}}}} \right)}^2}}} \cdot {f_{\text{c}}}\left( {{D_{ij}}} \right)} \end{array}\end{equation} \tag{ 1 }$

$\begin{equation}\begin{array}{*{20}{c}} {G_i^5 = {2^{1 - \zeta }}\mathop {{{\mathop \sum \nolimits }}}\limits_{j,k \ne i}^{{\text{all}}} {{\left( {1 + \lambda \cos {\theta _{ijk}}} \right)}^\zeta }{{\text{e}}^{ - \eta {{\left( {{D_{ij}} - {D_{\text{s}}}} \right)}^2}}} \cdot {f_{\text{c}}}\left( {{D_{ij}}} \right) \cdot {f_{\text{c}}}\left( {{D_{ik}}} \right)} \end{array}.\end{equation} \tag{ 2 }$

Here, D_ij refer to the distances between atoms i and j, and θ_ijk refer to the angles between atoms i, j, and k. A cutoff function f_c is used where values become zero beyond the cutoff radius. The parameters, η, D_s, $\lambda$, and ζ, for the symmetry functions used in this work are provided in the supplementary information.

The ACSFs are able to provide local representations of each atomic environment, however a permutationally invariant global representation is desired for use in the generative model. To address this limitation, we then aggregate the representations over each atom in the system using max and min pooling. This method was chosen over simple averages to allow for greater variations in the values of R between structures, though other aggregation methods may also be used which are not exhaustively explored in this work. The final expression for R is then:

$\begin{equation}\begin{array}{*{20}{c}} {R = \left[ {{\text{max}}_{i = 1}^NG_i^2,{\text{max}}_{i = 1}^NG_i^5,{\text{min}}_{i = 1}^NG_i^2,{\text{min}}_{i = 1}^NG_i^5} \right]} \end{array}.\end{equation} \tag{ 3 }$

The ACSF functions were written in Pytorch [40] to calculate the gradients via auto-differentiation, which was based on an existing python implementation [41]. For the loss function L(R*, R), we use the mean absolute error (MAE) loss. We then use a gradient descent optimizer, Adam, to update x using s_x. The Adam optimizer [42], which implements adaptive learning rates for each variable being optimized, has been highly effective in training neural network models and was found to be similarly effective for our problem.

It is important to note that with gradient descent alone, it is likely that converged solutions for x may fall into one of many local minima rather than the global optimum solution. Therefore, we incorporate multiple initializations with several 'hops' within each initialization to more effectively sample through the local minima. Finally, the solution with the lowest loss is chosen, and the reconstructed atomic structure is then defined using the optimized x, and the constant C.

For the generative model (figure 3), we use a conditional variational autoencoder (cVAE), which learns a Gaussian distributed latent representation z conditioned on y, following the approach by Sohn et al [43]. A loss function containing a reconstruction loss and a Kullback–Leibler loss is minimized in the training. Note this is not the same reconstruction loss referenced elsewhere in the work, which is associated with mapping the positions back from the representation. An additional parameter, beta, is used to adjust the weight of the two loss values [44], which was set to be 0.75 with respect to the Kullback–Leibler loss. Here we use a cVAE with three hidden layers, a hidden dimension of 128, and a latent dimension of 25. The model was trained with the Adam optimizer with an initial learning rate of 0.001 for 2000 epochs and with a batch size of 64.

Zoom In Zoom Out Reset image size

To quickly evaluate the generated samples for their formation energy, a surrogate model is trained which takes in R as the input, avoiding the need to first perform reconstruction to obtain x. The model is a feed-forward neural network with five layers and a hidden dimension of 256. The model was trained with the Adam optimizer with an initial learning rate of 5 × 10^−4 for 2000 epochs and with a batch size of 256. The training dataset was comprised of 100 000 structures described below, and the model obtains a MAE of 0.0028 eV/atom on the test split. We find this accuracy to be sufficiently high to provide initial estimates of the quality of the generated samples.

The training data for the cVAE consists of a set of Pt₁₀ nanoclusters with various structural configurations. The structures are obtained from density functional theory (DFT) optimizations of randomly initialized geometries; both the local minima as well as intermediate structures in the optimization are used in the dataset. The target property for this dataset is formation energy in units of eV/atom, which are also obtained from the DFT calculations. Structures with a formation energy from 0 to 3.7 eV/atom were chosen. The formation energy is defined as

$\begin{equation}\begin{array}{*{20}{c}} {{E_{{\text{form}}}} = { }\left( {{E_{{\text{total}}}} - N \cdot {E_{{\text{bulk}}}}} \right)/N} \end{array}\end{equation} \tag{ 4 }$

where E_bulk is the energy per atom of Pt in the bulk state and E_total is the total energy of the Pt₁₀ system. From this overall dataset of approximately 100 000 structures, three sub-datasets are created by sampling from the total dataset using furthest point sampling based on their computed R to yield datasets containing 30 000, 3000, and 300 structures (see figures S1–S3 in supplementary information). An 80:20 training to validation ratio is used.

The DFT calculations were performed using the Vienna Ab Initio Simulation Package [45, 46]. The Perdew–Burke–Ernzerhof [47] functional was used for electron exchange and correlation energies. The projector-augmented wave method was used to describe the electron-core interaction [45, 48]. A kinetic energy cutoff of 450 eV was used. The electronic convergence criterion was set to 1 × 10^−6 eV. All calculations were performed without spin polarization. The Brillouin zone was sampled at the Gamma point only. A cubic cell of 17 Å was used to prevent interactions between periodic images.

We first validate the effectiveness of the proposed reconstruction method outlined in algorithm 1 for the ACSF representation on a wide range of artificial and realistic atomic structures. Here, the aim is to obtain atomic positions which are as close to the ground truth material using only their ACSF representation. Starting with examples of reconstructing non-periodic geometric structures (figure 4), we find the reconstruction algorithm is generally able to find solutions which are nearly identical to the ground truth. In certain cases, the reconstruction loss can be reduced to within numerical error, while in others, minor distortions can still be perceptible (structure (D) in figure 4). Next, we move to reconstruction for periodic and multi-elemental structures. Periodicity is incorporated in the ACSF representation by following the minimum image convention when calculating distances and angles. For systems with more than one element, the canonical ACSF treats each atomic pair and triplet separately, and the system is then initialized with the same unit cell chemical formula as the ground truth structure. Here, we find the algorithm can obtain a nearly identical crystal structure (albeit translated in space) for the six examples in figure 5, ranging from simple binary and ternary compounds to more complex cases containing molecular species (structure (E)) and layered materials (structure (F)). As the algorithm cannot guarantee finding the global minimum solution, ten trials with five hops each were performed for each structure, and the solution with the lowest L(R*, R) was selected. We note that while for more complex structures containing large numbers of atoms, more trials may be necessary, in practice these settings were sufficient for all the examples shown in figures 4 and 5.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

To demonstrate a fully end-to-end generation process, we formulate a task of generating structures of Pt₁₀ nanoparticles based on a property metric of formation energy, a measure of their thermodynamic stability. Sub-nanometer Pt nanoparticles have been studied extensively and are known to exhibit a high degree of structural fluxionality with important implications for catalysis. The formation energy of these systems is highly sensitive to even minor changes in structure and is quantum mechanical in nature, which cannot be captured by simple pair-wise interaction models resulting in complex potential energy surfaces. This therefore serves as a challenging benchmark task for structure generation due to their complex structure-property relationship and high sensitivity to errors in the generation or reconstruction process. Three training datasets of 30 000, 3000 and 300 data points are compiled from a larger dataset containing various Pt configurations and their DFT-computed formation energies.

A cVAE is then trained to serve as a generative model for new samples within the representation space R, conditional on y, the formation energy. The generated samples are then evaluated with a surrogate model trained directly on R to predict y. Table 1 shows the MAE and percent within range performance of the cVAE over 10 000 generated samples for the three training dataset sizes. Starting with the largest 30 000 dataset, the cVAE does particularly well for generation, with nearly 100% of the samples being within ±0.1 eV/atom of the targeted formation energies ranging from 1.9 to 3.7 eV/atom. To investigate the data efficiency of this approach, smaller datasets of 3000 and 300 data points are also used to train the model and compared. Remarkably, the generation performance remains strong even using only 300 training samples, and on average 80% of the generated samples falls within ±0.1 eV/atom of the target value with respect to the surrogate model. This performance suggests ACSFs to be an effective and efficient representation for generative models, which can likely be attributed to the inclusion of the translational, rotational and permutational invariances within the representation.

In addition to being able to generate samples which are accurate with respect to the target property, the samples must also be sufficiently unique and diverse. In our present work for Pt nanoparticles, as they all share the same composition and number of atoms, the diversity should be evaluated solely from the variations in atomic structures. Since it is difficult to define a specific threshold for which two structures are considered unique from one another, we instead adopt a visual approach for qualitatively determining diversity using dimensionality reduction. Samples which are less diverse would appear to be localized in the same region of the reduced dimensional space, but be spread out for more diverse samples. We plot the generated R at several selected formation energy values, taking the first two principal components from a principal component analysis (PCA). We perform this analysis for the smallest dataset with 300 training data. As shown in figure 6, there is a large spread in the principal components of the generated samples (red squares) for y = 1.9, 2.8 and 3.7 eV/atom targets, demonstrating a high degree of diversity. These samples are also unique with respect to the training data (green circles), which only sparsely covers the space of the full 100 000 dataset (black circles). This suggests the cVAE model, having learned a continuous, low-dimensional latent representation of the chemical space, can now successfully populate that space with novel samples. We note for y = 1.9 many generated samples fall outside of the distribution of both the training data and the full dataset within the range of 1.8 ⩽ y ⩽ 2.0 (blue circles). These we hypothesize will result in 'invalid' representations and are likely due to the low number and proportion of training samples available in that range of y, which will be discussed further. Here, we define 'invalid' representations as those which cannot encode back into a structure in the Cartesian space; in other words no solution for the atomic positions exists which results in a zero, or sufficiently low, reconstruction loss.

Zoom In Zoom Out Reset image size

The generated representations are then reconstructed back to Cartesian positions using algorithm 1. The reconstructed atomic structures are subsequently validated with DFT to obtain their true formation energy (y) values. No further optimization, screening or down-sampling is applied here to provide the raw performance of this generation and reconstruction process. We perform this process on 100 generated samples for y = 1.9, 2.8 and 3.7 eV/atom for each of the three models, with the results shown in table 2. We find mean absolute errors for the structures to be consistently higher once reconstructed back to their positions, as compared to table 1. This discrepancy is likely a combination of several key factors. First, there is an inherent error or reconstruction loss when going from the representation to the structure due to finding the exact global minimum solution. Second, the MAEs obtained in table 1 are calculated from a surrogate model trained on R, which has its own inaccuracies. Third, it is possible to generate 'invalid' representations which do not encode into a valid structure. Of these factors, we hypothesize the third one to have the strongest impact, given the strong reconstruction performance (figures 4 and 5) and the reasonable accuracy of the surrogate model (figure S8). This was observed earlier in figure 6 where some of the generated samples fall outside of the range of principal component values of the full dataset of valid representations (black circles). A much higher proportion of generated samples fall outside this range for y = 1.9, and unsurprisingly the post-reconstruction performance is also much lower for y = 1.9 than the other two target values. We also find a convergence in performance from 3000 to 30 000 training data, where the MAE and percentage within range are roughly similar, or even slightly worse for the larger training data case. As the model, ACSF, and training hyperparameters are fixed for all three dataset sizes with an emphasis on the n = 300 case, it is likely that a better performance could be obtained. However, we leave this to a future study given the current focus on data efficiency in this work.

A random selection of some real and generated structures from the studies in table 2 are visualized in figures 7–9. For y = 1.9 eV/atom, the real samples show distorted globular structures which are largely captured in the generated structures but are also flatter with more planar motifs than the real structures. For y = 2.8 eV/atom, the real samples are more diffuse and line-like (being less stable and thus having higher formation energies) with partially dissociated atoms which is also found in the generated samples. For y = 3.7, the real samples are almost completely dissociated with many isolated atoms, and the generated samples follows a similar pattern. For all cases, the high diversity of the generated structures can be visually observed, with little discernible overlap in structure between the randomly chosen samples, and consistent with the PCA plots in figure 6. A limitation in the conditional generation performance can be readily seen for y = 1.9, where the generated structures lack the same degree of globularity and deviate noticeably in their computed energies from the target.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size