为自监督学习重构去噪扩散模型

在这项研究中，作者检验了最初用于图像生成的去噪扩散模型（DDM）的表示学习能力。其理念是解构DDM，逐渐将其转化为经典的去噪自动编码器（DAE）。这一解构过程让大家能够探索现代DDM的各个组成部分如何影响自监督的表征。观察到，只有极少数现代组件对学习好的表征至关重要，而其他许多组件则不重要。研究最终得出了一种高度简化的方法，在很大程度上类似于经典的DAE。

来自：Deconstructing Denoising Diffusion Models for Self-Supervised Learning

背景概述

去噪是当前计算机视觉和其他领域生成模型的核心。如今，这些方法通常被称为去噪扩散模型（DDM），学习了一种去噪自编码器（DAE），它可以去除由扩散过程驱动的多个级别的噪声。这些方法实现了令人印象深刻的图像生成质量，尤其是高分辨率、逼真的图像。这不禁让人想到，这些生成模型非常好，似乎对理解视觉内容有很强的表示。

虽然DAE是当今生成模型的强大力量，但它最初是为了以自监督的方式从数据中学习表示而提出的。在当今的表示学习社区中，DAE最成功的变体可以说是基于“masking noise”，例如预测语言中的缺失文本（例如，BERT）或图像中的缺失patch（例如，MAE）。然而，在概念上，这些基于mask的变体与去除加性噪声（例如，高斯噪声）存在显著不同：mask的token明确指定未知与已知内容，但在分离加性噪声的任务中没有明确的信号可用。然而，今天的DDM主要基于加性噪声，这意味着它们可以在不显式标记未知和已知内容的情况下学习表示。

最近，人们对DDM的表征学习能力越来越感兴趣。特别是，这些研究直接来自最初用于生成的预训练DDM，并评估其表示质量以进行识别。他们报告说，使用这些模型取得了令人鼓舞的结果。然而，这些开创性的研究显然留下了悬而未决的问题：这些现成的模型是为生成而设计的，而不是识别；目前还不清楚表示能力是通过去噪过程还是扩散过程获得的。

在这项工作中，作者对这些最近的探索相关初始化的方向进行了更深入的研究。不像之前使用面向生成的现成（off-the-shelf）DDM，作者训练面向识别的模型。核心是解构DDM，逐步将其转变为经典的DAE。通过这个解构性的研究过程，作者考察了现代DDM的每一个方面，目的是学习表征。这一研究过程使我们对DAE学习良好表示的关键组件有了新的理解。

令人惊讶的是，作者发现主要的关键组件是标记器（tokenizer），它创建了一个低维的潜在空间。有趣的是，这一观察结果在很大程度上独立于tokenizer的具体情况：作者探索了标准VAE、patch-wise VAE、patch-wise AE和patch-wise PCA编码器。发现使DAE能够实现良好的表示的是低维潜在空间，而不是tokenizer的细节。

由于PCA的有效性，解构轨迹最终达到了一个与经典DAE高度相似的简单架构（图1）。使用patch-wise PCA将图像投影到潜在空间上，添加噪声，然后通过逆PCA将其投影回来。然后训练一个自动编码器来预测去噪图像。作者将这种架构称为“潜在去噪自动编码器”（l-DAE，latent Denoising Autoencoder）。
fig1

图1：这种简单的架构在很大程度上类似于经典的DAE（主要区别是将噪声添加到潜在的DAE中），并实现了有竞争力的自监督学习性能。

解构轨迹还揭示了DDM和经典DAE之间的许多其他有趣的特性。例如，即使使用单个噪声水平（即，不使用DDM的噪声scheduling），也可以让l-DAE获得不错的结果。使用多个级别的噪声的作用类似于一种形式的数据增强，这可能是有益的，但不是一个促成因素。根据这一点和其他观察结果，作者认为DDM的表示能力主要是通过去噪驱动的过程获得的，而不是扩散驱动的过程。

最后，作者将结果与以前的基线进行比较。一方面，我们的结果比现有的结果要好得多：这正如预期的那样，因为这是解构的起点。另一方面，结果没有达到基线对比学习方法和基于mask的方法，但差距缩小了。研究表明，在DAE和DDM的方向上还有更多的研究空间。

去噪扩散模型

解构研究始于去噪扩散模型。简要描述了DDM如下。

扩散过程从干净的数据点 $z_{0}$ 开始，按顺序添加噪声。在指定的时间step $t$ ，加噪后的数据 $z_{t}$ 为： $z_{t}=\gamma_{t}z_{0}+\sigma_{t}\epsilon$ 其中， $\epsilon\sim N(0,I)$ 是一个从高斯分布采样的noise map， $\gamma_{t}$ 和 $\sigma_{t}$ 分别定义了信号和噪声的scaling因子。此外， $\gamma_{t}^{2}+\sigma_{t}^{2}=1$ 。

