何凯明: 扩散模型的解构研究

1 介绍

本研究对去噪扩散模型（DDM）进行了解构，发现其关键组件是分词器，而其他组件并非必要。DDM的表现能力主要来自去噪过程而非扩散过程。研究还发现，通过消除类标签条件化项和KL正则化项，使用补丁式分词器可获得与卷积VAE相当的表现。最后，将现代DDM推向经典DAE，通过消除输入缩放和直接定义噪声调度，可获得更好的结果。

本研究揭示了自监督学习中去噪扩散模型DDM的表示能力和生成质量之间的关系，以及分词器和潜在空间对DDM/DAE的影响。

2 相关工作

在机器学习和计算机视觉的历史中，图像生成与无监督或自监督学习的发展密切相关。模型在没有标记的数据上训练，学习捕捉输入数据的底层分布。生成对抗网络（GAN）和变分自编码器（VAEs）等模型生成高保真度数据的能力表明它有学习良好表示的潜力。图像修复和去噪扩散模型（DDM）的出色生成性能也引起了人们对表示学习潜力的关注。然而，虽然模型的生成能力表明了一定的理解水平，但这并不一定转化为对下游任务有用的表示。去噪自编码器（DAE）为基于自编码器的表示学习奠定了基础，但其成功主要局限于涉及基于遮罩的污染的场景。

图2 一个经典DAE和现代DDM的对比。(a)一个在图像空间上添加和预测噪声的经典DAE。(b)在潜在空间上操作的最新技术DDM（例如，LDM[33]，DIT[32]），其中噪声也是在该空间上被添加和预测的。

3 背景：去噪扩散模型

我们的解构研究从去噪扩散模型（DDM）出发，对其进行深入浅出的介绍。我们参照[11，32]对所使用的DDM进行详细的阐述。

扩散过程从一个干净的数据点（z0）开始，并逐步向其添加噪声。在指定的时间步长t，经过噪声处理后的数据点zt由以下公式给出：

其中ϵ∼N (0, I)是从高斯分布采样的噪声图，γt和σt分别定义了信号和噪声的缩放因子。默认情况下，γt2+σt2=1[29, 11]。

学习一个去噪扩散模型，基于时间步长t，用于去除噪声。与原始DAE不同，现代DDM通常预测噪声。我们最小化一个损失函数，其形式如下：

其中，net(zt)为网络输出，网络针对给定的噪声级别进行训练，训练条件为时间步长t，在生成过程中，训练后的模型迭代应用直到达到清洁信号z0。

DDMs有两种类型的输入空间：原始像素空间和由分隔器产生的潜在空间，如图2。在原始像素空间中，原始图像x0直接用作z0。而在潜在空间上构建的DDMs，则使用预训练的分隔器f将图像x0映射到其潜在空间z0=f(x0)。

扩散变换（DiT）。我们的研究基于扩散变换（DiT），选择基于Transformer的DDM的原因包括：（i）可与其他自监督学习基准公平比较；（ii）DiT的编码器和解码器有明确区分，无需额外评估编码器；（iii）DiT训练速度更快，生成质量更好。我们使用DiT-Large（DiT-L）变体作为DDM基准，其中编码器和解码器的大小与ViT-L相同（24个块）。我们评估编码器的表示质量（线性探针精度），包含12个块，称为“1/2 L”（一半大）。

分词器。分词器DiT是一种隐式扩散模型（LDM），使用VQGAN分词器。VQGAN分词器将输入图像的256×256×3转换为32×32×4的隐式映射，步长为8。

开始基准线。默认情况下，我们在ImageNet上使用分辨率为256×256像素的图像进行模型训练，训练周期为400轮。具体实现细节详见第A节。

我们的DiT基准结果在表1中有所报告。使用DiT-L，我们发现其12L编码器的线性探针精度为57.5%。该模型的生成质量（FID-50K）为11.6，这是我们研究的起点。尽管实现细节有所不同，但我们的起点与最近的研究（更具体地说，DDAE）在概念上是相符的，该研究评估了在线性探针协议下的现成DDMs。

表1 用于自监督学习的重新定向DDM。我们从DiT[32]基准开始，并评估其在ImageNet上的线性探针准确性。每行基于紧接的前一行的修改。在灰色条目中使用类标签的条目对于自监督学习来说不是合法的结果。请参阅第4.1节以获取说明。

4 解构去噪扩散模型

我们的解构轨迹分为三个阶段。首先，调整DiT的代际中心设置，强化自我监督学习（第4.1节）。其次，分解和简化分词器步骤（第4.2节）。最后，逆转DDM驱动设计，向经典DAE靠近（第4.3节）。在4.4节总结了解构过程的关键经验。

