Title: Uncertainty-aware Fine-tuning of Segmentation Foundation Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: Segmentation foundation model
Abstract: The Segment Anything Model (SAM) is a large-scale foundation model that has revolutionized segmentation methodology. Despite its impressive generalization ability, the segmentation accuracy of SAM on images with intricate structures is often unsatisfactory. Recent works have proposed lightweight fine-tuning using high-quality annotated data to improve accuracy on such images. However, here we provide extensive empirical evidence that this strategy leads to forgetting how to "segment anything": these models lose the original generalization abilities of SAM, in the sense that they perform worse for segmentation tasks not represented in the annotated fine-tuning set. To improve performance without forgetting, we introduce a novel framework that combines high-quality annotated data with a large unlabeled dataset. The framework relies on two methodological innovations. First, we quantify the uncertainty in the SAM pseudo labels associated with the unlabeled data and leverage it to perform uncertainty-aware fine-tuning. Second, we encode the type of segmentation task associated with each training example using a $\textit{task prompt}$ to reduce ambiguity. We evaluated the proposed Segmentation with Uncertainty Model (SUM) on a diverse test set consisting of 14 public benchmarks, where it achieves state-of-the-art results. Notably, our method consistently surpasses SAM by 3-6 points in mean IoU and 4-7 in mean boundary IoU across point-prompt interactive segmentation rounds. Code is available at https://github.com/Kangningthu/SUM


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces the Segmentation with Uncertainty Model (SUM) that combines high-quality annotated data with a large unlabeled dataset to improve performance without forgetting. First, the authors quantify the uncertainty in the SAM pseudo labels associated with the unlabeled data and leverage it to perform uncertainty-aware fine-tuning. Second, the task prompt encodes the type of segmentation task associated with each training example to reduce ambiguity. The proposed method is evaluated on different test sets consisting of 14 public benchmarks.\n\nSTRENGTHS:\n1. The motivation of the paper, which aims to improve the accuracy of SAM without affecting generalization, is interesting.\n2. The paper verifies the effectiveness of the proposed method through a large number of experiments.\n\nWEAKNESSES:\n1.The proposed uncertainty-aware fine-tuning method has limited innovation. The uncertainty-aware fine-tuning strategy has been widely used and has been proven to effectively purify pseudo-labels [1,2]. The proposed method aims to improve the initial prediction of SAM. Although the framework has many complex modules, it does not reveal any insightful views or lacks key theoretical analysis.\n\n2.The method section lacks theoretical introduction. Only the text introduction reads like a technical document. There should be more forms of introduction to help understand the principle of the method, such as formulas or figures.\n\n3.The presentation of the method is unclear and not very readable. For example, in the first sentence of section 3.2, it is introduced that \"SUM applies the same prompt-sampling strategy as SAM for the human-annotated data during interactive training\". In this case, should this module be placed in the Preliminary instead of a separate section in the method section? This makes it confusing to understand the proposed method.\n\n[1]Yuxi Wang, Junran Peng, and ZhaoXiang Zhang. Uncertainty-aware pseudo label refinery for domain adaptive semantic segmentation. In Proceedings of the IEEE/CVF international conference on computer vision, pages 9092–9101, 2021.\n[2]Zhedong Zheng and Yi Yang. Rectifying pseudo label learning via uncertainty estimation for domain adaptive semantic segmentation. International Journal of Computer Vision, 129(4):1106– 1120, 2021.\n\nQUESTIONS:\n1.The Related Work Section should be divided into subheadings. Summarizing the related work separately is conducive to improving the readability of the paper. In addition, emphasizing the insightful advantages of the proposed method in each section of related work is conducive to a better understanding of the proposed method.\n\n2. The introduction of the uncertainty-aware segmentation method is relatively broad, and a more detailed analysis of the related work should be provided.\n\nLIMITATIONS:\nAs above", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a novel framework for enhancing the accuracy of the Segment Anything Model (SAM) while maintaining its generalization capabilities. SAM is a foundational model for interactive binary segmentation, but it struggles with segmenting intricate structures accurately. Fine-tuning SAM with high-quality annotated data often leads to overfitting, degrading its generalization abilities.\n\nThe proposed framework, called Segmentation with Uncertainty Model (SUM), addresses these challenges by combining high-quality annotated data with a large unlabeled dataset. Key innovations include: Uncertainty-aware Fine-tuning: The framework quantifies uncertainty in SAM’s pseudo labels and incorporates this into the fine-tuning process to improve segmentation accuracy without losing generalization capabilities. Task Prompts: SUM uses task prompts to specify the segmentation task for each training example, reducing ambiguity and improving the model’s performance on diverse tasks. Uncertainty-aware Prompt Sampling: This technique avoids misleading prompt locations by focusing on regions with high confidence. Experiments demonstrate that SUM consistently outperforms SAM and other fine-tuning strategies across various datasets, achieving significant improvements in mean Intersection over Union (mIoU) and mean boundary IoU (mBIoU).\n\nSTRENGTHS:\nImproved Accuracy:\n\nSUM significantly enhances the segmentation accuracy of SAM, particularly for complex structures. The uncertainty-aware fine-tuning process focuses on regions with high confidence, leading to more precise segmentation results.\n\nMaintained Generalization:\n\nUnlike traditional fine-tuning methods that can lead to overfitting, SUM maintains SAM’s generalization abilities. This is achieved by effectively combining high-quality annotated data with a large, diverse set of unlabeled data.\n\nVersatility and Flexibility:\n\nThe use of task prompts allows SUM to handle various segmentation tasks, including salient-object, entity, and part segmentation. This flexibility makes SUM suitable for a wide range of applications.\n\nEfficiency in Handling Uncertainty:\n\nBy incorporating uncertainty maps and an uncertainty-aware loss function, SUM effectively manages the noise and inaccuracies in pseudo labels. This leads to more reliable training and improved overall performance.\n\nRobust Performance Across Datasets:\n\nSUM demonstrates robust performance on multiple public benchmarks and internal datasets, consistently outperforming existing methods in mIoU and mBIoU. This robustness underscores the model’s effectiveness in diverse settings.\n\nInnovative Use of Uncertainty Maps:\n\nThe generation and utilization of uncertainty maps to guide the training process is an innovative approach. It not only enhances segmentation accuracy but also improves the quality of the pseudo labels used in training.\n\nWEAKNESSES:\nTraining Complexity:\n\nThe SUM framework introduces significant complexity into the training process. Incorporating uncertainty maps, task prompts, and uncertainty-aware prompt sampling requires additional computation and fine-tuning, which may be challenging and time-consuming for implementation.\nDependency on Initial Model Performance:\n\nThe effectiveness of the SUM framework heavily depends on the initial performance of the SAM model. If SAM’s initial pseudo labels are highly inaccurate, the overall performance of SUM could be compromised, as the refinement process might not fully correct these inaccuracies.\nComputational Overhead:\n\nThe additional steps involved in generating and utilizing uncertainty maps, as well as the iterative prompt sampling, add computational overhead. This could be a barrier for practical deployment in resource-constrained environments where computational resources are limited.\nEvaluation on Specific Datasets:\n\nThe paper primarily evaluates SUM on a limited set of datasets focused on specific segmentation tasks. While results are promising, the generalizability of the framework to other datasets and segmentation tasks remains uncertain. Further validation on a broader range of datasets is necessary to confirm its robustness.\n\nQUESTIONS:\nsee the weakness\n\nLIMITATIONS:\nsee the weakness", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-03"}

