Unleashing the Potential of the Diffusion Model in Few-shot Semantic Segmentation

Authors: muzhi  zhu
yang  liu
zekai  luo
chenchen  jing
hao  chen
guangkai  xu
xinlong  wang
chunhua  shen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-11 05:40:52

Abstract: The Diffusion Model has not only garnered noteworthy achievements in the realm of image generation 
but has also demonstrated its potential as an effective pretraining method utilizing unlabeled data. 
Drawing from the extensive potential unveiled by the Diffusion Model in both semantic correspondence and open vocabulary segmentation, our work initiates an investigation into employing the Latent Diffusion Model for Few-shot Semantic Segmentation.
Recently, inspired by the in-context learning ability of large language models, Few-shot Semantic Segmentation has evolved into In-context Segmentation tasks, morphing into a crucial element in assessing generalist segmentation models.
In this context, we concentrate 
on Few-shot Semantic Segmentation, 
establishing a solid foundation for the future development of a Diffusion-based generalist model for segmentation. Our initial focus lies in understanding how to facilitate interaction between the query image and the support image, resulting in the proposal of a KV fusion method within the self-attention framework.
Subsequently, we delve deeper into optimizing the infusion of information from the support mask and simultaneously re-evaluating how to provide reasonable supervision from the query mask.
Based on our analysis, we establish a simple and effective framework named DiffewS, maximally retaining the original Latent Diffusion Model's generative framework and effectively utilizing the pre-training prior. Experimental results demonstrate that our method significantly outperforms the previous SOTA models in multiple settings.

Main paper: 
 introduction: 
 The Diffusion Model (DM) has demonstrated powerful capabilities in multiple visual generation tasks, including image generation [1, 2] , image editing [3, 4] , video generation [5] [6] [7] , etc. At the same time, DM has also been proven to be a powerful method for self-supervised pre-training [8, 9] employing unlabelled data. To exploit the representation ability of DM, there are currently two emerging topics in vision research: improving the learning paradigm [10, 11] and downstream task adaptation [12] [13] [14] . The latter often focuses on the Latent Diffusion Model [2] (LDM). By compressing images into latent space, it significantly decreases computational expenses and emerges as the first open-source Text-to-Image Diffusion Model scaled up to the LAION-5B [15] level. For example, ODISE [16] ,DVP [12] , DatasetDM [17] adapt LDM to multiple tasks such as depth estimation, semantic segmentation, but they all require training additional decoder heads, which increases training costs and may undermine the generalization ability and generation quality. Therefore, some works [13, 14] have emerged that attempt to repurpose the Diffusion Model's generative framework and apply it to visual perception tasks without adding extra decoder heads. Nonetheless, these paradigms still cannot uniformly adapt to all tasks.Let us reconsider the most fundamental question in using generative models for visual perception: how to design a fine-tuning framework that can guarantee both generalization ability and precise prediction of details? Unfortunately, existing methods do not sufficiently address this challenge. The demands of the FSS task for open-set generalization and high-quality segmentation results precisely align with this challenge. Thus, our first motivation is to further address the fundamental question posed above by exploring the Diffusion Model on the FSS task.FSS aims to segment query images given support samples. Traditional FSS methods [18] [19] [20] rely on a pre-trained backbone, achieving semantic matching and pixel-level prediction tasks through designing complex modules and long-term training. Recently, with the emergence of SAM [21] , some works are based on foundation models to complete FSS, such as Matcher [22] . It employs DINO [23] for semantic matching and SAM for segmentation. Similarly, other works [24, 25] combine SAM with CLIP [26] [27] [28] or MLLM [29, 30] to complete other open-set segmentation tasks. The current methods deal with matching(semantic) and segmentation as two distinct tasks through different modules. The Diffusion Model itself, however, exhibits significant potential in fine-grained pixel prediction tasks [13, 14, 16] and semantic correspondence tasks [31] [32] [33] . Hence, we seek to maximize the reuse of the generative framework by taking advantage of the innate priors within the Diffusion Model to accomplish the FSS task.Recently, inspired by the in-context learning ability of large language models, Few-shot Semantic Segmentation has further evolved into the In-context Segmentation [34, 35] task (see Section 2). In-context Segmentation requires the model to have in-context learning ability for few-shot samples, posing new challenges to model's generalization capabilities. Consequently, it's now recognized as a crucial component in the evaluation process for generalist segmentation models. Therefore, the second motivation of our work is to lay the groundwork for the development of diffusion-based generalist segmentation models.As a foundational work of Diffusion-based methods in the FSS field, we strive to achieve optimal performance with a simple and efficient design, while maximally preserving the generative framework of the Latent Diffusion Model. This minimal disruption to the original UNet structure allows us to better make use of pre-trained priors. We embark on a systematic exploration around the following four questions: 1) How to implement the interaction between the query image and the support image? 2) How to effectively inject information from the support mask? 3) What is a reasonable form of supervision from the query mask? 4) How to design effective generation process to transfer the pre-trained diffusion models to mask prediction task? Based on our observations, we ultimately establish the DiffewS framework and validate it in multiple settings, demonstrating the effectiveness of our method. Our main contributions include:• We systematically study four crucial elements of applying the Diffusion Model to Few-shot Semantic Segmentation. For each of these aspects, we propose several reasonable solutions and validate them through comprehensive experiments.• Building upon our observations, we establish the DiffewS framework, which maximally retains the generative framework and effectively utilizes the pre-training prior. Notably, we introduce the first diffusion-based model dedicated to Few-shot Semantic Segmentation, setting the groundwork for a diffusion-based generalist segmentation model.• We validate the effectiveness of the DiffewS framework under several experimental settings, demonstrating that our method not only achieves a performance comparable with the stateof-the-art (SOTA) model in a strict Few-shot Semantic Segmentation setting, but also significantly outperforms the current SOTA model in an 'in-context learning' setting, 