4.1 重新定位 DDM 以实现自我监督学习

虽然DDM在概念上属于DAE的一种，但其最初是为了图像生成而开发的。DDM中的设计大多针对生成任务，部分设计不适合自监督学习（如涉及类别标签的情况）。若不考虑视觉质量，其他一些设计也不再必要。因此，本节将调整DDM基准以适应自监督学习的目的，总结在表1中。

去除类别条件。在自监督学习研究中，高质量的DDM通常在类别标签上进行训练，以提高生成质量。然而，在我们的研究中，使用类别标签是非法的。因此，我们首先在我们的基准中去除类别条件。研究发现，去除类条件化后，线性探针精度从57.5%提高到62.1%，尽管生成质量受到很大影响（FID从11.6到34.2）。这可能是因为直接对模型进行类标签条件化会减少模型对编码类标签相关信息的需求。去除类条件化可以迫使模型学习更多的语义。

解构VQGAN。我们的基准中使用的VQGAN tokenizer是由LDM提出的，并由DiT继承。它使用多种损失项进行训练，包括自动编码重构损失、KL散度正则化损失、基于ImageNet分类训练的监督VGG网络的感知损失以及带有判别器的对抗性损失。我们通过删除感知损失训练了另一个标记器，该标记器的线性探测精度从62.5％显著降低到58.4％。这一比较表明，使用感知损失训练的标记器本身提供了语义表示。我们还注意到，去除任何一种损失都会对生成质量造成损害。因此，我们进一步训练下一代VQGAN tokenizer以进一步消除对抗损失，并略微提高了线性探测准确率。总的来说，我们的tokenizer基本上是一个VAE。

替换噪声图谱。在生成任务中，我们希望将噪声图谱转化为图像。原始的噪声图谱在嘈杂图像上花费大量时间（图3），但若模型非生成导向则此为不必要。为自我监督学习目的，我们研究了简化噪声时间表，使γt2在1>γt2≥0的范围内线性下降（图3）。这使模型能更多地在清晰图像上投入能力。此改动显著提升线性探测准确性，从59.0%提高至63.4%（表1），显示原始时间表过于关注嘈杂区。然而，此改变也损害生成能力，导致FID为93.2。

图3 噪声调度。原始调度[23, 32]，它通过线性调度β来设置γt 2=Πt s=1(1−βs)，在非常嘈杂的图像（小γ）上花费很多时间步骤。相反，我们使用一个简单的调度，它在γ 2上是线性的，这提供了不太嘈杂的图像。

总的来说，表1中的结果揭示了自监督学习性能与生成质量之间没有相关性。DDM的表示能力不一定是其生成能力的结果。

4.2 解构分词器

接下来，我们将通过实质性的简化来解构VAE的分词器。我们将自动编码器作为分词器，比较以下四种变体：

• 卷积VAE。卷积VAE的编码器f(·)和解码器g(·)是深度卷积神经网络。这种卷积VAE通过最小化以下损失函数来进行优化：

x是VAE的输入图像。第一项为重构损失，第二项为f(x)的潜在分布与单位高斯分布之间的Kullback-Leibler散度。

• Patch-wise VAE。这是一种简化的VAE模型，其中编码器和解码器都是线性投影，输入x是一个补丁。在训练过程中，该模型最小化损失。

这里x表示一个被展平为D维向量的补丁。U和V是d×D的矩阵，其中d是潜在空间的维度。补丁大小设置为16×16像素。

• 补丁级别的AE。移除正则化项对VAE简化：

这个分词器本质上是补丁的自动编码器（AE），其编码器和解码器都是线性投影。

• 补丁级别的PCA。探讨一种更简单的变体，即补丁空间上的主成分分析（PCA），这是自动编码器（AE）的一种特殊情况。

其中V满足VVT=I(d×d单位矩阵)。PCA基可以通过特征分解大量随机采样的补丁来简单计算，无需基于梯度的训练。

由于补丁的使用非常简单，对于这三个补丁式分词器，我们可以可视化它们在补丁空间中的滤波器（图4）。

图4 补丁式分词器的可视化。每个滤波器对应线性投影矩阵V的一行（d×D），为了可视化将其重塑为16×16×3。这里的d=16。

表2总结了使用这四种分词器变体对DiT的线性探测精度。我们展示了与“每个标记”的潜在维度相关的结果。我们得出以下观察结果。

表2 线性探针精度与潜在维度。使用DiT模型，我们研究了四种变体的令牌化器来计算潜在空间。我们改变潜在空间的维度d（每个令牌）。该表通过上面的图表可视化。四种变体令牌化器的趋势相似，尽管它们在架构和损失函数方面存在差异。“conv.VAE”的63.4%条目与表1最后一行的条目相同。

