Uncertainty-aware Fine-tuning of Segmentation Foundation Models

We thank the reviewer for the thoughtful feedback. We respond in detail below.

Q1 Difference with previous uncertainty-aware segmentation methods

Our approach differs fundamentally from previous approaches in both 1) the generation of uncertainty maps and 2) the utilization of these uncertainty maps, as detailed below. We will clarify this in the revised manuscript.

The uncertainty-aware fine-tuning strategies in [1][3] generate uncertainty via model prediction logit values. [2][4] use prediction differences from different heads to generate uncertainty. They are designed for domain-adaptive semantic segmentation or semi-supervised semantic segmentation. In contrast, our proposed strategy is designed for interactive binary segmentation, which has not been tackled by the previous method.

Uncertainty map generation: Our main insight is to utilize external supervision to detect and correct systematic biases that accumulate during training. This design logic is fundamentally different from relying on the model self-training itself, as most previous models including [1] and [2]. SAM undergoes several rounds of self-training, leading to the accumulation and overfitting of errors in pseudo-labels. In our evaluation, the model frequently predicts erroneous regions with high confidence, and different heads concur on these incorrect areas. Therefore, traditional methods to generate uncertainty maps, including [1][2][3][4], do not capture this uncertainty effectively.

Utilization of the uncertainty map: Our method is the first to introduce uncertainty-aware prompt sampling for training interactive segmentation models, representing a novel contribution to the field. In contrast, the semantic segmentation tasks addressed in [1] and [2] do not involve prompt sampling. Additionally, we have tailored the uncertainty-aware focal and Dice loss to train the binary segmentation foundation model. They are not proposed in [1] and [2]. We will report this in Section 2.

Quantitative comparison with previous methods: As reported in Table 1 of the paper, the proposed SUM method outperforms existing methods based on uncertainty quantification, including [1] and [2], "SUM Confidence-Uncer" corresponds to [1], and "SUM Discrepancy-Uncer" corresponds to [2]. Table 1 also provides a comparison with UP2L [3], an uncertainty-aware semi-supervised segmentation method that uses uncertainty maps in the feature space via contrastive learning, which also performs worse than SUM (see also Figure 5).

Q2 Key insights for module design

We will modify the introduction and method sections to clearly explain the three main novel components of our framework:

1. A module for uncertainty map generation, which is trained by external supervision to correct the systematic bias in the foundation model training. The module accurately quantifies the uncertainty in pseudo-labels, generalizing effectively across different tasks. We will also mention novel design considerations, including data-pair filtering, training mask generation, and model tuning design (now in Appendix Section D) in the main paper.
2. An uncertainty-aware cost function, which leverages the uncertainty map.
3. A strategy for uncertainty-aware prompt sampling during training, also leveraging the uncertainty map.

We will highlight these contributions more clearly in the introduction, comparing them in more detail to the existing uncertainty-aware segmentation methods, as explained below, modifying the paper structure as explained below. We will also provide a more mathematical formulation of the uncertainty map, as suggested by the reviewer, and highlight our contributions more clearly in the figures (see the provided modified figures).

Q3 Discussion about other uncertainty-aware methods in related work

We will expand the related work section by providing more detailed explanations of relevant techniques. Previous uncertainty quantification approaches such as [1][2][3][4] are not designed for training interactive foundation models. Our experiments validate that they are not effective in generating effective uncertainty maps in our setting. Additionally, the way these methods utilize uncertainty maps differs from our setting.

[1][2][4] are designed for domain-adaptive semantic segmentation, while [3] is for semi-supervised semantic segmentation. The concept of uncertain classes explored in [1] is not applicable to binary segmentation scenarios. Our strategy, however, is tailored for interactive binary segmentation. Please also see Q1.

We have already conducted comprehensive experiments and provided a detailed discussion of related work in Appendix E.1 and E.2, respectively. We will further expand the related work section to include these additional insights.

Q4 Paper organization.

• Related Work Section: We will reorganize the Related Work section into subheadings, summarizing related work in distinct categories to improve readability and highlight the novel aspects of our method.
• Method Section - Theoretical Introduction: We will add formulas and figures to the Method section to clarify the principles behind our approach.
• Method Section - Presentation: We will move the description of the prompt-sampling strategy to the Preliminary section and add figures to enhance clarity and readability.

[1]Wang et al. Uncertainty-aware pseudo label refinery for domain adaptive semantic segmentation. CVPR 2021.

[2]Zheng et al. Rectifying pseudo label learning via uncertainty estimation for domain adaptive semantic segmentation. IJCV 2021.

[3] Wang et al. Semi-supervised semantic segmentation using unreliable pseudo-labels. CVPR 2022.

[4] Wu et al. Upl-sfda: Uncertainty-aware pseudo label guided source-free domain adaptation for medical image segmentation. IEEE transactions on medical imaging 2023.

We thank the reviewer for the thoughtful feedback. It will help us improve the manuscript. We respond in detail below.

