DiT解读：当Diffusion遇上Transformer_dit模型

前置知识

ViT

Vision Transformer是一种基于Transformer架构的深度学习模型，专门用于处理计算机视觉任务。他的1出现给以往CNN base的图像工作带来了很多新的可能性

ViT的核心思想是将图像分割成均匀的图像块，然后将这些图像块转换为序列，并将序列输入Transformer模型进行处理。这使得ViT可以利用Transformer模型强大的序列建模能力来处理图像数据，避免了传统CNN需要手工设计、调整网络层次结构的缺点。ViT的工作流程大致如下：

1. 输入图像被分割成固定大小的图像块。
2. 每个图像块通过一个可学习的线性投影（projection）映射到特征空间，并与位置编码（position embeddings）结合，形成Transformer模型所需的输入序列。
3. 将这些输入序列输入到Transformer编码器中进行处理。
4. 最终，通过Transformer的输出进行任务相关的预测。

U-Net

U-Net是一种用于图像分割任务的深度学习架构，它的名字来源于其U形状的网络结构。U-Net在医学影像分割等领域取得了很大成功，也被广泛用于其他图像分割任务。

U-Net的结构包含两部分：编码器（Encoder）和解码器（Decoder）。整体结构如同一个U字形，因此得名U-Net。具体工作流程如下：

1. 编码器（Encoder）： 编码器由一系列卷积层和池化层组成，用于提取输入图像的特征并逐步减小特征图的尺寸。
2. 跳跃连接（Skip Connections）： 在每一层的编码器中，将特征图与对应层的解码器中的特征图进行连接，这种连接保留了高分辨率的特征信息。
3. 解码器（Decoder）： 解码器由一系列卷积层和上采样（或反卷积）层组成，用于将编码器提取的特征映射恢复为与输入图像相同大小的分割结果。
4. 输出层： 最后一层经过适当的激活函数（如sigmoid或softmax）产生最终的分割结果。

Diffusion Model

Diffusion model是一种用于建模数据分布的概率生成模型。它提供了一种对数据分布进行建模的方法，特别适用于处理具有复杂结构和高维度的数据。

Diffusion model的核心思想是通过一系列迭代步骤来逐渐将初始数据分布“扩散”（diffuse）到目标数据分布。在这些迭代步骤中，模型会逐渐将噪声注入原始数据，从而使原始数据在逐步迭代的过程中逼近目标数据分布。训练过程通常包括以下步骤：

1. 初始化： 从一个简单的初始分布开始，例如高斯分布。
2. 扩散过程： 在每个迭代步骤中，模型会通过引入逐渐增加的噪声来扩散数据分布。这个噪声可以是随机的，也可以是通过模型生成的。
3. 反向过程： 每个迭代步骤之后，模型将尝试逆向这个扩散过程，使得扩散后的数据逼近目标数据分布。
4. 训练目标： 模型的训练目标通常是最小化生成数据与真实数据分布之间的差异，例如通过最大似然估计或其他损失函数进行优化。

LDM(Latent Diffusion Model)

Latent Diffusion Model（潜变扩散模型）是一种生成模型，是对Diffusion Model的一种扩展和改进。它结合了变分自动编码器（VAE）的思想和概念，以更好地建模数据分布和学习高质量的数据表示。

Latent Diffusion Model的核心思想是在Diffusion Model的基础上引入潜在变量，并通过潜在变量来表示数据空间中的不确定性。这种方法使得模型更具灵活性和表达能力，能够更好地处理数据中的高维度、复杂性和多模态性。

模型结构

DiT的模型的基础架构如图，左边的*N可以看作是一个不断循环叠加的过程。可以发现这其实就是transformer的编码器架构，和ViT的区别是不大。DiT他的主要改进是运用到了扩散模型的流程当中

我们从LDM的设计图上可以看到他的核心其实就是U-Net结构的网络，上面的一段是先用VAE编码压缩图片，然后的模拟扩散过程得到ZT。接下来获取到相关的语义图，文本信息，以及各种潜在变量集合起来传入U-Net大网络，图片除了做卷积之外，还和其他的这些关键信息做注意力机制，以学习到图片的深层语义信息和类别信息等内容之间的关系。

在DiT中，作者创新的改造了LDM这个网络，把U-Net这一步换成了transformer的方式，其就是在中间的去噪步骤去除了CNN相关的操作，统一接入了transformer架构，新的大模型可以理解成下图的形式。这种方式的好处就是解决了CNN感受野限制的问题，图像可以通过自注意力机制来得到更加全面的信息

DiT Block的设计

DiT Block的设计是这个文章最关键的部分。对于各种相关的特征信息，比如图像的类别，扩散的步数这些信息，作者也尝试了很多种方式去把他们结合起来，并且做了很多的实验去对比，最后选择了N-Zero这种方式。

右边的两种方式其实都比较好理解。最右侧的In-Context方法就是把这些信息只是当作一个token去接入，然后forward训练，效果不好也是可以理解的。

中间的交叉注意力的方式就是在中间环节菜用上这些信息，其实效果提升也不大。

On-Zero的方式就是在每个环节都会结合上这些关键信息，模型能够更好地把深层信息和他们关联起来

这是一个形象化的效果图，原本在U-Net中的图片形式在DiT中被展开成了token的形式来进行学习。

规模效果

作者还实验对比了不同的模型规模以及patch_size对于模型性能的影响，实验结果如下，模型的参数量的提升都会提升模型的性能，但是后面边界效应也明显。更小的path_size会有更多的token，模型可以学习到更多的细节

Sora和DiT的联系

Sora的工作其实也是基于DiT，不过这里输入训练的是视频。首先把视频分成一帧一帧输进入做扩散以及ViT编码，每一帧的图片都要结合前后时间的编码，同时可还有另外一侧专门的文字prompt输入做注意力机制辅助可控生成，最终才可以达到我们锁看到的那种生成视频的效果

论文和源码

[2212.09748] Scalable Diffusion Models with Transformers (arxiv.org)

facebookresearch/DiT: Official PyTorch Implementation of "Scalable Diffusion Models with Transformers" (github.com)