以时间step $t$ 为条件，学习去噪扩散模型来去除噪声。不同于DAE预测干净的输入，DDM预测噪声 $\epsilon$ 。损失为最小化： $||\epsilon-net(z_{t})||^{2}$ 其中， $net(z_{t})$ 是网络输出。在给定以时间步长 $t$ 为条件的噪声调度下，针对多个噪声水平来训练网络。在生成过程中，迭代地应用经过训练的模型，直到它达到干净信号 $z_{0}$ 。

DDM可以在两种类型的输入空间上操作。一个是原始像素空间，其中原始图像 $x_{0}$ 直接用作 $z_{0}$ 。另一种选择是在tokenizer产生的潜在空间上构建DDM，见图2b。在这种情况下，预训练的tokenizer $f$ （通常是一个自编码器），比如VQVAE被用于 $z_{0}=f(x_{0})$ 。

DiT：Diffusion Transformer
该研究从DiT开始，选择这种基于Transformer的DDM有几个原因：i.与其他基于UNet的DDM不同，基于Transformer的架构可以与Transformer驱动的其他自监督学习基线进行更公平的比较；ii.DiT在编码器和解码器之间有更清晰的区别，而UNet的编码器和解码器通过skip连接，并且在评估编码器时可能需要在网络改动上付出额外的努力；iii.DiT的训练速度比其他基于UNet的DDM快得多，同时实现了更好的生成质量。

作者使用DiT-Large（DiT-L）变体作为DDM基线。在DiT-L中，编码器和解码器加在一起的大小为ViT-L（24个blocks）。作者评估编码器的表示质量（linear probe accuracy），编码器有12个blocks，称为 $\frac{1}{2}L$ （half large）。

Tokenizer
DiT是潜在扩散模型（LDM）的一种形式，它使用VQGAN tokenizer。VQGAN tokenizer将(256,256,3)的输入图像转换到 $(32, 32, 4)$ 的latent map。

开始的baseline
默认情况下，作者在ImageNet上以256×256像素的分辨率训练400个epochs的模型。DiT baseline的结果见表1（第一行）。使用DiT-L，作者报告使用其 $\frac{1}{2}L$ 编码器的线性探头精度为57.5%。该DiT-L模型的生成质量为11.6。

重构扩散模型

解构轨迹分为三个阶段。作者首先调整了DiT中以生成为中心的设置，使其更倾向于自监督学习。接下来，逐步解构和简化tokenizer。最后，作者试图扭转尽可能多的DDM驱动设计，将模型推向经典的DAE。

用于自监督学习的DDM重定向

虽然DDM在概念上是DAE的一种形式，但它最初是为了生成图像而开发的。DDM中的许多设计都面向生成任务。有些设计对于自监督学习是不合法的（例如，涉及类别标签）；如果不考虑视觉质量，则其他一些是不必要的。在本节中，作者为了自监督学习的目的重新调整了DDM基线，总结见表1。
tab1

表1：从DiT baseline开始，并在ImageNet上评估其线性探针精度（acc）。每一行都基于对前一行的修改。使用类标签的灰色条目不是自监督学习的结果。

移除类别-条件
高质量的DDM通常在类标签上进行条件训练，这可以在很大程度上提高生成质量。但是，在自监督学习中，标签的使用是不行的。作为第一步，作者在baseline中删除类条件。令人惊讶的是，尽管生成质量如预期的那样受到极大的损害（FID从11.6提高到34.2），但去除类别条件显著提高了线性探针的精度，从57.5%提高到62.1%（表1）。作者假设，直接将模型条件化为类标签会减少模型对编码的信息需求，删除类条件可以强制模型学习更多语义。

重构VQGAN
在baseline中，LDM的VQGAN tokenizer使用多个损失项进行训练：i.自动编码重建损失；ii.KL散度正则化损失；iii.基于针对ImageNet分类训练的VGGNet的感知损失；iv.具有判别器的对抗性损失。

由于感知损失涉及有监督的预训练网络，使用用这种损失训练的VQGAN是不合法的。相反，作者在其中消除了感知损失。使用这种tokenizer将线性探针的准确率从62.5%显著降低到58.4%（表1）。这种比较表明，用感知损失（带有类标签）训练的tokenizer本身提供了语义表示。

作者训练下一个VQGAN tokenizer，它可以进一步消除对抗性损失。它将线性探头的精度从58.4%略微提高到59.0%（表1）。tokenizer在这一点上本质上是一个VAE。还注意到，消除这两种损失都会损害生成质量。

