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

Reliable and accurate machine learning models are crucial for real-world applications. However, a significant challenge arises when the statistical properties of the training data differ from those of the test data, a phenomenon known as distribution shift. This is particularly problematic in dynamic systems where data distributions evolve over time and in transfer learning scenarios. Distribution shifts manifest in various ways, impacting model performance across diverse domains. For instance, in wind turbine performance monitoring, covariate shift occurs when environmental factors (wind speed, direction, temperature) differ between training and operational data. Target drift arises when the relationship between these factors and turbine performance changes due to aging equipment or operational modifications. Similarly, medical image analysis models trained on X-rays from one hospital may fail in another due to equipment differences, potentially leading to misdiagnoses. Autonomous vehicles can encounter unforeseen situations not present in training data, posing safety risks. Recommendation systems can deliver inaccurate suggestions as user preferences change. These examples highlight the critical need to address distribution shift's impact on model performance. Previous research by Shimodaira [18] and Liu et al. [10] has demonstrated the negative effects of distributional disparities. Various domain adaptation methods have been proposed to mitigate these disparities, including adversarial training [5], discrepancy-based approaches [11], and data augmentation techniques like mixup [22] and GANs [7]. Bayesian methods [6] and conformal prediction [1] offer approaches to uncertainty quantification. Ovadia et al. [14] emphasized the importance of quantifying distribution shift and its effect on model uncertainty. This study addresses the gap in understanding distribution shift by investigating the use of synthetic data to evaluate model robustness. Using the van der Waals equation for gases, we systematically quantify distributional disparities and assess their impact on model performance and uncertainty.

Existing literature extensively covers the detrimental effects of distribution shifts on machine learning model performance and generalization. Studies like Shimodaira [18] and Liu et al. [10] highlight the challenges posed by distributional disparities between training and test datasets. These findings spurred research into domain adaptation techniques to bridge the gap between source and target domains. Methods like adversarial training [5] aim to align source and target distributions, while discrepancy-based approaches [11] focus on minimizing the differences between them. Data augmentation techniques, including mixup [22] and generative adversarial networks (GANs) [7], have been explored to improve model robustness. Furthermore, Bayesian methods [6] and conformal prediction [1] provide frameworks for uncertainty quantification, offering probabilistic and frequentist approaches to assess prediction reliability. The work of Ovadia et al. [14] underscores the importance of quantifying distribution shifts and their impact on model uncertainty, advocating for the use of appropriate quantitative measures to evaluate model robustness and generalization capabilities. However, a comprehensive understanding of the role of synthetic data in evaluating model robustness under distribution shift remains an area needing further exploration.

This study employs two experiments using synthetic data generated from the van der Waals equation, providing control over distributional properties. **Experiment 1: Changes in Feature-Target Correlations:** This experiment explores how variations in feature-target correlations impact model accuracy and data similarity quantification. Synthetic training data is generated using the ideal gas approximation. Subsequently, datasets for different gases, deviating from ideal gas behavior (approximated by the van der Waals equation), are generated. Kullback-Leibler divergence and Jensen-Shannon distance quantify data similarity. A machine learning model is trained on the ideal gas data and used to predict properties of other gases, assessing the impact of distribution shift on accuracy. **Experiment 2: Feature Distribution Drift:** This experiment investigates the impact of feature distribution changes on model accuracy and uncertainty. It simulates covariate shift by varying feature distributions (pressure, temperature, etc.) between training and test datasets using the ideal gas approximation. The Mahalanobis distance measures each test set data point's deviation from the training data distribution. Model predictions are analyzed to determine how distribution shift affects accuracy and uncertainty. Monte Carlo Dropout estimates model uncertainty, correlated with the Mahalanobis distance to assess the relationship between distribution shift and uncertainty. **Data Similarity Quantification:** The study utilizes Kullback-Leibler (KL) divergence, Jensen-Shannon distance, and Mahalanobis distance to quantify data similarity and distributional shifts. KL divergence measures the relative entropy between two distributions; Jensen-Shannon distance is a symmetric measure derived from KL divergence; and Mahalanobis distance measures the distance of a point from a distribution, accounting for data covariance. **Model Architecture and Uncertainty Quantification:** A simple neural network architecture is employed to focus on distribution shifts and uncertainty, not complex modeling. The model includes input, dense, custom dropout, and output layers. Monte Carlo Dropout, keeping dropout active during inference, estimates model uncertainty by generating an ensemble of predictions for each input, calculating the mean and standard deviation.

