Uncertainty-Aware Multiple-Instance Learning for Reliable Classification: Application to Optical Coherence Tomography

TL;DR

This paper introduces UBIX, a method that improves the reliability of deep learning models in medical image classification by detecting and excluding corrupted instances, especially artifacts, to enhance generalization across different scanner vendors.

Contribution

The paper proposes UBIX, a novel uncertainty-based multiple-instance learning technique that detects out-of-distribution artifacts and reduces their impact without retraining, improving model robustness across diverse medical imaging data.

Findings

UBIX maintains high performance on new vendor data with artifacts (κ_w from 0.861 to 0.708).

State-of-the-art models significantly degrade on the same data (κ_w from 0.852 to 0.084).

Out-of-distribution artifacts can be effectively identified and mitigated.

Abstract

Deep learning classification models for medical image analysis often perform well on data from scanners that were used during training. However, when these models are applied to data from different vendors, their performance tends to drop substantially. Artifacts that only occur within scans from specific scanners are major causes of this poor generalizability. We aimed to improve the reliability of deep learning classification models by proposing Uncertainty-Based Instance eXclusion (UBIX). This technique, based on multiple-instance learning, reduces the effect of corrupted instances on the bag-classification by seamlessly integrating out-of-distribution (OOD) instance detection during inference. Although UBIX is generally applicable to different medical images and diverse classification tasks, we focused on staging of age-related macular degeneration in optical coherence tomography. After being trained using images from one vendor, UBIX showed a reliable behavior, with a slight decrease in performance (a decrease of the quadratic weighted kappa ($\kappa_w$) from 0.861 to 0.708), when applied to images from different vendors containing artifacts; while a state-of-the-art 3D neural network suffered from a significant detriment of performance ($\kappa_w$ from 0.852 to 0.084) on the same test set. We showed that instances with unseen artifacts can be identified with OOD detection and their contribution to the bag-level predictions can be reduced, improving reliability without the need for retraining on new data. This potentially increases the applicability of artificial intelligence models to data from other scanners than the ones for which they were developed.