Denoising Diffusion Autoencoders are Unified Self-supervised Learners (Paper reading)

Weilai Xiang, Beihang University, arXiv23, Code, Paper

1. 前言

受最近扩散模型进展的启发，这让人想起去噪自编码器，我们研究了它们是否可以通过生成预训练获得分类的判别表示。本文表明扩散模型中的网络，即去噪扩散自编码器(DDAE)，是统一的自监督学习器:通过对无条件图像生成的预训练，DDAE在没有辅助编码器的情况下已经在其中间层中学习了强线性可分表示，从而使扩散预训练成为生成和判别双重学习的一般方法。为了验证这一点，我们进行了线性探测和微调评估。我们基于扩散的方法在CIFAR10和Tiny-ImageNet上分别实现了95.9%和50.0%的线性评估精度，并且首次与对比学习和掩码自编码器相比较。来自ImageNet的迁移学习也证实了DDAE对Vision transformer的适用性，表明将DDAE扩展为统一基础模型的潜力。

2. 整体思想

预训练的生成模型是否可以做分类？本文通过实验验证了预训练的模型本身就学习到了图像的表示。主要的实验方法是使用线性探针对模型每层做分类，实验结果表示预训练网络内部是学习到了这种类别。

3. 介绍

在有限的人类监督下理解数据是机器学习中的一个关键挑战。为了处理具有稀缺注释的大量数据，深度学习范式正在从监督预训练转向自监督预训练。在自然语言处理(NLP)方面，BERT、gpt和T5等自监督模型在不同的任务中取得了出色的表现，ChatGPT等大型语言模型在机器学习社区之外也表现出了深远的影响。其中，BERT使用掩码语言建模(MLM)作为借口任务来训练编码器，但无法生成全文样本。GPTs和T5已显示出生成长段落自回归(AR)的能力。此外，他们证明了仅解码器或编码器-解码器模型可以通过生成式预训练获得深度语言理解，而无需有意训练编码器。随着人工智能生成内容(AI-Generated Content, AIGC)的兴起，gpt和T5与纯编码器相比得到了更多的关注，后者统一了生成(如翻译、摘要)和判别(如分类)任务。

在计算机视觉领域，自监督学习器在弥合生成和识别之间的差距方面尚未取得与gpt类似的成就。虽然生成对抗网络(GAN)和自回归变形器可以合成高保真图像，但它们对判别任务没有显著的好处。对于识别，对比方法通过在增强图像上的借口任务来构建特征。蒙面自编码器(MAE)为视觉变形器(Vision transformer)引入了类bert的蒙面图像建模(MIM)预训练，但对于卷积网络来说，它们似乎并不自然和实用。请注意，尽管基于MIM的方法可以恢复掩码图像令牌，但它们在合成完整图像时存在问题，主要是因为没有直接建模完整的数据分布。这些方法不被称为统一的生成和判别模型。相反，我们将它们表示为“半”生成，因为它们与完全生成模型部分相似。

从理论上讲，以判别为目的扩展生成模型并获得两种方法的好处更为实际。最近，由于扩散模型的出现，我们见证了人工生成视觉内容的蓬勃发展，在图像合成、图像编辑和文本到图像的合成方面都报道了最先进的结果。考虑到生成模型的这种能力、多功能性和可扩展性，我们问:扩散模型是否可以复制gpt和T5的成功，成为统一的生成和判别学习器?基于以下观察，我们认为它们是视觉领域非常有前途的替代方案:

1. 已经证明，判别图像或文本编码可以通过端到端生成预训练来学习。“综合分析”。直观地说，完整的生成任务应该包含，并且比半生成掩模建模更具挑战性，这表明图像(或语言)生成与视觉(或语言)理解兼容，而不是相反。生成式预训练范式支持扩散作为一种有意义的判别学习方法
2. 扩散网络被训练成多级去噪自编码器(DAE)。去噪自编码的思想已被广泛应用于判别视觉表征学习。最近，掩模自编码器(MAE)进一步强调了去噪预训练的有效性，这也可以由扩散网络继承-同样，恢复具有大，多尺度噪声的图像是非简单的，并且可能还需要对视觉概念有较高的理解。
3. 基于扩散的表示学习的好处得到了证明。DDPM-seg证实了扩散可以捕获逐像素的语义信息，表明了这种方法的可行性。此外，先前在gan和Transformers中的尝试，即BigBiGAN和iGPT，发现更好的图像生成能力可以转化为改进的特征质量，这表明扩散作为最先进的生成模型，更有能力进行表示学习。如果学习到的知识可以方便地从预训练的模型中转移，那么基于扩散的表示学习器也可以通过大型AIGC项目来促进。

在此分析的推动下，我们研究了融合了生成式预训练和去噪自动编码最佳实践的扩散模型(如表1所示)是否可以学习有效的图像分类表示。我们的方法很简单:我们通过测量中间激活的线性探针和微调精度来评估扩散预训练网络，即去噪扩散自编码器(DDAE)作为特征提取器(图1(b))。对于线性探测，我们将具有特定尺度(或时间步长)的噪声图像传递给DDAE，并检查不同层的激活情况。对于微调，我们在最佳表示层截断DDAE作为图像编码器，并在没有额外噪声的情况下对其进行微调。

