「兄弟們，DiT 是錯(cuò)的！」

最近一篇帖子在 X 上引發(fā)了很大的討論，有博主表示 DiT 存在架構(gòu)上的缺陷，并附上一張論文截圖。

圖 1. 我們引入了 TREAD，這是一種能夠顯著提升基于 token 的擴(kuò)散模型骨干網(wǎng)絡(luò)訓(xùn)練效率的訓(xùn)練策略。當(dāng)應(yīng)用于標(biāo)準(zhǔn)的 DiT 骨干網(wǎng)絡(luò)時(shí)，我們?cè)跓o(wú)引導(dǎo) FID 指標(biāo)上實(shí)現(xiàn)了 14/37 倍的訓(xùn)練速度提升，同時(shí)也收斂到了更好的生成質(zhì)量。

圖中橫軸代表訓(xùn)練時(shí)間（以 A100 GPU 的小時(shí)數(shù)為單位，log 尺度，從 100 小時(shí)到 10000 小時(shí)），縱軸代表 FID 分?jǐn)?shù)（越低越好，代表生成圖像質(zhì)量越高）。

博主認(rèn)為，這個(gè)圖的核心信息不是 TREAD 的速度優(yōu)勢(shì)，而是 DiT 的 FID 過(guò)早穩(wěn)定，暗示 DiT 可能存在「隱性架構(gòu)缺陷」，導(dǎo)致其無(wú)法繼續(xù)從數(shù)據(jù)中學(xué)習(xí)。

博主提到的論文發(fā)表于今年 1 月（3 月更新 v2），介紹了一種名為 TREAD 的新方法，該工作通過(guò)一種創(chuàng)新的「令牌路由」（token routing）機(jī)制，在不改變模型架構(gòu)的情況下，極大地提升了訓(xùn)練效率和生成圖像的質(zhì)量，從而在速度和性能上都顯著超越了 DiT 模型。

具體而言，TREAD 在訓(xùn)練過(guò)程中使用「部分令牌集」（partial token set） vs 「完整令牌集」（full token set），通過(guò)預(yù)定義路由保存信息并重新引入到更深層，跳過(guò)部分計(jì)算以減少成本，同時(shí)僅用于訓(xùn)練階段，推理時(shí)仍采用標(biāo)準(zhǔn)設(shè)置。這與 MaskDiT 等方法類(lèi)似，但更高效。

• 論文標(biāo)題：TREAD:Token Routing for Efficient Architecture-agnostic Diffusion Training
• 論文地址：https://arxiv.org/abs/2501.04765
• 代碼：https://github.com/CompVis/tread

博主在后續(xù)回復(fù)中逐步展開(kāi)了對(duì) DiT 的批判，并解釋 TREAD 如何暴露這些問(wèn)題。

博主指出，該論文揭示了 DiT 模型的設(shè)計(jì)缺陷。具體來(lái)說(shuō)，研究發(fā)現(xiàn)在訓(xùn)練過(guò)程中，如果將模型中的一部分計(jì)算單元替換為「恒等函數(shù)」（Identity Function）—— 也就是讓這些單元什么計(jì)算都不做，僅僅是「直通」數(shù)據(jù)，相當(dāng)于被臨時(shí)禁用了 —— 模型的最終評(píng)估分?jǐn)?shù)反而會(huì)提高。

接著博主指出 DiT 的兩個(gè)「可疑」的設(shè)計(jì)：

• 整個(gè)架構(gòu)都使用「后層歸一化」（Post-LayerNorm）

博主認(rèn)為 DiT 使用了一種已知不太穩(wěn)定的技術(shù)（后層歸一化），來(lái)處理一個(gè)數(shù)值范圍變化極其劇烈的任務(wù)（擴(kuò)散過(guò)程）。

• adaLN-zero

博主認(rèn)為，這個(gè)模型雖然整體上自稱(chēng)是「Transformer」架構(gòu)，但在處理最關(guān)鍵的「指導(dǎo)信息」（即條件數(shù)據(jù)）時(shí)，并沒(méi)有使用強(qiáng)大的 Transformer，而是用了一個(gè)非常簡(jiǎn)單的 MLP 網(wǎng)絡(luò)（多層感知機(jī)）。

更具體地，adaLN-zero 通過(guò)完全覆蓋注意力單元的輸入，并注入任意偏置來(lái)覆蓋輸出，這限制了模型的表達(dá)能力，相當(dāng)于「討厭注意力操作」（hate the attention operation），從而削弱了 DiT 的整體潛力。

博主還提到與早期論文相關(guān)的 LayerNorm 研究，指出 LayerNorm 的偏置和增益參數(shù)可能對(duì)梯度調(diào)整影響更大，而非真正改善模型性能。他認(rèn)為，adaLN-zero 正是利用了這一點(diǎn)，名為「梯度調(diào)節(jié)」，實(shí)則像是在「給小模型偷偷注入過(guò)擬合的偏置」。

• 論文標(biāo)題：Understanding and Improving Layer Normalization
• 論文地址：https://arxiv.org/abs/1911.07013

看了這篇帖子，DiT 的作者，紐約大學(xué)計(jì)算機(jī)科學(xué)助理教授謝賽寧有些忍不住了。

