Towards Integrating Uncertainty for Domain-Agnostic Segmentation

AI summary

Overview

• Research area: Computer vision, specifically foundation-model segmentation (the Segment Anything Model family) and uncertainty quantification (UQ).
• Technical level: Intermediate. The paper assumes familiarity with segmentation metrics, Bayesian approximations, and ensemble-style uncertainty estimation, though the core question is conceptually simple.
• Scope (one sentence): The paper curates a multi-domain segmentation benchmark called UncertSAM, compares four lightweight post-hoc uncertainty estimation methods for a frozen SAM-2 model, and tests whether those uncertainty maps can be used to refine predictions without any domain-specific fine-tuning.

Authors: Jesse Brouwers, Xiaoyan Xing, Alexander Timans, affiliated with the UvA-Bosch Delta Lab, University of Amsterdam. Preprint arXiv:2512.23427v1 [cs.CV], dated 29 Dec 2025.

What This Paper Is About

Foundation segmentation models such as the Segment Anything Model (SAM) work well zero-shot, but they break down on shifted or low-knowledge domains — shadows, transparent objects, camouflage, lighting artifacts, and medical scans. The authors ask whether attaching a pixel-level uncertainty estimate to a frozen SAM can flag where the model is likely wrong, and whether that signal can be used to improve the mask itself in a domain-agnostic way, without fine-tuning per dataset.

Key Contributions

1. UncertSAM benchmark. A curated, standardised benchmark of eight datasets (BIG, COIFT, COD, MSD Spleen, ISTD, SBU, Flare, Trans) covering fine-grained salient objects, camouflage, medical CT, shadows, lighting artifacts, and transparency. It contains over 23,000 images and 44,000 annotated masks, and is designed for domain-agnostic evaluation of segmentation, UQ, and uncertainty-guided refinement.
2. Systematic comparison of four pixel-level uncertainty methods for SAM — test-time augmentations (TTA), prompt perturbations, a last-layer Laplace approximation (LA), and a learnable variance network — all aimed at being lightweight and post-hoc, working on a frozen SAM encoder.
3. A preliminary uncertainty-guided refinement strategy, Dense Embedding Fusion, which fuses uncertainty maps into the model's internal spatial features, requiring no domain-specific fine-tuning.
4. Public release of the benchmark and code at https://github.com/JesseBrouw/UncertSAM.

Main Findings

• Uncertainty correlates strongly with error. Measured by the Pearson correlation coefficient ρ(U, E) between the uncertainty map U and the probabilistic model error E = |P − M|, correlation was generally high across all four methods, with ρ > 0.5 for all methods.
• Last-layer Laplace wins on error alignment. The LA produced the strongest correlation with segmentation errors; the variance network recorded the lowest values. The authors therefore focused on LA for the refinement experiments.
• All methods are comparably calibrated. A Brier score comparison (a proper scoring rule capturing both calibration and prediction properties) showed similarly low values across the UQ methods, indicating comparable probabilistic accuracy.
• Ensemble predictions from LA help on domain-shifted data, uncertainty does not (yet). Under mIoU, Dense Fusion w/ LA (Ones Map) improved over baseline SAM on ISTD (67.89 vs 65.71), SBU (67.37 vs 66.11), Flare (46.61 vs 42.82), and Trans (86.54 vs 84.87). Under mBIoU, the same configuration improved ISTD (48.92 vs 47.71), SBU (58.04 vs 57.49), Flare (36.71 vs 35.47), and Trans (72.31 vs 69.82).
• Explicit uncertainty fusion did not beat its control. Fusing the actual LA uncertainty map (Dense Fusion w/ LA) did not improve meaningfully over fusing a constant map of ones (Ones Map), suggesting the gains stem from the LA's improved ensemble predictions rather than from the uncertainty information passed through dense embedding fusion.
• Degradation on in-domain datasets. Results were mixed: on BIG, baseline SAM reached 89.95 mIoU and Dense Fusion w/ LA dropped to 88.58; on COIFT, 94.77 versus 94.73. The authors note these highly in-domain but fine-grained datasets already see strong baseline SAM performance.
• A large gap to the upper bound remains. The empirical upper bound using the ground truth mask for refinement was 93.96 (BIG), 95.65 (COIFT), 89.20 (COD), 91.35 (MSD), 92.52 (ISTD), 87.09 (SBU), 82.15 (Flare), and 94.92 (Trans) under mIoU — far above all refinement results.
• Two stated limitations of the approach. First, freezing SAM's image encoder, which contains most of the model's expressivity, prevents adaptation of internal representations. Second, all uncertainty fitting/training used an SA-1B subset containing predominantly high-confidence examples, limiting exposure to lower-confidence, domain-shifted uncertainty patterns.

Methodology in Plain English

The work proceeds in three stages.

1. Building the benchmark. The authors gathered eight public datasets spanning different visual challenge types and standardised them. Pre-processing included connected component analysis (morphological closing twice with a 3×3 kernel, retaining components larger than 1,000 pixels), splitting colour-coded masks into separate targets, and slicing 3D CT volumes along the axial plane, keeping only slices with foreground labels, with Z-score normalisation and clipping to the [0.05, 99.95] percentile range. All uncertainty fitting and training uses a subset of SAM's original SA-1B training data so that the evaluation remains domain-agnostic.

