Quantifying Distribution Shifts and Uncertainties for Enhanced Model Robustness in Machine Learning Applications

TL;DR

The paper tackles distribution shift in machine learning by leveraging synthetic data generated via the van der Waals equation to study model adaptation and uncertainty. It quantifies data similarity with $D_{ ext{KL}}(P\|Q)$, $D_{ ext{JS}}(P,Q)$, and Mahalanobis distance, and assesses uncertainty through Monte Carlo Dropout. Two controlled experiments reveal that increasing distributional dissimilarity correlates with reduced accuracy and higher uncertainty, with Mahalanobis distance offering per-point reliability signals and potential computational advantages. The findings support using distance-based monitoring alongside probabilistic uncertainty estimates to improve model robustness and generalization in real-world, shifting environments.

Abstract

Distribution shifts, where statistical properties differ between training and test datasets, present a significant challenge in real-world machine learning applications where they directly impact model generalization and robustness. In this study, we explore model adaptation and generalization by utilizing synthetic data to systematically address distributional disparities. Our investigation aims to identify the prerequisites for successful model adaptation across diverse data distributions, while quantifying the associated uncertainties. Specifically, we generate synthetic data using the Van der Waals equation for gases and employ quantitative measures such as Kullback-Leibler divergence, Jensen-Shannon distance, and Mahalanobis distance to assess data similarity. These metrics en able us to evaluate both model accuracy and quantify the associated uncertainty in predictions arising from data distribution shifts. Our findings suggest that utilizing statistical measures, such as the Mahalanobis distance, to determine whether model predictions fall within the low-error "interpolation regime" or the high-error "extrapolation regime" provides a complementary method for assessing distribution shift and model uncertainty. These insights hold significant value for enhancing model robustness and generalization, essential for the successful deployment of machine learning applications in real-world scenarios.

Figures (8)

• Figure 1: Neural Network with Dropout, as indicated by red nodes
• Figure 2: Pressure distribution and t-SNE plot for the various gases
• Figure 3: Predicted vs Real Pressure, Experiment 1
• Figure 4: Comparison of KL-Divergence and JS-Distance vs. Prediction Error, demonstrating a clear correlation between the degree of distribution shift and the consequent decline in model accuracy.
• Figure 5: Distribution of Volume, Moles, Temperature, and the resulting calculated Pressure for the train ("in distribution") vs. test ("out-of-distribution") datasets