扩散 Transformer 解释
原文：towardsdatascience.com/diffusion-transformer-explained-e603c4770f7e使用 DALL·E 生成的图像。
原文：
towardsdatascience.com/diffusion-transformer-explained-e603c4770f7e
使用 DALL·E 生成的图像。
引言
在使用 Vision Transformer (ViT)及其后续者将 NLP 摇身一变进入计算机视觉领域之后，transformers 现在正进入图像生成领域。它们正逐渐成为 U-Net 的替代品，而 U-Net 是所有早期扩散模型构建的基础卷积架构。本文探讨了威廉·皮布尔斯和赛宁·谢在他们的论文“Scalable Diffusion Models with Transformers”中引入的扩散 Transformer（DiT）。
DiT 影响了其他基于 transformer 的扩散模型的发展，如PIXART-α、Sora（OpenAI 令人惊叹的文本到视频模型），以及在我撰写本文时，Stable Diffusion 3。让我们开始探索这一新兴架构类别，它们正在推动扩散模型的演变。
预备知识
由于这是一个高级话题，我必须假设你对 AI 中的一些常见概念有一定的熟悉度，特别是在图像生成方面。如果你已经熟悉这个领域，本节将帮助你刷新这些概念，并提供进一步了解的参考资料。
如果在阅读本文之前你想对这个领域有一个广泛的概述，我推荐阅读我之前的文章，其中涵盖了众多扩散模型和相关技术，其中一些我们将在本文中重新探讨。
Comparing and Explaining Diffusion Models in HuggingFace Diffusers
扩散公式
图 2 来自Denoising Diffusion Probabilistic Models。
在直观层面上，扩散模型通过首先获取图像，引入噪声（通常是高斯噪声），然后训练一个神经网络来通过预测添加的噪声及其协方差矩阵（在某些情况下）来逆转这一添加噪声的过程。引入的噪声程度由时间步长变量t控制；在t=0 时，_x0代表原始图像，而在t=1000 时，_x1000几乎完全是噪声。
在实践中，对于每个时间步长t，我们从高斯分布中采样 _xt，条件是 _xt-1，__：
这可以通过以下方式重新参数化：
其中 ϵ_t ~ N(0, I). _βt 项表示一个预定的方差计划。要生成给定 _xt-1 的 _xt，我们从多元标准正态分布中采样 _ϵt 并应用上述方程。这一步一步添加噪声的过程被称为正向过程。
幸运的是，要生成 _xt，没有必要首先生成所有之前的 _xt-1；可以证明可以直接使用以下公式：
其中：
一旦我们建立了噪声添加的方法，我们就训练一个模型来预测添加的噪声。在训练过程中，我们为每张图像采样一个批次的图像和相应的 t 值，根据 t 添加噪声，并将这些带噪声的图像及其 t 值输入到模型中。该模型被训练来预测最小化损失函数的噪声，通常是实际添加的噪声与模型预测的噪声之间的均方误差。
注意，在这里，_ϵθ 是我们的神经网络，尽管没有显示，它也接受 t 的值作为输入。
对于图像生成，我们执行反向过程，从纯噪声开始，并使用以下条件分布进行迭代采样：
其中：
并且 Σ__θ(xθ, t) 被设置为对角矩阵。在去噪扩散概率模型（DDPM）的上下文中，这个矩阵是固定的，而对于改进的去噪扩散概率模型（iDDPM），这个矩阵是学习的。在我们正在分析的论文中，使用了 iDDPM。
使用训练好的模型，我们可以从一个纯噪声开始，尝试一次性移除所有噪声，以获得一个类似于训练数据集的样本。然而，一个逐步迭代的过程，其中部分噪声被移除，有时会重新引入一些“新鲜”的噪声（朗之万动力学），可以得到更好的结果。反向过程的细节由所采用的采样策略决定。
深入探讨技术细节将超出本总结的范围。对于那些想要深入了解该主题的人，我建议查阅以下关于扩散模型的奠基性论文：
无分类器指导
在实践中，我们很少旨在生成没有任何形式的控制的图像。引导最终结果最流行的方法是通过使用文本提示，本质上是我们想要图像描绘的内容的描述，称为文本到图像转换。这种方法的简化版本是类别条件，其中类别标签 c 作为提示。这种类型的提示被 Diffusion Transformer 使用。
不幸的是，模型有时往往会忽略我们的提示，无论是文本的还是其他形式的。为了解决这个问题，采用了一种称为“无分类器指导”的技术，以确保模型更紧密地遵循我们的指示。
你可能会好奇为什么它被称为“无分类器”指导。这种方法与分类器指导区分开来，后者依赖于相对于 _xt 的 log p(c | _xt, t) 的梯度，其中 p(c | _xt, t) 代表在时间步 t 的图像 _xt 属于类别 c 的可能性，并由分类器估计。函数的梯度指向最陡上升的方向。因此，通过添加这个梯度并赋予一定的权重，生成的图像与期望类别的对齐概率增加，至少根据我们的分类器来看是这样。
这种方法的一个主要挑战是需要一个分类器。如果没有为我们的数据集类别提供预训练的分类器，我们必须首先训练一个。更糟糕的是，即使存在合适的分类器，它也不太可能针对我们的图像类型进行定制，因为预训练的分类器是为了分类“干净”的图像而设计的，而不是那些被高斯噪声掩盖的图像。那么当我们想要根据文本提示“一个穿着 tutu 的婴儿萝卜在遛狗”或甚至非文本条件来生成时怎么办呢？
无分类器指导有效地解决了这些难题。它通过偶尔用表示训练期间没有提示的可学习嵌入替换提示嵌入来操作。在推理期间，这允许我们计算两个噪声估计：一个带有提示，一个不带提示。我们现在可以应用以下公式：
从公式中，我们可以看到将指导比例 s 设置为 1 维持了无指导过程（标准条件）的噪声估计。将 s 设置超过 1 将无条件噪声估计 _ϵ_θ(xt, ∅) 推向与特定提示更一致的值，朝着 _ϵ_θ(x_t, c) − ϵ_θ(xt, ∅) 的方向移动。
尽管无分类器引导通常会导致样本质量提高，但它也要求每次评估时都进行两次噪声预测——一次带有提示，一次不带——实际上加倍了计算工作量。此外，生成图像的质量和多样性之间存在权衡：随着s的增加，视觉保真度提高，但样本的多样性降低。在他们的论文中，作者在生成基于特定类标签的图像时使用了s=4 的引导比例。
潜在扩散模型
图 3. 来自使用潜在扩散模型进行高分辨率图像合成。
我已经写了一篇关于潜在扩散的完整文章：
既然如此，我这里只提供主要直觉。如果你想要深入了解，我建议阅读我上面的文章。
之前提到的扩散公式是在图像空间中操作的。与可以处理相对较低分辨率的分类模型不同，在我们的情况下，分辨率受我们打算生成的图像分辨率所限制，例如，1024 x 1024。此外，正如我们所看到的，反向过程是迭代的，这意味着要生成单个图像，我们必须多次使用计算量大的模型进行推理。使用类似于 U-Net 的卷积网络而不是 transformer 进一步加剧了这个问题，因为注意力机制是二次扩展而不是线性扩展。此外，正如我们将看到的，从图像生成的标记数量也是二次扩展的，这意味着如果我们从 256 x 256 的分辨率移动到 512 x 512，如果之前标记的数量是T，现在它变成 _T_²，如果注意力层的操作数量是 O(_T_²)，现在它变成 O(_T_⁴)。
解决这个问题的根本方法有两个：第一个，像DALL·E 2（OpenAI）或Imagen（Google）这样的模型所采用的方法，是使用相对较低的分辨率进行扩散过程，然后应用一个或多个超分辨率模型。第二个方法是放弃图像空间，在潜在空间中工作。
这个最后声明在实践中意味着什么？我们可以将潜在空间视为一种图像的压缩，它保持了其语义内容，可能只丢弃了一些边缘细节。能够以这种方式压缩和分解图像的模型是变分自编码器（VAE）。
在训练期间，每个图像通过 VAE 编码器进行压缩；在推理期间，简单地生成一个仅含噪声的张量z。如果所需的分辨率是，例如，256 x 256 x 3，那么z可能是 32 x 32 x 4。请注意，由于不再在图像空间中工作，通道数不再受限于保持为 3。在生成过程中，从z开始，逐渐从其中去除噪声，在最后一步，使用 VAE 解码器将“干净”的z解压缩成最终的图像。
这一部分只是一个回顾。如果你想了解更多，我建议阅读上面提到的我的文章。
扩展扩散变换器设计空间
图 3. 来自 Scalable Diffusion Models with Transformers。
Patchify
图 4. 来自 Scalable Diffusion Models with Transformers。
在本节中，我将假设你熟悉标准的自注意力机制和正弦位置编码的概念，这些概念与Transformer架构相关。如果你对这些概念不熟悉或希望复习它们，我推荐这篇文章：
变换器是接受序列作为输入的模型，或者更确切地说，是一个集合（序列中的顺序仅由位置编码/嵌入的添加给出）。因此，我们必须执行的第一步操作是将我们的潜在张量z（在这种情况下是 32 × 32 × 4）转换成一个标记序列。Patchify非常简单；它将z分成一个网格（32 / p) x (32 / p) x 4，其中每个网格元素p x p x 4 随后被线性投影成为 1 x d，其中d是一个超参数。在实践中，这可以通过应用一个具有等于p的内核和步长以及等于d的输出通道数的卷积来轻松完成，对于每个批处理元素，我们得到一个张量（32 / p) x (32 / p) x d。如果我们定义T=32² / p²，通过重新排列我们已有的元素，对于大小为 N 的批处理，我们得到一个张量N x T x d。将p减半将使T增加四倍，因此至少将总变换器 Gflops 增加四倍。在他们的工作中，作者尝试p=2,4,8；较小的p值会产生更好的结果，但计算需求更高。
在这一点上，回忆一下注意力机制本身并不知道元素顺序，我们向所有输入标记添加位置编码。
注意：尽管论文将它们称为嵌入，但这个术语通常保留用于训练期间学习的表示；在这种情况下，使用的是静态编码。
假设您已经熟悉 1D 正弦位置编码，我们将每个标记的网格的 x 坐标以标准 1D 方式进行编码。我们将相同的过程应用于 y 坐标，然后将这两个编码连接起来，为每个标记。这种方法比将标准 1D 位置编码应用于“展开”图像更受欢迎，因为后者会导致网格中相邻位置具有显著不同的编码。通过承认输入的 2D 性质，我们保持相邻标记具有相似的编码。
DiT 块设计
再次，我将假设您已经熟悉 Transformer 架构。因此，我不会深入探讨标准变换，如层归一化、多头自注意力或点积前馈，这些在链接的文章中有很好的记录。然而，DiT 有一个我们尚未讨论的特定方面：我们如何实现类条件化？
在这篇文章中，我将重点关注作者在实验中找到的最有效的方法：adaLN-Zero。
首先，我们需要将 t 和 c 转换为两个嵌入。对于类别标签 c，这涉及到简单地为训练期间学习的每个类别初始化一个嵌入向量。对于 t，则稍微复杂一些：对 t 应用正弦编码，然后通过一个小型 MLP 进行进一步转换，该 MLP 由一个线性层、一个 SiLU 激活和一个另一个线性层组成。然后，这两个嵌入被连接起来成为我们的条件向量。
现在，让我们描述 adaLN。首先，条件通过 SiLU 转换，然后进行线性投影；我将这个转换称为 adaLN_modulation，以与官方仓库的代码保持一致。adaLN_modulation 的输出必须分成六个不同的部分，每个部分的维度为 d（与我们的条件相同）。这些向量代表用于在不同部分的 DiT 块中缩放和偏移输入的 _γ1、_β1、_α1、_γ2、_β2、_α2。
变换的顺序如下：
def modulate(x, shift, scale):
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
class DiTBlock(nn.Module):
...
def forward(self, x, c):
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
return x
-
输入标记
x通过层归一化进行归一化：self.norm1(x)。 -
使用
modulate(...)，输出被缩放和偏移，其中shift_msa、scale_msa分别代表 _γ1、_β1。 -
应用多头自注意力：
self.attn(...)。 -
输出通过
gate_msa.unsqueeze(1)缩放，代表 _α1。 -
结果被添加回原始的
x，从而采用残差连接（类似于 ResNet）。 -
重复类似的过程，但用
self.mlp(...)替代self.attn(...)。
adaLN-Zero 与 adaLN 不同，它将 adaLN_modulation 的权重初始化为零。这意味着最初，正如我们可以在上面的 forward 方法中看到的那样，x 没有任何改变，网络逐渐学习最优缩放和偏移参数。
Transformer 解码器
最后，解码器由一个线性层加上与我们之前所见类似的归一化层组成。这个线性层将最终输出 x 的维度从 N x T x d 转换为 N x T x _p_² x 2C，其中 C 代表输入通道。需要注意的是，只有当我们旨在预测对角协方差矩阵 Σ 以及噪声时，它才是 2C；否则，它只是 C。
然后，输出被“去补丁化”（重塑），以产生预测的噪声或预测的 Σ 与噪声的组合。
结论
在这篇文章中，我们介绍了论文“Scalable Diffusion Models with Transformers”的必要部分，并提供了一些有用的背景信息。以一个有趣的故事作为总结，这项工作正变得越来越重要，因为它构成了最现代方法的基础，最初在 CVPR 2023 被拒绝，
提醒我们，即使是专家也很难预测哪些因素会产生影响，并且科学会议的评审过程远非完美。
感谢您花时间阅读这篇文章，请随时留下评论或与我联系，分享您的想法或提出任何问题。要了解我的最新文章，您可以在 Medium、LinkedIn 或 Twitter 上关注我。
有“AI”的1024 = 2048，欢迎大家加入2048 AI社区
更多推荐
3
0- 0
已为社区贡献360条内容
所有评论(0)