语义分割的解码器去噪预训练
Decoder Denoising Pretraining for Semantic Segmentation
https://arxiv.org/abs/2205.11423
23 May, 2022
Authors: Emmanuel Brempong Asiedu, Simon Kornblith, Ting Chen, Niki Parmar, Matthias Minderer, Mohammad Norouzi
摘要:语义分割标注的获取既昂贵又耗时。因此,通常使用预训练来提高分割模型的标注效率。通常,分割模型的编码器作为分类器进行预训练,解码器随机初始化。在这里,我们认为解码器的随机初始化可能是次优的,尤其是当标注样本很少时。我们提出了一种基于去噪的解码器预训练方法,该方法可以与编码器的有监督预训练相结合。我们发现,在ImageNet数据集上,解码器去噪预训练的效果明显优于编码器监督的预训练。尽管解码器去噪预训练非常简单,但它在标注高效语义分割方面取得了最先进的效果,并在城市景观、Pascal上下文和ADE20K数据集上获得了可观的收益。
1简介
计算机视觉中的许多重要问题,如语义分割和深度估计,都需要密集的像素级预测。为这些任务构建精确的监督模型是一项挑战,因为在所有图像像素上密集地收集真值标注成本高昂、耗时且容易出错。因此,最先进的技术通常求助于预训练,其中模型主干(即编码器)首先被训练为监督分类器(Sharif Razavian等人,2014;Radford等人,2021;Kolesnikov等人,2020)或自监督特征提取器(Oord等人,2018;Hjelm等人,2018;Bachman等人,2019;He等人,2020;Chen等人,2020b;c;Grill等人,2020)。主干架构(如Resnet(He et al.,2016))逐渐降低了特征图的分辨率。因此,为了进行像素级预测,需要解码器将采样恢复到像素级。大多数最先进的语义分割模型都不会预训练解码器引入的附加参数并随机初始化它们。在本文中,我们认为解码器的随机初始化远远不是最优的,并且使用简单但有效的去噪方法预训练解码器权重可以显著提高性能。
去噪自动编码器在机器学习方面有着悠久而丰富的历史(Vincent et al.,2008;2010)。一般的方法是向干净的数据中添加噪声,并训练模型将噪声数据分离回干净的数据和噪声分量,这需要模型学习数据分布。去噪目标非常适合于训练密集预测模型,因为它们可以在每像素级别上轻松定义。虽然去噪的想法由来已久,但去噪目标最近在去噪扩散概率模型(DPM)背景下吸引了新的兴趣;(Sohl Dickstein等人,2015;Song&Ermon,2019;Ho等人,2020))。DPMs通过学习通过一系列迭代去噪步骤将高斯噪声转换为目标分布,从而近似复杂的经验分布。这种方法在图像和音频合成方面取得了令人印象深刻的结果(Nichol&Dhariwal,2021;Dhariwal&Nichol,2021;Saharia et al.,2021b;Ho et al.,2021;Chen et al.,2021b),在样本质量分数方面优于强GAN和自回归基线。
受扩散模型中去噪的新兴趣和成功的启发,我们研究了通过去噪自动编码器学习的表示在语义分割中的有效性,特别是对于通常随机初始化的预训练解码器权重。
总之,本文研究了语义分割体系结构中解码器的预训练,发现通过随机初始化可以获得显著的收益,特别是在有限的标注数据设置中。我们提出将去噪用于解码器预训练,并将去噪自动编码器连接到扩散概率模型,以改进去噪预训练的各个方面,例如在去噪目标中预测噪声而不是图像,以及在添加高斯噪声之前缩放图像。这导致在三个数据集上对编码器的标准监督预训练有了显著的改进。
在第2节中,我们在深入研究常规去噪预训练的细节之前,先对其进行简要概述
第3节和第4节中的解码器对预训练进行去噪。
第5节介绍了与最新方法的实证比较。
2方法
我们的目标是学习能够很好地迁移到密集视觉预测任务的图像表示。我们考虑一种由编码器fθ和解码器gφ组成的体系结构,解码器gφ由两组参数θ和φ参数化。此模型将图像x作为输入∈ R H×W×C并将其转换为稠密表示y∈ R h×w×c,例如语义分割掩码。
我们希望找到一种初始化参数θ和φ的方法,这样模型就可以通过几个标注的样本有效地在语义分割上进行微调。对于编码器参数θ,我们可以按照标准实践,使用分类时预训练的权重初始化它们。我们的主要贡献涉及解码器参数φ,它通常是随机初始化的。我们建议将这些参数作为去噪自动编码器进行预训练(Vincent et al.,2008;2010):给定一幅未标注的图像x,我们通过将高斯噪声σc与固定的标准偏差σ相加到x,获得一幅带噪图像xe,然后将模型训练为自动编码器gφ◦ fθ以最小化重建误差kgφ(fθ(xe))− xk 2 2(仅优化φ并保持θ固定)。我们称这种方法为解码器去噪预训练(DDeP)。或者,φ和θ都可以通过去噪(去噪预训练;DeP)进行训练。下面,我们将讨论对标准自动编码器公式的几个重要修改,这些修改将显著提高表示的质量。
作为我们的实验装置,我们使用了Transune(Chen et al.(2021a));图2)。编码器根据ImageNet-21k(Deng等人,2009)分类进行预训练,而解码器则使用我们的去噪方法进行预训练,也使用ImageNet-21k图像,而不使用标注。预训练后,该模型根据城市景观、Pascal上下文或ADE20K语义分段数据集进行微调(Cordts等人,2016;Mottaghi等人,2014;Zhou等人,2018)。我们报告了所有语义类别的平均交并比(mIoU)。我们在第5.1节中描述了进一步的实施细节。
图1显示,我们的DDeP方法明显优于仅编码器的预训练,尤其是在少样本情况下。图6显示,即使是DeP,即在没有任何监督预训练的情况下对整个模型(编码器和解码器)进行去噪预训练,也可以与监督预训练相竞争。我们的结果表明,尽管去噪预训练简单,但它是学习语义分割表示的一种有效方法。
3编码器和解码器的去噪预训练
如上所述,我们的目标是学习能够很好地转换到语义分割和其他密集视觉预测任务的有效视觉表示。我们重新讨论去噪目标以实现这一目标。我们首先介绍标准去噪自动编码器公式(针对编码器和解码器)。然后,我们提出了对标准公式的一些修改,这些修改的动机是最近扩散模型在图像生成方面的成功(Ho等人,2020年;Nichol&Dhariwal,2021;Saharia等人,2021b)。
3.1标准去噪目标
在标准去噪自动编码器公式中,给定一个未标注的图像x,我们通过将高斯噪声σc与固定的标准偏差σx相加,得到一个带噪图像xe,
然后我们训练一个自动编码器gφ◦ fθ以最小化重建误差kgφ(fθ(xe))− xk 2。因此,目标函数采用以下形式
虽然这个目标函数已经产生了对语义分割有用的表示,但我们发现,一些关键的修改可以显著提高表示的质量。
3.2目标中去噪目标的选择
标准去噪自动编码器目标训练模型来预测无噪图像x。然而,扩散模型通常训练来预测噪声向量c(Vincent,2011;Ho et al.,2020):
对于具有从输入xe到输出的跳跃连接的模型,这两个公式的行为类似。在这种情况下,模型可以很容易地将其对c的估计与输入xe结合起来,以获得x。
然而,在没有明确的跳跃连接的情况下,我们的实验表明,预测噪声向量明显优于预测无噪声图像(表1)。
3.3作为预训练目标的去噪的可扩展性
无监督的预训练方法最终会受到由预训练目标学习的表征与最终目标任务所需表征之间不匹配的限制。对于任何无监督的目标,一个重要的“健全性检查”是,它不会很快达到这个极限,以确保它与目标任务很好地一致。我们发现,通过去噪学习的表示将继续提高到我们的最大可行预训练计算预算(图3)。这表明去噪是一种可扩展的方法,并且表示质量将随着计算预算的增加而继续提高。
3.4去噪与监督预训练
在标准去噪自动编码器公式中,使用去噪对整个模型(编码器和解码器)进行训练。然而,至少在微调数据丰富的情况下,全模型的去噪预训练性能不如编码器的标准监督预训练(表2)。在下一节中,我们将探索将去噪和有监督的预训练相结合,以获得两者的好处。
4仅对解码器进行去噪预训练
实际上,由于已经存在用于预训练编码器权重的强大且可伸缩的方法,因此去噪的主要潜力在于预训练解码器权重。为此,我们将编码器参数θ固定在通过在ImageNet-21k上进行监督预训练获得的值上,并且仅对解码器参数φ进行预训练并进行去噪,从而实现以下目标:
我们将这种预训练方案称为解码器去噪预训练(DDeP)。如下所示,在所有标注效率机制中,DDeP的性能优于纯监督或纯去噪预训练。在第5节给出基准结果之前,我们研究了DDeP的关键设计决策,如本节中的噪声公式和最佳噪声级。
4.1噪声大小和图像与噪声的相对比例
解码器去噪预训练的关键超参数是添加到图像中的噪声的大小。噪声方差σ必须足够大,网络必须学习有意义的图像表示才能将其移除,但不能太大,导致干净图像和噪声图像之间的分布过度偏移。对于目视检查,图4显示了σ的几个样本值。
除了噪声的绝对大小外,我们还发现干净图像和噪声图像的相对缩放也起着重要作用。不同的去噪方法在这方面有所不同。具体地说,DDPM生成一个有噪声的图像xe as
这不同于等式(1)中的标准去噪公式,因为x衰减为√γ和c衰减为√ 1.− γ,以确保如果x的方差为1,则随机变量xe的方差为1。使用此公式,我们的去噪预训练目标变为:
在图5中,我们将这种缩放的加性噪声公式与简单的加性噪声公式(公式(1))进行了比较,发现缩放图像可以显著提高下游语义分割性能。我们推测,噪声图像的方差与噪声大小的解耦减少了干净图像和噪声图像之间的分布偏移,从而改善了预训练表示到最终任务的传递。因此,本文其余部分将使用此公式。我们发现,对于缩放的加性噪声公式,最佳噪声幅值为0.22(图5),并将该值用于下面的实验。
4.2预训练数据集的选择
原则上,任何图像数据集都可以用于去噪预训练。理想情况下,我们希望使用大型数据集(如ImageNet)进行预训练,但这引发了一个潜在的担忧,即预训练数据和目标数据之间的分布变化可能会影响目标任务的性能。为了验证这一点,我们比较了解码器在几个数据集上进行去噪预训练,而编码器在ImageNet-21K上进行预训练,并保持分类目标不变。我们发现,对于所有测试数据集(城市景观、Pascal上下文和ADE20K;表3),在ImageNet-21K上预训练解码器比在目标数据上预训练解码器的效果更好。值得注意的是,这甚至适用于城市景观,城市景观在图像分布方面与ImageNet-21k存在显著差异。因此,在通用图像数据集上使用DDeP预训练的模型通常适用于范围广泛的目标数据集。
4.3解码器变体
考虑到解码器去噪预训练显著改善了解码器的随机初始化,我们假设该方法可以使解码器的大小扩大到使用随机初始化时效益减少的程度。我们通过在解码器的各个阶段改变特征映射的数量来测试这一点。我们所有实验的默认(1×)解码器配置为[1024、512、256、128、64],其中索引i处的值对应于第i个解码器块处的特征映射数。这反映在图2中。在Cityscapes上,我们尝试将所有解码器层的默认宽度加倍(2×),而在Pascal上下文和ADE20K上,我们尝试将宽度加倍(3×)。虽然较大的解码器通常即使在随机初始化时也能提高性能,但DDeP在所有情况下都会带来额外的增益。因此,DDeP可能会解锁新的解码器密集型架构。我们在第5节中给出了1×解码器和2×/3×解码器的主要结果。
4.4扩散过程的扩展
如上所述,我们发现,通过调整标准自动编码器公式的某些方面,如预测目标的选择以及图像和噪声的相对缩放,可以改进预训练表示,使其更类似于扩散模型。这就提出了一个问题,即是否可以通过使用全扩散过程进行预训练来进一步改进表征。在这里,我们研究了使该方法更接近DDPM中使用的完全扩散过程的扩展,但发现它们并没有改善上述简单方法的结果。
可变噪声表。
由于它使用单个固定噪声级(式(6)中的γ),因此我们的方法对应于扩散过程中的单个步骤。完整的DDPMs通过从每个训练样本的[0,1]中随机均匀采样噪声幅值γ,模拟从干净图像到纯噪声(及其相反)的完整扩散过程(Ho et al.,2020)。因此,我们也对随机抽样γ进行了实验,但发现固定γ表现最好(表4)。
调节噪声水平。
在扩散形式中,模型表示从一个噪声级到下一个噪声级的(反向)过渡函数,因此以当前噪声级为条件。在实践中,这是通过为每个训练样本提供所采样的γ作为额外的模型输入来实现的,例如,提供给归一化层。由于我们通常使用固定的噪声级,因此我们的方法不需要调节。当使用可变噪声时间表时,调节也不会提供任何改善。
噪声级加权。
在DDPM中,损失中不同噪声级的相对权重对样品质量有很大影响(Ho等人,2020年)。由于我们的实验表明,学习可迁移表征不需要多个噪声级,因此我们没有对不同噪声级的权重进行实验,但请注意,这可能是未来研究的一个有趣方向。
5基准结果
我们在几个语义分割数据集上评估了所提出的解码器去噪预训练(DDeP)的有效性,并进行了标注效率实验。
5.1实施细则
对于语义分割任务的预训练模型的下游微调,我们使用标准的逐像素交叉熵损失。我们使用Adam(Kingma&Ba,2015)优化器和余弦学习率衰减计划。对于解码器去噪预训练(DDeP),我们使用512的批量大小,并训练100个epoch。学习速度为6e−5用于1×和3×宽度解码器,以及1e−4用于2×宽度解码器。
在对目标语义分割任务中的预训练模型进行微调时,我们将权值衰减和学习率值扫过[1e−5,3e−4] 并为每个任务选择最佳组合。对于100%设置,我们报告所有数据集上10次运行的平均值。在Pascal上下文和ADE20K中,我们还报告了1%、5%和10%标注分数的10次运行(不同子集)的平均值,以及20%设置的5次运行的平均值。在城市景观上,我们报告了1/30设置下10次跑步的平均值,1/8设置下6次跑步,1/4设置下4次跑步的平均值。
在训练过程中,对图像及其相应的分割模板应用随机裁剪和随机左右翻转。对于城市景观,我们将图像随机裁剪为1024×1024的固定大小,对于ADE20K和Pascal上下文,我们将图像裁剪为512×512的固定大小。所有解码器去噪预训练运行均以224×224分辨率进行。
在城市景观推断过程中,我们通过将全分辨率1024×2048图像分割为两个1024×1024输入patch来评估它们。我们采用水平翻转,平均每一半的结果。将两半部分串联起来以产生全分辨率输出。对于Pascal Context和ADE20K,除了水平翻转外,我们还对图像的重缩放版本使用多尺度评估。使用的比例因子为(0.5、0.75、1.0、1.25、1.5、1.75)。
5.2解码器去噪预训练的性能增益
在城市景观方面,DDeP的表现优于DeP和有监督的预培训。在图6中,我们报告了城市景观上的DeP和DDeP结果,并将其与随机初始化或使用ImageNet-21K预训练编码器初始化的训练结果进行比较。DeP结果使用了缩放的加性噪声公式(方程式(5)),与标准降噪目标获得的结果相比,性能显著提高。
如图6所示,在1%和5%标注图像设置中,DeP优于监督基线。对于1×和2×解码器变体,解码器去噪预训练(DDeP)比DeP和ImageNet-21K监督预训练都有进一步改进(表6)。
如表5.2所示,DDeP在所有标注分数上都优于先前提出的城市景观标注有效语义分割方法。只有25%的训练数据,DDeP比最强的基线方法PC2Seg(Zhong等人,2021)在完整数据集上训练时产生的分割效果更好。与最近的工作不同,我们不对城市景观进行多尺度评估,这将导致进一步的改进。
DDeP还改进了Pascal上下文数据集上的有监督预训练。图1比较了DDeP在1%、5%、10%、20%和100%的训练数据上与有监督基线和Pascal上下文上随机初始化模型的性能。表5.2将这些结果与使用3×解码器获得的结果进行了比较。对于1×和3×解码器,DDeP的性能明显优于架构上相同的监督模型,在所有半监督设置中获得了4-12%的mIOU改进。值得注意的是,只有10%的标注,DDeP优于20%标注训练的监督模型。
图7显示了ADE20K数据集上DDeP的类似改进。我们再次看到,在5%和10%的设置中,收益超过10分,在1%的设置中,收益超过5分。这些一致的结果证明了DDeP在数据集和训练集大小方面的有效性。
我们的上述结果使用了Transune(Chen等人(2021a));图2)实现最大性能的体系结构,但DDeP与主干网无关,并且与更简单的主干网体系结构一起使用时也会带来好处。在表7中,我们训练了一个标准的U-Net,该U-Net带有一个ResNet50编码器,在Pascal上下文中具有DDeP(无多尺度评估)。DDeP在所有设置中都优于监督基线,表明我们的方法超越了transformer架构。
6相关工作
因为为语义分割收集详细的像素级标注成本高昂、耗时且容易出错,已经提出了许多方法来从较少的标注样本中实现语义分割(Tarvainen&Valpola,2017;Miyato et al.,2018;Hung et al.,2018;Mittal et al.,2021;French et al.,2019;Ouali et al.,2020;Zou et al.,2021;Feng et al.,2020b;Ke et al.,2020;Olsson et al.,2021;Zhong et al.,2021)。这些方法通常求助于半监督学习(SSL)(Chapelle et al.,2006;Van Engelen&Hoos,2020),在这种学习中,除了有标注的训练数据外,还假设可以访问大量的未标注图像数据集。在接下来的内容中,我们将讨论以前关于强数据增强、生成模型、自训练和自监督学习在标注有效语义分割中的作用的工作。虽然这项工作侧重于自监督的预训练,但我们相信,强大的数据增强和自训练可以与所提出的去噪预训练方法相结合,以进一步改善结果。
数据增强。
French et al.(French et al.,2019)证明,Cutout(DeVries&Taylor,2017)和CutMix(Yun et al.,2019)等强大的数据增强技术对于少数标注样本的语义分割特别有效。Ghiasi等人(2021)发现,简单的复制粘贴增强有助于实例分割。之前的工作(Remez et al.,2018;Chen et al.,2019;Bielski&Favaro,2019;Arandjelović&Zisserman,2019)也通过利用GANs(Goodfello et al.,2014)组合不同的前景和背景区域来生成新的可信图像,探索了完全无监督的语义分割。我们使用了相对简单的数据增强,包括水平翻转和随机起始式裁剪(Szegedy et al.,2015)。使用更强大的数据增强留给未来的工作。
生成模型。
标注有效语义分割的早期工作使用GANs生成合成训练数据(Souly et al.,2017),并区分真实和预测的分割掩码(Hung et al.,2018;Mittal et al.,2021)。DatasetGAN(Zhang等人,2021)表明,现代GAN架构(Karras等人,2019)在生成合成数据以帮助像素级图像理解方面是有效的,而此时只有少数标注图像可用。我们的方法与扩散和基于分数的生成模型高度相关(Sohl Dickstein et al.,2015;Song&Ermon,2019;Ho et al.,2020),这代表了一个新兴的生成模型家族,导致图像样本质量优于GANs(Dhariwal&Nichol,2021;Ho et al.,2021)。这些模型通过去噪分数匹配与去噪自动编码器相连(Vincent,2011),可以看作是训练基于能量的模型的方法(Hyvärinen&Dayan,2005)。去噪扩散模型(DDPM)最近已应用于超分辨率、着色和修复等条件生成任务(Li等人,2021;Saharia等人,2021b;Song等人,2021;Saharia等人,2021a),表明这些模型可能能够学习有用的图像表示。我们从DDPM的成功中得到了启发,但我们发现DDPM的许多组件都不是必需的,简单的去噪预训练效果很好。扩散模型已被用于迭代优化语义分段掩码(Amit等人,2021;Hoogeboom等人,2021)。Baranchuk et al.(Baranchuk et al.,2021)证明了通过扩散模型学习的特征在语义分割中的有效性,这些特征来自于极少数标注的样本。相比之下,我们利用简单的去噪预训练进行表示学习,并研究编码器-解码器体系结构的完全微调,而不是提取固定特征。此外,我们使用成熟的基准将我们的结果与之前的工作进行比较。
自训练,一致性规范化。
自训练(自学习或伪标注)是最古老的SSL算法之一(Scudder,1965;Fralick,1967;Agrawala,1970;Yarowsky,1995)。它的工作原理是使用一个初始的监督模型,用所谓的伪标注对未标注的数据进行注释,然后使用伪标注和人类标注数据的混合来训练改进的模型。这个迭代过程可以重复多次。自训练已被用于改进目标检测(Rosenberg等人,2005;Zoph等人,2020)和语义分割(Zhu等人,2020;Zou等人,2021;Feng等人,2020a;Chen等人,2020a)。一致性正则化与自训练密切相关,并在图像增强过程中加强预测的一致性(French et al.,2019;Kim et al.,2020;Ouali et al.,2020)。这些方法通常需要仔细调整超参数和合理的初始模型,以避免传播噪声。将自训练与去噪预训练相结合可能会进一步改善结果。
自监督学习。
自监督学习方法可以制定预测性的借口任务,这些任务很容易从未标注的数据中构建,并且有利于下游的区分性任务。在自然语言处理(NLP)中,掩蔽语言建模任务(Devlin等人,2019;Liu等人,2019;Raffel等人,2020)已成为事实上的标准,在NLP任务中显示出令人印象深刻的结果。在计算机视觉中,提出了用于自监督学习的不同借口任务,包括预测图像中相邻patch的相对位置的任务(Doersch et al.,2015),修复任务(Pathak et al.,2016),解决拼图(Noroozi&Favaro,2016),图像着色(Zhang et al.,2016;Larsson et al.,2016),轮换预测(Gidaris et al.,2018)和其他任务(Zhang et al.,2017;Caron et al.,2018;Kolesnikov et al.,2019)。最近,基于样本鉴别和对比学习的方法在图像分类方面显示出了有希望的结果(Oord et al.,2018;Hjelm et al.,2018;He et al.,2020;Chen et al.,2020b;c;Grill et al.,2020)。这些方法已被用于成功地预训练用于对象检测和分割的主干(He et al.,2020;Chen et al.,2020d),但与这项工作不同的是,它们通常随机初始化解码器参数。最近,也有一系列基于掩蔽自动编码的新兴方法,如BEIT(Bao等人,2021)、MAE(He等人,2021)和其他方法(Zhou等人,2021;Dong等人,2021;Chen等人,2022)。我们注意到,我们的方法是针对这一系列掩模图像建模同时开发的,我们的技术也是正交的,因为我们关注的是解码器预训练,这不是上述论文的重点。
用于密集预测的自监督学习。
Pinheiro et al.(2020)和Wang et al.(2021)提出了密集对比学习,这是一种针对密集预测任务的自监督预训练方法,对比学习应用于patch和像素级特征,而不是图像级特征。这让人想起AMDIM(Bachman et al.,2019)和CPC V2(Hénaff et al.,2019)。Zhong等人(2021)进一步考虑了这一想法,并将图像不同增强(可能未标注)的模型输出之间的分割掩码一致性与整个增强的像素级特征一致性相结合。
视觉Transformer。
受NLP中Transformer的成功(Vaswani et al.,2017)的启发,一些出版物研究了卷积和自注意相结合的目标检测(Carion et al.,2020)、语义分割(Wang et al.,2018;2020b)和全景分割(Wang et al.,2020a)。Vision Transformer(ViT)(Dosovitskiy et al.,2021)证明,当大量标注数据集可用时,无卷积方法可以产生令人印象深刻的结果。最近的研究探索了将ViT作为语义分割的主干(Zheng等人,2020年;Liu等人,2021;Strudel等人,2021)。这些方法在解码器的结构上有所不同,但它们显示了基于ViT的语义分割的能力。我们采用混合ViT(Dosovitskiy et al.,2021)作为主干,其中patch嵌入投影应用于从卷积特征图中提取的patch。我们研究了解码器的大小,发现更宽的解码器通常会改善语义分割结果。
7结论
受最近流行的图像合成扩散概率模型的启发,我们研究了这些模型在学习语义分割有用的可迁移表示方面的有效性。令人惊讶的是,我们发现,将语义分割模型作为去噪自动编码器进行预训练可以大大提高语义分割性能,尤其是在标注样本数量有限的情况下。我们基于这一观察结果,提出了一种两阶段预训练方法,其中监督预训练编码器与去噪预训练解码器相结合。这导致了数据集和训练集大小的一致收益,从而形成了一种切实可行的预训练方法。探索去噪预训练在其他密集预测任务中的应用也很有趣。