related work: 
 Diffusion models have shown impressive performance on visual generation tasks such as text-based image generation [1, 2] , image editing [3, 4] , and video generation [5] [6] [7] . Current research on leveraging Diffusion models to enhance visual perception tasks mainly focuses on two directions: one is the direct use of diffusion models to generate images, aiming to address the issue of insufficient data, such as instance segmentation [36] [37] [38] , semantic segmentation [39] , few-shot segmentation [40] and so on. Another direction is to transfer features from Diffusion models to other visual tasks, which aligns with the research direction of this paper.ODISE [16] uses frozen diffusion models for panoptic segmentation of any category in the wild. DVP [12] , DatasetDM [17] , GenPercept [41] , Geowizard [42] adapt LDM to multiple tasks such as depth estimation, semantic segmentation, and surface normal. Marigold [13] fine-tunes diffusion models on synthetic data for affine-invariant monocular depth estimation and achieves impressive performance. Different from the above methods, we focus on using diffusion models to model the visual correlations of multiple reference images and a target image for few-shot segmentation. The most related work to this paper is a concurrent study [43] , which focuses on utilizing diffusion models for in-context segmentation. However, it disrupts the original U-Net structure and the priors of the diffusion model to some extent. In contrast, our work offers a more comprehensive and systematic analysis of applying diffusion models to Few-Shot Semantic Segmentation tasks.Few-shot semantic segmentation [44, 45] aims to segment target objects in an input image given a few annotated support images. Traditional FSS methods either explore prototype learning [46] [47] [48] of support images to predict query images' masks or use pixel-level information [49, 50, 18] to exploit the support information. For example, some works [34, 35, 51] demonstrate powerful generalization ability by unifying various segmentation tasks in an in-context learning framework. SegGPT [35] can exactly segment any semantic conception by using one or a few support images, which motivates us to explore the potential of the diffusion model for the FSS task under the in-context setting [35] .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: In the appendix.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be analyzed. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
preliminary: 
 We first review the Latent Diffusion Model [2] used in our paper. It consists of an auto-encoder (VAE) and a UNet. The auto-encoder facilitates a two-way transformation between the RGB image I ∈ R H×W ×3 and the latent space z ∈ R h×w×c . Both the forward and backward processes of diffusion are carried out in the latent space, and we denote the noisy latent code at time t asz (t) = √ ᾱt z + √ 1 -ᾱt ϵ, where ᾱt = t s=1 (1 -β s )is the noise schedule. β s is the variance sampled from a variance schedule β t ∈ (0, 1) T t=1 . The UNet can be considered as a series of equally weighted denoiser ϵ θ (z (t) , t). The training objective L can be simplified as:L = E z,ϵ∼N(0,1),t∈U(T ) ϵ -ϵ θ z (t) , t 2 2 (1)Furthermore, to simplify comprehension and narration, we can reparametrize the output of UNet ϵ θ as the form of v-prediciton v θ . The training objective can be further elaborated as:L = E z,ϵ∼N(0,1),t∈U(T ) z -v θ z (t) , t 2 2 (2)This implies that the goal of every training round is to denoise z (t) to z for any time step t.Secondly, we present our task definition, using one-shot segmentation as an illustration. Given a data triplet (I s , M s , I q ), here I s and I q denote the support image and query image respectively, both sharing an overlapping category c. M s is the mask of category c in the support image. Our task is to predict the mask corresponding to category c in I q . In the strict one-shot segmentation setting, the category sets of the training set and the test set are disjoint.Our objective is to fully utilize the priors in the Latent Diffusion Model and equip it with Few-shot Semantic Segmentation capabilities. This leads us to reuse the original VAE to convert I s , I q and M q into latent variables z s , z q and z mq . Thus, our task is further simplified to explore how to improve the structure of UNet to v * θ so that it can accept z s , z q and M s as inputs, and use z mq as supervision. This supervised approach in the latent space has been certified effective in tasks such as depth estimation [13] and semantic segmentation [14] . Concretely, our training objective L FSS is transformed into:EQUATIONwhere D represents the constructed training dataset. In addition, we omitted the input of time t. Our early experiments revealed that performing multiple steps of noise addition and denoising during training did not bring performance improvement. 𝐗 ! 𝐗 !"# 𝐊 ! 𝐐 ! 𝐕 ! 𝐊 !" 𝐕 !" 𝐕 " 𝐊 " 𝐈 " 𝐌 " 𝐈 ! 𝐌 ! 𝐳 " 𝐳 ! 𝐳 #! 𝐳 #" 𝐳 " , 𝐳 #" 𝐳 " , 𝐳 #" 𝐳 ! 𝐳 #! ( 𝐳 #! ( 𝐳 #!Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
