OpenAI發(fā)布Sora，以及Stability.AI發(fā)布的SD3，根據(jù)其技術(shù)報(bào)告，使用了可擴(kuò)展的transformer擴(kuò)展模型，《Scalable Diffusion Models with Transformers》是其相關(guān)的一篇重要論文。
關(guān)于DIT作者進(jìn)階的論文SIT《SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers 》介紹，下一篇文章解析?。。ㄖ酰憾?；公眾號：二度簡并）

論文重點(diǎn)

文章提出，使用一種基于Transformer架構(gòu)的新型擴(kuò)散模型DiT，并訓(xùn)練了潛在擴(kuò)散模型，用對潛在patch進(jìn)行操作的Transformer替換常用的U-Net主干網(wǎng)絡(luò)。通過以Gflops衡量的前向傳遞復(fù)雜度來分析擴(kuò)散Transformer(DiT)的可擴(kuò)展性。通過增加Transformer的深度/寬度或增加輸入token的數(shù)量，始終具有較低的FID。

DiT架構(gòu)與標(biāo)準(zhǔn)Vision Transformer (ViT)非常相似，但有一些小但重要的調(diào)整。擴(kuò)散模型需要處理?xiàng)l件輸入，例如擴(kuò)散時(shí)間步長或類標(biāo)簽。最有效的是帶有adaptive layer norm layers (adaLN)的ViT塊。重要的是，這些adaLN層還會(huì)在塊內(nèi)任何殘差連接之前激活，并進(jìn)行初始化，以便每個(gè)ViT塊都是恒等函數(shù)。簡單地改變注入條件輸入的機(jī)制就會(huì)在FID方面產(chǎn)生巨大的差異，這一更改是獲得良好性能所需的唯一更改。除此之外，DiT是一個(gè)相當(dāng)標(biāo)準(zhǔn)的Transformer模型。

論文摘要

1. 探索了一種基于transformer架構(gòu)的新型擴(kuò)散模型。訓(xùn)練了圖像的潛在擴(kuò)散模型，將常用的 U-Net 骨干替換為在潛在補(bǔ)丁上操作的transformer。
2. 通過基于前向傳播復(fù)雜度的 Gflops 測量來分析擴(kuò)散transformer（DiTs）的可擴(kuò)展性。
3. 具有更高 Gflops 的 DiTs — 通過增加transformer的深度/寬度或增加輸入token的數(shù)量會(huì)具有較低的 FID。
4. 除了具有良好的可擴(kuò)展性屬性外，最大的 DiT-XL/2 模型在條件 ImageNet 512×512 和 256×256 基準(zhǔn)上均優(yōu)于所有先前的擴(kuò)散模型，后者實(shí)現(xiàn)了 2.27 的最新 FID。
5. 揭示擴(kuò)散模型中架構(gòu)選擇的重要性，并為未來的生成建模研究提供經(jīng)驗(yàn)基線。展示了 U-Net 的歸納偏差對擴(kuò)散模型的性能并不至關(guān)重要，它們可以輕松地被標(biāo)準(zhǔn)設(shè)計(jì)，如Transformer所替代。
6. 研究了transformer相對于網(wǎng)絡(luò)復(fù)雜度與樣本質(zhì)量的擴(kuò)展行為。通過在 Latent Diffusion Models（LDMs）框架下構(gòu)建和基準(zhǔn)測試 DiT 設(shè)計(jì)空間來展示，擴(kuò)散模型在 VAE 的潛在空間內(nèi)進(jìn)行訓(xùn)練，可以成功地用transformer替換 U-Net 骨干。
7. 進(jìn)一步展示了 DiTs 是擴(kuò)散模型的可擴(kuò)展架構(gòu)：網(wǎng)絡(luò)復(fù)雜度（通過 Gflops 測量）與樣本質(zhì)量（通過 FID 測量）之間存在著強(qiáng)烈的相關(guān)性。通過簡單地?cái)U(kuò)大 DiT 并訓(xùn)練具有高容量骨干（118.6 Gflops）的 LDM，能夠在條件 ImageNet 256×256 生成基準(zhǔn)上實(shí)現(xiàn) 2.27 FID 的最新結(jié)果。

簡介

