NeurIPS 2023 | 用扩散模型构建语义分割新数据，直接替代传统人工标注

随着AIGC逐渐进入大众视野，学术界也在考虑如何更加有效的利用生成式技术手段到具体的任务领域中。对于深度学习模型而言，第一要素便是准备有良好标注的数据集，这是一项非常耗时的工作。因此有研究者想到使用AIGC手段来合成可供模型训练的数据。目前最流行的生成方法莫过于大火的扩散模型Stable Diffusion（SD），本文提出了一种基于SD的从文本到图像的生成方法，同时提供与图像相对应的像素级语义分割标签。在原有扩散模型的基础上，本文作者提出了类提示叠加（class-prompt appending）、类提示交叉注意力（class-prompt cross-attention）和自注意力求幂（self-attention exponentiation）模块来构成同时生成图像和标签的pipeline。使用该框架生成的像素级标注图可以作为训练语义分割网络的伪标签。此外，为了提升生成数据的可信度，作者还引入了不确定性分析手段来解释生成伪标签中的缺陷区域，作者在PASCAL VOC和MSCOCO两个标准数据集上的实验证明了本文方法的有效性，本文的研究团队主要来自VinAI研究院，该文目前已被人工智能领域顶级会议NeurIPS 2023录用。

论文名称： Dataset Diffusion: Diffusion-based Synthetic Dataset Generation for Pixel-Level Semantic Segmentation 论文链接： https://arxiv.org/abs/2309.14303 代码仓库： https://github.com/VinAIResearch/Dataset-Diffusion （ 尚未开源）

一、介绍

本文主要关注的任务领域是语义分割，语义分割作为一项基础视觉任务，首先就需要非常精细的像素级标签，这种类型的标签标注势必会耗费大量的人力。为了解决这一问题，目前比较实用的方法是利用一些预训练生成模型（例如GAN）来生成图像标签，但是从GAN模型本身的特性来看，其主要集中于合成以对象为中心的图像，对于捕获真实场景中的背景语义存在一定的缺陷，此外GAN在接受语言提示方面的工作也非常有限。因此，本文作者选择了当前非常流行的文本到图像扩散模型，可以直接从文本描述出发来生成非常精细的逼真图像，例如DiffuMask[1]和DiffusionSeg[2]就可以根据简单的文本提示（例如“a photo of a [类名称][背景描述]”）来扩散生成语义标签，但是这些方法目前仅限于生成单个对象分割mask。

本文提出了一种称为数据集扩散（Dataset Diffusion）的新方法，整体的pipeline如上图所示，首先将输入的文本提示转换为扩散指导信息，其明确指定了模型扩散过程中所要生成的对象，然后根据扩散生成的语义分割标签以及不确定性感知分割损失训练分割模型，随后就可以使用该分割模型进行推理。作者在Dataset Diffusion中设计了一个类提示叠加模块，该模块直接以类名文本提示作为输入，可以促使模型在同一张图像中生成多种不同类别的对象。此外，作者还引入了一种全新的类提示交叉注意力模块，该模块可以在图像生成的过程中不断引导模型关注对象的空间线索，从而能够在生成的语义图中精确定位对象的具体位置。

二、本文方法

本文的整体框架可以视为一个三阶段方法，整体框架如下图所示，目标是合成一个新数据集 ，其由高分辨率图像和像素级语义Mask组成。这些图像和Mask需要捕获目标类别 的语义和位置信息 ，其中 表示类的数量，通过构建该数据集，可以在没有人工标注的情况下来训练语义分割器 ，这一流程可以称得上是研究界进入AIGC时代的一个典型的标志。

2.1 准备输入到扩散模型的文本提示

为了与现有的图像语义分割Benchmark保持一致，作者选择从COCO和PASCAL VOC数据集出发来生成， 在使用COCO数据集时，可以直接依靠数据集本身提供的标题来构造提示。对于PASCAL VOC数据集，作者使用BLIP模型为每个图像生成文本描述。 但在这一过程中作者发现，模型生成的文本提示与数据集中原本的类别名称 可能不完全相同，例如，模型可能会使用“男人”和“女人”等术语代替“人”，或者使用“bike”代替“bicycle”，从而目标类别不匹配。 其次，模型生成的描述并不 包含图像中实际存在的所有类别，如下图所示。这会导致某些类的文本提示不够充分，从而影响这些特定类的生成过程。

为了解决这些问题，作者提出了一种数据集类叠加提示方法，首先将生成的文本叠加到类标签上来创建新的文本提示 这可以保证其明确包含所有目标类 。类叠加提示的实现可以表示为： 。例如，在上图中的左侧图像中，最终的文本提示将是“房子内厨房的照片; 瓶,微波炉,水槽,冰箱”。这样就可以 确保新的文本提示包含所有目标类别，从而解决标题中类别名称不匹配或缺失的问题。

2.2 从自注意力图和交叉注意力图来生成语义分割图

本文作者提出可以直接使用扩散模型中的自注意力层和交叉注意力层构建分割图生成器， 给定文本提示 ，然后将其编码为文本嵌入 ，然后执行扩散过程。在每个去噪步骤 中，使用具有多层自注意力和交叉注意力的UNet架构将其嵌入编码 转换为 。其中，在层 和时间步 处，自注意力层可以有效的捕获内位置之间的相似性，以便通过中的全局上下文来增强局部特征。 同时，交叉注意力层可以对 的每个位置与文本嵌入 的每个token之间的关系进行建模。形式上，第 层和时间步 的自注意力图 和交叉注意力图 计算如下：

