UQ-ARMED: Uncertainty quantification of adversarially-regularized mixed effects deep learning for clustered non-iid data

TL;DR

This paper evaluates uncertainty quantification methods for a deep learning model in medical prognosis, demonstrating that ensemble approaches with subsampling offer reliable statistical metrics, good calibration, and maintain high predictive performance.

Contribution

It compares four epistemic UQ methods for ARMED MEDL models, highlighting ensemble approaches with subsampling as the most effective for uncertainty estimation and statistical inference.

Findings

Ensemble methods with 90% subsampling perform best across metrics.

UQ methods provide statistical significance for model fit and covariate effects.

Ensemble approaches maintain high predictive performance while improving uncertainty calibration.

Abstract

This work demonstrates the ability to produce readily interpretable statistical metrics for model fit, fixed effects covariance coefficients, and prediction confidence. Importantly, this work compares 4 suitable and commonly applied epistemic UQ approaches, BNN, SWAG, MC dropout, and ensemble approaches in their ability to calculate these statistical metrics for the ARMED MEDL models. In our experiment for AD prognosis, not only do the UQ methods provide these benefits, but several UQ methods maintain the high performance of the original ARMED method, some even provide a modest (but not statistically significant) performance improvement. The ensemble models, especially the ensemble method with a 90% subsampling, performed well across all metrics we tested with (1) high performance that was comparable to the non-UQ ARMED model, (2) properly deweights the confounds probes and assigns them statistically insignificant p-values, (3) attains relatively high calibration of the output prediction confidence. Based on the results, the ensemble approaches, especially with a subsampling of 90%, provided the best all-round performance for prediction and uncertainty estimation, and achieved our goals to provide statistical significance for model fit, statistical significance covariate coefficients, and confidence in prediction, while maintaining the baseline performance of MEDL using ARMED