method: 
 Our investigation into model design primarily adheres to two criteria: 1. Strive for the design to be as simple and efficient as possible, while optimizing performance in Few-shot Semantic Segmentation. 2. Maximize the preservation of the Latent Diffusion Model's generative schema, minimizing alteration to the original UNet structure, so as to better utilize the pre-training prior. Specifically, four key issues need to be addressed: 1) How to facilitate interaction between the query image and support image? 2) How to effectively incorporate information from the support mask? 3) What form of supervision from the query mask would be most reasonable? 4) How to design an effective generation process to transfer the pre-trained diffusion models to mask prediction task? In this section, we discuss the four issues mentioned above in detail. We engage in fair comparison tests and analysis on several feasible strategies. Drawing on our observations, we eventually settle on our framework, DiffewS (see Figure 1 ).We first decompose the block of the l-th layer in UNet into three components: a self-attention layer SelfAttn, a cross-attention layer CrossAttn, and a feedforward layer FFN. Given the feature map X l of the l-th image and the textual input t (which is an empty character in our task), we obtain:EQUATIONwhere CLIP text represents CLIP text encoder, and we have skipped over skip-connection in the formula.Before considering the incorporation of the support mask, two straightforward and intuitive methods can be leveraged to facilitate interaction between the query image and support image. One approach entails interaction within the self-attention module, while the other involves interaction within the cross-attention module. KV Fusion Self-Attention. We first propose a KV fusion method in self-attention layer to achieve interaction between query image and support image. For the input image feature X, the standard self-attention layer first maps it to query Q, key K and value V with a linear projection layer. . Therefore, SelfAttn(X) can be further represented as:EQUATION)where d is the dimension of query and key, while X * is the feature updated by self-attention. Back to our task, we can also map the features of the support image and query image X s and X q to Q s , K s , V s and Q q , K q , V q through the linear projection layer. We hope that the features of the query image can effectively utilize the information of the support image, so we need to let Q q access K s and V s . To achieve this, we can concatenate K q and K s to form K qs = [K q , K s ]. Similarly, we can get V qs = [V q , V s ]. Finally, our KV Fusion Self-Attention layer can be represented as:EQUATIONSince we only replaced K and V, we can fully reuse the weights of the original self-attention.Tokenized Interaction Cross-Attention The second alternative is to inject information originating from the support image via cross-attention. This strategy has been widely used in Customized Text-to-Image Generation [52] [53] [54] . In particular, the initial cross-attention is employed to introduce the text information, encoded using CLIP text encoder. We can encode the support image into a series of tokens using the CLIP image encoder and utilize it as the cross-attention input. At this point, the process can be represented as:EQUATIONwhere Flatten means flattening the token sequence after image encoding. CLIP img represents the CLIP image encoder corresponding to the CLIP text encoder used in the original UNet.Building upon the Self-attention kv fusion approach, we investigate methodologies for incorporating support mask information. We categorize the injection methods into four types: a. Concatenation The support mask M s can be converted into an RGB image, then directly encoded into a latent variable z ms using VAE, which is then concatenated with z s in the channel dimension. Due to the resulting mismatch in dimensionality from the concatenation, we adopt the approach of Marigold [13] , where the first layer weight tensor is duplicated and its values are halved. b. Multiplication We can directly multiply M s on the image I s to form the image I * s = I s •M s , and finally encode I * s into a latent variable z * s using VAE as the input of UNet. c. Attention Mask M s can serve as an attention mask to control self-attention so that only K s in the masked region can be accessed by Q q . Since the feature map sizes of different layers are different, we need to resize M s to fit the dimensions of each layer. d. Addition Alternatively, M s can be directly added to the image I s , generating the image I * s = 0.5I s + 0.5M s . Following that, I * s is encoded into a latent variable z * s using VAE, which is then used as the UNet input. For cross-attention tokenized interaction, information of the support mask can also be injected in the same four ways. There are just some slight differences in the implementation details (see the Appendix A.3).We carry out a comparison of two interaction methods (Section 4.1) paired with four injection methods (Section 4.2); these eight combinations are then verified experimentally, and the results are presented in Figure 2 . Overall, we observe that KV Fusion Self-Attention(FSA) outperforms Tokenized Interaction Cross-Attention(TCA). We attribute this mainly to the preservation and flexible utilization of information from the support image by FSA. Conversely, TCA, which only compresses support image to tokens via the CLIP image encoder, leads to some information loss. Notably, within the FSA, the Concatenation method surpassed the other three. It offered a more free-form handling of RGB images and MASK information via subsequent learnable convolutional layers, compared to other hard injection methods. In the case of TCA, the Attention Mask method seems more apt as other operations are actually constrained by the CLIP image encoder. The CLIP image encoder itself is not good at dealing with mask information. Of course, we believe that there is still room for further exploration here, referring to FGVP [55] . In Section 3, we mentioned that we encode the query mask M q into a latent variable z mq , and directly supervise in the latent space. However, M q ∈ [0, 1] H×W is a two-dimensional mask, while the input of VAE needs to be an RGB image. Consequently, conversion of M q into an RGB image becomes necessary, but it's unclear which form of conversion would yield optimal results as no research has delved into this as yet. A reasonable conversion method should satisfy the following two conditions:1.It is easier for UNet to learn 2. It is more convenient to get the final segmentation result through post-processing. In this section, we explore the following four forms of conversion.a. White foreground + black background Visualizing the segmentation annotation with a white mask and black background is a common way in the academic community. Specifically, we only need to replicate M q three thrice along the channel dimension to form the corresponding RGB image denoted by the mask. We employed this conversion approach as a default in Section 4.2.b. Real foreground + black background Considering LDM's original pre-training on real images, forcing the model to output purely black-and-white images that do not fit within real-image distribution might amplify the model's learning difficulty. Therefore, we also attempted to retain the real pixels of the foreground, while setting the background to black c. Black foreground + real background Following the same logic, we also try preserving the pixels of the real background but render the foreground pixel black.We also consider overlaying M q on the real image to form the mask on the real image, which is the Addition method mentioned in Section 4.2. This approach makes the output space of UNet closer to the distribution of real images, but it requires more complex post-processing to get the final segmentation results. That is, we need to subtract the original image from the model output to get the final segmentation result.As shown in Figure 3 , we assess the performance of the four forms of supervision, among which (a) method achieved the best performance in all experiments. Although (b) (c) (d) methods being closer to the real image distribution, the performance is lower. On the one hand, it is difficult to obtain the mask through simple post-processing, and on the other hand, it may increase the learning difficulty because the model needs to retain the ability to generate the original image. In conclusion, our results demonstrate that UNetUNet can effortlessly learn to output in forms such as 'white foreground + black background'. Therefore, we eventually chose this approach for Diffews.In this section, we further discuss how to design an effective generation process to transfer the pre-trained diffusion models to mask prediction tasks. Inspired by the success of transferring pretrained diffusion models to depth estimation task [13, 56] , we explore three different mask generation processes. The illustration of different mask generation processes is shown in Figure 4 .• Multi-step noise-to-mask generation (MN2M) MN2M follows the denoise pipeline of original diffusion models. The training and inference schemes of MN2M are similar to Marigold [13] . (b) Illustrations of different pipelines of mask generation process in latent space. concatenates with the mask latent ẑ(t) mq . The UNet takes it as input and predicts the new mask latent ẑ(t-1)𝐈 ! Prediction 𝐳 ! 𝐳 ! 𝐳 ! O I2 M M I2 M 𝛽 ! M I2 M 𝛽 " M N 2 M 𝛽 ! (c) Comparisons of different pipelines of mask generation process. # 𝐳 "! ($%&) # 𝐳 "! $ # 𝐳 "! ($%&) # 𝐳 "! U-Net U-Net U-Net U-Netmq . After T steps, the final mask latent ẑ(0) mq is decoded into mask prediction. The mask latent ẑ(T ) mq is initialized as random noise. We also use the annealed multi-resolution noise and test-time ensemble tricks [13] proposed in Marigold.• Multi-step image-to-mask generation (MI2M) MI2M formulates the diffusion denoising process as a deterministic multi-step conversion process from image to prediction, similar to DMP [56] . Figure 4 (b2) shows the illustration of inference process. The mask latent ẑ(T ) mq is initialized as image latent z q . Then similar to MN2M, the UNet takes ẑ(t) mq as input and predicts ẑ(t-1)mq . After T steps, the final mask latent ẑ(0) mq is decoded into mask prediction. • One-step image-to-mask generation (OI2M) OI2M further transforms MI2M's multi-step prediction into a one-step prediction, i.e., UNet takes z q as input and outputs the prediction ẑmq directly.We explore the mask generation pipeline starting from MN2M. As shown in Figure 4 (c), MN2N achieves 15.2% mIoU. Then, we change MN2M into MI2M keeping same variance β 1 = (0.00085, 0.012), respectively representing the initial and final values of β in the DDIM scheduler. The performance has improved by 4.7% mIoU. However, despite the improvement, both methods exhibit suboptimal performance. We hypothesize that this is because adding a very small noise or image to the binary mask during the training process and then predicting it does not lead to a challenging task compared with diffusion pre-training.We hypothesize that the suboptimal performance is due to the minimal noise or image added to the binary mask during training, which results in an insufficiently challenging task compared to diffusion pre-training. The binary mask is inherently simpler than natural images, and even after adding noise, the latent mask can still easily distinguish between the foreground and background. This simplicity causes significant information leakage during UNet training, ultimately leading to poor performance.To verify this hypothesis, we increase the variance of MI2M from β 1 = (0.00085, 0.012) to β 2 = (0.0272, 0.384). The performance has significantly improved by 23.3% mIoU. To fully increase the challenge of training, we convert MI2M into OI2M, which does not introduce any ground-truth information into the input of the UNet during training. Additionally, OI2M reduces the number of iterations to one, significantly boosting the network's predictive efficiency. As shown in Figure 4 (c), OI2M achieves the best performance, making it the preferred choice for the mask generation pipeline.So far, we have primarily explored the training and inference processes specifically designed for 1-shot scenarios. A natural question arises: can this framework be extended to n-shot settings? To address this, we first present the simplest and most straightforward method for adaptation, which requires only minor modifications during the inference phase to accommodate n-shot tasks.In the Section 4.1, we introduced how to inject the information of the support image into the features of the query image using the KV Fusion Self-Attention method. In inference, our support set S may contain more than one image, S = {I s1 , I s2 , ..., I sn }. We encode each image into the features X si . Correspondingly, after mapping, we can obtain a series of Q si , K si , V si and Q qi , K qi , V qi . We can concatenate K qi and K si to form K qs = [K qi , K s1 , K s2 , ..., K sn ], and similarly we can obtainV qs = [V qi , V s1 , V s2 , ..., V sn ].Finally, our kv fusion self attention layer can be represented as:EQUATIONWhile the aforementioned solutions enable N-shot inference, their performance does not match that of state-of-the-art (SOTA) models. This discrepancy primarily arises because the model receives only a single support image during the training phase, which leads to inconsistencies when transitioning to the inference phase with 5-shot or 10-shot configurations.To address this issue, we explore improvements from both the inference and training perspectives.From the perspective of inference, transitioning from 1-shot to N-shot involves concatenating the keys and values of additional support samples, which significantly increases the number of keys and values processed during inference. To address this, we implement random sampling of the keys and values from the support samples during inference, ensuring that their quantity matches that of the training phase (see Table 6 ). Another more straightforward idea is to introduce multiple support samples during the training phase. In this way, the model can learn how to utilize multiple support images during training. we randomly select 1 to N support samples as input using KV Fusion in Equation ( 8) during a single training iteration (see Table 7 ).Our experiments demonstrate that improvements during the training phase are more effective than those during the inference phase. Therefore, we include the results of the model with training phase improvements in Table 2 . 
