CVPR23 Highlight｜扩散模型也能用来做分割？UCSD与NVIDIA提出全景分割新SOTA

论文地址：https://arxiv.org/abs/2303.04803

简介

在这篇文章中，作者提出了ODISE：基于扩散模型的开放词汇全景分割（Open-vocabulary DIffusion-based panoptic SEgmentation），该模型将预训练的文本-图像扩散模型与判别模型统一起来，用于开放词汇全景分割任务。作者认为:文本-图像扩散模型具有出色的能力，可以根据各种开放词汇语言描述生成高质量图像。这表明它们内部的表征空间与现实世界中的开放概念高度相关。另一方面，像CLIP这样的文本-图像判别模型擅长将图像分类为开放词汇标签。因此作者利用冻结的扩散模型与判别模型来实现开放类别的全景分割。ODISE在开放词汇全景分割和语义分割任务上的表现优于以前的最先进方法。值得注意的是，该模型仅使用COCO训练就能在ADE20K数据集上达到23.4 PQ和30.0 mIoU，相对于以前的最先进方法，PQ和mIoU分别提高了8.3和7.9。

Introduction

人类可以看到世界并识别出无限的类别。例如在上图中，除了将每辆车辆都识别为“卡车”之外，我们也可以立即理解其中一辆卡车需要挂接拖车才能移动另一辆卡车。为了研究这种细粒度和近似无限理解能力的智能，开放词汇认知问题最近在计算机视觉领域引起了很大关注。然而，只有少数工作能够提供一个统一的框架，同时解析所有对象实例和场景语义，也即全景分割。ODISE是首个结合扩散模型与判别模型的全景分割方法。至于为什么要结合这两个模型，作者给出了以下观点：

首先，目前大多数开放词汇认知方法都应用了在大规模数据上训练的文本-图像判别模型（例如CLIP）。多个研究也证明了这样的模型在开放词汇任务中的潜力。然而，作者认为虽然这样的预训练模型擅长分类单独图像，但它们不一定是场景级别，或者像素级别结构理解的最佳选择。事实上，研究已经表明CLIP经常混淆对象之间的空间关系（参考Reclip）。在本文中作者也同样假设文本-图像判别模型中缺乏空间和关系理解是开放词汇全景分割的瓶颈。

基于这样的思考，作者尝试引入扩散模型。最近诞生的许多扩散模型提供了前所未有的图像质量、泛化能力、以及语义控制能力。为了将生成的图像限制在提供的文本上，扩散模型在文本的Embeddings和它们的视觉表示之间计算交叉注意力。这种设计意味着扩散模型潜在的特征可能与语言描述的高级语义概念有良好的区分和相关性。同时，由于需要生成正确的图片，扩散模型对于图片的空间信息也会有更好的理解。

基于以上的思考，作者提出了ODISE：基于扩散的开放词汇全景分割（Open-vocabulary DIffusion-based panoptic SEgmentation），该模型利用大规模文本-图像扩散和判别模型来完成全景分割。该方法概述如下图所示。

它包含一个预训练的冻结的文本-图像扩散模型，模型接受一个图像及其标题作为输入，并提取其扩散模型的内部特征。模型使用这些特征作为输入，其中的掩模生成器生成图像中所有可能对象的掩膜。作者使用训练集中带注释的掩模来训练掩模生成器。而掩膜分类器通过将每个预测掩模的Diffusion特征与几个对象类别名称的文本Embeddings相关联，将每个掩模分类为许多开放词汇类别之一。作者使用训练数据集中的掩模类别标签或图像级标题来训练分类模块。训练完成后，用文本-图像扩散和判别模型执行开放词汇全景推理以对预测掩模进行分类。实验证明，作者的方法大幅超过了现有的baseline方法。

总的来说，作者的贡献如下：

• 首次提出了将扩散模型用于开放词汇分割任务的方法。
• 提出了一种新的流程，有效利用文本-图像扩散模型和判别模型执行开放词汇全景分割。
• 作者的方法在许多开放词汇识别任务上超越了所有现有的baseline方法，从而达到了新的SOTA。

Method

1. 问题定义

首先是对问题的定义，这关系到训练与测试时的处理。全景分割要求测试集中的类别是开放的。作者在一组基础训练类别 训练基础模型，这些训练类别可能与测试类别 不同，即 。 可能包含在训练期间没有看到的新类别。同时在训练期间，对于图像中的每个类别都提供了二值的掩膜注释。此外，每个掩膜的类别标签或图像的文本标题都是可用的。在测试期间，对于任何图像，都不提供类别标签或标题，只提供测试类别 的名称。

2. 扩散模型的引入