**Experiment 1:** Synthetic datasets for various gases were generated using the van der Waals equation, with subtle differences in pressure distributions visualized using t-SNE. KL-divergence and Jensen-Shannon distance quantified data similarity. A machine learning model trained on ideal gas data was used to predict the properties of other gases. Results showed a clear correlation between smaller corrections to the van der Waals equation (indicating greater similarity to the training data), lower KL-divergence and JS distance, and lower prediction errors. Even a subset of the ideal gas data showed non-zero KL-divergence and JS distance due to finite sample size. A linear trend was observed between KL-divergence and mean absolute percentage error (MAPE), indicating that KL-divergence effectively predicts model accuracy degradation. **Experiment 2:** This experiment focused on quantifying prediction uncertainty in the presence of covariate shift. Two datasets were generated using the ideal gas approximation, differing in input variable distributions. KL-divergence and JS-distance were calculated. The model was trained on one dataset and tested on both in- and out-of-distribution datasets. The Mahalanobis distance was used to measure data points' deviation from the training distribution. Monte Carlo Dropout was used to estimate model uncertainty. Results showed that data points with low Mahalanobis distance (close to the training distribution) had lower prediction errors. As Mahalanobis distance increased (beyond the 95th percentile), prediction errors and uncertainty (standard deviation of predictions) increased. The Mahalanobis distance demonstrated a stronger correlation with prediction error than the standard deviation from Monte Carlo Dropout. The Mahalanobis distance offers per-data-point reliability assessment.

The findings confirm the significant impact of distribution shift on machine learning model performance. Experiment 1 demonstrated that KL-divergence and Jensen-Shannon distance are effective in quantifying the degree of distribution shift and predicting the subsequent decrease in model accuracy. Experiment 2 highlighted the value of the Mahalanobis distance in identifying regions of high uncertainty, specifically where model predictions deviate from the training data distribution. The combination of Monte Carlo Dropout and Mahalanobis distance provides a robust approach to uncertainty quantification, surpassing the performance of either method alone. The results suggest that focusing solely on the standard deviation from Monte Carlo Dropout may not fully capture prediction uncertainty; the Mahalanobis distance offers a more informative and computationally efficient alternative. This approach allows for a more nuanced understanding of prediction reliability, enabling more informed decision-making in real-world applications. The use of synthetic data generated from the van der Waals equation proved beneficial in controlling the distributional properties, providing valuable insights into distribution shift and its impact on model performance.

This study demonstrated the effectiveness of using KL-divergence, Jensen-Shannon distance, and Mahalanobis distance to quantify distribution shifts and uncertainties in machine learning models. Experiment 1 showed a clear correlation between distribution shift and prediction accuracy, while Experiment 2 highlighted the usefulness of Mahalanobis distance in assessing prediction reliability on a per-data-point basis. The findings emphasize the importance of monitoring distribution shift and utilizing appropriate uncertainty quantification methods for robust and reliable machine learning applications. Future research could explore more sophisticated models and different types of distribution shifts, as well as investigating other uncertainty quantification techniques to find optimal strategies for different application domains.

The study's use of synthetic data, while offering control over distributional properties, may not fully capture the complexities of real-world datasets. The specific neural network architecture used was deliberately simple; more complex architectures might exhibit different behaviors under distribution shifts. The study focused on specific metrics; exploring additional metrics could provide a more comprehensive understanding of distribution shifts and their impact on model performance. The focus was on regression problems; extending the findings to classification tasks would provide broader applicability.