experiment: 
 Datasets We test our method in two settings: 1. Strict few shot setting: Following the few-shot setting on COCO-20 i [57], we organize 80 classes from COCO2014 [58] into 4 folds. Each trial consists of 60 classes allocated for training and 20 classes designated for testing. For evaluation, we randomly sample 1000 reference-target pairs in each fold with the same seed used in HSNet [18] . 2. In-context setting: Following the setting in SegGPT [35] , COCO, ADE [59] , and PASCAL VOC [60] We first compare DiffewS with other generalist models such as Painter [34] , SegGPT [34] , PerSAM-F [63] , and Matcher [22] as well as specialist models like HSNet [18] , VAT [61] , FPTrans [62] .Regarding the specialist models, we directly refer to the results presented within the SegGPT [35] and Matcher [22] research papers. These specialist models are also trained on the test categories from COCO [58] and PASCAL VOC [60] . We employ COCO-20 [21] and DINO V2 [23] ), and SAM itself is pre-trained on an exhaustive, finely annotated segmentation dataset. On the other hand, DiffewS undergoes fine-tuning on a relatively smaller quantity of segmentation data for limited iterations, still delivering performance that rivals Matcher. This indicates that using the paradigm of DiffewS, there is potential to achieve significant breakthroughs in the segmentation field if further trained on larger-scale segmentation data. It should be noted that the improvement of DiffewS in 5-shot is not significant, with a 4.0 distinct improvement only on LVIS-92 i . This might be due to the presence of many small objects in the support images of LVIS, so increasing the number of support images can alleviate this problem. Conversely, the DiffewS 5-shot performance on PASCAL-5 i [44] is slightly deficient compared to the 1-shot. This could be ascribed to the presence of relatively larger and more simplistic objects within PASCAL VOC's support images, inputting more images might interfere with the original architecture of the model. In this case, we do not apply the improvement strategies discussed in Section 4.5, therefore, the relatively weaker performance in the 5-shot scenario is reasonable.We also undertake validation of DiffewS under the standard few-shot setting, comparing it with other specialist models such as HSNet [18] , CyCTR [64] , VAT [61] , BAM [65] , HDMNet [20] , and DCAMA [19] . For the one-shot setting, the average performance of DiffewS across all four folds attains 51.2, surpassing the current state-of-the-art (SOTA) model DCAMA, scoring 50.9 mIoU. Worth mentioning is that DCAMA relies on a highly complex additional block, whereas DiffewS entirely utilizes the generative framework of UNet. In terms of the efficiency of convergence, DiffewS necessitates just a 30000-iteration training, in contrast to both DCAMA and HSNet which require training spanning hundreds of epochs, typically costing several days. This demonstrates the successful employment of Stable Diffusion priors by DiffewS, thereby securing impressive performance without requiring extended periods of fine-tuning. In the five-shot setting, the average performance across four folds reaches 56.0, higher than all other models aside from DCAMA. Currently, DiffewS primarily focuses on the 1-shot situation lacking specific optimizations for the 5-shot scenario in its training and inference systems. This explains why DiffewS is at present marginally inferior to DCAMA. Furthermore, when employing our proposed training improvement strategy, DiffewS-n outperforms other models in both the 1-shot and 5-shot settings.As shown in Figure 5 , DiffewS effectively segments categories not in the training set, such as slippers and aprons. It also accurately segments objects of different styles and smaller items, demonstrating strong generalization capabilities. In some cases, DiffewS even achieves more accurate results than ground-truth human annotation.In addition, DiffewS demonstrates impressive results in various cross-style segmentation tasks and small object segmentation cases (see Figure 6 ). We hypothesize that DiffewS's exceptional generalization ability stems from its extensive utilization of prior knowledge from diffusion models. However, DiffewS also struggles with certain challenging cases, we also present several failure cases in Figure 7 and categorize the reasons for these failures. 
