Diffusion Models Beat GANs on Image Classification

使用扩散模型去做图像分类

摘要

虽然许多无监督学习模型关注于一个任务家族，无论是生成性的还是歧视性的，但我们探索了一个统一的表征学习者的可能性：一个使用单一的训练前阶段来同时处理两个任务家族的模型。我们确定扩散模型作为一个主要的候选模型。扩散模型作为一种最先进的图像生成、去噪、内画、超分辨率、操作等方法而日益突出。这种模型包括训练一个U-Net来迭代预测和去除噪声，所得到的模型可以合成高保真度、多样化、新颖的图像。U-Net体系结构，作为一种基于卷积的体系结构，以中间特征映射的形式生成一组不同的特征表示。我们提出了我们的发现，这些嵌入是有用的噪声预测任务之外，因为它们包含鉴别信息，也可以用于分类。我们探索了提取和使用这些嵌入进行分类任务的最佳方法，展示了在ImageNet分类任务上的有希望的结果。我们发现，通过仔细的特征选择和池化，扩散模型在分类任务中优于类似的生成鉴别BigBiGAN方法。我们研究了迁移学习机制中的扩散模型，检查了它们在几个细粒度视觉分类数据集上的性能。我们将这些嵌入与由竞争架构和分类任务的预训练生成的嵌入进行比较。

方法

        评估自我监督预训练有效性的两种最常见的方法是线性探测和微调，我们尽可能地与VISSL [71]记录的流行食谱进行匹配。虽然相关，但这些测试了训练前的不同属性。线性探测在冻结特征之上学习批归一化和线性层，测试学习到的特征表示的效用它显示了预训练前是否学习解纠缠表示，以及这些特征是否具有有意义的语义相关性。另一方面，微调学习批归一化和线性层，但没有冻结特征。在微调机制下，我们将预训练方法作为一种昂贵的权值初始化方法，并对整个体系结构进行分类再训练。在本文中，我们更多地关注冻结特征的代表能力，这在细粒度分类和少量镜头学习等领域特别感兴趣，这些领域的数据可能不足以进行微调。此外，这允许我们陈述关于学习的特性的效用，而不是学习的权重。我们注意到，扩散模型就像规则的卷积网，因为它们本身并不产生线性特征，而是在网络的不同点上生成一系列特征映射。因此，与其他cnn类似，

我们使用池化和扁平化的组合来生成每个图像的向量特征表示。通道的深度和特征图的大小自然都相当大，所以除了标准的池化之外，我们还尝试了其他的方法。

我们研究了多层感知器头。由于规模较大，

我们也尝试将cnn作为一种学习过的池化机制，并在附录中给出了更完整的设计细节。

我们还研究了注意力头对空间和通道信息进行适当聚合的能力，详见附录。

结果：Main Results: ImageNet Classification

首先，我们在表1中展示了扩散的线性探测性能，使用了我们通过第4.1.1节中描述的消融所选择的设置。作为一个基线，我们与扩散预训练的分类器进行比较，因为它使用相同的U-Net编码器。我们还提供了与其他统一模型的比较： BigBiGAN [3]和MAGE [8]。我们在生成和分类方面都优于BigBiGAN，特别是当BigBiGAN被迫处理更高分辨率的256×256张图像时。因此，扩散模型在图像分类（和生成）方面优于GANs。我们承认，与仅分类模型相比，扩散还不是最先进的，有超过10%与顶级模型精度的差距，或与强大的统一MAGE模型相比。然而，我们注意到，我们无法完全匹配模拟其他方法的线性探测设置所需的资源。例如，MAE [10]训练他们的线性层为100个时代，每批有16,384张图像。因此，很难与这种方法进行“公平”的比较。

我们在类似的条件下进行微调。如表2所示，引导扩散滞后于其他使用分类特定调整的方法。无论如何，这是一个比U-Net编码器更好的结果（+9.38%），这表明引导扩散是一个有用的分类预训练。

如前所述，我们还提出了几种方法来处理U-Net表示的大空间和信道维度。天真地，我们可以使用具有不同初始池化的单一线性层，并展示了不同池化维度的结果。

或者，我们也可以使用一个更强大的MLP、CNN或注意力头，来解决特征图的高度、宽度和深度的不同方面。