根据观察，设置不同幅度的时间步长只会对最终的生成结果产生轻微影响，因此作者直接对这些交叉和自注意力图在层和时间步长上进行平均扩散， 扩散过程如下：

下图展示了给定文本提示 “一辆自行车停在房间里：自行车”的生成效果，其中“自行车”为目标类别，从图中我们可以观察到，尽管交叉注意力图已经可以展示图像中目标类的位置，但它们仍然是粗粒度和有噪声的。

因此，作者提出使用自注意力图（如下图-左所示）来增强交叉注意力图使其获得更精确的物体位置， 这是由于捕获嵌入特征 内位置之间的成对相关性可以帮助将初始交叉注意力图传播到空间上非常相似的位置， 尤其是图像与对象背景中的非显著区域，从而提高整体的图像生成质量。

在具体操作时，作者提出了一种自注意力求幂运算，如下所示：

接下来，模型需要对一个特殊矩阵 中每个位置计算对象值， 对象性越高，该位置包含对象的可能性越大，同时还需要确定指示向量中每个位置最有可能是类标签中的哪些对象，直接执行像素级argmax和max运算符可以获得上述矩阵和指示向量：

随后，通过对矩阵 中的每个位置 进行阈值比较，如果其值小于阈值，则可以将其标签设置为背景类0。作者对背景决策使用较低的阈值 ，而对某些前景目标的决策使用较高的阈值，因此落在范围内的任何值都表示不确定的Mask预测，其值直接设置为255。综上所述，最终Mask 的计算公式为：

三、实验效果

本文的实验在两个标准图像语义分割数据集PASCAL VOC 2012和COCO 2017上进行，为了实现数据集扩充，作者首先对PASCAL VOC和COCO构建了文本提示集，并对其生成扩充后的验证数据集，分别称为synth-VOC和synth-COCO，为了使生成的数据集在类别之间达到分布平衡，作者为PASCAL VOC的每个对象类生成2k个图像，总共产生40k个图像Mask对，为COCO的每个对象类生成约1k个图像，总共产生80k个图像Mask对。如果其中某个类别的文本提示数量不足，就对其设置更多的随机种子来生成更多的图像。作者首先对整体的Dataset Diffusion框架进行了定量的实验对比，实验结果如下表所示。

上表中展示了DeepLabV3和Mask2Former在真实训练集、DiffuMask合成数据集和DatasetDiffusion 合成数据集上训练的结果。在PASCAL VOC上，本文方法实现了64.8 mIoU的语义分割效果，这相比先前的工作DiffuMask提升了4.2 mIoU。而在COCO数据集上，本文方法达到了34.2的分割效果，使用真实COCO数据集的结果为54.9 mIoU，可以看到在COCO数据集上，本文方法与真实数据集训练相比仍然具有一定的差距，作者分析，这可能是因为COCO中的图像内容比VOC的图像内容更复杂，从而降低了StableDiffusion生成具有相同复杂程度的图像的能力。

此外，作者还在上图中进行了本文方法的生成效果展示，在上图（a）中，作者直接使用合成图像及其对应的Mask训练语义分类器，图中前两行展示了两种较好的分割示例，后两行为失败示例。作者发现，在处理小尺寸目标时，扩散生成的Mask会对模型学习产生不利的影响。在上图（b）中展示了将上一步训练好的模型在真实VOC验证集上的预测结果，其中前三行展示出了令人满意的分割结果，而后三行也表明模型仍然无法对含有多个小目标和遮挡目标进行很好的分割。

四、结论

本文针对深度学习数据集构建问题提出了一种全新的框架DatasetDiffusion，通过引入目前流行的扩散模型，本文方法实现了从指定的对象类别生成高质量高语义的图像Mask对，为了实现这一目标，作者首先通过类提示叠加操作来将目标类别嵌入到扩散模型自身的隐藏空间中，随后从模型自注意力和交叉注意力图出发，通过自注意力求幂计算来提高合成图像的目标显著定位和位置感知效果。作者通过一系列的实验证明，使用DatasetDiffusion来生成数据并辅助训练的方法可以有效的提升模型的语义分割能力。这也表明，巧妙的利用AIGC技术完全可以为深度学习领域构建具有精细标注的大规模数据集，而减少对人力标注的依赖。

参考

[1] Weijia Wu, Yuzhong Zhao, Mike Zheng Shou, Hong Zhou, and Chunhua Shen. Diffumask: Synthesizing images with pixel-level annotations for semantic segmentation using diffusion models. arXiv preprint arXiv:2303.11681, 2023.

[2] Chaofan Ma, Yuhuan Yang, Chen Ju, Fei Zhang, Jinxiang Liu, Yu Wang, Ya Zhang, and Yanfeng Wang. Diffusionseg: Adapting diffusion towards unsupervised object discovery. arXiv preprint arXiv:2303.09813, 2023.

Illustration by IconScout From Pixel True

-The End-

扫码观看！

本周上新！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（www.techbeat.net）。社区上线500+期talk视频，3000+篇技术干货文章，方向覆盖CV/NLP/ML/Robotis等；每月定期举办顶会及其他线上交流活动，不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台，希望为AI人才打造更专业的服务和体验，加速并陪伴其成长。

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

chenhongyuan@thejiangmen.com

或添加工作人员微信（chemn493）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈

⤵一键送你进入TechBeat快乐星球