conclusion: 
 In this work, we have presented DiffewS, a simple and efficient framework for few-shot semantic segmentation. By directly generating the target mask, DiffewS is capable of retaining the original latent diffusion models' generative framework and effectively utilizing the visual prior of pre-trained diffusion models. By introducing several designs about multi-image interaction, information injection, and supervision signals, DiffewS outperforms SOTA models in the in-context learning setting, and reaches comparable performance to specialist models in the strict few-shot setting.Limitation and more Discussions are provided in Appendix A.1.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We clearly described the training and test details in the setup section and in the training details of appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: Like the previous works we follow, we do not report error bars. The overhead of retraining with different random seeds is significant. If computational resources permit, we will supplement this in the future.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: Yes, we discuss these in the discussion. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: All the thing in this paper credited and are the license and terms of use explicitly mentioned and properly respected Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: A.1 Discussion Broader Impacts We do not foresee any obvious undesirable ethical or social impacts now.Limitations Our method, as the first diffusion-based FSS model, proposes a simple and intuitive design, which maximizes the retention of the generative framework of LDM. There is still a lot of room for improvement in performance (especially in the n-shot setting), including more sophisticated model design and more optimized training strategies. We hope that our method can serve as a diffusion-based FSS baseline to inspire more researchers to invest in this field.On the other hand, we believe that our method is not limited to FSS. Our framework has the potential to unify few-shot segmentation and open vocabulary segmentation by leveraging prompts from different modalities, as some work [66] [67] [68] has already proven the possibility.In the above section we have discussed the generation process of DiffewS. In addition to the final choice of OI2M, we also tried MN2M and MI2M. Here we detail the training objectives of these three generation processes.OI2M We directly input the image and let the UNet output the mask. This process can be described as:EQUATIONMN2M We add noise to query mask z mq , z(t) mq = √ ᾱt z mq + √ 1 -ᾱt ϵ,and during inference we use z (0) mq as the mask prediction. The supervised form is as follows:EQUATIONMI2M We add image(as noise) to the query mask z mq , z(t) mq = √ ᾱt z mq + √ 1 -ᾱt z q .The supervised form is as follows:EQUATIONIn the Section 4.2, we only discussed how to inject information from the support mask based on the Self-attention kv fusion method. Here we discuss how to inject information from the support mask based on the Tokenized Interaction Cross-Attention method. There are also the following four ways.a. Concatenation We can convert the support mask M s into an RGB image, encode I s and M s into token sequences using CLIP image encoder respectively, concatenate them on the sequence, and finally use them as the input of cross-attention. b. Multiplication We can directly multiply M s on the image I s to form the image I * s = I s •M s , and finally encode I * s into a token sequence using CLIP image encoder as the input of crossattention. c. Addition We can also directly add M s to the image I s to form the image I * s = 0.5I s + 0.5M s . Similarly, we encode I * s into a token sequence using CLIP image encoder as the input of cross-attention. d. Attention Mask We can use M s as an attention mask to control self-attention, so that only K s in the masked area can be accessed by Q q .The original prediction of the model is an RGB three-channel image. We first average over the channel dimension to obtain a single-channel Mq ∈ [0, 1] H×W . Then we tried two thresholding methods, absolute threshold τ a and relative threshold τ r . The absolute threshold is a fixed value, and the final binary mask M q can be represented as:EQUATIONUsing relative threshold, we have: Our experiments (see Table 3 ) have shown that the relative threshold method achieved better results on COCO-20 i [57] fold-0. The optimal τ r is 0.25.EQUATIONMultiplication We found in the experiment that Multiplication can be directly applied to RGB images, and another choice is to apply it to the latent space. As shown in Table 4 , the Multiplication method directly applied to RGB images achieved better results. However, the overall disparity is not significant.In previous sections, we mentioned that we use a KV fusion strategy. An alternative is to use a QKV fusion strategy, in which we also concatenate Q q and Q s to formQ qs = [Q q , Q s ].This strategies means the support image can also access the query image information. As shown in 5 , KV fusion is slightly better than QKV fusion, and KV fusion has lower computational complexity, which can effectively reduce memory usage and inference time. Therefore, we choose KV fusion as our default strategy.To better explore the capabilities of DiffewS, we visualize its performance on COCO-20 i [57] LVIS-92 i [22] and several cases from Internet. Figure 8 shows the remarkable results of DiffewS on COCO-20 i . Figure 6 demonstrates the impressive generalization capabilities of DiffewS. For some categories not present in the training set, such as apron and violin, DiffewS is able to perform accurate segmentation. In addition, DiffewS is demonstrated effective results in some cross-style segmentation and small object segmentation cases. For abstract concepts, such as Western dragons and Chinese dragons, DiffewS links them together to achieve accurate results. We speculate that the impressive generalization ability of DiffewS stems from its effective utilization of prior knowledge from the diffusion model. As shown in Figure 7 , DiffewS also fails to segment some challenging cases. When there is a significant appearance disparity between the reference image and the target image (Appearance disparity), DiffewS may encounter segmentation errors. Additionally, if there are other objects with similar appearances in the target image (Look-alike interference) or if the objects in the image are severely occluded (Occlusion interference), DiffewS struggles to produce accurate results.Reference Prediction Reference Prediction Reference Prediction As mentioned in Section 4.5, we conduct Inference Time Improvement on COCO and PASCAL, Training Time Improvement on COCO fold-0, see Table 6 and Table 7 . Justification: In the abstract and the end of the introduction.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: The paper introduces a latent diffusion model (DM)-based method for few-shot segmentation (FSS). It develops several effective solutions from four key aspects of applying DM to FSS and achieves strong performance on several benchmarks. 