機(jī)器學(xué)習(xí)正在經(jīng)歷由 Transformer 驅(qū)動(dòng)的復(fù)興。在過去的五年里，自然語言處理、視覺和其他幾個(gè)領(lǐng)域的神經(jīng)架構(gòu)很大程度上已被 Transformer 所涵蓋。 然而，許多類別的圖像級生成模型仍然抵制這一趨勢——雖然 Transformer 在自回歸模型中得到廣泛使用，但在其他生成模型框架中的采用較少。 例如，擴(kuò)散模型一直處于圖像級生成模型最新進(jìn)展的前沿； 然而，它們都采用卷積 U-Net 架構(gòu)作為骨干網(wǎng)絡(luò)的實(shí)際選擇。U-Net 最初在像素級自回歸模型和條件 GAN 中取得了成功，它繼承自 PixelCNN++，并進(jìn)行了一些更改。該模型是卷積模型，主要由ResNet塊組成。與標(biāo)準(zhǔn)U-Net相比，額外的空間自注意力塊（變壓器中的重要組成部分）以較低的分辨率散布。Dhariwal和Nichol取消了 UNet 的幾種架構(gòu)選擇，例如使用自適應(yīng)歸一化層為卷積層注入條件信息和通道計(jì)數(shù)。
通過這項(xiàng)工作，文章的目標(biāo)是揭開擴(kuò)散模型中架構(gòu)選擇的重要性，并為未來的生成建模研究提供經(jīng)驗(yàn)基線。 表明，U-Net 歸納偏差對于擴(kuò)散模型的性能并不重要，并且可以很容易地用Transformer等標(biāo)準(zhǔn)設(shè)計(jì)代替。 因此，擴(kuò)散模型完全可以從最近的架構(gòu)統(tǒng)一趨勢中受益，例如，通過繼承其他領(lǐng)域的最佳實(shí)踐和訓(xùn)練方法，以及保留可擴(kuò)展性、穩(wěn)健性和效率等有利的屬性。標(biāo)準(zhǔn)化架構(gòu)也將為跨領(lǐng)域研究開辟新的可能性。
在文章中，關(guān)注一類基于Transformer的新型擴(kuò)散模型。 稱它們?yōu)镈iffusion Transformer，簡稱 DiT。 DiT 遵循Vision Transformers (ViTs) 的最佳實(shí)踐，事實(shí)證明，與傳統(tǒng)的卷積網(wǎng)絡(luò)（例如 ResNet）相比，它可以更有效地進(jìn)行視覺識別擴(kuò)展。更具體地說，文章研究了 Transformer 在網(wǎng)絡(luò)復(fù)雜性與樣本質(zhì)量方面的縮放行為。 表明，通過在潛在擴(kuò)散模型（LDM）框架下構(gòu)建和基準(zhǔn)化 DiT 設(shè)計(jì)空間，其中擴(kuò)散模型在 VAE 的潛在空間中進(jìn)行訓(xùn)練，可以成功地用Transformer替換 U-Net 主干。 進(jìn)一步證明 DiT 是擴(kuò)散模型的可擴(kuò)展架構(gòu)：網(wǎng)絡(luò)復(fù)雜性（以 Gflops 衡量）與樣本質(zhì)量（以 FID 衡量）之間存在很強(qiáng)的相關(guān)性。通過簡單地?cái)U(kuò)展 DiT 并訓(xùn)練具有高容量主干（118.6 Gflops）的 LDM，能夠在類條件 256 × 256 ImageNet 生成基準(zhǔn)上實(shí)現(xiàn) 2.27 FID 的最先進(jìn)結(jié)果。
文章發(fā)現(xiàn)U-Net并非不可替代，并且很容易使用諸如Transformers的結(jié)構(gòu)替代U-Net，使用Transformers可以很好地保持原有的優(yōu)秀特性，比如可伸縮性、魯棒性、高效性等，并且使用新的標(biāo)準(zhǔn)化架構(gòu)可能在跨領(lǐng)域研究上展現(xiàn)出更多的可能。文章從網(wǎng)絡(luò)復(fù)雜度和采樣質(zhì)量兩個(gè)方面對DiTs方法進(jìn)行評估。

相關(guān)工作

Transformer

