稳定提升！CAN和Diffusion Transformer才是最佳拍档(MIT&清华&英伟达)

论文链接：https://arxiv.org/pdf/2404.01143.pdf

项目链接：​​https://github.com/mit-han-lab/efficientvit​​

今天和大家一起学习条件感知神经网络（CAN），这是一种向图像生成模型添加控制的新方法。与先前的条件控制方法并行，CAN通过动态调整神经网络的权重来控制图像生成过程。这是通过引入一个条件感知权重生成模块来实现的，该模块根据输入条件生成卷积/线性层的条件权重。

本文在ImageNet上进行了类别条件图像生成以及在COCO上进行了文本到图像的生成测试。CAN始终为扩散Transformer模型（包括DiT和UViT）提供显著改进。特别是，结合EfficientViT（CaT）的CAN在ImageNet 512×512上实现了2.78的FID，超过了DiT-XL/2，同时每个采样步骤需要的MAC数量减少了52倍。

部分成果展示

介绍

大规模图像和视频生成模型已经展示了在合成逼真图像和视频方面的惊人能力。要将这些模型转化为人类的有效工具，一个关键步骤是添加控制。本文不希望模型随机生成数据样本，而是希望生成模型能够按照本文的指示生成图像（例如，类别标签、文本、姿态）。

为了实现这一目标，已经进行了大量研究。例如，在GANs中，一个广泛采用的解决方案是使用自适应归一化，根据输入条件动态调整和偏移中间特征图。此外，另一种广泛采用的技术是使用交叉注意力或自注意力将条件特征与图像特征融合。尽管在使用的操作上有所不同，但这些方法分享着相同的基础机制，即通过特征空间操作添加控制。与此同时，神经网络权重（卷积/线性层）对不同条件保持不变。

本文旨在回答以下问题：

• 可以通过操纵权重来控制图像生成模型吗？
• 受控图像生成模型能否从这种新的条件控制方法中获益？

为此，本文引入了条件感知神经网络（CAN），这是一种基于权重空间操作的新的条件控制方法。与普通神经网络不同，CAN引入了一个额外的权重生成模块（如下图2）。该模块的输入是条件embedding，包括用户指令（例如，类别标签）和扩散模型的时间步长。模块的输出是用于调整卷积/线性层的静态权重的条件权重。本文进行了大量消融研究实验，探究了CAN在diffusion transformers上的实际应用。本文的研究揭示了CAN的两个关键见解。

首先，与使所有层都具有条件感知能力相比，精心选择一小部分模块具有条件感知能力（下图3）对于效率和性能（下表1）都是有益的。其次，本文发现，直接生成条件权重比自适应地合并一组基本静态层以进行条件控制（下图4）要更有效。

本文在两种代表性的扩散transformer模型上评估了CAN，包括DiT和UViT。CAN为所有这些扩散transformer模型实现了显著的性能提升，同时增加的计算成本可以忽略不计（下图7）。本文还发现，CAN单独为图像生成模型提供了有效的条件控制，其FID较低，CLIP分数较高，优于先前的条件控制方法（下表3）。除了将CAN应用于现有的扩散transformer模型外，本文还通过将CAN与EfficientViT结合，构建了一个名为CaT的新型扩散transformer模型家族（下图6）。本文总结本文的贡献如下：

• 本文引入了一种用于控制图像生成模型的新机制。据本文所知，本文的工作是第一个展示权重操作在条件控制中有效性的工作。
• 本文提出了Condition-Aware Neural Network（CAN），这是一种用于控制图像生成的新条件控制方法。本文还提供了设计见解，使CAN在实践中可用。
• 本文的CAN在图像生成模型上提高性能，在很大程度上优于先前的条件控制方法。此外，CAN还可以有助于图像生成模型的部署。在ImageNet 512×512上实现更好的FID，本文的CAN模型在每个采样步骤中需要的MAC数量比DiT-XL/2少52倍（上图1），为扩散模型在边缘设备上的应用铺平了道路。

方法

Condition-Aware Neural Network