The paper received mostly positive reviews. The three reviewers considered the proposed approach novel and effective (WLLP, zc26), the paper well-written (WLLP, zc26), the study of its solutions comprehensive (WLLP, uveh) and the experimental results strong/convincing (WLLP, uveh, zc26). However, they initially also raised several important concerns regarding its mixed performance under multi-shot settings (WLLP, uveh, zc26), insufficient motivation (uveh), and the lack of clarity in technical details and settings (uveh, zc26). The author's rebuttal provided thorough clarifications on its motivation and technical details, and additional results on 5 and 10-shot settings, which largely addressed the reviewers' concerns. During the rebuttal discussion, the three reviewers reached a consensus with positive recommendations on this work.  

After considering the paper, the accompanying reviews, and the rebuttal discussion, despite some limitations in the multi-shot setting, the AC concurs with the reviewers regarding the substantive contributions of this work and therefore recommends it for acceptance. The author should take into account the rebuttal and the reviewers' feedback, especially regarding the refined 5 and 10-shot results, when preparing the final version.

Reviews: review: paper_summary: paper_summary: This paper introduces DiffewS, a novel Diffusion-based generalist model designed for few-shot segmentation. It systematically examines four key components involved in applying Diffusion Models to Few-shot Semantic Segmentation. For each component, the work proposes several viable solutions, which are validated through extensive experiments. Notably, the work introduces Key-Value Fusion Self-Attention (FSA) to facilitate the interaction between the query and support images.
strength_weakness: strength_weakness: The method of incorporating support masks, adapted from Marigold, may have limitations in a few-shot setting. The experiments detailed in Table 1 indicate that DiffewS achieves only marginal improvements in the COCO few-shot scenario compared to the one-shot scenario and even exhibits a decline in performance on the PASCAL dataset. Further evaluation of DiffewS’s capabilities with additional support samples, such as through experiments in a ten-shot setting, would be beneficial.
questions: questions: I am highly intrigued by the possibility of whether DiffewS can perform open-vocabulary segmentation by substituting the support mask with text.
limitations: limitations: The limitation is discussed in the paper supplementary material.
paper_summary: paper_summary: This paper explores the potential of diffusion model in Few-Shot Segmentation (FSS) tasks. For achieve better performance, it examines four critical elements of applying the diffusion model to FSS. Building on this research, the paper introduces the DiffewS framework, with experimental results validating its effectiveness. The authors assert that this is the first diffusion-based model specifically designed for Few-shot Semantic Segmentation.
strength_weakness: strength_weakness: The weakness of this article is as follows:

