Information Bottleneck-based Feature Weighting for Enhanced Medical Image Out-of-Distribution Detection

Brayden Schott, Zan Klanecek, Alison Deatsch, Victor Santoro-Fernandes, Thomas Francken, Scott Perlman, Robert Jeraj

Keywords: Out of distribution detection, uncertainty quantification, tumor segmentation

Abstract: Deep learning models are subject to failure when inferring upon out-of-distribution (OOD) data, i.e., data that differs from the models’ train data. Within medical image settings, OOD data can be subtle and non-obvious to the human observer. Thus, developing highly sensitive algorithms is critical to automatically detect medical image OOD data. Previous works have demonstrated the utility of using the distance between embedded train and test features as an OOD measure. These methods, however, do not consider variations in feature importance to the prediction task, treating all features equally. In this work, we propose a method to enhance distance-based OOD measures via feature importance weighting, which is determined through an information bottleneck optimization process. We demonstrate the utility of the weighted OOD measure within the metastatic liver tumor segmentation task and compare its performance to its non-weighted counterpart in two as-sessments. The weighted OOD measure enhanced the detection of artificially perturbed data, where greater benefit was observed for smaller perturbations (e.g., AUC = 0.8 vs. AUC = 0.72). In addition, the weighted OOD measure achieved better correlation to liver tumor segmentation Dice coefficient (e.g., ρ = -0.76 vs ρ = -0.21). In summary, this work demonstrates the bene-fit of feature importance weighting for distance-based OOD detection.

Submission Number: 12