Transformers 已經(jīng)取代了跨語言、視覺、強(qiáng)化學(xué)習(xí)和元學(xué)習(xí)的領(lǐng)域特定架構(gòu)。作為通用自回歸模型和ViT，它們在增加模型大小、訓(xùn)練計(jì)算和語言領(lǐng)域數(shù)據(jù)的情況下表現(xiàn)出了顯著的擴(kuò)展特性。除了語言之外，Transformer 還經(jīng)過訓(xùn)練可以自回歸預(yù)測像素。 他們還接受了離散codebook本作為自回歸模型和掩蔽生成模型的訓(xùn)練； 前者在高達(dá) 20B 參數(shù)時(shí)表現(xiàn)出出色的縮放行為。 最后，在 DDPM 中探索了 Transformer 來合成非空間數(shù)據(jù)； 例如，在 DALL·E 2 中生成 CLIP 圖像嵌入。 文章主要研究了Transformer用作圖像擴(kuò)散模型的骨干時(shí)的縮放特性。

DDPMs

擴(kuò)散和基于分?jǐn)?shù)的生成模型作為圖像生成模型特別成功，在許多情況下優(yōu)于生成對抗網(wǎng)絡(luò)（GAN）以前是最先進(jìn)的。過去兩年 DDPM 的改進(jìn)很大程度上是由改進(jìn)的采樣技術(shù)推動(dòng)的，最顯著的是無分類器指導(dǎo)、重新制定擴(kuò)散模型來預(yù)測噪聲而不是像素以及使用級聯(lián) DDPM 管道，其中 低分辨率基礎(chǔ)擴(kuò)散模型與上采樣器并行訓(xùn)練 。對于上面列出的所有擴(kuò)散模型，卷積 U-Nets 是主干架構(gòu)的事實(shí)上的選擇。 并行工作引入了一種基于 DDPM 注意力的新穎、高效的架構(gòu)。
擴(kuò)散模型是借鑒了物理學(xué)上的擴(kuò)散過程，在生成模型上，分為正向和逆向的過程。正向過程是向信號中逐漸每步加少量噪聲，當(dāng)步數(shù)足夠大時(shí)可以認(rèn)為信號符合一個(gè)高斯分布。所以逆向過程就是從隨機(jī)噪聲出發(fā)逐漸的去噪，最終還原成原有的信號。去噪過程一般采用UNet或者ViT，使用t步的結(jié)果和條件輸入預(yù)測t-1步增加的噪聲，然后使用DDPM可以得到t-1步的分布，經(jīng)過多步迭代就可以從隨機(jī)噪聲還原到有實(shí)際意義的信號。如果使用原始DDPM速度會(huì)慢很多，所以很多工作如DDIM、FastDPM等工作實(shí)現(xiàn)了解碼加速。在圖像的無條件生成任務(wù)上，擴(kuò)散模型的性能已經(jīng)超過了GANs，并且在有條件生成如文圖生成任務(wù)上大放異彩。

架構(gòu)的復(fù)雜度

在評估圖像生成文獻(xiàn)中的架構(gòu)復(fù)雜性時(shí)，使用參數(shù)計(jì)數(shù)是相當(dāng)常見的做法。 一般來說，參數(shù)計(jì)數(shù)可能不能很好地代表圖像模型的復(fù)雜性，因?yàn)樗鼈儧]有考慮到顯著影響性能的圖像分辨率等因素。相反，作者的大部分模型復(fù)雜性分析都是通過理論Gflops的視角進(jìn)行的。 這使與架構(gòu)設(shè)計(jì)文獻(xiàn)保持一致，其中Gflop被廣泛用于衡量復(fù)雜性。 在實(shí)踐中，黃金復(fù)雜度指標(biāo)仍然存在爭議，因?yàn)樗?jīng)常取決于特定的應(yīng)用場景。Nichol和Dhariwal改進(jìn)擴(kuò)散模型的開創(chuàng)性工作與我們最相關(guān)——他們分析了U-Net架構(gòu)類的可擴(kuò)展性和Gflop屬性。文章重點(diǎn)關(guān)注Transformer類。

方法

基礎(chǔ)

Diffusion基礎(chǔ)

高斯擴(kuò)散模型假設(shè)前向噪聲過程逐漸將噪聲應(yīng)用于實(shí)際數(shù)據(jù)，前向過程是一個(gè)T步逐漸加噪的馬爾科夫鏈。所以，給定前向擴(kuò)散過程作為先驗(yàn)，擴(kuò)散模型訓(xùn)練反轉(zhuǎn)的過程，可以通過去除所加噪聲從XT恢復(fù)成X0，并且每步的擴(kuò)散過程都采樣自特定的高斯分布。

