Soft Checksums to Flag Untrustworthy Machine Learning Surrogate Predictions and Application to Atomic Physics Simulations

TL;DR

This paper introduces soft checksums, a novel method for neural networks to identify when their predictions are trustworthy or untrustworthy, especially in complex atomic physics simulations, by learning a checksum function that signals prediction errors.

Contribution

The paper proposes a new soft checksum technique that differentiates between in-distribution and out-of-distribution predictions with minimal computational overhead.

Findings

Soft checksums effectively identify high-error OOD predictions.

Incorporating checksum loss improves ID and OOD separation.

Method applied successfully to complex atomic physics dataset.

Abstract

Trained neural networks (NN) are attractive as surrogate models to replace costly calculations in physical simulations, but are often unknowingly applied to states not adequately represented in the training dataset. We present the novel technique of soft checksums for scientific machine learning, a general-purpose method to differentiate between trustworthy predictions with small errors on in-distribution (ID) data points, and untrustworthy predictions with large errors on out-of-distribution (OOD) data points. By adding a check node to the existing output layer, we train the model to learn the chosen checksum function encoded within the NN predictions and show that violations of this function correlate with high prediction errors. As the checksum function depends only on the NN predictions, we can calculate the checksum error for any prediction with a single forward pass, incurring negligible time and memory costs. Additionally, we find that incorporating the checksum function into the loss function and exposing the NN to OOD data points during the training process improves separation between ID and OOD predictions. By applying soft checksums to a physically complex and high-dimensional non-local thermodynamic equilibrium atomic physics dataset, we show that a well-chosen threshold checksum error can effectively separate ID and OOD predictions.