聊聊 DiT 和 GenTron

作者丨無惡不作

來源丨 https://zhuanlan.zhihu.com/p/683930767

編輯丨 GiantPandaCV

1. 介紹

隨着 OpenAI 最新 Sora 的火爆出圈，Diffusion Transformer 無疑將成爲繼 Stable Diffusion 後最出圈的模型結構。Diffusion Transformer Model 是一種結合了 Transformer 架構和擴散模型的圖像、視頻及其他數據的生成模型。也就是說 Diffusion Transformer 繼承了 Transformer 這種具有 scaling laws 優勢的模型結構，具有知識隨參數增加帶來迸發效果的先天優勢。

DiT 模型：基於 transformer 架構的擴散模型，不僅將 transformer 成功應用在擴散模型，還探究了 transformer 架構在擴散模型上的 scalability 能力。
GenTron 模型：基於 Transformer 的擴散模型，在針對 SDXL 的人類評估中，GenTron 在視覺質量方面取得了 51.1% 的勝率（19.8% 的平局率），在文本對齊方面取得了 42.3% 的勝率（42.9% 的平局率）。

2. DiT

論文：https://https://arxiv.org/pdf/2212.09748.pdf8.pdf

代碼：https://github.com/facebookresearch/DiT

DiT 基於 Vision Transformer（ViT）架構，該架構操作於圖像的 patch 序列。DiT 保留了 ViTs 的許多最佳實踐。下圖顯示了完整的 DiT 架構概述，描述了 DiT 的前向傳播，以及 DiT 類的設計空間組件。

2.1 Patchify

DiT 的第一層是 “patchify”，它通過卷積和線性操作將輸入的圖像 patch 化，BCHW -> BNC。在 patchify 之後，對所有輸入標記應用標準 ViT 的基於頻率的位置編碼（正弦 - 餘弦版本）。如下圖所示。

2.2 Timestep t

Timestep t 是一個時間維度的 embedding，在 Timestep 中會根據預製的 frequency_embedding_size，對於預製的時間步數進行相對位置編碼。再將相對位置編碼後的 embedding 映射到隱藏層維度。

2.3 Label y

Label y 類似於 NLP 模型的詞表映射，因爲 DiT 模型只有 1000 類，所以 num_classes=1000，在實際代碼中，如下代碼發現還會增加 class_dropout_prob 類別。（這部分對於類別不清晰的部分做歸納）

| self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size) | | --- |

Timestep t （N, D）和 Label y （N, D）的結果會進行相加得到 conditioning。

2.4 DiTBlock

conditioning 經過線性層生成，，，，，γ1，β1，α1，γ2，β2，α2 。如下代碼所示。

| shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1) | | --- |

如上圖 1 所示，經過 depth 層 DiTBlock 迭代，(N, T, D) -> (N, T, D)。

2.5 FinalLayer

同上 DiTBlock 有 adaLN_modulation 模塊，將 adaLN_modulation 的結果經過線性層得到最後的輸出結果。(N, T, D) -> (N, T, patch_size ** 2 * out_channels)

| shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) | | --- |

2.6 unpatchify

將 FinalLayer 的結果 unpatchify （此處使用 reshape 和 einsum 做轉換）到最後的輸出噪聲。

(N, T, patch_size ** 2 * out_channels) -> (N, out_channels, H, W)

2.7 vae 解碼

通過 vae 解碼得到最後的輸出圖片。