Classifier-free guidance

條件擴(kuò)散模型采用額外信息作為輸入，例如類標(biāo)簽 c。 在這種情況下，逆過程變?yōu)閜θ(xt-1|xt, c)，其中θ和Σθ以c為條件。 在這種情況下，可以使用無分類器指導(dǎo)來鼓勵(lì)采樣過程找到 x，使得 log p(c|x) 較高。所以在想要條件的概率較大，就可以將條件的梯度增加到優(yōu)化目標(biāo)里。模型在訓(xùn)練時(shí)，使用一個(gè)網(wǎng)絡(luò)架構(gòu)優(yōu)化兩個(gè)模型（uncond，cond）。眾所周知，與通用采樣技術(shù)相比，無分類器指導(dǎo)可以產(chǎn)生顯著改進(jìn)的樣本，并且這種趨勢也適用于DiT模型。

潛在擴(kuò)散模型

直接在高分辨率像素空間中訓(xùn)練擴(kuò)散模型在計(jì)算上可能會(huì)令人望而卻步。 潛在擴(kuò)散模型（LDM）通過兩階段方法解決這個(gè)問題：（1）學(xué)習(xí)一個(gè)自動(dòng)編碼器，使用學(xué)習(xí)的編碼器 E 將圖像壓縮成更小的空間表示； (2) 訓(xùn)練表示 z = E(x) 的擴(kuò)散模型，而不是圖像 x 的擴(kuò)散模型（E 被凍結(jié)）。 然后可以通過從擴(kuò)散模型中采樣表示 z 并隨后使用學(xué)習(xí)的解碼器 x = D(z) 將其解碼為圖像來生成新圖像。
簡單來說就是，模型使用VAE（固定權(quán)重）將圖片encoder到隱空間，生成結(jié)果同樣也是通過VAE解碼成原始大小的圖片。

DIT架構(gòu)與設(shè)計(jì)思路

Patchify

DiT 的輸入是通過VAE后的一個(gè)稀疏的空間表示z（對于256×256×3圖像，z的形狀為32×32×4）。DiT 的第一層是“patchify”，它通過將每個(gè) patch 線性嵌入到輸入中，將空間輸入轉(zhuǎn)換為一系列 T 標(biāo)記，每個(gè)標(biāo)記的維度為 d。 patchify 之后，將標(biāo)準(zhǔn) ViT 基于頻率的位置嵌入（正弦余弦版本）應(yīng)用于所有輸入token。patchify創(chuàng)建的token T 的數(shù)量由patch大小超參數(shù) p 決定。如圖 4 所示，將 p 減半將使 T 增加四倍，從而使Transformer總 Gflop 至少增加四倍。 盡管它對 Gflops 有重大影響，但更改 p 對下游參數(shù)計(jì)數(shù)沒有有意義的影響。 文章采用將 p = 2, 4, 8 添加到 DiT 設(shè)計(jì)空間，進(jìn)行相關(guān)對比試驗(yàn)。

DIT模塊設(shè)計(jì)

在Patch化之后，輸入token由一系列Transformer塊處理。除了噪聲圖像輸入之外，擴(kuò)散模型有時(shí)還會(huì)處理其他條件信息，例如噪聲時(shí)間步 t、類標(biāo)簽 c、自然語言等。文章探索了以不同方式處理?xiàng)l件輸入的 Transformer 塊的四種變體。 這些設(shè)計(jì)對標(biāo)準(zhǔn) ViT 模塊設(shè)計(jì)進(jìn)行了微小但重要的修改。

In-context 條件

將 t 和 c 的向量emb作為輸入序列中的兩個(gè)附加token（與輸入的圖像token做拼接），將它們與圖像token沒有區(qū)別。 這類似于 ViT 中的 cls token，它允許無需修改即可使用標(biāo)準(zhǔn) ViT 塊。 在最后一個(gè)塊之后，從序列中刪除條件token。 這種方法向模型引入的新 Gflop 可以忽略不計(jì)。

Cross-attention模塊