在 2022 年，謝賽寧發(fā)表了 DiT 的論文，這是擴(kuò)散模式首次和 Transformer 相結(jié)合。

• 論文標(biāo)題：Scalable Diffusion Models with Transformers
• 論文鏈接：https://arxiv.org/pdf/2212.09748

在 DiT 問(wèn)世之后，Transformer 逐步代替原始擴(kuò)散模型中的 U-Net，在圖像和視頻生成任務(wù)中生成高質(zhì)量的結(jié)果。

其核心思想是采用 Transformer 代替?zhèn)鹘y(tǒng)的卷積神經(jīng)網(wǎng)絡(luò)作為擴(kuò)散模型的主干網(wǎng)絡(luò)。

這一方法業(yè)已成為 Sora 和 Stable Diffusion 3 的基礎(chǔ)架構(gòu)，同時(shí)也確定了 DiT 的學(xué)術(shù)地位。

在 DiT 論文剛剛問(wèn)世時(shí)，就已接連受到質(zhì)疑，甚至以「缺乏創(chuàng)新」為由被 CVPR 2023 拒稿。

這一次面對(duì) DiT 在數(shù)學(xué)和形式上都「是錯(cuò)的」的論調(diào)，謝賽寧發(fā)推做出了幾點(diǎn)回應(yīng)。

從字里行間來(lái)看，謝賽寧對(duì)這個(gè)帖子多少有些情緒：

我知道原帖是在釣魚(yú)騙點(diǎn)擊率，但我還是咬一下鉤……

坦白講，每個(gè)研究者的夢(mèng)想其實(shí)就是發(fā)現(xiàn)自己的架構(gòu)是錯(cuò)的。如果它永遠(yuǎn)都沒(méi)問(wèn)題，那才是真正的大問(wèn)題。

我們每天都在用 SiT、REPA、REPA-E 等方法試圖打破 DiT，但這需要基于假設(shè)、做實(shí)驗(yàn)、進(jìn)行驗(yàn)證，而不是只在腦子里扮演式地做科學(xué)…… 否則，你得出的結(jié)論不僅僅是錯(cuò)的，而是根本連錯(cuò)都談不上。

也難怪謝賽寧語(yǔ)氣有些不善，原帖博主的一些說(shuō)法可能有些拱火的嫌疑：

謝賽寧也從技術(shù)角度對(duì)于原帖子提出的一些問(wèn)題進(jìn)行了回復(fù)，在對(duì)原帖的部分問(wèn)題進(jìn)行了反駁后，他也同樣說(shuō)明了 DiT 架構(gòu)目前存在一些硬傷。

截至今天，DiT 的問(wèn)題：

• tread 更接近于 stochastic depth，我認(rèn)為它的收斂性來(lái)自正則化效應(yīng)，這讓表示能力更強(qiáng)（注意推理過(guò)程是標(biāo)準(zhǔn)的 —— 所有模塊都會(huì)處理所有 token）；這是非常有意思的工作，但和原帖說(shuō)的完全不是一回事。
• Lightning DiT 已經(jīng)是經(jīng)過(guò)驗(yàn)證的穩(wěn)健升級(jí)版（結(jié)合了 swiglu、rmsnorm、rope、patch size=1），有條件就應(yīng)該優(yōu)先使用它。
• 沒(méi)有任何證據(jù)表明 post-norm 會(huì)帶來(lái)負(fù)面影響。
• 過(guò)去一年最大的改進(jìn)點(diǎn)在于內(nèi)部表示學(xué)習(xí)：最早是 REPA，但現(xiàn)在有很多方法（例如 tokenizer 層面的修正：VA-VAE / REPA-E，把語(yǔ)義 token 拼接進(jìn)噪聲潛變量、解耦式架構(gòu)如 DDT，或者通過(guò)分散損失、自表示對(duì)齊等正則化手段）。
• 始終優(yōu)先采用隨機(jī)插值 / 流匹配（SiT 在這里應(yīng)該是基線）。
• 對(duì)于時(shí)間嵌入，使用 AdaLN-zero；但遇到更復(fù)雜的分布（如文本嵌入）時(shí)，應(yīng)采用交叉注意力。
• 不過(guò)要用對(duì)方式 —— 采用 PixArt 風(fēng)格的共享 AdaLN，否則會(huì)白白浪費(fèi) 30% 的參數(shù)。
• 真正的「硬傷」其實(shí)是 DiT 里的 sd-vae：這是顯而易見(jiàn)卻長(zhǎng)期被忽視的問(wèn)題 —— 它臃腫低效（處理 256×256 圖像竟需要 445.87 GFlops？）、不是端到端的。像 VA-VAE 和 REPA-E 只是部分修復(fù)，更多進(jìn)展還在路上。

評(píng)論網(wǎng)友也對(duì)回應(yīng)中提到的技術(shù)細(xì)節(jié)感興趣，謝賽寧也都對(duì)相關(guān)疑惑做出了回復(fù)：

算法的迭代進(jìn)步總是伴隨著對(duì)現(xiàn)有算法的質(zhì)疑，雖說(shuō)所謂「不破不立」，但 DiT 仍然在擂臺(tái)中央，不是么？