Uncertainty quantification in machine learning and nonlinear least squares regression models
Ni Zhan and John R. Kitchin, “Uncertainty quantification in machine learning and nonlinear least squares regression models”, AIChE Journal 68(6), e17516 (2022; published online November 2021). https://doi.org/10.1002/aic.17516
This summary was written by a language model from the paper and has not been checked line by line. Rely on the published paper for details.
Question
Machine learning models, and neural networks in particular, tend to extrapolate unreliably outside their training data, and in high-dimensional input spaces it is hard to tell when a prediction is an extrapolation. Common remedies (bootstrap, ensembles, Gaussian processes, dropout or Bayesian networks) either require training many models or restrict the model form. The paper asks whether a simple, post hoc uncertainty estimate can be obtained for an already trained model whose parameters were fitted by least squares, and whether that estimate flags extrapolation in a realistic machine-learned interatomic potential (Zhan 2021).
Approach
The authors use the multiparameter delta method. A first-order Taylor expansion of the model prediction in its parameters gives the standard error of a prediction at a point x as the square root of a quadratic form: the gradient of the prediction with respect to the parameters, contracted with the inverse Fisher information matrix, which for a summed squared error loss is the loss Hessian up to a scale factor. Residuals are assumed independent, identically and normally distributed about zero. Both the gradient and the Hessian come from automatic differentiation, and a short autograd listing fitting a quadratic to noisy data shows the whole procedure in a few lines. Because neural network Hessians are often nearly singular, with a few very large eigenvalues, and the optimum may be a saddle point, they propose a recipe: add a small shift to the diagonal (the larger of 1e-5 and 1.05 times the magnitude of the smallest eigenvalue), take a pseudoinverse, scale by a factor alpha chosen so uncertainties match training residuals in magnitude (usually the mean squared error, or nbeta times it), and clip negative eigenvalues to obtain the nearest positive semidefinite matrix. The t-value uses the number of data points rather than an ill-defined effective degrees of freedom (Zhan 2021).
Two examples follow. A [1, 4, 1] network (13 parameters) was fitted to noisy one-dimensional Lennard-Jones data, once with 23 points and once with a gap in the middle. Then a SingleNN neural network potential in PyTorch, with four weighted G2 symmetry functions and a [4, 11, 11, 1] architecture (211 parameters), was trained on energies and forces of AuPd slabs from the Boes and Kitchin 2017 dataset (3907 unique slabs), initially using only 572 configurations at a lattice constant of 3.934 Å split 64/16/20 into train, validation and test sets (Zhan 2021).
Key findings
In the one-dimensional case the confidence band widens beyond the ends of the data and inside the data gap, so the estimate depends on training data coverage in the intended way. For the AuPd potential, energy RMSEs were 0.003, 0.0023 and 0.003 eV/atom for train, validation and test, the uncertainty distributions of the three sets overlapped, and 98% of test energies fell inside the 95% prediction intervals. Applying the same potential to slabs at 4.034 and 4.134 Å lattice constants, whose fingerprints lie partly outside the training range, raised the average standard error of confidence from about 0.002 eV/atom (test) to 0.023 and 0.034 eV/atom, an order of magnitude larger. Coverage of the 95% intervals fell to 75% and 0%, so the intervals become less quantitative under strong extrapolation, but their size still signals that the model should not be trusted. Larger standard errors broadly tracked larger absolute errors. After adding 365 points from each new lattice constant and retraining, the predictions returned to parity, coverage was 98.7% (train) and 98.5% (predict), and average standard errors fell back to about 0.002 (confidence) and 0.003 eV/atom (prediction), illustrating an active learning loop driven by the delta method (Zhan 2021).
Practical notes: computing the inverse Fisher information for an 861-parameter model with 1900 training points took about 5 minutes on a laptop CPU with autograd, it is done once per trained model, and its cost scales quadratically in parameters and linearly in data. Using only the energy loss in the Fisher information worked well even though training used forces, force uncertainties follow by differentiating the force with respect to parameters, and the Fisher information can be approximately updated by adding contributions from new points when retraining changes the parameters little. Very large networks need Hessian approximations or other methods (Zhan 2021).
Significance
The paper offers a cheap alternative to ensembles for any parametric least squares model, from simple nonlinear regression to neural network potentials, that needs no retraining and uses tools already present in machine learning frameworks. Its main use is detecting extrapolation and choosing which new data to compute, for example on the fly during molecular dynamics. It is the basis for the group’s later comparison of model-specific and model-general uncertainty (Zhan 2021).