DIRECT PREDICTION OF TENSORIAL PROPERTIES WITH EQUIVARIANT MESSAGE-PASSING: APPLICATIONS TO NONLINEAR OPTICS

Accurate machine-learned property prediction enables data-driven design and discovery of a wide range of materials. While prediction of scalar quantum mechanical properties like energies have recently reached unprecedented levels of accuracy, prediction of higher-order polar tensors remains relatively difficult and uncommon, despite their ubiquity in fields such as nonlinear optics. The ability to perform accurate ML-based predictions of optical tensors could greatly expedite the discovery of nonlinear optical media. Here, we present a study on the performance of a simple equivariant message-passing neural network for the prediction of molecular hyperpolarizability tensors. Our key findings demonstrate the ability for a modest architecture to perform highly accurate, direct prediction of the full 27-element hyperpolarizability tensor, which we attribute to the network respecting the natural transformation properties of polar tensors, and also the ability of the network to recognize the global symmetries of the input molecules. To provide a mechanistic understanding of these results, we employ dimensionality reduction techniques on the learned equivariant representations to visualize and reason about their latent structure.