为了公平起见，我们在各种池化设置下，用与线性层相当的参数计数来训练cnn、mlp和注意头。我们在图1（右）中展示了ImageNet-50上这些头部的结果，

附录中显示了完整的数值结果和模型细节。我们注意到注意力头的表现最好。

在表3中，我们在ImageNet（所有类）上尝试了表现最好的注意头，并发现它的性能明显优于简单的线性探针。不管池。这表明分类头是从扩散模型中提取有用表示的重要机制，它可以扩展到其他生成模型中。

如图1所示，从扩散模型中提取良好的特征需要仔细考虑噪声步长、块数和池化大小。我们开始搜索ImageNet的那个超参数空间。我们为噪声设置了一个大致对数等距时间步长的搜索空间。我们在U-Net瓶颈周围的均匀间隔尝试几个块。我们还讨论了特征的高度和宽度（池化）。从我们的线性搜索中，如图2所示，我们发现t应该设置为90，b应该设置为24。然而，正如我们在第4.2节中讨论的，我们发现这些设置至少在某种程度上依赖于数据。因此，虽然在这项工作中，我们提取了一些一般的设置和原则，但自动选择和组合的特征可以在未来的工作中进行探索。

为了进一步的消融，我们探索了我们的想法在多大程度上适用于其他扩散模型。我们专门研究了稳定扩散，在15个时期的冻结特征上训练一个分类器，t固定在150。

因此，在表4中，我们表明，稳定的扩散特征也很适合于分类。重要的是，这不仅意味着我们的方法是灵活的，而且我们未来开发的性能更好的轻扩散模型可以通过我们的方法立即作为统一的表示模型。

结果：细粒度视觉分类（FGVC）

在这里，我们给出了在传输设置中应用我们的方法到表5中定义的数据集的结果。我们同时使用标准的线性探测，以及我们的每个分类头（具有它们最好的ImageNet-50配置）。我们在图4中显示了这些结果。注意，无论使用分类头，扩散模型和SimCLR之间存在性能差距。一个明显的例外是飞机，4个头部中的3个的扩散性能优于SimCLR；这表明了它良好的性能。

此外，我们发现特征选择并不简单，而且通常针对各种FGVC数据集工作的设置并不符合理想的ImageNet设置。例如，考虑到注意力，对于ImageNet-50的最佳头部，往往对FGVC的表现最差。这可能是由于他们依赖于需要正确学习的数据量。此外，当我们在图5中探索CUB的特征选择问题时，我们发现与b = 19相比，ImageNet（b = 24）的理想块数明显低于b = 19。超参数变化的效果更柔和，在ImageNet上，如池化大小，可以导致在幼崽的精度上发生多达3×的性能变化。因此，在扩散训练过程中确定一个更鲁棒的特征选择过程或引入一些正则化可能是未来使转移更可靠的重要工作。

我们使用线性中心核对齐（CKA）[13]来寻找扩散模型的不同块的表示之间的相似程度。根据之前使用CKA [12,73]样本的工作的约定，我们使用了ImageNet-50的2,500个图像测试集（见表5）。我们首先在图6中检查了扩散方法中不同时间步长和特征维度（池化大小）中引导扩散块之间的表示差异。我们还比较了我们的标准设置（t = 90和d = 4096）与ResNet-50和ViT表示，以及一组代表性流行的预训练方法，以及稳定扩散。对于ResNet-50，我们从每个瓶颈块中提取特征，而对于ViT，我们从每个隐藏层中提取特征。

我们注意到，在所有情况下，早期层往往具有更高的相似性，这表明扩散模型可能在前几个块中捕获了类似的低水平细节。还要注意时间步长的影响：当使用来自不同噪声时间步长的图像计算表示时，在后面的层中，表示非常不同。然而，有趣的是，我们发现在瓶颈附近，GD的层倾向于具有与ResNets和vit相似的表示，这表明GD的后期层自然地学习了区别属性。这进一步支持了我们在表1和表3中的发现，其中我们展示了GD特征的良好分类性能。

表示分析

图6：通过中心内核对齐（CKA）进行的特征表示比较。在前2行中，我们比较了在不同时间步长和特征大小下的自身层之间的引导扩散（GD）表示。在底部的2行中，我们将标准t = 90和d = 4096与各种训练前方法的ResNets和ViTs进行了比较。在右下角，我们与稳定扩散（SD）进行了比较，b=18，大小为= 512。