Q1 Training complexity and Q3 computational overhead

We acknowledge the reviewer's concern about the additional training phase for obtaining uncertainty maps in the SUM framework. While the framework adds steps, they remain manageable within practical constraints, as detailed below. We will mention this in the paper and add a subsection to the appendix with a more detailed explanation.

Training the Uncertainty Map Generation Module: Training the uncertainty map generation module is relatively efficient on the human-labeled samples. It involves tuning a small number of parameters and can be completed within 4 hours using 8 A100 GPUs.

Generating uncertainty map for the unlabeled image The computational overhead introduced by the SUM framework is minimal. The process of generating pseudo labels and uncertainty maps for unlabeled data utilizes the existing SAM encoder, thereby sharing the primary computational workload. Beyond pseudo-label generation, the additional time required to produce refined masks using a lightweight decoder is negligible compared to the time taken by the image encoder. ViT-H image encoder in SAM has 632M parameters while the prompt-based decoder only has 3.87M parameters. As per the reference [1], on an RTX4090, the decoder adds only approximately 2% more computational time compared to the encoder.

Fine-Tuning the SAM Model: Once the uncertainty map is generated, the fine-tuning SAM model using the uncertainty map is similar to the standard training used in SAM.

• Iterative Point Prompt Sampling: This method is used in standard SAM training as well, and replacing the original uniform sampling with weighted (uncertainty-aware) sampling results in a negligible training burden. Sampling a point from a 1024x1024 candidate pool using both methods can be completed within 0.006-0.009 seconds (average of 1000 runs) on the CPU of a MacBook.

• Task Prompt: The task prompt, a learnable single vector, is combined via element-wise addition with the embeddings from the SAM image encoder and is used only in the first round of interactive segmentation. The element-wise addition of two tensors is relatively fast.

• Uncertainty-aware Loss Computation: This involves thresholding and a weighted loss computation, which requires a similar running time to the original loss.

Inference Phase: Uncertainty maps are utilized only during training. Once our model is trained, inference operates similarly to the SAM model, without additional computational burden. This results in negligible computational overhead.

Q2 Dependence on the initial model performance

As shown in Figure 9, although the quality of refined results depends on the initial model performance, the gains are positive for most examples. We provide Figure 5 in the rebuttal PDF to demonstrate that even when the initial SAM input quality is suboptimal, our mask refinement module still enhances the input and improves performance.

Q4 Evaluation set

We organize our evaluation tasks according to the hierarchical level of granularity, covering various levels (part, entity, multiple instances). Each task is evaluated using several diverse datasets designed to encompass a wide range of images and subtasks.

For instance, our part segmentation evaluation utilizes five diverse datasets: Fashionpedia, Fashionpedia Subpart, Paco, Multi-Human Parsing, and Portrait. The first two include a comprehensive ontology of fashion elements and different levels of part granularity. Paco covers 75 object classes and the last two focus on different human-specific part segmentation subtasks.

That said, we acknowledge that evaluating a broader range of segmentation tasks would further strengthen the robustness of our proposed methods. To address this, we have extended our evaluation by testing SUM and SAM on additional image types. For reproducibility, SUM is fine-tuned on the Public dataset FT-Medium.

We selected 7 datasets from the evaluation sets used in SAM to complement our existing 14 public evaluation sets. Additionally, we have included part one of a synthetic dataset, GTAV [2]. These additional evaluation sets encompass various image types e.g. driving, synthetic, egocentric, irregular shapes, paintings, underwater animals, drones, and underwater trash.

The mIoU comparison results, reported in the following tables, confirm that SUM consistently outperforms SAM. We appreciate the reviewer’s suggestion and will include these additional results in the Appendix of the final version.

[1] Songa et al. SAM-Lightening: A Lightweight Segment Anything Model with Dilated Flash Attention to Achieve 30 times Acceleration." arXiv preprint arXiv:2403.09195 2024.

[2] Richter et al. Playing for data: Ground truth from computer games. ECCV 2016.

We thank the reviewer for the thoughtful feedback. We respond in detail below.

Q1 Clarification on the figures

We appreciate the reviewer’s suggestion to make the figures clearer. We have included updated figures in the rebuttal PDF (see Figure 1, 2, 3 in the rebuttal PDF) and will incorporate them into the final version of our paper.

In Figure 1, we have allocated individual legends to each corresponding sub-figure to clarify the relationship between the figures and their respective legends.

In Figure 2, we highlight the novel components of our proposed framework, making them easier to identify.

In Figure 3, we have emphasized the uncertainty map and toned down the colors to avoid distraction.

Q2 Highlight the context in intro

We appreciate the reviewer's suggestion. We will highlight key concepts in the introduction and method sections in our revised draft.

Q3 Definition of the uncertainty map

