Neural network retrained to predict how crystals absorb light, from few examples

The work

Technical · from the record

The authors computed random-phase-approximation (RPA) optical spectra for the materials in an existing independent-particle-approximation (IPA) database with at most 8 atoms per unit cell, about 6000 spectra in total, and used them to test whether a graph attention network (OPTIMATE) trained on cheaper IPA spectra can be retrained to predict RPA spectra. Retraining on 300 RPA spectra gave median similarity-coefficient and mean-squared-error values on the 639-material test set comparable to training a randomly initialised model directly on about 3000 RPA spectra, and all trained models scored better than using the DFT IPA spectrum as a stand-in for the RPA spectrum. Retraining only on materials with up to 4 atoms per primitive unit cell, around 1500 materials, reached median similarity coefficients above 0.85 and median mean squared errors below 0.1, with similar errors across all unit-cell sizes in the test set. A separate network with the same architecture and a scalar output predicted the similarity coefficient between the IPA and RPA spectra with a mean absolute error of 0.026 and a median of 0.018.

How AI was used

Crystal structures taken from the Alexandria database were converted into multigraphs with element identity on nodes and interatomic distances on edges, and fed to the OPTIMATE graph attention network: a per-node multilayer perceptron, three message-passing layers using an improved graph attention operator, softmax vector attention pooling, and a final multilayer perceptron emitting the target optical property as a 2001-element vector sampled in 10 meV steps from 0 to 20 eV. Target spectra, the imaginary part of the trace of the dielectric tensor at 300 meV broadening, were computed ab initio with QUANTUM ESPRESSO for the ground state and YAMBO for the RPA response. Two training strategies were compared on the same reused training, validation and test splits of 4610, 603 and 639 materials: direct learning, where weights were randomly initialised and fitted to RPA spectra, and transfer learning, where all learnable parameters of the published IPA OPTIMATE model were retrained on RPA spectra. Both were run on random training subsets of 100, 300, 1000, 3000 and 4610 materials, and transfer learning was also run on subsets restricted to materials with at most 2, 3, 4, 5 or 6 atoms per primitive unit cell. All models used an L1 loss and the Adam optimizer, with hyperparameters, and for direct learning also the architecture, selected by 5-fold cross-validation on the relevant subset and evaluation on the validation set. A further model with the same architecture but a one-dimensional output was trained to predict the similarity coefficient between the ab initio IPA and RPA spectra, and its latent embeddings after the pooling step were reduced to two dimensions with UMAP using 30 neighbours and a minimum distance of 0.1.