结论

在本文中，我们提出了一种利用扩散模型学习到的表示来进行分类任务的方法。这种重新定位的扩散模型作为潜在的最先进的统一自监督表示学习者。我们解释了识别这些表示的最佳实践，并为从扩散过程中提取高效用的鉴别嵌入提供了初始指导。我们展示了很有前途的迁移学习特性，并研究了不同的数据集如何需要不同的方法来提取特征。我们比较了CKA的扩散表示，以显示扩散模型在不同层次学习什么，以及扩散表示与其他方法的比较。

更广泛的影响。通过我们的论文，我们分析了算法；我们不提供新的现实世界的应用程序。然而，我们的工作涉及到图像生成，这涉及到潜在的错误信息生成的伦理问题。然而，我们并没有改进现有的发电方法，所以潜在的危害似乎可以忽略不计。

训练扩散模型，即使仅仅用于线性探测，也是非常密集的。因此，我们不能在这项工作中提供对可变性的分析。然而，我们的工作是利用扩散模型的能力来区分任务的重要的第一步。

具体结果：

卷积头。在标准线性探测的背景下，我们将卷积头作为特征池步骤的替代品。也就是说，我们不是训练一个自适应的平均池和一个线性层，而是训练一个卷积头，然后是一个线性层。当我们探索不同的通道维度时，该架构由2个块组成，每个块都有2×2的卷积，然后是2×2的最大池化操作。我们执行一个2×2的自适应平均池，将结果变平，以便它可以被一个线性层使用。我们指出了我们在表8中训练的卷积头，其中第一个（输入）通道总是1024（在所选的U-Net块上的特征映射的通道数），但两个学习到的卷积操作的输出通道都被视为一个超参数

注意头。除了应用卷积头之外，我们还实验了一个更复杂的架构来应用变压器块。具体来说，我们首先使用2×2池化层来将特征映射的空间分辨率降低到8×8。每个令牌的信道维度为1024，对应于从U-Net块中提取的特征映射的信道数。然后，我们将这些特性变平，为我们的变压器生成64个令牌。我们向这些标记集添加一个CLS标记，然后使用它来进行最终的分类。我们遵循标准的变压器块体系结构，包括层归一化和QKV-附接子。我们将注意块的数量作为一个超参数，我们的选择和相应的结果如表9所示。

B Experiment Details

我们为我们的实验提供了额外的澄清细节。首先，图中由于资源限制，图2所示的烧蚀是在ImageNet上训练15个时代的结果，而不是我们完整的线性探测配方。其次，我们比较的稳定扩散模型在训练前是文本引导的，因此，与我们的引导扩散检查点不同，它不是完全无监督的。最后，我们在表10中提供了我们的方法和其他统一表示学习者的参数计数比较。

C Ablations

如前所述，在主论文中，我们只对线性探针使用随机翻转，而不像标准设置，后者也使用随机调整大小的作物。这在一定程度上是因为，标准的无监督预训练会使用随机裁剪，而引导扩散则不会。然而，为了进行彻底的比较，我们也尝试了FGVC线性探测与标准的，随机调整大小1裁剪设置的精度（%）。我们在图7中提供了这些结果。注意，扩散头与额外的增强往往表现更差，可能是由于缺乏学习的不变性。我们还包括了BigBiGAN。请注意，该模型是在张量流中实现的，我们也只在张量流中进行线性探测

D Figure 1 Details

现在，我们提供了表7、表8和表9中关于图1所示的信息。我们显示了精确的准确性和参数计数。我们还为每个头部指定了超参数选择。

表7：线性和MLP的结果。对于线性的，1k、4k、16k和65k表示池化和扁平化后的特征的大小。对于MLP，第一个数字是池化和扁平化后的特征的大小，随后的数字是在最后一个（分类）层之前隐藏的层的大小。

表9：注意力头结果。超参数在破折号后表示。第一个超参数类似地是输入特征映射的信道大小，下一个是使用的变压器块的数量。

表10：主要统一无监督表示学习方法的参数计数。对于每一个系统，我们考虑整个系统，而不仅仅是编码网络。