Neural Structure Fields with Application to Crystal Structure Autoencoders
Abstract
Representing crystal structures of materials to facilitate determining them via neural networks is crucial for enabling machine-learning applications involving crystal structure estimation.
Here we propose neural structure ﬁelds (NeSF) as an accurate and practical approach for representing crystal structures using neural networks.
Inspired by the concepts of vector ﬁelds in physics and implicit neural representations in computer vision, the proposed NeSF considers a crystal structure as a continuous ﬁeld rather than a discrete set of atoms.
Unlike existing grid-based discretized spatial representations, the NeSF overcomes the tradeoff between spatial resolution and computational complexity and can represent any crystal structure.
To evaluate the NeSF, we propose an autoencoder of crystal structures. Quantitative results demonstrate the superior performance of the NeSF compared with the existing grid-based approach.
Introduction
Representing crystal structures in a neural-network-familiar expression has promising utility in machine-learning (ML)-assisted material development and materials informatics (MI).
For example, graph-based representations of crystal structures have allowed the encoding of crystal structures via graph neural networks [1, 2, 3, 4] for high-throughput material property predictions.
When this type of task is viewed as the discovery of structure-to-property relationships among materials, there is yet another important type of task, namely, the discovery of property-to-structure relationships , which constitutes an inverse design problem [5, 6, 7, 8, 9]. Despite the great potential utility of this inverse approach in developing materials, few studies have addressed this [5, 6, 7, 8, 9] or its underlying problem [10, 11, 12], that is, the estimation of crystal structures under given conditions. Regarding MI and ML, whether to input or output crystal structures ( i.e. , encode or decode crystal structures in MI and ML terms) induces a crucial difference.
Although encoding crystal structures is suitably established using graph neural networks [1, 2, 3, 4], a technical bottleneck persists for decoding crystal structures. We addressed the bottleneck in this study.
The crystal structure of an inorganic material is a regular and periodic arrangement of atoms in a three-dimensional (3D) space.
This arrangement is usually described by the 3D positions and species of atoms in a unit cell and the lattice constants deﬁning the translations of the unit cell in 3D space.
The atoms in a unit cell have no explicit order, and their quantity varies from one to hundreds in number.
Because ML models, including neural networks, generally accept ﬁxed-dimensional
Position field z x y v nbr = f P ( x,y,z ) = (0.41, 0.12, 0.36)
Species = f S ( x,y,z ) =
(0,..., 0.9, 0.1,...) =
Fe H
Fe Co f P f S Species field v nbr query
[Figure]
and consistently ordered tensors for processing, treating crystal structures with ML models is not straightforward [1], and determining crystal structures via the models is even more difﬁcult. In this paper, we propose a continuous vector-ﬁeld-based crystal structure representation named neural structure ﬁelds (NeSF). The proposed NeSF is inspired by the concepts of vector ﬁelds in classical physics and implicit neural representations [13, 14, 15, 16, 17] in computer vision. Implicit neural representations have recently been proposed to handle some representation issues in 3D computer vision applications, such as 3D shape estimation [14, 15, 16] and free-viewpoint image synthesis [13, 17].
The NeSF follows the basic idea of implicit neural representations and further extends it to the estimation of crystal structures described by atomic positions and species.
Our idea of representing crystal structures as continuous vector ﬁelds has been partially and implicitly explored by grid-based discretization ( i.e. , voxelization) in recent MI studies [5, 6, 12, 8, 10].
In those studies, the 3D space within the unit cell is discretized into voxels, and each voxel is then assigned an electron density, which essentially represents the presence or absence of an atom around the voxel.
However, the discretization of 3D data considerably suffers from the tradeoff between spatial resolution and computational complexity in terms of both computational time and memory space.
For instance, the ICSG3D method [10] based on 3D convolutional neural networks (CNNs) uses a limited resolution of 32 ×
32 × 32 voxels.
The NeSF overcomes the limitation of voxelization.
Theoretically, NeSF can achieve inﬁnitely high spatial resolution with compact (memory- and parameter-efﬁcient) neural networks in place of costly 3D CNNs.
Method
In this section, we ﬁrst describe the key concept of the NeSF and the procedure for estimating crystal structures with the NeSF. We then propose an autoencoder of crystal structures using NeSF.
2.1 Crystal structure representation via NeSF
We propose the NeSF as a general representation of crystal structures that enables neural networks to decode or determine such structures.
The key concept underlying our approach is illustrated in Fig. 1a, where a crystal structure is represented as a continuous vector ﬁeld tied to 3D space rather than as a discrete set of atoms.
The NeSF uses two types of vector ﬁelds, namely, the position and species ﬁelds, to implicitly represent the positions and species of the atoms in the unit cell of the crystal structure, respectively.
To illustrate the concept of the NeSF, assume that we are given target material information as a ﬁxed-dimensional vector, z , and consider the problem of recovering a crystal structure from z .
In the NeSF, we use neural network f as an implicit function to indirectly represent the crystal structure embedded in z instead of letting neural network f directly output the crystal structure as f ( z ) .
Speciﬁcally, we treat f as a vector ﬁeld on the Cartesian 3D coordinates, p , conditioned on the target material information, z :
[Equation]
In the position ﬁeld, network f is trained to output a 3D vector pointing from query point p to its nearest atom position, a , in the crystal structure of interest.
Thus, we expect output s to be a − p .
If the position ﬁeld is ideally trained, we can retrieve position a of the nearest atom at any query point p as p + f ( p , z ) .
The precise description of atomic positions in crystal structures is of crucial interest in materials science.
Thus, the NeSF outputs vectors pointing to the nearest atoms to represent atomic positions more directly than existing implicit neural representations of 3D geometries [14, 15, 16].
Analogously, the species ﬁeld is trained to output a categorical probability distribution that indicates the species of the nearest atom.
Thus, the output dimension of the species ﬁeld is the number of candidate atomic species.
The training procedure of NeSF is simple and similar to those for existing implicit neural representations [14, 18].
We randomly sample 3D query points { p si } in the unit cell and compute loss values for the ﬁeld outputs at these points, thus supervising f p ( p si , z ) and f s ( p si , z ) to indicate the position and species of the nearest atom, respectively.
To efﬁciently train the NeSF, we introduce two sampling methods.
Global grid sampling considers 3D grid points that uniformly cover the entire unit cell and samples the points with perturbations that follow a Gaussian distribution.
Local grid sampling considers local 3D grid points centered at each atomic position and samples the points with perturbations that follow a Gaussian distribution.
To train the position ﬁeld, we combine both sampling methods.
Hence, query points are sampled uniformly over the entire unit cell and densely around the atoms.
To train the species ﬁeld, we use local grid sampling to concentrate training query points in the neighborhood of atoms.
See Appendix A for more details.
2.2 Crystal structure estimation using NeSF
Given the target material information as vector z , the estimation of the crystal structure from z amounts to estimating the positions and species of the atoms in the unit cell along with the lattice constants.
The lattice constants are modeled as lengths a , b , and c and angles α , β , and γ , and are estimated via simple multilayer perceptrons (MLPs) with input z .
On the other hand, the atomic positions and species are estimated by position ﬁeld f p and species ﬁeld f s of the NeSF, respectively.
These ﬁelds are also implemented as simple MLPs, each taking query position p and vector z as input and predicting a ﬁeld value ( i.e. , 3D pointing vector or categorical probability distribution).
An overview of the NeSF network architecture is illustrated in the right part of Fig.
1b.
Given vector z from the encoder, the estimation of the atomic positions and species using the NeSF is summarized in the ﬁve steps:
Initialize particles.
We ﬁrst estimate the lattice constants via MLPs.
Then we regularly spread initial query points { p 0 i } , which we call particles , at 3D grid points within a bounding box.
The bounding box is common to each dataset and is given to loosely encompass the atoms of all training samples.
Move particles.
We update the position of each particle p ti using the position ﬁeld according to p t +1 i
= p ti + f p
( p ti , z ) .
We iterate this process for all particles to obtain { p i } as candidate atomic positions.
Through this process, the particles travel toward their nearest atoms.
Score particles.
We score each particle p i and ﬁlter outliers.
Since the norm of the output of the position ﬁeld, ∥ f p ( p i , z ) ∥ , indicates an estimated distance from p i to its nearest atom, we score each particle p i by ∥ f p ( p i , z ) ∥ and discard the particle if the score is above a certain threshold (set to 0.9 Å in this study).
Detect atoms.
Until this point, the particles are expected to form clusters around atoms.
Thus, we apply a simple
clustering algorithm to detect each cluster position as an atomic position, determining the number of atoms in the crystal structure.
Here we employ a well-known clustering algorithm in object detection, non-max suppression, and set its suppressing radius to 0.5 Å.
Estimate species.
Finally, we use species ﬁeld f s ( p , z ) to estimate the atomic species at each atomic position a i .
For robust estimation against errors in a i , we spread new particles intensively around each a i as queries to the species ﬁeld.
We select the most frequent atomic species among them as the ﬁnal estimate.
2.3 Crystal structure autoencoder
To demonstrate and evaluate the expressive power of the NeSF, we propose an autoencoder of crystal structures.
Similar to other common autoencoders, the proposed NeSF-based autoencoder consists of an encoder and decoder.
The encoder is a neural network that transforms an input crystal structure ( i.e. , positions and species of the atoms in the unit cell and lattice constants) into abstract latent vector z .
The decoder, for which we use the NeSF, reconstructs the input crystal structure from latent vector z .
Autoencoders are typically used to learn latent vector representations of data via self-supervised learning, in which the input data can supervise the learning via a reconstruction loss.
While we focus on decoding crystal structures, their encoding requires a speciﬁc type of neural network called set encoders .
Among them, the family of graph neural networks [1, 2, 3, 4] serves as popular crystal-structure encoders.
However, these networks implicitly represent atomic positions as edges, encoding distances between atoms while discarding the exact coordinates.
Although this distance-based graph representation is key to ensure the invariance to coordinate systems, its information loss in input may unintentionally hinder the reconstruction performance. Thus, we adopt the basic encoder architecture from PointNet [19] and DeepSets [20]. This architecture not only represents the simplest type of set encoders but can also preserve the information of input crystal structures, thus being appropriate for the performance evaluation of the NeSF decoder. The proposed autoencoder architecture is detailed in Appendix B.
During training, the four types of outputs from the decoder ( i.e. , atomic positions, species, lattice lengths, and angles) were evaluated using the following loss function:
[Equation]
Here, L pos is the mean squared error (MSE) between estimated and true atomic positions, L spe is the cross-entropy loss function for the atomic species distributions, and L len and L ang are the mean square error for the lattice lengths and angles, respectively.
The total loss function, L , is given by the weighted sum of these loss functions with weights ( λ pos , λ spe , λ len , λ ang ) =
We optimized the loss function using stochastic gradient descent with a batch size of 128 .
We used Adam [21] as the optimizer with an initial learning rate of 10 − 3 decaying every 640 epochs by a factor of 0 . 5 .
We conducted iterative training for 3200 epochs for all training samples in the dataset with early stopping.
These hyperparameters were chosen through the procedure described in Appendix C. To reduce the performance variation owing to randomness, we repeated training and evaluation 10 times with different random seeds and evaluated the average performance.
Results and Discussions
We compared the performance of the proposed NeSF-based autoencoder with the grid-based ICSG3D baseline [10].
We reproduced the ICSG3D dataset [10] for training and evaluation by curating 7897 material samples from the Materials Project [22] (see Appendix D for details).
The reconstruction performance was measured in terms of errors in the number of atoms, position, and species.
The error in the number of atoms is the rate of materials for which the number of atoms in the unit cell is incorrectly estimated.
The position error is the average error of the reconstructed atomic positions.
Depending on the denominator of the metric, we evaluated the position error in two ways. An actual metric was used to evaluate the mean position errors at the actual atomic sites of the crystal structure by computing their shortest distances to estimated atomic sites. By contrast, a detected metric was used to evaluate the errors at the estimated sites by computing their shortest distances to the actual atomic sites. The actual metric is more sensitive to errors related to underestimation of the number of atoms, whereas the detected metric is more sensitive to errors related to overestimation.
The species error is the average rate of atoms with incorrectly estimated species.
Analogous to the position error, the species error was evaluated using actual and detected metrics. Lower values of these metrics indicate better performance.
[Table]
Table 1 lists the reconstruction errors of the proposed NeSF-based autoencoder and ICSG3D baseline on the test set.
Overall, the proposed method consistently outperforms ICSG3D in all the evaluation metrics, with substantial performance improvements for the species error.
Meanwhile, ICSG3D achieves good performance for the error in number of atoms and the position.
The evaluated dataset [10] consists of materials with limited crystal systems ( i.e. , cubic) and prototypes ( i.e. , AB, ABX2, and ABX3) and may thus offer relatively simple problem instances.
Evaluating the reconstruction performance on more diverse crystal structures is future work.
We further discuss potential negative societal impacts of this work in Appendix E.
We believe that the notable high performance of the proposed method is attributable to two reasons.
First, our method does not use discretization, being advantageous over the grid-based ICSG3D for estimating crystal structures. In a grid-based method, the spatial resolution is limited by the cubically increasing computations and memory usage. The proposed NeSF is free from such a tradeoff between resolution and computational complexity, and it can thus effectively represent complex structures. Second, the model size of the proposed NeSF using MLPs is much smaller than that of the 3D CNN architecture of ICSG3D. In general, the number of training samples required for an ML model is correlated with the number of trainable parameters.
Grid-based methods use layers of 3D convolution ﬁlters that involve many trainable parameters. On the other hand, NeSF employs implicit neural representations to indirectly describe the 3D space as a ﬁeld instead of voxels. Thus, it is efﬁciently implemented by MLPs with fewer parameters than a 3D CNN.
Speciﬁcally, the NeSFbased autoencoder has 0 . 76 million parameters, which is only 2 . 24 % of the number of parameters in the 3D CNN-based ICSG3D ( 34 million parameters).
Conclusion
We propose the NeSF to estimate crystal structures using neural networks.
Determining crystal structures directly using neural networks is challenging because these structures are essentially represented as an unordered set with varying numbers of atoms.
The NeSF overcomes this difﬁculty by treating the crystal structure as a continuous vector ﬁeld rather than as a discrete set of atoms.
Unlike existing grid-based approaches for representing crystal structures [10], the NeSF is free from the tradeoff between spatial resolution and computational complexity and can represent any crystal structure.
The NeSF was applied as an autoencoder of crystal structures and demonstrated its expressive power.
A quantitative performance analysis showed a clear advantage of the NeSF-based autoencoder over an existing grid-based method.
Meanwhile, the present study is subject to the following limitations.
First, the NeSF does not explicitly consider space-group symmetry.
Symmetry is an important concept in crystallography and incorporating the constraint of space-group symmetry into NeSF is thus an important direction of future work.
Second, we adopted the autoencoder architecture rather than generative models, such as variational autoencoders [11, 23, 1] and generative adversarial networks [24, 23, 25, 12], for an evaluation purpose.
While these generative models are more suitable for novel structure discovery, the lack of ground-truth structures prevents quantitative and reliable performance analysis.
The NeSF should be applied to generative models in future work with appropriate performance analysis.
Such generative models for crystal structures will be important for inverse design of materials, which is a major challenge in MI.
Acknowledgments
This work is partly supported by JST-Mirai Program, Grant Number JPMJMI19G1 and JST ACT-I grant number JPMJPR18UE
Appendix A Sampling parameters
For the training of the NeSF, we use the global grid sampling and the local grid sampling.
Sampling ranges and parameters are shown in followings.
• For position ﬁeld training: – Global grid sampling: * Interval: 1 Å * σ : 0 .
5 Å – Local grid sampling: * Interval: 0 .
3 Å * Range for each axis: 1 . 2 Å
* σ : 0 .
5 Å • For species ﬁeld training: – Local grid sampling: * Interval: 0 .
1 Å * Range for each axis:
5 Å * σ : 0 .
5 Å
B Network architecture
The proposed network architecture is speciﬁed as follows, where N is the number of atoms in the input crystal structure and A is the number of atomic species in the dataset.
• Encoder – Atom position encoder *
Input: Atom positions: [Nx3] * Shared MLP: [Nx3] -> [Nx128] · Layers: 3, 64, 128,· Activation function: ReLU · Using batch normalization · No dropout * Output: Atom position features: [Nx128] – Atom species encoder *
Input: Atom species with one-hot encodings: [NxA] * Shared MLP: [NxA] -> [Nx128] · Layers: A, 128,· Activation function: ReLU · Using batch normalization · No dropout * Output: Atom species features: [Nx128] – Crystal structure encoder *
Input: Atom position features and atom species features: [Nx128], [Nx128] * Concatenation of atom position features and atom species features: [Nx128], [Nx128] -> [Nx256] * Shared MLP: [Nx256] -> [Nx512] · Layers: 256,· Activation function: ReLU · Using dropout with 0.3 dropout probability * Max-Pooling (along the atoms): [Nx512] -> [512] * Concatenation with lattice parameters (length and angles): [512], [
] -> [518]
* MLP: [518] -> [192] ·
Layers
: 518, 256, 256,· Activation function: ReLU · Using batch normalization · Using dropout with 0.3 dropout probability *
Output: Latent vector: [192] • Decoder – Lattice lenght parameter decoder
Input: Latent vector: [192] * MLP: [192] -> [3] · Layers: 192, 128,· Activation function: ReLU · No batch normalization · No dropout * Output: Lattice length parameter: [3] – Lattice angle parameter decoder *
Input: Latent vector: [192] * MLP: [192] -> [3] · Layers: 192, 128,· Activation function: ReLU · No batch normalization · No dropout * Output: Lattice angle parameter: [3] – Position ﬁeld decoder: * Input: Latent vector and query position: [192], [3] * Concatenation of Latent vector and query position: [192], [3] -> [195] * FC: [195] -> [256] · Activation function: ReLU · Using batch normalization · No dropout * Concatenation of above output and query position: [256], [3] -> [259] * FC: [259] -> [128] · Activation function: ReLU · Using batch normalization · No dropout * Concatenation of above output and query position: [128], [3] -> [131] * FC:
[131] -> [128] · Activation function: ReLU · Using batch normalization · No dropout * Concatenation of above output and query position: [128], [3] -> [131] * FC:
[131] -> [128] · Activation function: ReLU · Using batch normalization · No dropout * Concatenation of above output and query position: [128], [3] -> [131] * FC:
[131] -> [128] · Activation function: ReLU · Using batch normalization · No dropout * Concatenation of above output and
[Table]
· No dropout * Concatenation of above output and query position: [128], [3] -> [131] * FC:
[131] -> [128] · Activation function: ReLU · Using batch normalization · No dropout * Concatenation of above output and query position: [128], [3] -> [131] * FC:
[131] -> [128] · Activation function: ReLU · Using batch normalization · No dropout * Concatenation of above output and query position: [128], [3] -> [131] * FC:
[131] -> [3] · Activation function:
None · No batch normalization · No dropout * Output: Position ﬁeld value of the query position: [3] – Species ﬁeld decoder: * Input: Latent vector and query position: [192], [3] * Concatenation of Latent vector and query position: [192], [3] -> [195] * FC: [195] -> [256] · Activation function: ReLU · Using batch normalization · No dropout * Concatenation of above output and query position: [256], [3] -> [259] * FC: [259] -> [128] · Activation function: ReLU · Using batch normalization · No dropout * Concatenation of above output and query position: [128], [3] -> [131] * FC:
[131] -> [A] · Activation function:
None · No batch normalization · No dropout * Output: Species ﬁeld value (class likelihood) of the query position: [A]
C Hyperparameter search
We searched hyperparameters by the procedure summarized as the eight steps below.
In each step, we evaluated multiple conﬁgurations of certain hyperparameters by training and validating models on a dataset and chose the best conﬁguration among them for the next step.
We evaluated the Number of Atom Accuracy as the primal performance indicator.
When the performance was the same by this metric, the Position Error (Actual) was used as the secondary performance indicator.
Since the ICSG3D dataset [10] is restricted to certain material classes and prototypes, conducting the hyperparameter search on it may result in a model with limited generalizability despite the huge computational effort.
We thus prepared a more general dataset, named Limited Cell Size (6 Å) (LCS6Å), for the hyperparameter search.
The LCS6Å dataset is a collection of the material entries in the Materials Project whose unit cell sizes are 6 Å or less along the x , y , and z axes.
The dataset amounts to 6005 material entries ( 4 . 31 % of the Materials Project database) and was split into the training (5418 entries), validation (286 entries), and test (301 entries) sets.
The training of a NeSF model on the LCS6Å dataset took approximately 9 hours using a single
NVIDIA RTX
GPU (48GB VRAM).
The hyperparameter search was carried out on a single
NVIDIA DGX-1 system (NVIDIA V100 GPU 32GB VRAM x8), taking approximately 10 days in total.
[Table]
3. Learning rate
[Table]
[Table]
[Table]
[Table]
For ICSG3D
[Table]
we also tuned the default hyperparameters provided by the authors for the LCS6Å dataset
[Table]
by searching for the best conﬁguration from the list below.
• batch_size=10, latent_size=
[Table]
• batch_size=10, latent_size=
[Table]
• batch_size=10, latent_size=
[Table]
• batch_size=10, latent_size=
[Table]
• batch_size=10, latent_size=
[Table]
• batch_size=
[Table]
• batch_size=20, latent_size= • batch_size=20, latent_size= • batch_size=20, latent_size= • batch_size=20, latent_size= • batch_size=20, latent_size= • batch_size=40, latent_size= • batch_size=40, latent_size= • batch_size=40, latent_size= • batch_size=40, latent_size= • batch_size=40, latent_size= • batch_size=5, latent_size= • batch_size=5, latent_size= • batch_size=5, latent_size= • batch_size=5, latent_size= • batch_size=5, latent_
[Table]
=
[Table]
D ICSG3D Dataset Reproduction
We reproduced the dataset used for ICSG3D [10] by curating crystal structure data from the Materials Project [22], a public online database distributed under Creative Commons Attribution [-4.
-]4.
License.
Unfortunately, the authors of ICSG3D do not provide speciﬁc material IDs used for their experiments.
Therefore, by following the procedures in the paper, we crawled the three material classes having cubic structures, i.e. , binary alloys (AB), ternary perovskites (ABX 3 ), and Heusler compounds (ABX 2 ), from the Materials Project.
We consequently obtained 7896 material entries, which were split into the training (7194 entries), validation (342 entries), and test (360 entries) sets.
E Potential Negative Societal Impacts
We are aware of some potential negative societal impacts of our work.
Since our work proposes a fundamental idea for the decoding of 3D structures consisting of atoms, it may be broadly related to ML-based approaches for the discovery of novel materials and chemical substances.
Therefore, misuses of our idea may lead to the designing of substances that are harmful to humans.