我们证实，通过端到端扩散预训练，DDAEs确实学习了强线性可分特征，这些特征位于上采样的中间，可以在图像受到噪声干扰时提取出来。此外，我们通过对噪声配置、训练步骤和数学模型的研究，验证了DDAEs的生成性能和判别性能之间的相关性。对CIFAR-10和Tiny-ImageNet的评估表明，基于扩散的方法首次可与有监督的WideResNet、对比SimCLRs和MAE相媲美。在ImageNet预训练模型上也验证了转移能力，包括由像素空间UNets和潜在空间视觉变形器构建的模型。

我们的研究强调了扩散模型作为统一视觉基础模型的本质。DDAEs作为最先进的生成模型和竞争识别模型的二元性可能会启发视觉预训练及其在这两个领域的应用。通过本文提出的富有洞察力的阐明和观察，在不久的将来，它很有可能从大规模预训练的AIGC模型(如Stable Diffusion)中转移出强大的判别知识。

4. 方法

4.1 将DDAEs评估为判别学习器

从DDAEs中提取有意义和有区别的表示并非易事。尽管像DDIM这样的确定性推理方法能够推导出唯一可识别的编码，但所包含的信息不够紧凑，无法进行分类。还有一种趋势是使用额外的编码器来学习属性操作的表示或分类。相比之下，受iGPT和DDPM-seg的启发，它们在自回归编码器和扩散unet中评估学习到的特征，我们建议直接在预训练的DDAEs中取中间激活。这种方法不需要修改常见的扩散框架，并且与所有现有模型兼容。

从与去噪自编码器的连接中，DDAEs可以在一些隐式编码器-解码器接口上产生线性可分表示，类似于MAE。在此驱动下，我们将之前对GPTs和DDPM的研究扩展到不同框架(DDPM和EDM)下的各种网络骨干网(UNets和DiT)。考虑到具有跳过连接的UNet已经是事实上的设计，我们避免显式地拆分编码器-解码器以防止降低生成性能。然而，提取特征的最佳层仍然是未知的。此外，为了防止预训练和部署之间的差距，图像必须通过一定的尺度进行线性评估。考虑到上述两个事实，我们通过网格搜索研究特征质量和层噪声组合之间的关系，遵循DDPM-seg。

为了应用噪声，我们随机采样ε并使用Eq. 1来获得xt，因为随机噪声和确定性噪声之间没有明显的差异。检查全局平均池化后特征上的线性探针精度，如图1(b)所示。从图2可以看出，层深和噪声尺度作为一个凹函数共同影响特征质量，可以经验地找到其全局最大值点。对于322像素的分辨率，最好的特征位于上采样的中间，而不是像通常的做法那样在最低分辨率处。此外，我们发现用相对较小的噪声扰动图像可以提高线性探头的性能，特别是在EDM训练的DDPM++上，达到95%以上的线性探头精度，超过了经典的AE或VAEs。

这些特性已经在不同的数据集和模型中得到了验证，但在不同的设置下，最佳的层-噪声组合可能会有所不同。在本文的其余部分中，线性探针精度被认为是网格搜索中最高的。为了进行微调，干净的图像被传递给DDAE编码器，这些编码器在最佳层被截断。时间步长输入也固定为最优值。请注意，这对于微调可能不是最好的，因为我们发现如果使用更多的层，微调精度可以进一步提高，并且它对时间步长的输入不那么敏感。然而，我们保持层噪声设置与线性探测一致，以减少符号开销。

4.2 无标签监控层选择

尽管DDAE为判别任务学习了高质量的表示，但它可能依赖于探测带注释的数据来搜索适当的配置。因此，如果无标签度量可以指示训练中表现最好的层，这将是有价值的。受对比学习中的特征分布分析的启发，我们将每一层的特征归一化到单位超球上，并研究它们的对齐性和均匀性。在对比学习中，对准损失直接评价正对之间的距离，而均匀性损失旨在度量单位超球上分布的均匀性。在对比学习中，对准损失直接评价正对之间的距离，而均匀性损失旨在度量单位超球上分布的均匀性:
$$\begin{gathered} L_{align}={\mathbb{E}}_{(x,y)\sim p_{pos}}[||f(x)-f(y)|{|}_2^2] \\ {L}_{uniform}=log{\mathbb{E}}_{x,y\sim p_{data}}[e^{-2||f(x)-f(y)|{|}_2^2}] \end{gathered}$$
为了将它们应用于DDAEs，我们需要正确地定义正对$(x,y)p_{pos}$。考虑到第3.1节中提出的特征提取方法不依赖于确定性噪声，我们假设该表示应该独立于使用的噪声样本ε。为此，我们将这两个指标调整为:
$$\begin{gathered} L_{align}(t)={\mathbb{E}}_{(x,y)\sim p_{data,{ϵ}_1,{ϵ}_2}}[||f_t(x_t^{{ϵ}_1})-f_t(x_t^{{ϵ}_2})|{|}_2^2] \\ {L}_{uniform}(t)=log{\mathbb{E}}_{x,y\sim p_{data,ϵ}}[e^{-2||f_t(x_t^{{ϵ}_1})-f_t(x_t^{{ϵ}_2})|{|}_2^2}] \end{gathered}$$
式中$f_t(\cdot )$为DDAE编码器，$x_t^{\epsilon }$为$\epsilon$处理后的噪声图像，${\epsilon }_1、{\epsilon }_2$为独立采样噪声。这里，第一个损失的含义是不同噪声强度中学习正样本的特征，第二个意思和原本的损失区别在于加入了噪声，这是为了适应扩散模型。