將 t 和 c 的emb連接成一個(gè)長度為 2 的序列，與圖像token序列分開。Transformer塊被修改為在多頭自注意力塊之后包含一個(gè)額外的多頭交叉注意力層，類似于 LDM 用于調(diào)節(jié)類標(biāo)簽的方法（DiT結(jié)構(gòu)與Condition交互的方式，與原來U-Net結(jié)構(gòu)類似）。 交叉注意力為模型增加了最多的 Gflops，大約增加了 15% 的開銷。

Adaptive layer norm（adaLM）模塊

探索用自適應(yīng)layer norm (adaLN) 替換Transformer塊中的標(biāo)準(zhǔn)layer norm。不是直接學(xué)習(xí)維度尺度和平移參數(shù) γ 和 β，而是根據(jù) t 和 c 的嵌入向量之和對它們進(jìn)行回歸。 這種方式，adaLN 添加的 Gflops 最少，因此計(jì)算效率最高。 它也是唯一限制對所有token應(yīng)用相同功能的調(diào)節(jié)機(jī)制。（使用adaLN替換原生LayerNorm（NeurIPS2019的文章，LN 模塊中的某些參數(shù)不起作用，甚至?xí)黾舆^擬合的風(fēng)險(xiǎn)。所以提出一種沒有可學(xué)習(xí)參數(shù)的歸一化技術(shù)））

adaLN-zero 模塊

之前關(guān)于 ResNets 的工作發(fā)現(xiàn)，將每個(gè)殘差塊初始化為恒等函數(shù)是有益的。擴(kuò)散 U-Net 模型使用類似的初始化策略，在任何殘差連接之前對每個(gè)塊中的最終卷積層進(jìn)行零初始化。作者探索了 adaLN DiT 塊的修改，它具有相同的功能。 除了回歸 γ 和 β 之外，還回歸維度縮放參數(shù) α，這些參數(shù)在 DiT 塊內(nèi)的任何剩余連接之前立即應(yīng)用。（對殘差塊的γ, β, α進(jìn)行衰減）

模型大小

模型應(yīng)用一系列 N DiT 塊，每個(gè)塊都以隱藏維度大小 d 進(jìn)行操作。 遵循 ViT，使用標(biāo)準(zhǔn)Transformer配置來聯(lián)合縮放 N、d 和注意力頭。使用四種配置：DiT-S、DiT-B、DiT-L 和 DiT-XL。 它們涵蓋了廣泛的模型大小和觸發(fā)器分配，從 0.3 到 118.6 Gflops，使得能夠衡量擴(kuò)展性能。這些模型配置范圍從 33M 到 675M 參數(shù)

Transformer Decoder

在最終的 DiT 塊之后，需要將圖像token序列解碼為輸出噪聲預(yù)測和輸出對角協(xié)方差預(yù)測。這兩個(gè)輸出的形狀都等于原始空間輸入。使用標(biāo)準(zhǔn)線性解碼器來做到這一點(diǎn)； 應(yīng)用最終層Layer Norm（如果使用 adaLN 則為自適應(yīng)）并將每個(gè)token線性解碼為 p×p×2C 張量，其中 C 是 DiT 空間輸入中的通道數(shù)。 最后，將解碼后的token重新排列為其原始空間布局，以獲得預(yù)測的噪聲和協(xié)方差。文章主要探索的完整 DiT 設(shè)計(jì)空間是patch大小、Transformer塊架構(gòu)和模型大小。

實(shí)驗(yàn)設(shè)置

文章重點(diǎn)探索 DiT 設(shè)計(jì)空間并研究模型類的縮放屬性。模型根據(jù)其配置和潛在patch大小p命名；例如，DiT-XL/2 指的是 XLarge 配置且 p = 2。

訓(xùn)練

