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

📄 arXiv: 2410.02369v3 📥 PDF

作者: Muzhi Zhu, Yang Liu, Zekai Luo, Chenchen Jing, Hao Chen, Guangkai Xu, Xinlong Wang, Chunhua Shen

分类: cs.CV

发布日期: 2024-10-03 (更新: 2024-10-29)

备注: Accepted to Proc. Annual Conference on Neural Information Processing Systems (NeurIPS) 2024. Webpage: https://github.com/aim-uofa/DiffewS

💡 一句话要点

提出DiffewS，利用扩散模型解决少样本语义分割问题

🎯 匹配领域: 支柱三：空间感知与语义 (Perception & Semantics) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 少样本学习 语义分割 扩散模型 潜在扩散模型 自注意力机制

📋 核心要点

  1. 现有少样本语义分割方法缺乏对扩散模型预训练先验的有效利用，限制了性能。
  2. 提出DiffewS框架，通过KV融合和优化信息注入，有效利用潜在扩散模型的生成框架和预训练先验。
  3. 实验表明，DiffewS在多个少样本语义分割设置下，显著超越了现有最优方法。

📝 摘要（中文）

扩散模型不仅在图像生成领域取得了显著成就，而且作为一种利用无标签数据的有效预训练方法也展现了潜力。受扩散模型在语义对应和开放词汇分割方面巨大潜力的启发，本文着手研究将潜在扩散模型应用于少样本语义分割。最近，受到大型语言模型上下文学习能力的启发，少样本语义分割已经演变为上下文分割任务，成为评估通用分割模型的关键要素。在此背景下，我们专注于少样本语义分割，为未来开发基于扩散的通用分割模型奠定坚实的基础。我们的初步重点是理解如何促进查询图像和支持图像之间的交互，从而在自注意力框架内提出一种KV融合方法。随后，我们深入研究如何优化支持掩码的信息注入，并同时重新评估如何提供来自查询掩码的合理监督。基于我们的分析，我们建立了一个简单有效的框架DiffewS，最大限度地保留了原始潜在扩散模型的生成框架，并有效地利用了预训练先验。实验结果表明，我们的方法在多种设置下显著优于之前的SOTA模型。

🔬 方法详解

问题定义：论文旨在解决少样本语义分割问题，即在只有少量标注样本的情况下，对图像进行像素级别的语义分割。现有方法通常难以充分利用无标签数据提供的先验知识，并且在查询图像和支持图像之间建立有效的关联仍然是一个挑战。

核心思路：论文的核心思路是利用潜在扩散模型（Latent Diffusion Model）强大的生成能力和预训练先验知识，将其应用于少样本语义分割任务。通过在扩散模型的框架内，有效地融合查询图像和支持图像的信息，并优化监督信号的注入方式，从而提升分割性能。

技术框架：DiffewS框架主要基于潜在扩散模型，包含以下几个关键模块：1) 图像编码器：将查询图像和支持图像编码到潜在空间；2) KV融合模块：在自注意力机制中融合查询图像和支持图像的特征，建立关联；3) 信息注入模块：优化支持掩码的信息注入方式，指导分割；4) 扩散模型解码器：将潜在空间表示解码为分割结果。整个流程是在扩散模型的生成框架内进行的，最大限度地保留了其预训练先验。

关键创新：论文的关键创新在于：1) 提出了KV融合方法，用于在自注意力框架内有效地融合查询图像和支持图像的特征；2) 优化了支持掩码的信息注入方式，使其能够更好地指导分割过程；3) 构建了DiffewS框架，该框架能够最大限度地保留原始潜在扩散模型的生成框架，并有效地利用预训练先验。

关键设计：KV融合模块通过修改自注意力机制中的Key和Value，将支持图像的信息融入到查询图像的特征中。信息注入模块通过调整损失函数的权重，平衡来自查询掩码和支持掩码的监督信号。扩散模型的具体参数设置和训练策略沿用了原始论文的设置，以保证预训练先验的有效性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DiffewS在多个少样本语义分割数据集上取得了显著的性能提升。例如，在COCO数据集上，DiffewS相比之前的SOTA模型，在1-shot和5-shot设置下分别提升了超过5%和3%的mIoU。实验结果表明，DiffewS能够有效地利用扩散模型的预训练先验，并实现更好的分割性能。

🎯 应用场景

该研究成果可应用于医疗图像分析、遥感图像处理、自动驾驶等领域。在这些领域中，标注数据通常非常稀缺，而少样本语义分割技术可以有效地利用少量标注数据和大量无标注数据，提高分割精度和效率。未来，该研究有望推动通用分割模型的发展，使其能够适应各种不同的分割任务。

📄 摘要（原文）

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.