1. Why should the diffusion model be applied to FSS tasks? What advantages does DM offer over Multi-modal fundamental models (like SAM or LLaVA) for FSS tasks? This article's research motivation requires stronger support through theoretical explanations or experimental evidence.
2. The default generation process (OI2M) transforms "gradually approaching ground truth in denoising" to "single-step prediction in segmentation", yielding optimal performance. Does this suggest that DM's multi-step denoising process is unsuitable for FSS? Does this paper not unleash the potential of DM, instead transforming it into a segmentation model? Section 4.4 should include more theoretical explanations.
3. The comparison of inference speed and training cost is important and should be supplemented to illustrate the advantages of DiffewS.
4. Other studies that apply the Diffusion Model to Few-Shot Segmentation (e.g., DifFSS[1], Ref LDM-Seg[2]) need to be compared. 

[1] DifFSS: Diffusion Model for Few-Shot Semantic Segmentation. Github: https://github.com/donglongzi/DifFSS-BAM/tree/main

[2] Explore In-Context Segmentation via Latent Diffusion Models. Project Page: https://wang-chaoyang.github.io/project/refldmseg/
questions: questions: 1. Did the pretrained Stable Diffusion (SD) weights utilized in the model training? How much parameters were updated? 
2. What are the experimental settings of the comparing results in Fig2/3/4?
3. In Section 4.4, this paper chose OI2M as the generation pipeline. Since it is a one-step process, what is the timestep setting in this pipeline during training? And did any accelerated generation schedule (such as DDIM) is utilized?
4. Did the few-shot support images being processed as the same as one-shot support image in support image encoding branch? What is the difference of model settings between using few-shot support images and one-shot support image? How about the comparison of computational cost?
5. In Section 5.2, the “In-context” setting means “these specialist models are also trained on the test categories from COCO and PASCAL VOC”. Is it a general experimental dataset setting? And is it a fair comparison for the methods in Table 1?
limitations: limitations: No
paper_summary: paper_summary: This paper explores the potential of the Latent Diffusion Model for Few-shot Semantic Segmentation. The authors study four crucial elements of applying the Diffusion Model to Few-shot Semantic Segmentation and propose several reasonable solutions for each aspect. 
Based on their observations, the authors establish the DiffewS framework, which maximally retains the generative framework and effectively utilizes the pre-training prior. 
Experimental results demonstrate that the proposed method outperforms the previous SOTA models in multiple settings.
strength_weakness: strength_weakness: - Some places lack proper references. For example, in L110, "v-prediction," we typically cite Tim Salimans and Jonathan Ho's work: "Progressive distillation for fast sampling of diffusion models" (ICLR, 2022).
In L241, I noticed that the appendix contains an analysis of the training objectives for different generation processes, and a reference should be added here.

