Uncertainty-Aware Image Classification In Biomedical Imaging Using Spectral-normalized Neural Gaussian Processes

TL;DR

This work addresses unreliable, overconfident predictions in digital pathology by employing Spectral-normalized Neural Gaussian Processes (SNGP), which enable single-pass, uncertainty-aware classification via spectral normalization and a Gaussian-process final layer. The authors evaluate SNGP across six datasets and three biomedical tasks, comparing against deterministic and MC dropout baselines, and demonstrate substantially improved out-of-distribution detection and calibration while preserving in-distribution accuracy. They provide an open-source PyTorch implementation to facilitate adoption. The approach promises safer deployment of AI in pathology by delivering reliable uncertainty estimates that can flag unfamiliar inputs for human review.

Abstract

Accurate histopathologic interpretation is key for clinical decision-making; however, current deep learning models for digital pathology are often overconfident and poorly calibrated in out-of-distribution (OOD) settings, which limit trust and clinical adoption. Safety-critical medical imaging workflows benefit from intrinsic uncertainty-aware properties that can accurately reject OOD input. We implement the Spectral-normalized Neural Gaussian Process (SNGP), a set of lightweight modifications that apply spectral normalization and replace the final dense layer with a Gaussian process layer to improve single-model uncertainty estimation and OOD detection. We evaluate SNGP vs. deterministic and MonteCarlo dropout on six datasets across three biomedical classification tasks: white blood cells, amyloid plaques, and colorectal histopathology. SNGP has comparable in-distribution performance while significantly improving uncertainty estimation and OOD detection. Thus, SNGP or related models offer a useful framework for uncertainty-aware classification in digital pathology, supporting safe deployment and building trust with pathologists.

Figures (3)

• Figure 1: Overview of the predictions of deterministic DNN vs SNGP on In-Distribution (blood cell) vs OOD (cardiac) data.
• Figure 2: Sample images from paired datasets: (a,d) white blood cells, (b,e) amyloid plaques, and (c,f) colorectal pathology.
• Figure 3: Class logits uncertainty distribution