圖像生成模型是目前業(yè)內研究的焦點，而目前諸如Sora等前沿生成模型，其所基于的主體架構都是Diffusion Transformers（DiT）。Diffusion Transformers（DiT）是論文Scalable Diffusion Models with Transformers（ICCV 2023）中提出的，是擴散模型和Transformer的結合，也是Sora使用的底層生成模型架構，將Diffusion Transformers從圖像生成擴展到了視頻生成。這篇文章給大家總結了目前主要的幾個DiT模型結構，帶大家梳理DiT系列模型的核心。?

更加完整的多模態(tài)生成模型技術總結，可以加入我的星球【圓圓的算法筆記】，獲取Sora底層原理解析專欄。

1.DiT

在之前的圖像生成擴散模型中，底層的網絡結構一般都是U-Net。而本文基于Vision Transformer（ViT）中的Transformer圖像分類模型結構，替代擴散模型中的U-Net，得到DiT模型，實現(xiàn)了更優(yōu)的生成效果。

在輸入部分，基本采用了和ViT相同的方法。對輸入的圖像分成多個patch，并轉換成一個token序列，每個token拼接上相應的position embedding。這個底層的embedding序列作為后續(xù)DiT模塊的輸入。

在擴散模型中，Transformer除了像ViT那樣輸入圖像patch token序列，往往還要輸入一些額外的信息，包括擴散模型中當前的生成時間步、文本信息的輸入等，如何將這些信息輸入到DiT中，文中嘗試了幾種方案。最簡單的方法是將這些額外的embedding直接拼接到原始的序列上。第二種是將外部的embedding單獨拼接成一個序列，和原始的圖像patch序列額外做一個cross attention。第三種方法是修改Transformer中的layer normalization模塊，將其替換成adaptive layer normalization，LN的均值和方差由外部embedding的加和生成。第四種是在第三種的基礎上，引入了基于外部embedding生成的縮放因子，對multi-head attention的輸出進行縮放。

在經過多層的DiT模型后，需要將預測的噪聲結果還原出來，這里使用一個MLP作為Decoder，將DiT生成的結果映射到噪聲預測結果。

上述就是DiT的整體結構，主要還是Vision Transformer。用這個DiT結構，替代擴散模型中的去噪模塊，也就是噪聲預測網絡，就是DiT模型

從實驗對比中可以看出，DiT的生成效果是超過基于U-Net等之前的SOTA模型的。

2.U-ViT

U-ViT是另一個基于ViT的擴散模型網絡。U-ViT也是將擴散模型中的噪聲預測網絡替換成Transformer結構，并且借鑒了U-Net等傳統(tǒng)CV模型中的殘差網絡思路，每一層的輸出都會通過龍skip connection加到更深層的網絡中。此外，文中對一些模型結構也進行了嘗試，包括殘差網絡怎么加，是直接拼接到深層+MLP還是add到生成；擴散步驟embedding怎么加入到U-ViT中；以及Transformer之后的卷積網絡怎么加。

3.MDT

MDT發(fā)表于論文Masked diffusion transformer is a strong image synthesizer（ICCV 2023），在DiT的基礎上，引入了mask latent modeling，進一步提升了DiT的收斂速度和生成效果。

文中分析發(fā)現(xiàn)，DiT在學習過程中，并不能很好的學習各個語義單元之間的關系。為了解決這個問題，MDT引入了一個重構任務，對輸入的圖像的部分patch進行mask，然后使用一個Transformer模型在生成過程中，對這部分被mask掉的patch進行還原。在擴散模型中，每一層MDT輸入被mask掉一部分的token序列，只根據(jù)這部分序列進行噪聲預測。同時，使用一個Transformer網絡來還原被mask掉的部分。通過這種方式，讓模型在學習過程中強行學習patch之間的關系。同時通過position embedding的引入提升對mask token的還原能力。

由于在生成階段，decoder在處理token的時候都是沒有mask的，訓練的時候是mask的，這種不一致會影響效果。因此文中采用side-interpolater，對被mask掉的部分使用side-interpolater的預測結果，融合上沒被mask的結果，保證訓練和預測階段decoder的輸入都是沒有mask掉的。

4.Diffit

Diffit是英偉達發(fā)表于論文Diffit: Diffusion vision transformers for image generation（2023）中的一種方法，也是Diffusion Transformer的一個變體，在模型結構上進行了改進。整體的結構類似于U-Net和Transformer的結合，通過增加downsample和upsample實現(xiàn)層次性的建模。

Diffit在引入擴散步驟embedding的時候，采用了一種Time-dependent Self-Attention的方式，即將步驟embedding直接加入到輸入token序列上，讓self-attention在計算的過程中就考慮到擴散步驟的信息。在模型結構上，采用U-Shape的形式，Encoder部分每一層Transformer后做downsample，來提取不同分辨率下的圖像信息，Decoder部分再逐漸upsample。

本文轉載自 ??圓圓的算法筆記??，作者： Fareise