- In Section 4.4, L257, the variances $\beta_1$ and $\beta_2$  are not clearly defined. Since $\beta_t$ in L106 represents the variance at time t, it would be better to use $\beta^1$ and $\beta^2$  to denote these two different settings here. Additionally, the meaning of the two values within the tuple is not explained. In my view,  $\beta^1 = (\beta^{start},\beta^{end})$ should represent the initial and final values of the DDIM scheduler's $\beta$ .
Furthermore, it is necessary to explain the intuitive differences between these two settings.

- Other typos:
L273,  "we tested" -> "we test"
L151, "linear projection layer. ." ->"linear projection layer"
L279, "function as"-> "functions as"
questions: questions: I've noticed that the current method doesn't perform well enough on 5-shot, can this be mitigated by making adjustments during training?
limitations: limitations: NaN

Comments: comment: We sincerely appreciate your positive feedback and are grateful for your recognition of the merits of our work. Your constructive comments have been invaluable in enhancing the quality of our paper.

Thank you again for your time and expertise.
Thank the authors for the detailed rebuttal, which has satisfactorily addressed my concerns. I believe that the paper has merit and should be accepted. I will be keeping my original rating.
Thanks again for your thoughtful review and for your recognition of our efforts to address your concerns. We deeply appreciate your professional insights and the constructive feedback that has helped us improve our work.
Thanks to the author's thorough response, most of my concerns have been addressed. Therefore, I am upholding the original rating.
Thank you very much for your thoughtful review and for acknowledging our work. We are delighted that our responses addressed your concerns, and we greatly appreciate your recognition of the insights and contributions our work offers to the field. Your positive feedback is invaluable to us.

We also want to assure you that we will incorporate your feedback and any potential improvements into future versions of our paper.
I would like to thank the authors for their detailed responses. All my concerns were thoroughly addressed. Overall,  this work provides clear insights and significant contributions to the field. Therefore, I will maintain my original rating.