We will add the following mathematical definition of the uncertainty map to Section 3.3. The uncertainty map is obtained from the absolute difference between the original SAM input and the refined output by the Mask-refinement Module. Both the sigmoid-transformed probabilities of the SAM logits and the refined prediction have values in the range $[0,1]$ and share the same spatial dimensions. The uncertainty map is calculated as the absolute difference between them for each pixel. Let $u_i$ represent the uncertainty value for the $i$-th pixel in the uncertainty map, it is equal to:

$u_i = \left| \sigma(\mathbf{S}_i) - \sigma(\mathbf{R}_i )\right|$

where $\sigma$ is the sigmoid function, $\mathbf{S}$ denotes the SAM logits, and $\mathbf{R}$ denotes the refined prediction. This yields values between 0 (no difference, indicating low uncertainty) and 1 (large difference, indicating high uncertainty).

Q4 The similarity of uncertainty map with respect to the attention in the transformer

The attention mechanism in transformers serves multiple crucial functions, such as capturing long-range dependencies and assigning importance weights to tokens. Similarly, the uncertainty map in our method assigns importance to regions of the pseudo labels, which is somewhat analogous in spirit to attention. However, there are important differences.

The uncertainty map is specifically used to guide the training process of the segmentation model, modifying the training cost function and the prompt sampling process. In contrast, attention is a component of the function implemented by transformers, which enables the model to weigh the relevance of different tokens within the sequence, facilitating better context understanding and representation. It does not modify the cost function (or the prompt sampling process).

We thank the reviewer for their encouraging feedback. We believe it will improve the manuscript. We respond in detail below.

Q1 Extra training burden

The reviewer is correct that the additional training phase for obtaining uncertainty maps in the SUM framework adds some computational overhead, but it is manageable (relative to the large size of the backbone model). It involves tuning a small number of parameters and can be completed within 4 hours using 8 A100 GPUs. We will mention this in the revised version of the paper.

Q2 Scaling up with parameter numbers

The reviewer is correct that the performance gains eventually saturate when scaling the backbone, but this already occurs in the original SAM framework, as was noted in the SAM paper [1] (arXiv version, page 12): “ViT-H improves substantially over ViT-B, but has only marginal gains over ViT-L. Further image encoder scaling does not appear fruitful at this time.” Our results in Table 6 align with this observation, showing similar performance trends.

We reorganized Table 6 to highlight the boundary IoU performance of the original SAM model, confirming that our results are consistent with the conclusions of the SAM paper:

Our SUM framework consistently improves SAM performance across different tasks and backbones, we display the Gains of SUM with respect to SAM in the following table:

In summary, while the SAM framework provides limited gains from scaling the backbone, our SUM framework consistently enhances performance across various tasks and backbones.

Q3 Applying uncertainty-aware training for the labeled data

We designed the uncertainty-aware training strategy to mitigate the influence of noisy annotations in pseudo-labels during model training, as our labeled data are of high quality. However, the reviewer's suggestion is very intriguing. In scenarios where the quality of the available annotations is not assured, our uncertainty quantification module could be applied to these labels to enhance training. We will mention this in Section 5. In the rebuttal pdf, we provide a proof of concept that the uncertainty map can be applied to human-labeled data. The uncertainty map corresponding to an image from Cityscapes [2] training set with coarse mask is provided in Figure 4 of the rebuttal PDF. This map accurately highlights the boundary regions where the human annotations tend to be less precise.

Q4 continuous TP in SUM

This is a good point. SUM only adds the task prompt in the first round (i.e. single prompt) during training and inference. Our results indeed indicate that adding the task prompt in all rounds (i.e. both single prompt and multi-prompt scenario) during the interactive training and inference may be counterproductive. A possible explanation is that the task prompt provides an implicit prior to the model regarding the desired output mask. When only one prompt is provided, this prior is useful and improves performance. However, when multiple prompts are provided, this already sufficiently constrains the desired mask, in a way that may slightly contradict the prior associated with the task prompt for some images. We will mention this in the revised manuscript, and point out that correctly balancing different user prompts is an important topic for future research.

In order to illustrate the effect of providing different task prompts, we report the results of applying SUM (FT-Medium) to the salient object segmentation task (on the same evaluation sets as the main paper) for three different round-1 task prompts. The results show that task prompt 1 enables the model to achieve the best performance in round 1, which makes sense since task prompt 1 is associated with the salient object task. However, for later rounds, the difference in performance is very small.

[1] Kirillov et al. Segment Anything. arXiv preprint arXiv:2304.02643 (2023).

[2] Cordts et al. The cityscapes dataset for semantic urban scene understanding. CVPR. 2016.

Thank you for your response. We're pleased that we addressed your concern. With respect to the code, we will make the inference code and the novel components of the SUM training code publicly available upon acceptance. We are currently in the process of obtaining approval from our organization for the release of the full training code.

As an update regarding code release, we are confident in securing approval for the full release of the code, given that it is based on public implementations and does not contain proprietary information. In the meantime, upon acceptance, we will publish a website for the paper, where we will release the inference code and the novel components of the SUM training code.