2. Four ways to estimate uncertainty. Given an image X and a bounding box prompt q, SAM outputs pixel-wise logits, converted with a sigmoid to a probability map P. Uncertainty is then obtained by perturbing different parts of the pipeline:

• Test-Time Augmentations (TTA): apply random augmentations (flips, resizes, jitter, hue shifts), run several forward passes, and compute the entropy of the averaged probability map.
• Prompt Perturbations: jitter the bounding box coordinates using SAM's training-time perturbation schedule and ensemble the results the same way.
• Last-layer Laplace Approximation (LA): fit a diagonal Hessian over the final linear decoder layer, treat the pretrained weights as a maximum a posteriori estimate, and sample weights from the resulting Gaussian posterior to build an ensemble.
• Learnable variance network: train an auxiliary head on top of SAM's decoder features to predict a pixel-wise log variance, using the heteroscedastic loss of Kendall and Gal (2017). Unlike the other three, this is deterministic rather than ensemble-based.

3. Refinement. Dense Embedding Fusion extends SAM's prompt encoder with a parallel uncertainty encoder. Its dense features are concatenated with the mask prompt embeddings and fused through a 1×1 convolution, so uncertainty is injected into the spatial features the decoder uses. A second forward pass through SAM then produces the refined mask.

Implementation details. Experiments ran on NVIDIA H100 GPUs (80GB VRAM) with CUDA 12.6.0. Training followed an 8-step schedule; the first prompt is a bounding box or a single foreground point with equal probability, and each subsequent step adds one point sampled uniformly from foreground or background. Optimisation used AdamW with a constant learning rate of 1e-4, gradient clip norm 0.1, weight decay 1e-4 for the variance network and 0.01 otherwise, batch size of 1 image, 17,500 steps, 195 warmup steps, and 972 cooldown steps. The LA was fit for one epoch with batch size 1, sampling one of the eight prompts per optimisation step and computing the loss on predictions and masks downsampled to 128×128 due to memory constraints. Bounding box noise was up to 10% of box dimensions (maximum 20 pixels), and augmentation was limited to random horizontal flips to mirror SAM's pretraining setup.

Why This Matters

Impact on research. The paper tests a specific hypothesis that has been under-explored: that uncertainty, estimated post-hoc from a frozen foundation model, can act as a domain-agnostic substitute for per-domain fine-tuning. Its negative result on explicit uncertainty fusion — no meaningful gain over a constant map control — is as informative as its positive one, and it points toward deeper architectural or training-time integration rather than post-hoc add-ons. The UncertSAM benchmark also provides a standardised, over-23,000-image testbed for future UQ work on segmentation.

Real-world applications (drawn from the domains the benchmark covers):

• Medical imaging, where CT scan segmentation (MSD Spleen) benefits from knowing which regions the model is unsure about before a clinician reviews the mask.
• Shadow and lighting correction in photography and computational imaging (ISTD, SBU, Flare), where the paper's example shows SAM missing shadow regions despite an accurate bounding box.
• Robotics and scene understanding for transparent or camouflaged objects (Trans, COD), where failure to segment an unseen or blending object has safety consequences.
• Interactive annotation tools and dataset construction, where confidence signals could tell a human annotator which regions need correction.

Industry relevance. Any product built on promptable segmentation — annotation platforms, medical imaging pipelines, photo editing, autonomous perception — inherits SAM's brittleness under domain shift. A lightweight, post-hoc confidence layer that runs on a frozen encoder is attractive industrially because it avoids retraining large models for every new domain. The authors' finding that this promise is only partly realised sets realistic expectations and identifies where investment is needed.

Future Directions

1. Move beyond post-hoc integration. The authors argue that fusing uncertainty into the model architecture, or into the learning process itself, should give better leverage than freezing SAM's image encoder and attaching modules downstream.
2. Improve exposure to low-confidence patterns. Because all fitting used an SA-1B subset dominated by high-confidence examples, the refinement module saw few of the uncertain patterns that arise under domain shift. A more balanced regime or cross-domain fine-tuning exposure is proposed to build a richer representation of uncertainty.
3. Improve the refinement mechanism. The current Dense Embedding Fusion is explicitly described as a preliminary first step; more elaborate strategies for turning an uncertainty map into a corrected mask remain open.
4. Explain why ensemble predictions help. Gains on ISTD, SBU, Flare, and Trans came from the LA's ensemble predictions rather than the uncertainty map — the authors note this may partly stem from the fusion layer's ability to re-weigh latent features, an effect that is not fully characterised.

Target Audience

Researchers and practitioners working on segmentation foundation models, uncertainty quantification, or robustness under domain shift will get the most from this paper. It is also relevant to applied computer-vision engineers who need to decide whether adding a confidence estimate to a frozen SAM deployment is worth the effort, and to benchmark builders interested in domain-agnostic evaluation across shadows, transparency, camouflage, and medical imaging. Readers looking for a fully solved refinement method will find the results preliminary; readers looking for an honest assessment of where post-hoc uncertainty helps and where it does not will find the paper useful.

Authors’ 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.