标记器的隐含维度对于DDM在自监督学习中发挥良好的作用至关重要。

如表2所示，四种分词器变体虽有架构和损失函数的差异，但表现趋势相似。有趣的是，最优维度较低（d=16或32），而每个补丁的全维度要高得多（768）。令人惊讶的是，卷积变分自编码器（ConvVAE）既非必要也不利；所有基于补丁的转换器表现相似且优于ConvVAE。此外，KL正则化项不必要，因为AE和PCA变体表现良好。更令人惊讶的是，PCA tokenizer也表现良好，它不需要基于梯度的训练。使用预计算的PCA基，PCA tokenizer类似于图像预处理而非网络架构。其有效性有助于将现代DDM推向经典的DAE，我们将在下一节展示。

高分辨率、基于像素的DDMs对于自监督学习来说是不好的。

我们考虑了一个“朴素的分词器”，它在缩放图像中提取的补丁上进行身份映射，提取的标记是扁平向量，由补丁中的所有像素组成。如图5，在不同图像大小下，基于像素的标记器呈现出类似的趋势，最优维度为d=48，这相当于图像大小为64，补丁大小为4时的图像大小。在自监督学习场景中，分词器和潜在空间对于DDM/DAE具有竞争力。传统DAE在像素空间上应用带有加性高斯噪声会导致结果不佳。

4.3 朝向经典型去噪自编码器

我们将继续解构轨迹，以更接近经典的DAE问题为目标。我们将努力消除PCA基础上的DDM和经典DAE之间的差异，并讨论现代设计如何影响经典DAE。表3将总结我们的讨论结果。

表3 从补丁式的PCA分词器开始，逐步转向经典的有向无环图（DAE）。每行基于紧接的前一行的修改。请参阅第4.3节以获取描述。

预测干净数据（而不是噪声）。

虽然现代DDM通常预测噪声ε（参见式（2）），但经典DAE预测的是干净数据。我们通过最小化以下损失函数来考察这种差异：

z0是干净数据（在潜在空间中），net(zt)是网络预测。λt是随时间变化的损失权重，用于平衡不同噪声水平的影响[34]。建议根据[34]设置λt = γt2/σt2。我们发现设置λt = γt2在我们的场景中效果更好。直观地说，它只是给清洁数据（γt更大）的损失项赋予更多的权重。

修改对干净数据的预测（而不是噪声）导致线性探针精度从65.1%降至62.4%（表3）。这显示预测目标的选取对表示质量有影响。尽管精度有所下降，我们仍坚持这一修改，因为我们的目标是迈向经典 DAE。

移除输入缩放。在DDMs中，输入通常被γt因子缩放，但在经典DAE中并不常见。研究移除输入缩放，即令γt≡1，需要在σt上直接定义一个噪声调度。简单地设定σt为从0到√2的线性调度，并经验性地设定式（3）中的权重λt。在确定γt≡1之后，达到了63.6%的精度，这表明按γt缩放数据在我们的场景中是不必要的。在图像空间上执行逆主成分分析（PCA），通过PCA基将输入图像投影到潜在空间，在潜在空间中添加噪声，并通过逆PCA基将带噪声的潜在空间图像投影回图像空间。在输入侧应用这种修改具有63.6％的准确性，进一步将其应用于输出侧具有63.9％的准确性。两种结果都表明，使用反PCA在图像空间中进行操作可以获得与在潜在空间中进行操作类似的结果。

预测原始图像。虽然反PCA可以在图像空间中产生一个预测目标，但这个目标并不是原始图像，这是因为PCA对于任何降低的维度d来说是一个有损编码器。直接预测原始图像是一个更自然的解决方案。引入的“噪声”包括两部分：高斯噪声和PCA重建误差。我们对这两部分的损失赋予不同的权重。形式上，我们使用原始图像x0和网络预测net(xt)来计算投影到主成分分析(PCA)空间的残差r：r ＝V(x0−net(xt))，其中V是一个D×D的矩阵，表示PCA基的完整表示。然后我们最小化以下损失函数：

i表示向量r的第i维。对于前d维，权重wi设为1；对于d到D维，权重wi设为0.1。这意味着PCA重建误差的损失在wi中的权重较低。采用这种方法，线性探查预测原始图像的准确率达到了64.5%（表3）。

该变体在概念上非常简单：其输入是在PCA隐空间中添加噪声的图像，其预测是原始清洁图像（图1）。