理論解釋：1. 變分自編碼器（Variational Autoencoder）[https://www.zhangzhenhu.com/aigc/%E5%8F%98%E5%88%86%E8%87%AA%E7%BC%96%E7%A0%81%E5%99%A8.html]
代碼解釋：Pytorch 實現: VAE | DaNing 的博客 [https://adaning.github.io/posts/9047.html]

2.8 實驗

配置：

256 × 256 and 512 × 512 image resolution
ImageNet dataset
Initialize the final linear layer with zeros and otherwise use standard weight initialization techniques from ViT
AdamW
Learning rate of 1 × 10−4
No weight decay and a batch size of 256
Horizontal flips
Moving average (EMA): 0.9999

Scaling the DiT model improves FID at all stages of training

Transformer Gflops are strongly correlated with FID

Larger DiT models use large compute more efficiently

Benchmarking class-conditional image generation on ImageNet 256×256

3. GenTron

論文：https://arxiv.org/pdf/2312.04557.pdf7.pdf

代碼：暫未找到

GenTron 是繼 DiT 後的 Transformer 擴散生成模型。將 Diffusion Transformers 擴展到文本輸入，增加 cross attention 機制。GenTron 將參數規模從約 900M 擴展到超過 3B，生成圖片質量顯著提升。此外，GenTron 從文本擴展到到視頻生成，引入了無運動引導以提高視頻質量。GenTron 在視覺質量方面達到了 51.1% 的勝率（平局率爲 19.8%），在文本對齊方面達到了 42.3% 的勝率（平局率爲 42.9%）。

主要特點：

基於 Transformer 的 T2I（文本到圖像）生成與擴散模型進行了徹底而系統的調查。
探索了一種純基於 Transformer 的擴散模型，用於 T2V（文本到視頻）生成。引入了無動作引導（motion-free guidance），這是一種創新技術，可以有效地對 T2I 生成模型進行微調，以產生高質量的視頻。
實驗結果表明，在人類評估中，GenTron 相比於 SDXL 有明顯的優勢。此外，GenTron 在 T2I-CompBench 評估中相比現有方法展現出了更優秀的性能。

3.1 主體結構

GenTron 基於 DiT-XL/2 構建，它將形狀爲 32×32×4 的 Latent 特徵轉換爲非重疊 tokens 序列，採用 2×2 的裁剪層。然後，這些 tokens 被送入一系列的 transformer 塊。最後，應用標準的線性解碼器將這些圖像 tokens 轉換爲 Latent 空間。

文本的 cross attention
文本 pooling 後與 timestamp 進行相加

3.2 Embedding integration

自適應層歸一化（Adaptive layernorm (adaLN)）
主體和 DiT 的 adaLN 類似
跨注意力（Cross-attention）

圖像特徵充當 query，而文本 embedding 則充當 key 和 value。該設置允許通過注意力機制在圖像特徵和文本 embedding 之間進行直接交互

3.3 Scaling Up GenTron

爲了探索大幅擴展模型規模的影響，作者開發了 GenTron 的高級版本，稱爲 GenTron-G/2。該模型擴展了三個關鍵方面：transformer 塊的數量（深度）、patch embeddings 的維度（寬度）和 MLP 的隱藏維度（MLP 寬度）。GenTron 模型的規格和配置詳見下表 1。

3.4 GenTron-T2V 結構

本文方法僅在每個 transformer 塊中集成輕量級的時間自注意力（TempSelfAttn）層。如下圖 9 所示，TempSelfAttn 層位於交叉注意力層之後和 MLP 層之前。

作者通過在進入 TempSelfAttn 層之前更改交 cross attention 層的輸出，然後在通過後將其重新改回其原始格式，修改了交叉注意力層的輸出。這個過程可以正式表示爲：

其中，b、t、n、d 分別表示 batch 大小、幀數、每幀的 patch 數量和通道維度。

3.4.1 初始化

使用預訓練的 T2I 模型作爲初始化 T2I 和 T2V 模型之間共享層的基礎。此外，對於新增的 TempSelfAttn 層，將輸出 project 層的權重和偏差初始化爲零。這確保在 T2V 微調階段的開始，這些層產生零輸出，與 shortcut 連接一起有效地作爲恆等映射。

3.5 無運動引導（Motion-Free Guidance）

3.5.1 問題和挑戰

作者觀察到當前 T2V 擴散模型中存在一個顯著現象，即每幀的視覺質量明顯滯後於 T2I 模型。此外，分析還顯示，在 T2V 模型經過微調後，其視覺質量明顯下降，特別是與其原始的 T2I 對應物相比。注意到了這些問題通常存在於當前的 T2V 擴散模型中，不僅限於基於 transformer 的 T2V。

3.5.2 問題分析

作者推測 T2V 視覺質量滯後的現象主要源於兩個因素：

公開可用的視頻數據集在質量和數量上通常不及圖像數據集，許多視頻幀受到運動模糊和水印的影響，進一步降低了它們的視覺質量。
在視頻微調期間專注於優化時間方面可能會無意中損害空間視覺質量，導致生成的視頻總體質量下降。

3.5.3 解決方案

從數據方面考慮，採用聯合圖像 - 視頻訓練策略來緩解視頻數據不足的問題。
將視頻片段內的時間運動視爲一種特殊的條件信號，類似於 T2I/T2V 擴散模型中的文本條件。基於這個類比，並提出了一種新方法，即無運動引導（MFG），靈感來自於無分類器引導，用於條件生成視頻中運動信息的權重。作者使用一個單位矩陣以概率 pmotion free 抵消時間注意力。這個單位矩陣在圖 9（無運動掩碼）中呈現，其對角線填充爲 1，而所有其他位置爲 0。這種配置將時間自注意力限制在單個幀內。

在推理過程中，有文本和運動的條件。可以修改分數估計爲：

可以將解決方案 I 和 II 結合在一起，具體方法如下：當在訓練步驟中省略運動時，加載一對圖像文本，並將圖像重複 T-1 次以創建一個僞視頻。相反，如果包含運動，則加載視頻片段並將其提取成 T 幀。

3.6 實驗