我们在图3中给出了与所选层相关的度量，以验证它们是否与特征质量一致。我们在CIFAR-10训练集上评估了三种DDAE实现:DDPM、DDPM训练的DDPM++和EDM训练的DDPM++。对于每个模型，t被固定到各自的最优值，我们在训练过程中监控指标的轨迹。图3显示，总体而言，Lalign和Luniform与线性探测结果一致(接近左下角)。此外，来自表现良好的层的特征在训练期间显示了对这两个指标的一致改进。这些结果表明，扩散训练可能与对比表示学习有相似之处，并且通过度量选择层可以缓解层搜索问题。

5. 实验

我们首先通过烧蚀研究考察了扩散模型中一些核心设计对生成性能和判别性能的影响。然后，我们在线性探测、微调和ImageNet传输设置下，将结果与CIFAR-10和TinyImageNet数据集上的对应结果进行比较。

所有模型都是从官方(或同等)代码库中检索或训练的。对于图像分类，我们不使用mixup或cutmix等正则化方法，只使用轻量级增强方法。代码库和附录中提供了实现细节以及最佳层噪声设置。

5.1 去噪自编码器

扩散模型可以看作是多级去噪自编码器。基于图2中DDAEs以无监督方式学习强线性可分特征的发现，我们探索了多层次去噪中的哪些设计使DDAEs比经典AEs、VAEs和DAEs更强的表示学习器。我们特别考虑了DDPM中可能有助于改进去噪预训练的两个关键因素:(1)噪声级的数量$(T)$和(2)噪声尺度的范围$({\beta }_{1\dots T})$。

为了研究噪声水平的影响，我们将默认T = 1000降低到T = 512、256和64。我们选择不再降低它，因为当T = 64时，DDPM不能生成有意义的图像。对于每种配置，噪声表${\beta }_{1\dots }$在DDPM默认后，T在$[{\beta }_{min}，{\beta }_{max}]$范围内呈线性间隔，其中${\beta }_{min}=10-4，{\beta }_{max}=0.02$。此外，我们通过使用默认$[{\beta }_{min}，{\beta }_{max}]$调度范围的大一半和小一半来检验噪声尺度的影响。为了公平比较，我们在这两种情况下保持T = 1000。图4(a)显示了噪声配置和训练步骤的影响，图4(b)提供了消噪噪声配置的说明。

降低噪声水平或缩小噪声范围会削弱生成性能和判别性能。训练T = 256或更大尺度的残障扩散模型会导致更显著的性能下降。有趣的是，我们观察到识别能力似乎更少地依赖于密集和广泛的噪声配置，而不是生成。如图4(a)所示，虽然生成FID可能会有很大的下降，但最高线性探头精度的下降幅度小于3%。T = 64的模型甚至比其他残障模型产生更强的特征，尽管它不能产生有意义的图像。

还有一些不关痛痒的实验请看原文。。。

6. 结论

我们提出扩散预训练作为一种统一的方法，同时获得卓越的生成能力和深刻的视觉理解，这可能导致统一的视觉基础模型的发展。然而，作为第一个在规模上研究扩散识别的研究，仍然存在一些局限性和开放性问题。

对于编码器来说，在中间截断DDAEs并不是一种优雅和最佳的实践，我们的方法依赖于探测来找到最佳层。尽管基于度量的方法在CIFAR-10上工作得很好，但在特征线性可分性较差的更复杂的数据集上可能会失败。相反，理想的DDAE主干网可能具有明确的编码器-解码器解纠缠。此外，基于潜在的网络是否能够在更多的识别任务上与像素空间模型竞争，还需要更多的探索。

虽然DDAEs可以达到与一些纯识别模型相当的精度，但它们依赖于更大的模型规模，效率不高。此外，扩散模型需要更长的训练时间才能达到最优的生成性能，这使得它们的扩展成本很高。与其他自我监督方法的关系。我们假设，不同图像之间沿噪声轨迹的对齐可能隐含地有助于判别特性，其操作类似于仅正面对比学习。同时进行的一项名为Consistency Models的研究[52]与基于自蒸馏的对比学习[21,9]有相似之处，该研究已经将这一思想用于去噪输出。此外，我们认为存在另一种将这种一致性或自预测约束整合到DDAE特征中的可能性，类似于将MAEs与对比方法相结合的研究[28]。它还可以通过在指定的编码器-解码器接口上聚合判别特征来缓解先前的问题，并提高学习效率。