Diffscaler: Enhancing the Generative Prowess of Diffusion Transformers
公和众和号：EDPJ（进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群）
目录
0. 摘要
最近，随着 OpenAI 发布的 SORA 引起了广泛关注，扩散 transformer（DiT）成为了备受瞩目的话题，强调了将 transformer 作为扩散模型骨干的必要性。与一般视觉任务的 CNN 模型相比，基于 transformer 的模型表现出更好的泛化能力。然而，在现有文献中对基于 transformer 的扩散骨干的能力以及将其生成能力扩展到其他数据集方面进行的研究要少得多。本文侧重于使单个预训练的 DiT 模型能够迅速适用于多个数据集，从而仅使用一个模型就可以完成多样的生成任务。为此，我们提出了 DiffScaler，一种用于扩散模型的高效缩放策略，在这种策略中，我们训练了最少量的参数以适应不同的任务。特别地，我们通过结合利用预训练模型学到的子空间以及学习额外的任务特定子空间的能力，在每一层学习任务特定的转换。由于这些参数是相互独立的，一个具有这些任务特定参数的单个扩散模型可以同时执行多个任务。此外，我们发现基于 transformer 的扩散模型在对较小数据集进行微调时显著优于基于 CNN 的扩散模型方法。我们在四个无条件图像生成数据集上进行实验。我们展示了使用我们提出的方法，单个预训练模型可以扩展到执行这些条件和无条件任务，同时进行最少的参数调整，同时表现得几乎与为特定任务微调整个扩散模型一样好。
3. 方法
在本节中，我们解释我们提出的方法 DiffScaler。无论是 GAN，VAE 还是扩散模型，通过仔细观察扩散模型的推理过程，可以看到在采样过程中学到的上下文是不同的。特别是，推理过程的初始步骤调整了数据集的广泛上下文，然后随着时间步骤的推进逐渐变为去噪任务。最近的工作利用了这一特性来集成多个模型以学习特定于时间步骤的上下文。网络对不同噪声水平的去噪能力主要归因于权重层，正如以前的工作所示。生成模型中的数据集特定信息归因于网络中的偏置。因此，为了使模型能够学习上下文信息并生成逼真的输出，需要额外的权重和偏置来促进特定于数据集的上下文生成。
3.1 解锁增量数据集添加
给定一个能够生成 N 个类别的扩散模型，我们希望将该模型的能力扩展到包含 M 个类别的额外数据集。同时，我们还需要确保在添加新类别后，模型生成旧的 N 个类别的能力不会下降。我们通过冻结 N 个类别的原始模型参数，并为学习新的 M 个类别引入少量新参数来解决这个问题。此外，为了将模型映射到新数据集，我们需要添加新的嵌入层，这些层可以映射到新的类别。考虑一个具有 N 个类别的条件扩散模型的嵌入矩阵 W_emb ∈ R^((N+1)×d)，其中 d 是矩阵的嵌入维度。额外的一个类别表示无条件类别，模拟所有其他类别的分布。我们用 W_uc 表示这个。嵌入矩阵 W_emb 可以表示为：
我们引入一种新的训练方案，逐步添加这些新的类别，并确保训练过程稳定，因为我们希望网络在训练开始时能够很好地工作并去噪。为了适应这种设置，我们将 W_(N+2), ..,W_(N+M+1) 初始化为无条件类别嵌入 W_uc。这个设置确保了当训练开始时，模型能够满足去噪的目标，即使它不能生成新的类别。同样的原理也适用于 ControlNet，其中在训练开始时添加的初始特征嵌入是零，从而确保了在训练过程开始时满足去噪目标的稳定预测。仅改变嵌入矩阵的权重只能给出网络中所有先前见过的类别的组合。我们在实验部分验证了这一说法，在那里我们看到一个在 ImageNet 类别上预训练的模型仅通过训练嵌入矩阵无法生成面部图像。为了使网络能够学习新类别，需要网络学习这些新类别的基函数（basis functions），而在新类别与原始预训练模型类别有很大差异的情况下，需要添加新的参数来适应这种生成能力。在下一节中，我们描述了我们提出的新的权重添加机制，以实现现有模型对新数据集的扩展。
3.2 Affiner
让我们考虑神经网络中的任何可训练层。对于输入 x ∈ R_n，输出 y 可以被建模为
其中 W ∈ R^(m×n) 和 b ∈ R_m 分别是该层中的可训练权重和偏置。在大多数情况下，训练过程中的内存消耗发生在存储权重矩阵 W ∈ R^(m×n) 及其相应梯度时。鉴于这一公式，LORA 作为增强网络生成能力的理想候选者，因为它可以向权重矩阵添加新的基向量。LORA 模块可以被视为，通过学习附加基函数，来学习较低维度的微调子空间。在 W 的秩 R = min(m, n) （满秩）的情况下，LORA 块变得多余。此外，LORA 无法有效利用 W 的基函数，并且可能不得不学习与 W 中相同的列向量，因为 LORA 中不存在 W 的缩放。因此，该模型仍然缺乏对与训练数据中的数据完全不同的数据进行建模的能力，并且根据我们在多个情况下的实验结果，这导致了梯度爆炸。因此，我们添加了额外的权重重参数化来解决这些额外的基空间。
可以注意到，在一个带有权重 W 的预训练扩散模型中，模型能够执行去噪。因此，手头的任务是学习描述数据集分布的上下文信息。受到 StyleGAN [19] 和 BitFit [3] 中样式嵌入作用的启发，我们添加了两个轻量级可训练参数，一个缩放参数 a 和一个偏移参数 b，分别初始化为 a 和 b。具体来说，为了能够将新的样式融入预训练网络，我们重新参数化网络层为：
其中 a 和 b 是可学习的参数，分别缩放和偏移权重。在许多扩散模型的情况下，需要考虑的一个关键方面是其依赖于分布的去噪能力。虽然上述公式有助于有效地参数化网络的权重，但这里需要注意的一个关键点是，如果权重矩阵 W ∈ R^(m×n) 具有秩 R，使得 R < min{m, n}，那么所提出的偏移和缩放无法对基向量进行建模，如满秩的权重矩阵 R = min{m, n} 那样。因此，我们尝试为网络提供建模这些额外基函数的能力。我们对原始权重矩阵进行低秩分解，表示如下：
其中 W_up ∈ R^(m×d)，W_down ∈ R^(d×n)，d ≪ {m, n}，s 是一个可训练的缩放参数，以零初始化。我们使用 Kaiming 正态初始化 [11] 初始化 W_down，W_up 以满足平凡条件。我们的初始化策略不同于 [6, 16] 中的方法，其中权重初始化为零。
将新的权重公式包含到我们之前的可学习缩放和偏移参数中，每个可训练层的整体参数化变为：
这是我们引入的插件层，称为 Affiner。总结一下，a ∈ R^n 被初始化为 1，b ∈ R^n 被初始化为0，W_down 和 W_up 被初始化为 Kaiming 正态分布，s 被初始化为 0。请注意，Affiner 可以插入卷积和 transformer 层。在二维卷积层的情况下，我们用 1 × 1 卷积层模拟 {W_down，W_up}，并且额外的偏移和缩放只是可学习的参数。对于一维和三维卷积的情况也是如此。在 ViT [8, 41] 的情况下，transformer 块包括多头注意力和带有 MLP 和层归一化的前馈块。在自注意力和 MLP 层中，权重转换通过线性层进行。因此，我们的公式仍然成立。对于 transformer 模型，我们为每个键、查询、值权重和偏差参数以及 MLP 块添加我们的 Affiner 块。Affiner 已在图 2 中说明。
3.3 将 Affiner 扩展到多个数据集
请注意，DiffScaler 可以用于跨多个数据集或任务扩展模型。在处理特定数据集的无条件生成时，我们使用术语 “数据集”，在处理有条件生成任务时，我们使用术语 “任务”。为简单起见，在本节中，我们以多个数据集的规模为例来解释缩放，但请注意，相同的策略也适用于跨多个任务的缩放。要将扩散模型扩展到 N 个不同的数据集，我们首先为 N 个数据集的每一个创建新的参数集
我们为每个任务单独训练这些新参数。然后，这些参数被堆叠在一起，使用单个扩散模型执行生成任务。也可以简单地同时在所有数据集上训练模型以节省时间。在这种情况下，将会有 N 组这些新参数，它们将同时跨所有数据集进行训练。另外，要将已经在 N 个数据集上训练过的扩散模型添加到新数据集中，只需简单地添加一个新的参数集并为该新任务训练该模型。进一步地，一旦以这种建议的方式训练了模型，我们提出了一种新的权重合并方案，以在下一节的 zero-shot 任务泛化中融合这些权重。
3.4 无零卷积的条件生成
最近，ControlNet 模型通过调节（conditioning）来增加对稳定扩散微调的能力。它首先获得 {图像，标题，条件} 对，并根据图像标题对稳定扩散进行微调。并行编码器通过零初始化的卷积层连接到基础网络，这些卷积层占据了大量参数。通常，条件是独立的并与文本标题对齐。
作为 ControlNet 的升级，DiffScaler 可以通过使用所提出的重新参数化技术和修改每个可训练层并为手头的新任务训练新添加的参数，来用于跨多个任务的条件生成。通过我们的公式，无需单独的编码器或像 ControlNet 中那样的零卷积。而且，我们获得了可比较的性能，7M 可训练参数，而 ControlNet 的可训练参数约为 300M。
4. 实验