替换噪声调度
在生成任务中，目标是逐步将噪声图转换为图像。因此，原始噪声调度在噪声很大的图像上花费了许多时间步长（图3）。如果模型不是面向生成的，那么这是没有必要的。

为了进行自监督学习，作者研究了一种更简单的噪声调度。具体来说，让 $γ^{2}_{t}$ 在[0,1]的范围内线性衰减（图3）。这一变化大大提高了线性探头的精度，从59.0%提高到63.4%（表1），表明原始调度过于关注噪声较大的状态。另一方面，正如预期的那样，这样做会进一步损害生成能力，导致FID为93.2。
fig3

图3：原始调度器设置 $\gamma_{t}^{2}=\prod_{s=1}^{t}(1-\beta_{s})$ ，如果使用线性调度器，可以减少step。让模型一开始就关注比较干净的数据，而不是纯纯的噪声。

总体而言，表1中的结果表明，自监督学习表现与生成质量无关。DDM的表示能力不一定是其生成能力的结果。

重构Tokenizer

接下来，通过进行实质性的简化来进一步解构VAE tokenizer。作者将以下四种自动编码器变体作为tokenizer进行比较，每种变体都是前一种的简化版本。

卷积VAE
到目前为止，解构被引导到VAE tokenizer，现在将VAE的编码器和解码器设置为卷积网络 $f(\cdot)$ 和 $g(\cdot)$ 。卷积VAE被如下损失最小化： $x-g(f(x))||^{2}+KL[f(x)|N(0,I)]$ 其中， $x$ 是VAE的输入图像。

Patch-wise VAE
接下来考虑简化情况，VAE的编码器和解码器都是线性投影，并且VAE的输入 $x$ 是一个patch，最小化损失为： $x-U^{T}Vx||^{2}+KL[Vx|N(0,I)]$ 这里， $x$ 表示flatten为 $D$ 维向量的patch。 $U$ 和 $V$ 都是 $d \times D$ 矩阵，其中 $d$ 是潜在空间的维数。patch大小设置为16×16像素。

Patch-wise AE
作者通过去除正则化项对VAE进行进一步简化： $x-U^{T}Vx||^{2}$ 因此，这个tokenizer本质上是patch上的自动编码器（AE），编码器和解码器都是线性投影。

Patch-wise PCA
最后，考虑一个更简单的变体，它在patch空间上执行主成分分析（PCA）。很容易证明PCA等效于AE的一个特殊情况： $x-V^{T}Vx||^{2}$ 其中， $V$ 符合 $VV^{T}=I\in R^{d\times d}$ ，PCA可以简单地通过在一大组随机采样的patch上进行特征分解来计算，不需要基于梯度的训练。

由于使用patch的简单性，对于三个patch tokenizer，我们可以在patch空间中可视化它们的滤波器（图4）。
fig4

图4：patch tokenizer的可视化。每个滤波器对应于线性投影矩阵 $V : d \times D$ 的一行，为了可视化，将其整形为16×16×3。

表2总结了使用这四种tokenizer变体的DiT的线性探针准确性。作者展示了关于潜在维度“per token”的结果。PCA tokenizer的有效性在很大程度上将现代DDM推向经典DAE。
tab2

表2：线性探针精度与latent维度的权衡。利用DiT模型，作者研究了用于计算潜在空间的tokenizer的四种变体。作者改变潜在空间的维度 $d$ （per token）。通过上面的图显示出来。尽管在体系结构和损失函数方面存在差异，但tokenizer的所有四种变体都表现出相似的趋势。

高分辨率、基于像素的DDM在自监督学习方面较差
在继续之前，报告了一个与上述观察结果一致的额外的消融实验。具体来说，考虑一种naive tokenizer，它对从resized image中提取的patch执行identity mapping。在这种情况下，token是由patch的所有像素组成的展平向量。在图5中，作者展示了这种“pixel-based” tokenizer的结果，该tokenizer分别对256、128、64和32的图像大小进行操作，patch大小为16、8、4、2。这些token化空间的“潜在”维度分别为768、192、48和12。在所有情况下，Transformer的序列长度都保持不变（256）。
fig5

图5：pixel-based的tokenizer的线性探测结果，分别在256、128、64和32的图像大小上操作，补丁大小为16、8、4、2。这些token化空间的“潜在”维度分别为768、192、48和12。与之前研究的其他tokenizer类似，这种基于像素的标记器呈现出类似的趋势：相对较小的潜在空间维度是最优的。

这些比较表明，tokenizer和由此产生的潜在空间对于DDM或DAE在自监督学习场景中至关重要。特别是，在像素空间上应用具有加性高斯噪声的经典DAE会导致较差的结果。

走向经典的DAE