1. ImageNet 數(shù)據(jù)集，以256 × 256 和 512 × 512 圖像分辨率訓(xùn)練
2. 零初始化最終的線性層，否則使用 ViT 的標(biāo)準(zhǔn)權(quán)重初始化技術(shù)。
3. AdamW 優(yōu)化器，1×10?4 的恒定學(xué)習(xí)率，沒有權(quán)重衰減，批量大小為 256
4. 數(shù)據(jù)增強(qiáng)是水平翻轉(zhuǎn)
5. 與之前的 ViT 工作不同，沒有發(fā)現(xiàn)學(xué)習(xí)率預(yù)熱或正則化對于訓(xùn)練 DiT 達(dá)到高性能是必要的。
6. 訓(xùn)練過程中保持 DiT 權(quán)重的指數(shù)移動(dòng)平均值 (EMA)，衰減為 0.9999。所有結(jié)果均使用 EMA 模型。
7. 在所有 DiT 模型大小和 patch 大小上使用相同的訓(xùn)練超參數(shù)。訓(xùn)練超參數(shù)幾乎完全保留自 ADM。
8. 沒有調(diào)整學(xué)習(xí)率、衰減/熱身計(jì)劃、Adam β1/β2 或權(quán)重衰減

Diffusion

1. 使用來自Stable Diffusion的現(xiàn)成預(yù)訓(xùn)練變分自動(dòng)編碼器（VAE）模型
2. VAE編碼器的下采樣因子為 8，給定RGB圖像x的形狀為256×256×3，z=E(x)的形狀為32×32×4。
3. 從擴(kuò)散模型中采樣新的潛在變量后，使用 VAE 解碼器x=D(z)將其解碼為像素。保留來自ADM的擴(kuò)散超參數(shù)

評價(jià)指標(biāo)

1. 使用 Frechet Inception Distance (FID)來衡量縮放性能，這是評估圖像生成模型的標(biāo)準(zhǔn)指標(biāo)。
2. 使用250個(gè)DDPM采樣步驟報(bào)告FID-50K。FID 對小的實(shí)施細(xì)節(jié)很敏感，為了確保準(zhǔn)確的比較，文章所有值都是通過導(dǎo)出樣本并使用 ADM 的 TensorFlow 評估套件獲得的。
3. 文章的FID編號不使用無分類器指導(dǎo)。
4. 另外對比Inception Score、sFID和Precision/Recall 作為次要指標(biāo)

實(shí)驗(yàn)結(jié)果

1. DIT結(jié)構(gòu)：adaLN-Zero效果最好

訓(xùn)練了四個(gè)最高 Gflop DiT-XL/2 模型，每個(gè)模型都使用不同的塊設(shè)計(jì)：

1. in-context（119.4 Gflops）
2. cross-attention（137.6 Gflops）
3. adptive-LN（adaLN，118.6 Gflops）
4. adaLN-zero （118.6 Gflops）
   圖5為FID結(jié)果。adaLN-Zero塊產(chǎn)生的FID低于corss-attention和in-content，同時(shí)計(jì)算效率最高。adaLNZero將每個(gè) DiT 塊初始化為恒等函數(shù)，顯著優(yōu)于普通的 adaLN。 對于文章的其余部分，所有模型都將使用 adaLN-Zero DiT 模塊。
   

2. 縮放模型大小和patch大?。Ｐ驮酱蟆atch size越小生成圖像質(zhì)量越好）

訓(xùn)練了12個(gè) DiT 模型，涵蓋模型配置（S、B、L、XL）和patch大?。?、4、2）。與其他配置相比，DiT-L 和 DiT-XL 在相對 Gflop 方面明顯更接近?？s放模型尺寸，通過增加 Transformer 深度 / 寬度或增加輸入 token 數(shù)量，具有較高 Gflops 的 DiT 始終具有較低的 FID

1. 圖 2（左）概述了每個(gè)模型的 Gflops 及其在 400K 訓(xùn)練迭代時(shí)的 FID。 在所有情況下，增加模型大小和patch大小可以顯著改進(jìn)擴(kuò)散模型。
2. 圖 6（上）展示了 FID 如何隨著模型尺寸的增加而變化并且patch尺寸保持不變。在所有四種配置中，通過使Transformer更深更寬，F(xiàn)ID 在訓(xùn)練的所有階段都獲得了顯著改進(jìn)。
3. 圖 6（下）顯示了隨著patch尺寸減小而模型尺寸保持不變的 FID。通過簡單地?cái)U(kuò)展 DiT 處理的token數(shù)量并保持參數(shù)大致固定，整個(gè)訓(xùn)練過程中 FID 有了相當(dāng)大的改進(jìn)。
   
   

3. DiT Gflops 對于提高性能至關(guān)重要（計(jì)算開銷與模型效果關(guān)系：Gflops越大模型效果越好）

