基于扩散模型的标签高效语义分割
摘要：
去噪扩散概率模型最近受到了很多研究的关注，因为它们优于gan等替代方法，并且目前提供了最先进的生成性能。扩散模型的优越性能使其成为一些应用程序的吸引人的工具，包括绘图，超分辨率和语义编辑。在本文中，我们证明了扩散模型也可以作为语义分割的工具，特别是在标记数据稀缺的情况下。特别地，对于几个预训练的扩散模型，我们研究了执行反向扩散过程的马尔可夫步骤的网络的中间激活。我们表明，这些激活有效地从输入图像中捕获语义信息，并且似乎是分割问题的优秀像素级表示。基于这些观察，我们描述了一种简单的分割方法，即使只提供少量训练图像也可以工作。对于相同数量的人类监督，我们的方法在几个数据集上显著优于现有的替代方法。该项目的源代码是公开的。
1 介绍
去噪扩散概率模型(DDPM) (Sohl-Dickstein等，2015;Ho等人，2020)最近在个体样本的真实性及其多样性方面优于其他方法来模拟自然图像的分布(Dhariwal和Nichol, 2021)。DDPM的这些优点已成功地应用于着色(Song et al .， 2021)、涂漆(Song et al .， 2021)、超分辨率(sahara et al .， 2021;Li et al .， 2021b)和语义编辑(Meng et al .， 2021)，其中DDPM通常比gan取得更令人印象深刻的结果。
然而，到目前为止，DDPM还没有被用来作为鉴别计算机视觉问题的有效图像表示的来源。虽然之前的文献已经证明了各种生成范式，如gan (Donahue & Simonyan, 2019)或自回归模型(Chen等人，2020a)，可以用于提取常见视觉任务的表示，但目前尚不清楚DDPM是否也可以作为表示学习器。在本文中，我们从语义分割的角度对这个问题给出了肯定的答案。
特别地，我们研究了来自U-Net网络的中间激活，该网络近似于DDPM中反向扩散过程的马尔可夫步长。直观地说，这个网络学习去噪它的输入，并且不清楚为什么中间激活应该捕获高级视觉问题所需的语义信息。然而，我们表明，在某些扩散步骤中，这些激活确实捕获了这些信息，因此，可以潜在地用作下游任务的图像表示。鉴于这些观察结果，我们提出了一种简单的语义分割方法，该方法利用这些表示并成功地工作，即使只提供了少数标记图像。在几个数据集上，我们表明我们基于ddpm的分割方法在相同监督量下优于现有基线。
综上所述，本文的贡献有:1。我们研究了最先进的DDPM学习的表征，并表明它们捕获了对下游视觉任务有价值的高级语义信息。
2. 我们设计了一种简单的语义分割方法，利用这些表示，并在少数几个操作点上优于替代方法。
3. 我们在相同的数据集上比较了基于ddpm的表示和基于gan的表示，并展示了前者在语义分割方面的优势。
2 相关工作
在本节中，我们简要地描述了与我们的工作相关的现有研究方向。
扩散模型(Sohl-Dickstein等，2015;Ho等人，2020)是一类生成模型，它通过源自简单参数分布(通常是标准高斯分布)的马尔可夫链的端点近似真实图像的分布。每个马尔可夫步骤都由一个深度神经网络建模，该网络有效地学习用已知的高斯核反转扩散过程。Ho等人强调了扩散模型和分数匹配的等价性(Song & Ermon, 2019;2020)，表明它们是通过迭代去噪过程将简单已知分布逐渐转换为目标分布的两种不同视角。V近期作品(尼科尔，2021;Dhariwal & Nichol, 2021)开发了更强大的模型架构以及不同的高级目标，这导致DDPM在生成质量和多样性方面优于gan。DDPM已广泛应用于图像着色(Song et al .， 2021)、超分辨率(sahara et al .， 2021;Li et al .， 2021b)， inpainting (Song et al .， 2021)和semantic editing (Meng et al .， 2021)。在我们的工作中，我们证明了人们也可以成功地将它们用于语义分割。
基于生成模型的图像分割是目前一个活跃的研究方向，但现有的方法主要是基于gan的。第一行作品(沃诺夫和巴班科，2020;Voynov et al, 2021;Melas-Kyriazi等人，2021)基于以下证据:最先进的gan的潜在空间具有与影响前景/背景像素不同的效果相对应的方向，这允许生成合成数据来训练分割模型。然而，这些方法目前只能进行二值分割，是否可以用于一般的语义分割设置还不清楚。第二行作品(Zhang et al .， 2021;Tritrong等，2021;徐,2021;Galeev等人，2020)与我们的研究更相关，因为它们是基于gan中获得的中间表示。特别是，(Zhang et al .， 2021)中提出的方法在这些表示上训练像素类预测模型，并确认其标签效率。在实验部分，我们将(Zhang et al .， 2021)的方法与我们基于ddpm的方法进行了比较，并展示了我们的解决方案的几个独特优势。
判别任务的生成模型表示。作为表征学习器，生成模型的使用已被广泛研究用于全局预测(Donahue & Simonyan, 2019;Chen et al .， 2020a)和密集预测问题(Zhang et al .， 2021;Tritrong等，2021;徐,2021;Xu et al, 2021)。虽然以前的工作强调了这些表示的实际优势，例如分布外鲁棒性(Li等人，2021a)，但与其他无监督方法(例如基于对比学习的Chen等人，2020b)相比，生成模型作为表示学习器受到的关注较少。主要原因可能是在复杂多样的数据集上训练高质量的生成模型很困难。然而，考虑到DDPM最近在Imagenet上的成功(Deng et al .， 2009)，可以预期，这个方向将在未来吸引更多的关注。
3 来自扩散模型的3种表示
在下一节中，我们将研究通过扩散模型学习的图像表示。首先，我们提供DDPM框架的简要概述。然后，我们描述了如何使用DDPM提取特征，并研究了这些特征可能捕获什么样的语义信息。背景。扩散模型通过将xT逐渐去噪到噪声较小的样本xT，将噪声xT ~ N(0, I)转换为样本x0。形式上，我们得到一个正向扩散过程:
对于某些固定方差表β1，…,βt。
图1:建议方法的概述。(1) x0−→xt，按q(xt|x0)加噪声。
(2)从噪声预测器θ(xt, t)中提取特征映射。(3)通过将特征映射上采样到图像分辨率并将它们连接起来，收集像素级表示。(4)使用逐像素特征向量训练mlp集合来预测每个像素的类标签。
重要的是，可以直接从数据x0中得到一个带噪声的样本xt:
预训练的DDPM近似于一个相反的过程:
在实践中，噪声预测网络θ(xt, t)不是预测式(3)中分布的均值，而是预测步长t处的噪声分量;那么均值就是这个噪声分量和xt的线性组合。协方差预测器Σθ(xt, t)可以是一组固定的标量协方差，也可以是学习的(后者被证明可以提高模型质量(Nichol, 2021))。
去噪模型θ(xt, t)通常由UNet架构的不同变体参数化(Ronneberger等人，2015)，在我们的实验中，我们研究了(Dhariwal & Nichol, 2021)中提出的最先进的模型。
提取表征。对于给定的真实图像x0∈RH×W ×3，可以从噪声预测网络∊θ(xt, T)中计算T组激活张量。时间步长T的总体方案如图1所示。首先，根据式(2)，我们通过添加高斯噪声来破坏x0。噪声xt用作由UNet模型参数化的∊θ(xt, t)的输入。然后用双线性插值将UNet的中间激活上采样到H × W。这允许将它们视为x0的像素级表示。
3.1 表征分析
我们分析了噪声预测器θ(xt, t)对不同t产生的表示。我们考虑了在LSUN-Horse和FFHQ-256数据集上训练的最先进的DDPM检查点1。
来自噪声预测器的中间激活捕获语义信息。在这个实验中，我们从LSUN-Horse和FFHQ数据集中取了一些图像，并手动将每个像素分别分配给21个和34个语义类中的一个。我们的目标是了解DDPM生成的像素级表示是否有效地捕获了关于语义的信息。为此，我们训练了一个多层感知器(MLP)，从特定扩散步骤t上18个UNet解码器块之一产生的特征中预测像素语义标签。请注意，我们只考虑解码器激活，因为它们还通过跳过连接聚合编码器激活。mlp在20张图片上进行训练，并在20张保留的图片上进行评估。
预测性能是根据平均借据来衡量的。
图2:基于ddpm的逐像素表示在不同UNet解码器块和扩散步骤下预测性能的演变。这些街区从深到浅依次编号。信息量最大的特征通常对应于反向扩散过程的后期步骤和UNet解码器的中间层。前面的步骤对应于无信息表示。其他数据集的图见附录A
图3:对于平均面积最小(左)和最大(右)的类，在LSUN-Horse数据集上基于ddpm的逐像素表示的预测性能演变。小型对象的预测性能在相反的过程中开始增长较晚。对于较大的对象，较深的块具有更多的信息，对于较小的对象，较浅的块具有更多的信息。附录A提供了对其他数据集的类似评估。
预测性能在不同区块和扩散步骤t上的演变如图2所示。这些街区从深到浅依次编号。图2显示了噪声预测器?θ(xt, t)产生的特征的可判别性因不同的块和扩散步骤而异。特别是，与反向扩散过程的后期步骤相对应的特征通常更有效地捕获语义信息。相比之下，与早期步骤相对应的那些通常是没有信息的。在不同的块中，UNet解码器中间的层产生的特征在所有扩散步骤中似乎是最具信息量的。
此外，我们还根据标注数据集的平均面积分别考虑小型和大型语义类。然后，我们在不同的UNet块和扩散步骤中独立评估这些类的平均IoU。LSUN-Horse的结果如图3所示。正如预期的那样，大型对象的预测性能在相反的过程中开始增长得更早。对于较小的对象，较浅的块提供的信息更多，而对于较大的对象，较深的块提供的信息更多。在这两种情况下，最具判别性的特征仍然对应于中间块。
图2表明，对于某些UNet块和扩散步骤，类似的基于ddpm的表示对应于相同语义的像素。图4显示了FFHQ检查点从扩散步骤{50,200,400,600,800}的块{6,8,10,12}中提取的特征形成的k-means聚类(k=5)，并证实了聚类可以跨越连贯的语义对象和对象部分。在block B=6中，特征对应于粗语义掩码。在另一个极端，B=12的特征可以区分细粒度的面部部分，但对粗碎片表现出较少的语义意义。在不同的扩散步骤中，最有意义的特征对应于后面的特征。我们将这种行为归因于这样一个事实，即在逆向过程的早期步骤中，DDPM样本的全局结构尚未出现，因此，在此阶段几乎不可能预测分割掩码。图4中的遮罩定性地证实了这种直觉。当t=800时，掩码较差地反映了实际图像的内容，而当t值较小时，掩码和图像在语义上是一致的。
图4:从UNet解码器块{6,8,10,12}中提取的特征在扩散步骤{50,200,400,600,800}上形成的k-means聚类示例(k=5)。中间块的聚类在空间上跨越连贯的语义对象和部分。
3.2基于ddpm的少镜头语义分割表示
上面观察到的中间DDPM激活的潜在有效性意味着它们用作密集预测任务的图像表示。图1简要地展示了我们对图像分割的总体方法，它利用了这些表示的可辨别性。更详细地说，我们考虑少量半监督设置，当大量未标记的图像{X1，…XN}⊂RH×W ×3来自特定域的图像是可用的，并且只对n个训练图像{X1，…Xn}⊂RH×W ×3真k类语义掩码{Y1，…y}⊂RH×W x{1，…，K}。
作为第一步，我们在整个{X1，…上训练扩散模型。以无监督的方式。
然后，使用该扩散模型提取标记图像的像素级表示，使用UNet块的子集和扩散步骤t。在这项工作中，我们使用来自UNet解码器的中间块B={5,6,7,8,12}和反向扩散过程的后期步骤t={50,150,250}的表示。这些块和时间步是由3.1节的见解激发的，但有意没有针对每个数据集进行调整。
虽然在特定时间步长的特征提取是随机的，但我们固定了所有时间步长的噪声，并在4.1节中消除了噪声。从所有块B和步骤t中提取的表示向上采样到图像大小并连接，形成训练图像的所有像素的特征向量。像素级表示的总维度为8448。
然后，接下来(Zhang等人，2021)，我们在这些特征向量上训练一组独立的多层感知器(mlp)，其目的是预测可用于训练图像的每个像素的语义标签。我们采用了(Zhang et al .， 2021)中的集成配置和训练设置，并在我们的实验中将它们应用于所有其他方法，详见附录C。
为了分割测试图像，我们提取其基于ddpm的逐像素表示，并使用它们来预测集成的像素标签。最终的预测是通过多数投票获得的。
4 实验
本节通过实验证实了基于ddpm的表示在语义分割问题上的优势。我们从对现有替代方案的全面比较开始，然后通过额外的分析剖析DDPM成功的原因。
数据集。在我们的评估中，我们主要使用LSUN (Y u et al .， 2015)和FFHQ-256 (Karras et al .， 2019)中的“卧室”、“猫”和“马”类别。作为每个数据集的训练集，我们考虑了几个图像，这些图像的细粒度语义掩码是按照(Zhang et al .， 2021)的协议收集的。对于每个数据集，聘请专业评估员对训练和测试样本进行注释。我们将收集到的数据集表示为Bedroom-28, FFHQ-34, Cat-15, Horse21，其中数字对应语义类的数量。
表1:在我们的评估中使用的每个数据集的注释图像的数量。
此外，我们考虑了两个数据集，与其他数据集相比，它们具有公开可用的注释和相当大的评估集:
•ADE-Bedroom-30是ADE20K数据集的一个子集(Zhou等人，2018)，其中我们仅提取具有30个最频繁类别的卧室场景图像。我们将每个图像的大小调整为较小的一面256，然后裁剪它们以获得256×256样本。
•CelebA-19是CelebAMask-HQ数据集的一个子集(Lee et al, 2020)，它提供了19个面部属性的注释。所有图像都被调整为256分辨率。
每个数据集带注释的图像数量见表1。其他细节见附录E。
方法。在评估中，我们将我们的方法(表示为DDPM)与先前的几种处理少量语义分割设置的方法进行了比较。首先，我们描述了生成大量带注释的合成图像以训练分割模型的基线:
•DatasetGAN (Zhang et al .， 2021)——该方法利用了gan产生的像素级特征的可判别性。更详细地说，评估人员注释了一些gan生成的图像。
然后，使用这些图像的潜在代码来获得中间生成器的激活，并将其视为像素级表示。给定这些表示，训练分类器来预测每个像素的语义标签。然后使用该分类器标记新的合成GAN图像，这些图像作为DeepLabV3分割模型的训练集(Chen等人，2017)。对于每个数据集，我们增加合成图像的数量，直到验证集上的性能不饱和。根据(Zhang et al .， 2021)，我们还删除了10%的预测最不确定的合成样本。
•DatasetDDPM反映了DatasetGAN基线，唯一的区别是gan被ddpm取代。我们包含这个基线是为了比较同一场景中基于gan和基于ddpm的表示。
请注意，与DatasetGAN和DatasetDDPM相比，我们在3.2节中描述的分割方法更直接，因为它不需要合成数据集生成和在其上训练分割模型的辅助步骤。
然后，我们考虑一组基线，这些基线允许直接从真实图像中提取中间激活，并将它们用作像素级表示，类似于我们的方法。与DatasetGAN和DatasetDDPM相比，由于真实图像和合成图像之间没有域间隙，这些方法可能是有益的。
•MAE (He et al .， 2021)——最先进的自监督方法之一，它学习去噪自编码器来重建缺失的补丁。我们使用ViT-Large (Dosovitskiy et al .， 2021)作为主干模型，并将patch大小减小到8×8，以增加特征图的空间维度。我们使用官方代码在与DDPM相同的数据集上预训练所有模型2。附录F描述了该方法的特征提取。
•SwA V (Caron等人，2020)——一种最新的自我监督方法。我们考虑使用两倍宽的ResNet-50模型进行评估。所有模型都是在与DDPM相同的数据集上进行预训练的，也使用官方源代码3。输入图像分辨率为256。
•GAN反演采用最先进的方法(Tov et al, 2021)获取真实图像的潜在代码。我们将标注的真实图像映射到GAN潜在空间，这允许计算中间生成器激活并将其用作像素级表示。
表2:平均欠条分割方法对比。(*)在CelebA-19和ADE Bedroom-30上，我们分别对FFHQ-256和LSUN Bedroom预训练的模型进行了评估。
•GAN编码器-当GAN反演努力从LSUN域重建图像时，我们还考虑了用于GAN反演的预训练GAN编码器的激活。
•VDV AE (Child, 2021) -最先进的自动编码器模型。中间激活从编码器和解码器中提取并连接起来。虽然在LSUN数据集上没有预训练的模型，但我们仅在FFHQ-256上公开可用的检查点4上评估该模型。值得注意的是，在LSUN上，V ae仍然明显低于gan和ddpm。
•ALAE (Pidhorskyi et al .， 2020)采用StyleGANv1生成器，并在对抗性训练中加入编码器网络。我们从编码器模型中提取特征。在我们的评估中，我们使用了LSUN-Bedroom和FFHQ-10245的公开可用模型。
生成式预训练模型。在我们的实验中，我们将最先进的StyleGAN2 (Karras等人，2020)模型用于基于gan的基线，将最先进的预训练adm (Dhariwal & Nichol, 2021)用于基于ddpm的方法。由于没有预先训练的FFHQ-256模型，我们使用官方实现自己训练它。为了对ADEBedroom-30数据集进行评估，我们使用了在LSUN-Bedroom上预训练的模型(包括基线)。
对于Celeba-19，我们评估了在FFHQ-256上训练的模型。
主要的结果。表2给出了这些方法在平均欠条度量方面的比较。对于不同的数据分割，结果在5次独立运行中取平均值。我们还在附录d中报告了每个类的借据。此外，我们在图5中提供了几个使用我们的方法进行分割的定性示例。下面我们重点介绍几个关键的观察结果:
•提出的基于DDPM表示的方法在大多数数据集上显著优于替代方法。
•MAE基线是基于ddpm的分割的最强竞争者，并在FFHQ-34和Cat-15数据集上展示了可比较的结果。
•与基于ddpm的细分相比，SwA V基线表现不佳。我们将这种行为归因于该基线以判别方式训练，并且可以抑制细粒度语义分割所需的细节。这一结果与(Cole et al, 2021)中的最新发现一致，该发现表明，最先进的对比方法产生的表征对于细粒度问题来说是次优的。
•DatasetDDPM在大多数基准测试中优于其对应的DatasetGAN。注意，这两种方法都使用DeepLabV3网络。我们将这一优势归因于DDPM合成物的高质量，因此合成物与实际数据之间的域差距较小。
•在大多数数据集上，DDPM优于DatasetDDPM的竞争对手。在下面的讨论部分中，我们将提供一个额外的实验来对此进行研究。
总体而言，所提出的基于ddpm的分割优于利用替代生成模型的基线以及以自监督方式训练的基线。这个结果突出了使用最先进的ddpm作为强无监督表示学习器的潜力。
图5:我们的方法在测试图像上预测的分割掩码示例以及groundtruth注释掩码。
表3:在真实图像和合成图像上训练时基于ddpm的分割性能。
当使用DDPM生成的数据进行训练时，DDPM显示出与DatasetDDPM相当的性能。当在gan生成的数据上进行训练时，DDPM仍然明显优于DatasetGAN，但两者之间的差距减小了。
4.1 讨论
训练对真实数据的影响。所提出的DDPM方法是在带注释的真实图像上进行训练的，而DatasetDDPM和DatasetGAN是在合成图像上进行训练的，这些合成图像通常不那么自然、多样化，并且可能缺乏特定类的对象。此外，人工注释合成图像更困难，因为它们可能有一些扭曲的对象，难以将其分配给特定的类。在接下来的实验中，我们量化了在真实数据或合成数据上训练导致的性能下降。具体来说，表3报告了DDPM方法在真实、DDPM生成和gan生成的注释图像上训练的性能。可以看出，在生成模型保真度仍然相对较低的领域，例如LSUN-Cat，对真实图像进行训练是非常有益的，这表明带注释的真实图像是更可靠的监督来源。此外，如果DDPM方法在合成图像上进行训练，其性能与DatasetDDPM相当。另一方面，当对gan生成的样本进行训练时，DDPM显著优于DatasetGAN。我们将此归因于与gan相比，ddpm提供了更有语义价值的像素表示。
Sample-efficiency。在这个实验中，我们评估了我们的方法在使用较少注释数据时的性能。我们为表4中的四个数据集提供了mIoU。重要的是，DDPM仍然能够优于表2中的大多数基线，使用更少的监督。
随机特征提取的效果。在这里，我们研究我们的方法是否可以从3.2节中描述的随机特征提取中获益。我们考虑确定性情况，当噪声∊ ~ N(0, I)被采样一次，并在(2)中使用，以获得训练和评估期间所有时间步t的xt。然后，我们将其与以下随机选项进行比较:首先，在不同的时间步长t上采样不同的∊t，并在训练和评估过程中共享。
其次，在每次训练迭代中对所有时间步采样不同的噪声;在评估过程中，该方法还使用了看不见的噪声样本。
图6:在Bedroom-28和Horse-21数据集上不同图像损坏水平下的mIoU退化。DDPM表现出更高的鲁棒性，并保留了其在所有失真水平下的优势。
表4:不同标记训练数据个数对所提方法的评价。即使使用较少注释的数据，DDPM仍然优于表2中的大多数基线。
表5:基于ddpm的方法在不同特征提取变化下的性能。所有考虑的随机选项都提供了与确定性选项相似的mIoU。
结果如表5所示。可以看到，性能上的差异是微乎其微的。
我们将这种行为归因于以下原因:•我们的方法使用了噪声强度较低的反向扩散过程的后期t。
•由于我们利用了UNet模型的深层，噪声可能不会显著影响这些层的激活。
对输入损坏的鲁棒性。在这个实验中，我们研究了基于ddpm的表示的鲁棒性。首先，我们使用Bedroom-28和Horse-21数据集上的DDPM、SwAV和MAE表示来学习干净图像上的像素分类器。然后，采用(Hendrycks & Dietterich, 2019)的18种不同的腐败类型应用于测试图像。每种腐败都有五个严重级别。在图6中，我们提供了针对1、3、5个严重级别的所有损坏类型计算的平均欠条，分别表示为“弱”、“中”和“强”。
可以观察到，所提出的基于ddpm的方法具有更高的鲁棒性，并且即使在严重的图像失真情况下也保持了其优于SwA V和MAE模型的优势。
5 结论
本文证明了ddpm可以作为判别性计算机视觉问题的表征学习器。与gan相比，扩散模型允许对真实图像的这些表示进行直接计算，并且不需要学习将图像映射到潜在空间的额外编码器。该DDPM的优势和优越的生成质量为少镜头语义分割任务提供了最先进的性能。基于ddpm的分割的显著限制是需要在手头的数据集上训练高质量的扩散模型，这对于像ImageNet或MSCOCO这样的复杂领域来说可能是具有挑战性的。然而，考虑到DDPM的快速研究进展，我们预计它们将在不久的将来达到这些里程碑，从而扩大相应表示的适用性范围。
