A prediction rigidity formalism for low-cost uncertainties in trained neural networks

Regression methods are fundamental for scientific and technological applications. However, fitted models can be highly unreliable outside of their training domain, and hence the quantification of their uncertainty is crucial in many of their applications. Based on the solution of a constrained optimization problem, we propose "prediction rigidities" as a method to obtain uncertainties of arbitrary pre-trained regressors. We establish a strong connection between our framework and Bayesian inference, and we develop a last-layer approximation that allows the new method to be applied to neural networks. This extension affords cheap uncertainties without any modification to the neural network itself or its training procedure. We show the effectiveness of our method on a wide range of regression tasks, ranging from simple toy models to applications in chemistry and meteorology.

Wigner kernels: body-ordered equivariant machine learning without a basis

Machine-learning models based on a point-cloud representation of a physical object are ubiquitous in scientific applications and particularly well-suited to the atomic-scale description of molecules and materials. Among the many different approaches that have been pursued, the description of local atomic environments in terms of their neighbor densities has been used widely and very succesfully. We propose a novel density-based method which involves computing ``Wigner kernels''. These are fully equivariant and body-ordered kernels that can be computed iteratively with a cost that is independent of the radial-chemical basis and grows only linearly with the maximum body-order considered. This is in marked contrast to feature-space models, which comprise an exponentially-growing number of terms with increasing order of correlations. We present several examples of the accuracy of models based on Wigner kernels in chemical applications, for both scalar and tensorial targets, reaching state-of-the-art accuracy on the popular QM9 benchmark dataset, and we discuss the broader relevance of these ideas to equivariant geometric machine-learning.

Fast evaluation of real spherical harmonics and their derivatives in Cartesian coordinates

Spherical harmonics provide a smooth, orthogonal, and symmetry-adapted basis to expand functions on a sphere, and they are used routinely in computer graphics, signal processing and different fields of science, from geology to quantum chemistry. More recently, spherical harmonics have become a key component of rotationally equivariant models for geometric deep learning, where they are used in combination with distance-dependent functions to describe the distribution of neighbors within local spherical environments within a point cloud. We present a fast and elegant algorithm for the evaluation of the real-valued spherical harmonics. Our construction integrates many of the desirable features of existing schemes and allows to compute Cartesian derivatives in a numerically stable and computationally efficient manner. We provide an efficient C implementation of the proposed algorithm, along with easy-to-use Python bindings.

A smooth basis for atomistic machine learning

Machine learning frameworks based on correlations of interatomic positions begin with a discretized description of the density of other atoms in the neighbourhood of each atom in the system. Symmetry considerations support the use of spherical harmonics to expand the angular dependence of this density, but there is as yet no clear rationale to choose one radial basis over another. Here we investigate the basis that results from the solution of the Laplacian eigenvalue problem within a sphere around the atom of interest. We show that this generates the smoothest possible basis of a given size within the sphere, and that a tensor product of Laplacian eigenstates also provides the smoothest possible basis for expanding any higher-order correlation of the atomic density within the appropriate hypersphere. We consider several unsupervised metrics of the quality of a basis for a given dataset, and show that the Laplacian eigenstate basis has a performance that is much better than some widely used basis sets and is competitive with data-driven bases that numerically optimize each metric. In supervised machine learning tests, we find that the optimal function smoothness of the Laplacian eigenstates leads to comparable or better performance than can be obtained from a data-driven basis of a similar size that has been optimized to describe the atom-density correlation for the specific dataset. We conclude that the smoothness of the basis functions is a key and hitherto largely overlooked aspect of successful atomic density representations.