单噪声水平。我们研究了单噪声水平变体，发现噪声调度是DDM中扩散过程驱动力的一种属性，但在经典DAE中并不必要。我们将噪声水平设定为常数，使用单级噪声达到了61.5％的准确度，相对于多级噪声的对等物下降了3％。这表明DDM的表示能力主要通过去噪驱动的过程获得，而不是扩散驱动的过程。由于多级噪声既有用且概念上简单，我们将其保留在参赛作品中。

4.4 总结

我们对现代DDM进行了分析，提炼出经典DAE。在DDM中，我们保留了两个核心设计概念：一是低维潜在空间中的噪声添加；二是多级噪声。我们基于表3中的信息，创建了名为“潜在去噪自编码器”的DAE实例，简称l-DAE。

5 分析与比较

可视化潜在噪声。l-DAE是另一种形式的DAE，它的目标是学习消除潜在空间中的噪声。由于PCA的简单性，我们能够轻松地通过逆PCA可视化潜在的噪声。图7比较了添加到像素中的噪声与潜在的噪声。与像素噪声不同，潜在噪声与图像分辨率的关系不大。使用PCA作为标记器，潜在噪声的模式主要由补丁大小决定。这可以理解为使用补丁而不是像素来解析图像，类似于MAE中的做法，其中掩蔽的是补丁而不是单个像素。

图7 可视化：像素噪声 vs. 潜在噪声。左图：清洁图像，256×256像素。中间图：向像素空间中添加高斯噪声。右图：通过使用逆PCA将PCA分词器生成的潜在噪声添加到图像空间中，以可视化潜在噪声。两种情况下的σ=√1/3。

去噪结果。图8展示了基于l-DAE的去噪结果示例。我们的方法在噪声大量存在的情况下也能得出合理预测，这是由于图像恢复的神经网络研究已深入进行。可视化可能帮助我们更好地理解L-DAE如何学习良好表示。在潜在空间添加的大量噪声为模型设定了高挑战任务；基于一或几个局部噪声补丁来预测内容不易（甚至对人类也难）；模型被迫学习更高层次、更整体的语义以理解底层物体和场景。

图8 对ImageNet验证图像进行评估的l-DAE去噪结果。这个去噪问题作为一个伪任务，鼓励网络以自我监督的方式学习有意义的表示。对于每个案例，我们展示：(左)清洁图像；(中)网络输入的噪声图像，其中噪声添加到潜在空间；(右)去噪输出。

数据扩充。我们目前所提出的所有模型都没有进行数据扩充，只使用了图像的中心裁剪，遵循[11, 32]。我们进一步探索了针对我们的最终l-DAE模型的温和数据扩充（随机调整大小的裁剪）。该扩充方法有助于提高模型的泛化能力。

这表明，l-DAE的表示学习能力与其对数据增强的依赖关系不大，这与MAE的行为相似，但与对比学习方法的行为完全不同。

训练周期。我们目前所有的实验都是基于400周期的训练。我们也研究了800周期和1600周期的训练，参考了MAE的研究

MAE在400到800个epochs之间有显著增益（4%），而MoCo v3在300到600个epochs之间增益很小（0.2%）。

模型大小。我们目前所有的模型都是基于DiT-L变体，编码器和解码器都是“ViT-1/2 L”（ViT-L的一半深度）。我们训练了不同大小的模型，其中编码器是ViT-B或ViT-L，解码器大小与编码器相同。

我们目前的模型都是基于DiT-L变体，其编码器和解码器都是“ViT-1/2 L”。在此基础上，我们还训练了不同大小的模型，其中编码器是ViT-B或ViT-L，解码器始终与编码器大小相同。

我们观察到模型规模与良好的缩放行为：从ViT-B到ViT-L的缩放具有10.6%的巨大收益。在MAE中也观察到了类似的缩放行为[21]，从ViT-B到ViT-L的缩放具有7.8%的收益。

与之前的基线进行比较。为了对比MoCo v3和MAE的性能，我们在表4中对比了这两种自监督学习算法。MoCo v3是对比学习算法，MAE则是基于掩码的算法。实验结果显示，l-DAE相比MAE表现良好，仅出现1.4%（ViT-B）或0.8%（ViT-L）的退化。训练设置公平，两者都训练1600个周期，并使用随机裁剪作为数据增强。MAE训练效率更高，仅对未遮罩的补丁操作。尽管我们已经大幅缩小MAE与DAE驱动的方法间的准确性差距，但基于自编码器的方法（MAE和l-DAE）在这种协议下仍然无法与对比学习方法相媲美，特别是在模型规模较小的情况下。我们希望我们的研究能够引起更多对基于自编码器自监督学习方法的研究的关注。