图像生成过程可以看作是从源域（噪声或嘈杂图像）到目标域（真实图像）的映射。对于受控图像生成，不同条件（例如，猫图像的数据分布与城堡图像的数据分布）给定时，目标数据分布是不同的。此外，扩散模型在不同时间步的输入数据分布也是不同的。尽管存在这些差异，先前的模型对所有情况使用相同的静态卷积/线性层，由于不同子任务之间的负迁移，限制了整体性能。为了缓解这个问题，一个可能的解决方案是为每个子任务都有一个专家模型。然而，由于巨大的成本，这种方法在实际使用中是不可行的。本文的条件感知神经网络（CAN）通过使神经网络能够根据给定条件动态调整其权重来解决这个问题，而不是显式地拥有专家模型。

上图2展示了CAN的一般思想。与常规神经网络的关键区别在于CAN具有额外的条件权重生成模块。该模块以条件embeddingc作为输入，并输出条件权重。除了条件权重外，每一层都有静态权重W。在训练和推理过程中，通过对权重值进行求和，将和W融合成单个核调用。这相当于在输入图像特征上独立应用和W，然后将它们的输出相加。

实用设计

哪些模块应该是条件感知的？ 从理论上讲，可以使神经网络中的所有层都是条件感知的。然而，在实践中，这可能不是一个好的设计。首先，从性能的角度来看，具有太多条件感知层可能会使模型的优化变得困难。其次，从效率的角度来看，虽然为所有层生成条件权重的计算开销可以忽略不计，但会产生显著的参数开销。例如，假设将条件embedding的维度表示为d（例如，384、512、1024等），将模型的静态参数大小表示为＃params。使用单个线性层将条件embedding映射到条件权重需要＃params×d个参数，这在实际应用中是不切实际的。在这项工作中，本文仔细选择了一组模块应用CAN来解决这个问题。

上图3提供了将CAN应用于扩散transformer的概述。深度卷积的参数大小比常规卷积小得多，使其成为一个低成本的条件感知候选模块。因此，本文在FFN中间添加了一个深度卷积，遵循先前的设计。本文使用UViT-S/2在ImageNet 256×256上进行消融研究实验，以确定要成为条件感知模块的一组模块。所有模型，包括基线模型，都具有相同的架构。唯一的区别是条件感知模块的集合不同。

在上表1中总结了结果。在消融研究实验中，有以下观察：

• 使模块成为条件感知并不总是会提高性能。例如，使用静态头部比使用条件感知头部的FID更低，CLIP分数更高（上表1中第2行与第4行对比）。
• 使深度卷积层、block embedding层和输出投影层成为条件感知可以显著提高性能。它将FID从28.32提高到8.82，将CLIP分数从30.09提高到31.74。

基于这些结果，本文选择了这种CAN的设计。具体细节如上图3所示。对于深度卷积层和block embedding层，本文为每个层使用单独的条件权重生成模块，因为它们的参数大小较小。相反，对于输出投影层，本文使用一个共享的条件权重生成模块，因为它们的参数大小较大。由于不同的输出投影层具有不同的静态权重，因此本文仍然对不同的输出投影层使用不同的权重。

CAN与自适应核选择。与直接生成条件权重不同，另一种可能的方法是维护一组基本卷积核，并动态生成缩放参数来组合这些基本核。这种方法的参数开销小于CAN。然而，这种自适应核选择策略无法达到CAN的性能（上图4）。这表明，单纯的动态参数化并不是获得更好性能的关键；更好的条件感知适应能力是关键。

实现。 由于条件感知层在不同样本下具有不同的权重，本文不能进行批处理训练和推理。相反，本文必须独立地为每个样本运行核调用，如下图5（左）所示。这将显著减慢GPU上的训练过程。为了解决这个问题，本文采用了一种高效的CAN实现方法（下图5右）。其核心思想是将所有卷积核调用融合成一个分组卷积，其中#Groups是批处理大小B。本文在运行分组卷积之前进行批到通道转换以保留功能。操作后，本文添加通道到批次转换，将特征图转换为原始格式。

理论上，使用这种高效的实现方式，与运行静态模型相比，训练开销将可以忽略不计。在实践中，由于NVIDIA GPU对常规卷积的支持要比对分组卷积好得多，本文仍然观察到30%-40%的训练开销。这个问题可以通过编写定制的CUDA核来解决。本文将其留给未来的工作。

实验

setups