comments: SUMMARY:
In this paper, the authors proposed the Segmentation with Uncertainty Model (SUM) which combines high-quality annotated data with a large unlabeled dataset. This novel framework improves the performance of the large-scale foundation model without forgetting.

STRENGTHS:
Paper clarity. The paper is overall well-written and structured. The appendix is

Good results. The method improves over the original SAM method and achieves SoTA results. 

Adequate appendix. The quantitative and qualitative results in supplementary material support the paper's idea and concept.

WEAKNESSES:
Although the paper is overall well-written and structured, the figures are confusing. 

1. The legend in Figure 1 looks confusing to me.

2. Figure 2 misses the emphasized part to me.

3. Figure 3 missed the uncertainty map.


Besides, I would suggest the authors use bold to highlight the context they want to emphasize in the Intro Sec.

QUESTIONS:
Please see the weaknesses parts and the following questions. 

1. Can the authors provide the uncertainty map mathematically?

2. Is the uncertainty map similar to the attention in the transformer?

LIMITATIONS:
The authors adequately addressed their work's limitations in the appendix.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes the Segment with Uncertainty Model (SUM), a fine-tuning framework for existing foundational segmentation models like SAM. Specifically, SUM consists of two main components: an uncertainty-aware training pipeline and the task prompt concept to reduce ambiguity.\n\nTechnically, the uncertainty-aware training pipeline comprises three distinct techniques. First, the uncertainty maps generation module is fine-tuned from SAM using filtered high-quality images from human-annotated datasets. With uncertainty maps generated from this module, the uncertainty-aware prompt sampling strategy is then proposed to increase the probability of selecting prompts from high-confidence regions. Finally, the SUM is trained in a semi-supervised manner, and the uncertainty-aware focal and dice loss is applied to the unlabeled branch. Additionally, the task prompt concept is introduced to differentiate data from various sources, reducing ambiguity related to different segmentation types of interest across datasets.\n\nThe experimental evaluation in this paper assessed the performance of the proposed strategy on various datasets. The results indicate that the proposed SUM achieved superior performance across multiple benchmarks compared to the other methods.\n\nSTRENGTHS:\n- The paper provides a clear explanation and detailed analysis of the proposed method. Also, it is good that even the supplementary material was faithfully written, allowing me to check most of the things I was curious about while reading the main text.\n- The ideas behind uncertainty-based training strategy are considered intuitive and novel.\n- The experimental analysis is comprehensive. In addition to the evaluation of multiple benchmarks across various datasets, the ablation studies and quantitative results analyses presented in the supplementary materials further enhance credibility.\n- The discussion on semi-supervised baselines makes the approach by which SUM utilizes unlabeled data more convincing.\n\nWEAKNESSES:\nThere are a few concerns about this paper:\n- Compared to SAM fine-tuning methods, SUM requires an additional training phase to obtain uncertainty maps. This introduces extra training burden.\n- SUM demonstrates outstanding performance on numerous benchmarks. However, as shown in Table 6, the performance improvement is not substantial with the increase in the number of parameters of the backbone. Does this indicate that SUM may face challenges in terms of scaling up the model?\n\nQUESTIONS:\nI thank the author for their detailed experiments and analysis. Nevertheless, there are a few questions I wonder about.\n\n- The uncertainty-based strategy is mainly applied to unlabeled data. What would happen if a similar strategy were applied to the labeled data as well?\n- It appears that the task prompt is used to distinguish the type of input images. In Table 2, the SUM Continuous TP setting employs the task prompt during inference as well. I am curious why the performance is better without using the task prompt during inference and how the performance would be affected by using different task prompts.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}