圖 6 的結(jié)果表明參數(shù)計(jì)數(shù)并不能唯一地確定 DiT 模型的質(zhì)量。當(dāng)模型大小保持不變并且patch大小減小時(shí)，Transformer的總參數(shù)實(shí)際上沒有改變（實(shí)際上，總參數(shù)略有減少），并且僅 Gflops 增加。這些結(jié)果表明，縮放模型 Gflops 實(shí)際上是提高性能的關(guān)鍵。 為了進(jìn)一步研究這一點(diǎn)，在圖 8 中針對模型 Gflops 繪制了 400K 訓(xùn)練步驟下的 FID-50K。結(jié)果表明，當(dāng)總 Gflops 相似時(shí)，不同的 DiT 配置會(huì)獲得相似的 FID 值（例如，DiT-S/2 和 DiT- B/4）。發(fā)現(xiàn)模型 Gflops 和 FID-50K 之間存在很強(qiáng)的負(fù)相關(guān)性，這表明額外的模型計(jì)算是改進(jìn) DiT 模型的關(guān)鍵因素。 在圖 12（附錄）中，發(fā)現(xiàn)這種趨勢也適用于其他指標(biāo)，例如 Inception Score。

4. 較大的 DiT 模型的計(jì)算效率更高（計(jì)算開銷與模型效果關(guān)系：模型越大計(jì)算約高效）

在圖 9 中，將 FID 繪制為所有 DiT 模型的總訓(xùn)練計(jì)算的函數(shù)。將訓(xùn)練計(jì)算估計(jì)為模型 Gflops · 批量大小 · 訓(xùn)練步驟 · 3，其中因子 3 大致近似為后向傳遞的計(jì)算量是前向傳遞的兩倍。發(fā)現(xiàn)相對于訓(xùn)練步驟較少的大型 DiT 模型而言，小型 DiT 模型即使訓(xùn)練時(shí)間較長，最終也會(huì)變得計(jì)算效率低下。同樣，發(fā)現(xiàn)除了patch大小之外相同的模型即使在控制訓(xùn)練 Gflop 時(shí)也具有不同的性能配置文件。 例如，在大約 1010 Gflops 后，XL/4 的性能優(yōu)于 XL/2。

5. 可視化
   在圖 7 中可視化縮放對樣本質(zhì)量的影響。在 400K 訓(xùn)練步驟中，使用相同的起始噪聲 xtmax 、采樣噪聲和類標(biāo)簽從 12 個(gè) DiT 模型中的每一個(gè)中采樣圖像。可以直觀地解釋縮放如何影響 DiT 樣本質(zhì)量。 事實(shí)上，縮放模型大小和token數(shù)量可以顯著提高視覺質(zhì)量。
   

6. SOTA Diffusion Models

256×256 ImageNet

在縮放分析之后，在表 2 中展示最終結(jié)果。當(dāng)使用無分類器引導(dǎo)時(shí)，DiT-XL/2 優(yōu)于所有先前的擴(kuò)散模型，將 LDM 實(shí)現(xiàn)的先前最佳 FID-50K 3.60 降低到 2.27。圖 2（右）顯示 DiT-XL/2 (118.6 Gflops) 相對于 LDM-4 (103.6 Gflops) 等潛在空間 U-Net 模型具有計(jì)算效率，并且比 ADM 等像素空間 U-Net 模型更高效 (1120 Gflops) 或 ADM-U (742 Gflops)。文章方法實(shí)現(xiàn)了所有現(xiàn)有生成模型中最低的 FID，包括之前最先進(jìn)的 StyleGANXL 。 最后，還觀察到，與 LDM-4 和 LDM-8 相比，DiT-XL/2 在所有測試的無分類器指導(dǎo)量表中實(shí)現(xiàn)了更高的召回值。當(dāng)僅訓(xùn)練 2.35M 步驟（類似于 ADM）時(shí)，XL/2 的 FID 為 2.55，仍然優(yōu)于所有先前的擴(kuò)散模型。

512x512 ImageNet

