Inadequacy of common stochastic neural networks for reliable clinical decision support

TL;DR

This work critically evaluates the reliability of stochastic neural networks for clinical decision support by applying an Encoder-Only Transformer to ICU mortality prediction on MIMIC-III. Despite strong in-distribution performance and calibration, both ensemble and Bayesian neural network approaches substantially underestimate epistemic uncertainty, with evidence of posterior collapse that hampers reliable OoD detection. The study demonstrates that current stochastic DL methods fail to provide trustworthy uncertainty in safety-critical medical contexts and advocates for distance-aware or kernel-based uncertainty quantification. The findings have direct implications for deploying AI in healthcare, highlighting the need for uncertainty mechanisms that robustly reflect evidence and data coverage beyond the training distribution.

Abstract

Widespread adoption of AI for medical decision making is still hindered due to ethical and safety-related concerns. For AI-based decision support systems in healthcare settings it is paramount to be reliable and trustworthy. Common deep learning approaches, however, have the tendency towards overconfidence under data shift. Such inappropriate extrapolation beyond evidence-based scenarios may have dire consequences. This highlights the importance of reliable estimation of local uncertainty and its communication to the end user. While stochastic neural networks have been heralded as a potential solution to these issues, this study investigates their actual reliability in clinical applications. We centered our analysis on the exemplary use case of mortality prediction for ICU hospitalizations using EHR from MIMIC3 study. For predictions on the EHR time series, Encoder-Only Transformer models were employed. Stochasticity of model functions was achieved by incorporating common methods such as Bayesian neural network layers and model ensembles. Our models achieve state of the art performance in terms of discrimination performance (AUC ROC: 0.868+-0.011, AUC PR: 0.554+-0.034) and calibration on the mortality prediction benchmark. However, epistemic uncertainty is critically underestimated by the selected stochastic deep learning methods. A heuristic proof for the responsible collapse of the posterior distribution is provided. Our findings reveal the inadequacy of commonly used stochastic deep learning approaches to reliably recognize OoD samples. In both methods, unsubstantiated model confidence is not prevented due to strongly biased functional posteriors, rendering them inappropriate for reliable clinical decision support. This highlights the need for approaches with more strictly enforced or inherent distance-awareness to known data points, e.g., using kernel-based techniques.

Figures (6)

• Figure 1: Visualisation of the embedding procedure.The parts of $\textit{token}^{(v/b)}$ tuples are encoded, embedded and combined. The concept part $c$ is encoded to a 1-hot vector $\boldsymbol{c}$. The value part $v$ is encoded to a 1-value vector $\boldsymbol{v}$. Both parts are concatenated ($||$) and embedded ($[.]_{EL}$). Linear time $t$ is embedded to circular time $\boldsymbol{t}$ via multiple sine/cosine waves of different frequencies. Time and concept embeddings are combined additively resulting in the final embedding $\textbf{token}^{(v/b)}$.
• Figure 2: Exemplary remapping for MI: \ref{['subfig:pre_remapping']} MI is binned by predicted probabilities. \ref{['subfig:post_remapping']} Within each bin the MI of the predictions is remapped to a uniform distribution. The successive cut-off procedure is indicated by the threshold $n$. By successively discarding predictions above the threshold $n$ only predictions with less model disagreement remain.
• Figure 3: Correlations between discrimination performance and $\text{m}_{\text{U}}$. By the described procedure the most uncertain predictions, as signaled by naive MI, $\text{PE}(\mathcal{L})$, $\sigma(\delta)$ and isolated MI$_{\text{EU}}$ and $\sigma(\delta)_{\text{EU}}$, are successively discarded (Cut-off Quantile) for \ref{['subfig:aucrocte']} TE and \ref{['subfig:aucrocbt']} BT. The table shows the performance metrics for TE, BT and two benchmark models (S-LSTM and MTCW-LSTM) HARUTYUNYANMultitaskLearningAndBenchmarkingWithClinicalTimeSeriesData ($\pm$ = 1 standard deviation, $(.,.)$ = 0.05/0.95 quantiles). The top two lines show discrimination performance on the entire data set in terms of AUC ROC and AUC PR. The lower portion shows the change in performance (+/-) when removing the 50% most uncertain predictions. Scores shown in the table are the result of a cross validation using 4 splits and swapping evaluation and test set for each split resulting in a total of 8 results for each model. The plots shown in the figures show only the mean results.
• Figure 4: Calibration results for both TE and BT models. Calibration curves shows close to optimal behaviour (indicated by the diagonal) after calibration for both models.
• Figure 5: Normed MI, $\sigma(\delta)$ and $\text{PE}(\mathcal{L})$ over pred. probabilities for \ref{['subfig:pvcsbt']} BT and \ref{['subfig:pvcste']} TE. Measures show secondary correlations evident by the correlations between pred. probabilities and the distribution of uncertainty. A correlation similar to $\text{PE}(\mathcal{L})$ (middle) is shown by MI and an inverted correlation is shown by $\sigma(\delta)$. These allow for an selection of samples statistically targeted at certain probabilities and as such AU.