目前的扩散模型通常使用UNet架构来学习去噪过程。UNet由卷积块、上采样和下采样块、跳跃连接和注意力块组成，它们在文本Embeddings和UNet特征之间执行进行交叉注意力计算。在去噪过程中，扩散模型使用文本输入来推断噪声输入图像的去噪方向。文本是通过交叉注意力层注入到模型中的，作者这样的设计是为了让视觉特征与丰富的语义描述相关联。因此，UNet块输出的特征图可以为全景分割提供丰富而密集的特征。

作者的方法只需要通过扩散模型的单次前向传播来提取其视觉表示。形式上，给定一个输入图像-文本对 ，首先在时间 处采样噪声图像 ，如下所示：

其中 是diffusion step，其中 。作者使用预训练的文本编码器 对标题 进行编码，并通过将其输入到UNet中提取内部特征 ：

作者同样也指出，扩散模型对于 的视觉表示 是依赖于其对应的标题 的。当配对的图像-文本数据可用时，例如在文本到图像扩散模型的预训练过程中，可以正确提取它。然而，当我们想要提取没有配对标题的图像的视觉表示时，原来的方法就不行了，但是这是全景分割中常见的问题。因此对于没有标题的图像，作者引入了一个可训练网络来从输入图像本身生成一个隐式的文本Embeddings。这个Embeddings会被输入到扩散模型中用于产生特征。作者将这个模块命名为Implicit Captioner。

最后，文本-图像扩散模型的UNet和Implicit Captioner一起形成ODISE的特征提取器，用于计算输入图像 的视觉表示 。最后，计算的视觉表示 如下：

3. 掩膜生成器

掩膜生成器以视觉表示 为输入，输出 个没有类别的二进制掩膜 及其相应的 个掩膜特征 。掩膜生成器的架构可以是任何能够生成整个图像掩码预测的全景分割网络。作者使用了直接分割的架构。按照Maskrcnn的方法，通过像素级二元交叉熵损失以及相应的真实掩膜来监督预测的没有类别信息的二进制掩膜。

4. 掩膜分类器

为了对每个掩膜进行分类，作者使用文本-图像判别模型（例如CLIP，ALIGN等等）。这些模型是在大规模的图像-文本对上训练的，表现出了强大的开放词汇表分类能力。它们由图像编码器 和文本编码器 组成。

在训练过程中，模型可以访问每个掩膜的真实类别标签。因此，训练过程类似于传统的语义分割任务。假设训练集中有 个类别。每个掩膜特征 的真实类别称为 。作者使用了冻结的文本编码器 对 中所有类别的名称进行编码得到Text Embeddings，并将所有训练类别名称的Embeddings集合定义为

其中类别名称 。然后通过分类损失计算掩码嵌入特征 属于 个类别中的一个的概率，如下所示：

其中

其中 是可学习的参数。

5. 前向过程

模型的前向过程如下图，测试类别名称的集合 是可用的。测试类别可能与训练类别不同。此外，测试图像没有可用的标题/标签。因此，作者选择将其通过Implicit Captioner模块以获取隐式标题；将其输入扩散模型以获取UNet的特征；并使用掩码生成器来预测图像中语义类别的所有可能二值掩膜。要将每个预测掩码 分类为测试类别之一，作者使用上述公式计算 ，最后预测具有最大概率的类别。

在实验中，作者发现扩散模型的内部表示在空间上是明确区分的，以产生许多对象级的可信掩膜。而它的对象分类能力可以通过再次与文本-图像判别模型相结合来进一步增强，特别是对于开放式词汇。因此，在此作者利用文本-图像判别模型的图像编码器 将原始输入图像的每个预测掩膜区域分类为测试类别之一。

作者使用 用于计算来自文本-图像判别模型的最终分类概率，并取扩散模型和判别模型的类别预测的几何平均值作为最终分类预测，即

其中 是一个固定的平衡因子。

讨论

本文的出发点很有意思，目前常见的开放式语义分割都使用多模态大模型处理未见类，而CLIP，ALIGH这类模型，训练时都是以图片为单位的，很可能缺少分割任务所需的空间信息。在这里，作者给出的解决方案是：通过引入扩散模型来弥补多模态大模型的缺陷。因此就有了这样的工作：同时使用Diffusion和CLIP全景分割模型。具体来说，作者的工作可以总结为以下几点。

• 第一次利用大规模文本-图像扩散模型来完成分割任务

• 在全景分割任务中超越以往的模型，实现了新的SOTA

• 该架构的成功证明了扩散模型不仅仅能用于做生成，同时也能学习到良好的语义表示，对图片的空间信息把握更到位。

公众号后台回复“极市直播”获取100+期极市技术直播回放+PPT