Towards Integrating Uncertainty for Domain-Agnostic Segmentation

TL;DR

This work tackles the robustness of segmentation models like SAM under domain shifts by introducing UncertSAM, a multi-domain benchmark, and systematically evaluating four post-hoc uncertainty methods plus a simple uncertainty-guided refinement. The last-layer Laplace approximation yields the strongest correlation between uncertainty and segmentation error, indicating a meaningful signal for potential refinement. However, a lightweight Dense Embedding Fusion refinement provides limited gains, suggesting that deeper integration of uncertainty into the model architecture or training is necessary for consistent domain-agnostic improvements. Overall, the study demonstrates meaningful uncertainty signals and outlines a clear direction for future uncertainty-aware segmentation approaches that generalize across domains.

Abstract

Foundation models for segmentation such as the Segment Anything Model (SAM) family exhibit strong zero-shot performance, but remain vulnerable in shifted or limited-knowledge domains. This work investigates whether uncertainty quantification can mitigate such challenges and enhance model generalisability in a domain-agnostic manner. To this end, we (1) curate UncertSAM, a benchmark comprising eight datasets designed to stress-test SAM under challenging segmentation conditions including shadows, transparency, and camouflage; (2) evaluate a suite of lightweight, post-hoc uncertainty estimation methods; and (3) assess a preliminary uncertainty-guided prediction refinement step. Among evaluated approaches, a last-layer Laplace approximation yields uncertainty estimates that correlate well with segmentation errors, indicating a meaningful signal. While refinement benefits are preliminary, our findings underscore the potential of incorporating uncertainty into segmentation models to support robust, domain-agnostic performance. Our benchmark and code are made publicly available.

Figures (8)

• Figure 2: Overview of our four uncertainty strategies.From left to right: We quantify pixel-level uncertainty by targeting the input image (via test-time augmentations, TTA), prompts, model parameters (via the Laplace approximation, LA), and an additional variance prediction head. The first three methods leverage stochasticity to generate ensembles, whereas the latter is deterministic.
• Figure 3: Pearson correlation coefficient$\rho$ between error maps and uncertainty maps for each method, averaged across samples per dataset. Higher positive values indicate better correlation, with the LA giving strongest error alignment.
• Figure 4: Modified SAM mask decoder architecture with an auxiliary variance network.
• Figure 5: Architecture of the Dense Embedding Fusion Network. The prompt encoder is extended with a CNN-based uncertainty encoder, operating in parallel to the mask prompt encoder used in SAM models. The resulting dense spatial embeddings are concatenated along the channel dimension and fused via a $1 \times 1$ convolution layer.
• Figure 6: Radar plot illustrating Brier scores of the UQ methods, averaged across samples per dataset. Lower scores indicate better probabilistic performance.