数据集。 由于资源限制，本文使用ImageNet数据集进行类别条件的图像生成实验，并使用COCO进行文本到图像生成实验。对于大规模文本到图像生成实验，本文将其留给未来的工作。

评估指标。 按照常见做法，本文使用FID作为图像质量的评估指标。此外，本文使用CLIP分数作为可控性的指标。本文使用公开的CLIP ViT-B/32来衡量CLIP分数，遵循的做法。文本prompt是根据CLIP的zero-shot图像分类设置构建的。

实现细节。 本文将CAN应用于最近的扩散transformer模型，包括DiT和UViT。本文遵循官方论文或GitHub存储库中建议的训练设置。除非另有说明，默认情况下，所有模型均使用无分类器引导。基线模型的架构与CAN模型相同，都在FFN层中使用深度卷积。本文使用Pytorch实现本文的模型，并使用A6000 GPU进行训练。训练过程中使用自动混合精度。除了将CAN应用于现有模型外，本文还通过将CAN与EfficientViT结合，构建了一个名为CaT的新型扩散transformer家族。CaT的宏观架构如上图6所示。

消融研究

本文对所有模型进行了80个epoch的训练，批量大小为1024（大约10万次迭代），用于消融研究实验，除非另有明确说明。所有模型使用DPM-Solver进行图像采样，步数为50。

CAN的有效性。 上图7总结了CAN在各种UViT和DiT变体上的结果。对于所有变体，CAN显著提高了图像质量和可控性，超过了基线。此外，这些改进带来的计算成本增加可以忽略不计。因此，CAN通过以更低成本的模型实现相同的FID和CLIP分数，也提高了效率。

下图8比较了UViT-S/2和DiT-S/2上CAN和基线模型的训练曲线。本文可以看到，即使对这两个模型进行更长时间的训练，绝对改善仍然显著。这表明改善不是由于更快的收敛而来，而是通过添加CAN提高了模型的性能上限。

分析。 对于扩散模型，条件embedding包含类别标签和时间步长。为了分析哪一个对条件权重生成过程更重要，本文使用UViT-S/2进行消融研究实验，并在下表2中总结结果。本文发现：

• 在权重生成过程中，类别标签信息比时间步信息更重要。仅添加类别标签比仅添加时间步长可以使FID降低5.15，CLIP分数提高0.33。
• 包括类别标签和时间步长在条件embedding中可以获得最佳结果。因此，在接下来的实验中，本文坚持采用这种设计。

​

与先前条件控制方法的比较在之前的实验中，本文保留了原来的条件控制添加 CAN 时，DiT（自适应归一化）和 UViT（以条件为标记的注意力）方法保持不变。为了看看CAN是否可以单独工作以及CAN与以前的条件控制方法的比较，进行实验并提供上表3中的结果，有以下发现

• CAN单独可以作为一种有效的条件控制方法。例如，CAN单独在DiT-S/2上比自适应归一化实现了13.00更好的FID和0.97更高的CLIP分数。此外，CAN单独在UViT-S/2上比注意力（条件作为标记）实现了19.53更低的FID和1.66更高的CLIP分数。
• CAN可以与其他条件控制方法结合以获得更好的结果。例如，将CAN与自适应归一化结合，可以为DiT-S/2提供最佳结果。
• 对于UViT模型，将CAN与注意力（条件作为tokens）结合会略微降低性能。因此，在接下来的实验中，本文将在UViT模型上使用单独的CAN。

与最先进模型的比较

将本文的最终模型与ImageNet和COCO上的其他扩散模型进行比较。结果总结在下表4和下表6中。对于CaT模型，本文使用UniPC来采样图像，以减少采样步骤的数量。

在ImageNet 256×256上进行类别条件生成。 如上表4（底部）所示，使用无分类器引导（cfg），本文的CaT-B0在ImageNet 256×256上实现了2.09的FID，优于DiT-XL/2和UViT-H/2。更重要的是，本文的CaT-B0比这些模型要高效得多：比DiT-XL/2少了9.9×的MAC，比UViT-H/2少了11.1×的MAC。在没有分类器引导的情况下，本文的CaT-B0在所有比较的模型中也实现了最低的FID（8.81 vs. 9.62 vs. 10.56）。