在 ImageNet 上以 512 × 512 分辨率訓(xùn)練新的 DiT-XL/2 模型，進(jìn)行 3M 次迭代，其超參數(shù)與 256 × 256 模型相同。 patch大小為 2 時(shí)，該 XL/2 模型在patch 64 × 64 × 4 輸入潛在變量 (524.6 Gflops) 后總共處理 1024 個(gè)token。 表 3 顯示了與最先進(jìn)方法的比較。 在此分辨率下，XL/2 再次優(yōu)于所有先前的擴(kuò)散模型，將 ADM 實(shí)現(xiàn)的先前最佳 FID 3.85 提高到 3.04。 即使token數(shù)量增加，XL/2 仍然保持計(jì)算效率。 例如，ADM使用1983 Gflops，ADM-U使用2813 Gflops； XL/2 使用 524.6 Gflops。 在圖 1 和附錄中展示了來自高分辨率 XL/2 模型的樣本。

7. 縮放模型與采樣計(jì)算（擴(kuò)大采樣計(jì)算不能彌補(bǔ)模型計(jì)算的不足）

擴(kuò)散模型的獨(dú)特之處在于，它們可以通過在生成圖像時(shí)增加采樣步驟數(shù)來在訓(xùn)練后使用額外的計(jì)算。 考慮到模型 Gflops 對樣本質(zhì)量的影響，研究較小模型計(jì)算 DiT 是否可以通過使用更多采樣計(jì)算來優(yōu)于較大模型？？？ 在 400K 訓(xùn)練步驟后計(jì)算所有 12 個(gè) DiT 模型的 FID，每個(gè)圖像使用 [16, 32, 64, 128, 256, 1000] 個(gè)采樣步驟。主要結(jié)果如圖 10 所示?？紤]使用 1000 個(gè)采樣步驟的 DiT-L/2 與使用 128 個(gè)采樣步驟的 DiT-XL/2。 在這種情況下，L/2 使用 80.7 Tflops 來采樣每個(gè)圖像； XL/2 使用減少 5 倍的計(jì)算量（15.2 Tflops）來對每個(gè)進(jìn)行采樣圖像。 盡管如此，XL/2 的 FID-10K 更好（23.7 vs 25.9）。 一般來說，擴(kuò)大采樣計(jì)算不能彌補(bǔ)模型計(jì)算的不足。

結(jié)論

文章主要引入了 Diffusion Transformers (DiTs)，一種基于 Transformer 的簡單擴(kuò)散模型主干，其性能優(yōu)于之前的 U-Net 模型，并繼承了 Transformer 模型類出色的縮放特性。在Gflops與Stable Diffusion相當(dāng)?shù)腄iTs-XL/2的結(jié)構(gòu)上，把ImageNet 256×256數(shù)據(jù)集上的FID指標(biāo)優(yōu)化到了2.27，達(dá)到了SOTA的水平。未來的工作應(yīng)該繼續(xù)將 DiT 擴(kuò)展到更大的模型和token數(shù)量。 DiT 還可以作為 DALL·E 2 和穩(wěn)定擴(kuò)散等文本到圖像模型的直接骨干進(jìn)行探索。

1. DiT的本質(zhì)思路就是用 Transformer 替換掉 UNet，并驗(yàn)證了不同scaling 能力和條件注入的結(jié)構(gòu)。
2. 在 Diffusion 成功之后，以 DiT 為代表的，有很多用 Transformer 替換 UNet 的工作，但都沒有受到很大的關(guān)注。究其原因，應(yīng)該是UNet 已經(jīng)做的足夠好，計(jì)算開銷也低，使用 Transformer scaling 上去，雖然生成質(zhì)量有所提升，但也沒那么顯著。
3. 在最近 OpenAI 的 Soray以及Stability.AI 的SD3 炸裂登場之后，借助 Transformer 實(shí)現(xiàn)建模的方式再次步入人們視線。

關(guān)于DIT作者進(jìn)階的論文SIT《SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers 》介紹，下一篇文章解析?。。ㄖ酰憾?；公眾號：二度簡并）

參考

論文地址：https://arxiv.org/pdf/2212.09748.pdf
GitHub地址：https://github.com/facebookresearch/DiT
項(xiàng)目地址：https://www.wpeebles.com/DiT.html文章來源地址http://www.zghlxwxcb.cn/news/detail-840950.html

到了這里，關(guān)于DIT: Scalable Diffusion Models with Transformers--Sora/SD3相關(guān)DIT技術(shù)論文閱讀的文章就介紹完了。如果您還想了解更多內(nèi)容，請?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！