接下来，继续解构轨迹，目标是尽可能接近经典的DAE。作者试图消除当前基于PCA的DDM和经典DAE之间仍然存在的每一个方面。通过这个解构过程，作者更好地理解了每一个现代设计如何影响经典DAE。表3给出了下面讨论的结果。
tab3

表3：从Patch-wise PCA标记器开始，迈向经典DAE。每一行都基于对前一行的修改。

预测干净的数据（而不是噪声）
虽然现代DDM通常预测噪声 $ε$ ，但经典DAE预测的是干净数据。作者通过最小化以下损失函数来检验这种差异： $\lambda_{t}||z_{0}-net(z_{t})||^{2}$ 其中， $z_{0}$ 是干净数据（latent空间中）， $\lambda_{t}=\gamma_{t}^{2}$ 是基于 $t$ 的loss权重，为不同等级的噪声引入去噪贡献。直观地说，它只是对更干净的数据的损失项给予了更多的重视（large $\gamma_{t}^{2}$ ）。

随着预测干净数据（而不是噪声）的修改，线性探头的精度从65.1%下降到62.4%（表3）。这表明预测目标的选择会影响表示质量。尽管在这一步中受到了下降影响，但从现在起，作者依然坚持这种修改，因为目标是向经典DAE迈进。

移除输入scaling
在现代DDM中，输入按 $γ_t$ 因子缩放。这在传统DAE中并不常见。接下来，作者研究去除输入缩放因子，即设置 $γ_t=1$ 。由于 $γ_t$ 是固定的，我们需要直接在 $σ_t$ 上定义一个噪声调度。作者简单地将 $σ_t$ 设置为从0到 $\sqrt{2}$ 的线性调度。此外，根据经验将等式中的权重设置为 $λ_t=1/(1+σ^{2}_{t})$ 。此后获得了63.6%的准确率（表3），这有缩放的62.4%相比是有利的。这表明缩放数据是不必要的。

使用逆PCA对图像空间进行运算
到目前为止，对于探索过的所有条目，该模型在tokenizer产生的潜在空间上运行。理想情况下，作者希望可以直接在图像空间上工作，同时仍然具有良好的准确性。利用主成分分析，可以通过逆主成分分析来实现这一目标。这个想法如图1所示。特别地，通过PCA基（即 $V$ ）将输入图像投影到潜在空间中，在潜在空间中添加噪声，并通过逆PCA基（ $V^{T}$ ）将带噪的潜在空间投影回图像空间。图1（底部）显示了在潜在空间中添加了噪声的示例图像。有了这个有噪声的图像作为网络的输入，我们可以应用直接对图像进行操作的标准ViT网络，就好像没有tokenizer一样。

进一步将其应用于输出侧（即，用逆PCA预测图像空间上的输出）具有63.9%的准确度。结果表明，用逆PCA在图像空间上操作可以获得与在潜在空间上操作类似的结果。

预测原始图像
虽然逆PCA可以在图像空间中产生预测目标，但该目标不是原始图像。这是因为PCA对于任何降维都是有损编码器。相比之下，直接预测原始图像是一种更自然的解决方案。当让网络预测原始图像时，引入的“噪声”包括两部分：i.加性高斯噪声，其固有维数为 $d$ ，ii.PCA重建误差，其固有维数为 $D - d$ （ $D$ 为768）。作者对这两部分的损失进行了不同的加权。

形式上，使用干净的原始图像 $x_{0}$ 和网络预测 $net(x_{t})$ ，可以计算投影到全PCA空间上的残差 $r$ ： $r=V(x_{0}-net(x_{t}))$ ，这里 $V$ 是PCA基 $D\times D$ 。然后最小化： $\lambda_{t}\sum_{i=1}^{D}w_{i}r_{i}^{2}$ 其中， $i$ 表示向量 $r$ 的第 $i$ 维。 $w_{i}$ 对PCA重建误差的损失进行加权。采用该公式，预测原始图像的线性探头精度达到64.5%（表3）。

这个变体在概念上非常简单：它的输入是一个噪声图像，其噪声来自PCA潜在空间中，它的预测是原始的干净图像（图1）。

单级噪声
最后，出于好奇，作者进一步研究了具有单级噪声的变体。并注意到，由噪声调度给出的多级噪声是由DDM中的扩散过程驱动的；在传统DAE中，这在概念上是不必要的。

作者将噪声水平 $σ$ 固定为常数（ $\sqrt{1/3}$ ）。使用这种单级噪声实现了61.5%的良好精度，与多级噪声对应（64.5%）相比，降低了3%。使用多级别的噪声类似于DAE中的一种数据增强形式：这是有益的，但不是一个促成因素。这也意味着DDM的表示能力主要是通过去噪驱动过程获得的，而不是扩散驱动过程。