Dimensionality Reduction and Nearest Neighbors for Improving Out-of-Distribution Detection in Medical Image Segmentation

TL;DR

This work applied the Mahalanobis distance post hoc to the bottleneck features of four Swin UNETR and nnU-net models that segmented the liver on T1-weighted magnetic resonance imaging and computed tomography to detect out-of-distribution images at inference.

Abstract

Clinically deployed deep learning-based segmentation models are known to fail on data outside of their training distributions. While clinicians review the segmentations, these models tend to perform well in most instances, which could exacerbate automation bias. Therefore, detecting out-of-distribution images at inference is critical to warn the clinicians that the model likely failed. This work applied the Mahalanobis distance (MD) post hoc to the bottleneck features of four Swin UNETR and nnU-net models that segmented the liver on T1-weighted magnetic resonance imaging and computed tomography. By reducing the dimensions of the bottleneck features with either principal component analysis or uniform manifold approximation and projection, images the models failed on were detected with high performance and minimal computational load. In addition, this work explored a non-parametric alternative to the MD, a k-th nearest neighbors distance (KNN). KNN drastically improved scalability and performance over MD when both were applied to raw and average-pooled bottleneck features.

Figures (6)

• Figure 1: MD and KNN pipelines with dimensionality-reduced features using either PCA, t-SNE, UMAP, or average pooling (Pool). The encoder is a trained encoder from a U-Net architecture. k is the k-th nearest neighbor.
• Figure 2: Segmentations with high (top) and low (bottom) DSCs along with their corresponding MDs, calculated in conjunction with PCA with two components. Pink is the ground truth segmentation; teal is the MRI UNETR segmentation.
• Figure 3: Visualization of 2D projections of MRI UNETR embeddings. (Top Row) PCA projections. (Middle Row) t-SNE projections. (Bottom Row) UMAP projections. (Left Column) Test projections split into ID and OOD by 95% DSC. (Middle Column) Test projections by DSC. (Right Column) Projections for the training data by source. The gray ellipses are the covariance ellipses (one and two standard deviations) for the training distribution.
• Figure 4: OOD scores plotted against DSC for MRI$_{\text{Te}}$. Horizontal lines represent 95% DSC. Vertical lines represent the 90% TPR.
• Figure 5: Sample images from the AMOS dataset Ji2022data that represent the baseline and low perceptual resolution images that were clustered separately by the dimensionality reduction techniques.