ImageNet 512×512 上的类条件生成 在更具挑战性的512×512图像生成任务上，本文观察到CAN的优点变得更加显著。例如，本文的CAN（UViT-S-Deep/4）可以与UViT-H的性能相匹配（4.04 vs. 4.05），而每个扩散步骤只需要UViT-H计算成本的12%。此外，本文的CaT-L0在ImageNet 512×512上实现了2.78的FID，优于每个扩散步骤需要52×更高计算成本的DiT-XL/2（3.04 FID）。此外，通过略微扩大模型，本文的CaT-L1将FID从2.78提高到2.48。

除了计算成本比较外，下表5还比较了CaT-L0和DiT-XL/2在NVIDIA Jetson AGX Orin上的性能。延迟是使用TensorRT，fp16测量的。在ImageNet 512×512上提供更好的FID，结合训练无关的快速采样方法（UniPC），CaT-L0在Orin上比DiT-XL/2快229倍。通过将CaT与基于训练的少步骤方法结合起来，有可能进一步推动效率的前沿，展示在边缘设备上实现实时扩散模型应用的潜力。

除了定量结果外，上图9展示了由CAN模型随机生成的图像样本，展示了本文模型在生成高质量图像方面的能力。

在COCO 256×256上的文本到图像 

相关工作

受控图像生成。 受控图像生成要求模型将条件信息纳入计算过程中，以生成相关的图像。社区中已经开发了各种技术用于受控图像生成。一个典型的例子是自适应归一化，它从条件信息中回归出尺度和偏移参数，并应用特征级别的仿射变换来影响输出。除了自适应归一化，另一个典型的方法是将条件信息视为标记，并使用交叉注意力或自注意力来融合条件信息。ControlNet是另一种代表性技术，它使用特征级别的加法来为预训练的文本到图像扩散模型添加额外的控制。与这些技术并行，本工作探索了另一种机制，用于将条件控制添加到图像生成模型中，即使神经网络层（conv/linear）的权重具有条件感知性。

动态神经网络。 本文的工作可以被视为一种新型的动态神经网络。除了在本文中探讨的添加条件控制外，动态调整神经网络可以应用于许多深度学习应用。例如，CondConv提出根据输入图像特征动态组合一组基础卷积核，以增加模型容量。类似地，专家混合技术使用门控网络动态地将输入路由到不同的专家。为了实现高效部署，一次性全网络和可调整神经网络根据给定的效率约束动态调整神经网络架构，以实现效率和准确性之间更好的权衡。

权重生成网络。 本文的条件权重生成模块可以被视为一种专门设计用于向生成模型添加条件控制的新型权重生成网络。在其他情况下，一些先前的工作利用权重生成网络。例如，[33]提出使用一个小网络为一个更大的网络生成权重。这些权重对于数据集中的每个示例都是相同的，以提高参数效率。此外，权重生成网络已被应用于神经结构搜索，以预测给定其架构的神经网络的权重，以降低神经结构搜索的训练和搜索成本。

高效深度学习计算。 本文的工作也与高效深度学习计算联系在一起，旨在提高深度学习模型的效率，使它们适合在硬件上部署。最先进的图像生成模型具有巨大的计算和内存成本，这使得在资源受限的边缘设备上部署它们并保持高质量具有挑战性。本文的工作可以通过提供相同性能但使用更少的扩散步骤和更低成本模型来提高受控生成模型的效率。对于未来的工作，本文将探索将本文的工作与高效深度学习计算技术结合起来，以进一步提高效率。

结论

这项工作研究了通过操纵权重向图像生成模型添加控制。本文引入了一种新的条件控制方法，称为条件感知神经网络（CAN），并提供了高效实用的设计，使CAN在实践中可用。本文进行了广泛的实验，使用ImageNet进行类别条件生成和使用COCO进行文本到图像生成，以评估CAN的有效性。CAN相对于先前的条件控制方法实现了一致且显著的改进。本文还通过将CAN与EfficientViT结合，构建了一个新的扩散transformer模型系列。对于未来的工作，本文将把CAN应用于更具挑战性的任务，如大规模文本到图像生成、视频生成等。

本文转自 AI生成未来 ，作者：Han Cai等

原文链接:​​https://mp.weixin.qq.com/s/ROX6yyh2pfOg0pcz778GZw​​