Laplacian Segmentation Networks Improve Epistemic Uncertainty Quantification

TL;DR

The paper addresses the problem that segmentation models can be overconfident on out-of-distribution medical images and proposes Laplacian Segmentation Networks (LSN) that use a Laplace approximation to the weight posterior with a Gaussian $p(y|x,D)$ expressed as $p(y\mid x,\theta) p(\theta\mid D)$ integrated over $\theta$, and a low-rank covariance for scalable uncertainty. It demonstrates a practical MAP estimation scheme for segmentation by sharing parameters across mean and variance networks and employing a fast Hessian diagonal backpropagation to scale to networks with skip connections. Evaluations on ISIC19, BRATS, and a QUBIQ Prostate task show that LSN with uncertainty measures like the Expected Pairwise KL ($\mathrm{EPKL}$) and Pixel Variance ($\mathrm{PV}$) achieves stronger OOD detection (AUROC) than baselines, while mutual information (MI) can be less reliable. The results underscore the importance of targeted uncertainty metrics for segmentation and provide a scalable framework for jointly quantifying epistemic and aleatoric uncertainty to improve safe deployment in clinical settings.

Abstract

Image segmentation relies heavily on neural networks which are known to be overconfident, especially when making predictions on out-of-distribution (OOD) images. This is a common scenario in the medical domain due to variations in equipment, acquisition sites, or image corruptions. This work addresses the challenge of OOD detection by proposing Laplacian Segmentation Networks (LSN): methods which jointly model epistemic (model) and aleatoric (data) uncertainty for OOD detection. In doing so, we propose the first Laplace approximation of the weight posterior that scales to large neural networks with skip connections that have high-dimensional outputs. We demonstrate on three datasets that the LSN-modeled parameter distributions, in combination with suitable uncertainty measures, gives superior OOD detection.

Figures (3)

• Figure 1: Model overview - uncertainty measures are calculated by approximating expectations by Monte Carlo-sampling mean networks from the Laplace approximation $q(\theta^*)$ and predicting the respective logit distributions $p(\eta \vert x, \theta)$ for $x$.
• Figure 2: Left - OOD performance measured by AUROC across epistemic components for Mutual Information (MI), Expected Pairwise KL (EPKL) and Pixel Variance (PV). Models using Laplace Approximations with EPKL and PV reach highest AUROC values on average. Right - Uncertainty Measures for predictive and aleatoric uncertainty perform on par indicating weak disentanglement.
• Figure 3: OOD performance measured by AUROC across models, marginalized over aggregation strategies, for Mutual Information (MI), Expected Pairwise KL (EPKL) and Pixel Variance (PV). LSN models with EPKL and PV reach highest AUROC values for Prostate and ISIC respectively.

Theorems & Definitions (2)

• proof
• proof