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

TL;DR

This work tackles the problem of trusting ML surrogates for physics simulations by distinguishing in-distribution (ID) from out-of-distribution (OOD) predictions. It introduces soft checksums, a mechanism that augments neural networks with a checksum output and a user-specified checksum function $\mathbb{C}(\bm{y})$, enabling a checksum error computed in a single forward pass to signal prediction reliability. A composite loss combining $\mathcal{L}_{\text{prediction}}$, $\mathcal{L}_{\text{checksum}}$, $\mathcal{L}_{\text{ID}}$, and $\mathcal{L}_{\text{OOD}}$ is proposed to shape the checksum response, including exposing the model to OOD examples outside the training hypercube. Empirical results on a NLTE atomic physics surrogate (88-dimensional input/output) show that appropriate checksum thresholds (e.g., achieving a 99% true-negative rate on validation data) and OOD-focused training markedly improve OOD detection (FNR99 down to the low single digits for certain checksum forms), while revealing a correlation between checksum error and prediction error for OOD data. The approach offers a lightweight, generalizable reliability signal for scientific ML surrogates with potential broad impact for physics-informed modeling and high-stakes simulations.

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.

Figures (3)

• Figure 1: Adding a check node to the neural network allows the user to encode a checksum function into the output layer. We can then use the degree of violation of this function as a metric for determining prediction reliability.
• Figure 2: We generated the training, validation and out-of-distribution (OOD) datasets from trusted Cretin simulations scott_cretinradiative_2001. While the data has 87 dimensions, we split the OOD data points with an arbitrary dividing line in the density-temperature plane to create a set of data points not shown to the surrogate model in training by construction.
• Figure 3: Relationship between checksum error and prediction error with an optimized loss function, and either a summation (\ref{['fig:results_sum']}) or sinusoid (\ref{['fig:results_sine']}) checksum function. We determine reliability based on a threshold checksum error with 99% of the validation data below this value. With respect to out-of-distribution data points, we see a positively correlated relationship between the checksum and prediction errors.