作者：Yongsheng Yu等

解讀：AI生成未來

亮點(diǎn)直擊

PixelDiT，一種單階段、完全基于Transformer的像素空間擴(kuò)散模型，無(wú)需單獨(dú)的自編碼器即可進(jìn)行端到端訓(xùn)練。

證明了高效的像素建模是實(shí)現(xiàn)實(shí)用像素空間擴(kuò)散的關(guān)鍵因素，并提出了一種雙層級(jí)DiT架構(gòu)，該架構(gòu)將全局語(yǔ)義學(xué)習(xí)與像素級(jí)紋理細(xì)節(jié)解耦。

引入了像素級(jí)AdaLN調(diào)制機(jī)制和像素token壓縮機(jī)制，共同實(shí)現(xiàn)了密集的每像素token建模。

PixelDiT在類別條件圖像生成和文本到圖像生成中均實(shí)現(xiàn)了高圖像質(zhì)量，顯著優(yōu)于現(xiàn)有的像素空間生成模型，并接近最先進(jìn)的隱空間擴(kuò)散模型。

總結(jié)速覽

解決的問題

兩階段流程的缺陷：傳統(tǒng)的隱空間擴(kuò)散模型依賴于預(yù)訓(xùn)練的自編碼器，這導(dǎo)致：

有損重建：自編碼器的重建過程會(huì)丟失高頻細(xì)節(jié)，限制生成樣本的保真度。

誤差累積：擴(kuò)散過程和自編碼器重建的誤差會(huì)累積。

聯(lián)合優(yōu)化困難：兩階段 pipeline 阻礙了端到端的聯(lián)合優(yōu)化。

像素空間建模的挑戰(zhàn)：直接在像素空間進(jìn)行擴(kuò)散面臨計(jì)算效率與生成質(zhì)量的權(quán)衡：

計(jì)算開銷大：對(duì)密集的像素級(jí)token進(jìn)行全局注意力計(jì)算，復(fù)雜度高，成本高昂。

細(xì)節(jié)丟失：為降低計(jì)算量而采用大尺寸圖像塊會(huì)削弱像素級(jí)建模，導(dǎo)致紋理細(xì)節(jié)生成不佳。

提出的方案

PixelDiT模型：一個(gè)單階段、端到端的完全基于Transformer的擴(kuò)散模型。

核心設(shè)計(jì)：采用雙層級(jí)架構(gòu)來解耦圖像語(yǔ)義和像素細(xì)節(jié)的學(xué)習(xí)：

塊層級(jí)DiT：使用大尺寸圖像塊，在較短的token序列上執(zhí)行遠(yuǎn)程注意力，以捕獲全局語(yǔ)義和布局。

像素層級(jí)DiT：進(jìn)行密集的像素級(jí)token建模，以細(xì)化局部紋理細(xì)節(jié)。

應(yīng)用的技術(shù)

像素級(jí)AdaLN調(diào)制：利用來自塊層級(jí)的語(yǔ)義token對(duì)每個(gè)像素token進(jìn)行條件調(diào)制，使像素級(jí)更新與全局上下文對(duì)齊。

像素token壓縮機(jī)制：在執(zhí)行全局注意力之前壓縮像素token，之后再進(jìn)行解壓縮。這使得像素級(jí)token建模得以實(shí)現(xiàn)，同時(shí)保持了全局注意力的計(jì)算效率。

達(dá)到的效果

圖像生成質(zhì)量：在ImageNet 256×256上達(dá)到FID 1.61，大幅超越了之前的像素空間生成模型。

可擴(kuò)展性：成功擴(kuò)展至文本到圖像生成，并在1024×1024像素空間直接進(jìn)行預(yù)訓(xùn)練。在GenEval上達(dá)到0.74，在DPG-bench上達(dá)到83.5，性能接近最佳的隱空間擴(kuò)散模型。

優(yōu)勢(shì)體現(xiàn)：避免了VAE偽影：由于直接在像素空間操作，在圖像編輯任務(wù)中能更好地保留內(nèi)容細(xì)節(jié)。高效與細(xì)節(jié)并存：雙層級(jí)設(shè)計(jì)結(jié)合高效像素建模，實(shí)現(xiàn)了高訓(xùn)練效率和快速收斂，同時(shí)保留了精細(xì)細(xì)節(jié)。

方法

PixelDiT，這是一種基于 Transformer 的擴(kuò)散模型，直接在像素空間執(zhí)行去噪。本工作的目標(biāo)是在保持潛空間（latent space）方法所具備的收斂行為和樣本質(zhì)量的同時(shí)，使像素 Token 建模在計(jì)算上更加高效。

雙層級(jí) DiT 架構(gòu)

如圖 2 所示，本文采用了一種雙層級(jí) Transformer 組織結(jié)構(gòu)，將語(yǔ)義學(xué)習(xí)集中在粗粒度的Patch 級(jí)路徑（patch-level pathway）上，并在像素級(jí)路徑（pixel-level pathway）中利用專門的 Pixel Transformer (PiT) 模塊進(jìn)行細(xì)節(jié)精修。這種組織方式允許大部分語(yǔ)義推理在低分辨率網(wǎng)格上進(jìn)行，從而減輕了像素級(jí)路徑的負(fù)擔(dān)并加速學(xué)習(xí)，這與文獻(xiàn) [11, 28, 29] 中的觀察一致。

Patch 級(jí)架構(gòu)：設(shè)輸入圖像為 。本文構(gòu)建非重疊的  patch tokens ，其中  為 token 數(shù)量，并將它們投影到隱藏層維度 ：

遵循文獻(xiàn) [7]，本文通過將 LayerNorm 替換為 RMSNorm 并在所有注意力層應(yīng)用 2D RoPE 來增強(qiáng) DiT 模塊。Patch 級(jí)路徑由  個(gè)增強(qiáng)的 DiT 模塊組成；對(duì)于第  個(gè)模塊，有：

其中 AdaLN 調(diào)制參數(shù)由全局條件向量  生成，然后廣播到  個(gè) patch tokens 上。這種全局到 Patch 的廣播（global-to-patch broadcasting）對(duì)所有 patch tokens 應(yīng)用相同的逐特征（per-feature）AdaLN 參數(shù)（即在 patch 級(jí)別是 token 無(wú)關(guān)的），這與隨后在像素級(jí)路徑中使用的逐像素 AdaLN 形成對(duì)比。

經(jīng)過  個(gè)模塊后，得到語(yǔ)義 tokens 。本著文獻(xiàn) [11, 28] 的設(shè)計(jì)精神，本文將像素級(jí)路徑的條件信號(hào)定義為 ，其中  是時(shí)間步嵌入。這些 tokens 通過逐像素 AdaLN 為 PiT 模塊提供語(yǔ)義上下文。

像素級(jí)架構(gòu)：像素級(jí) DiT 由  層 PiT Blocks 組成。它接收像素 tokens 和 Patch 級(jí) DiT 的輸出  作為輸入，以執(zhí)行像素 token 建模并生成最終結(jié)果。每個(gè) PiT 模塊的詳細(xì)信息如下所述。

設(shè)計(jì)要點(diǎn)。Patch 級(jí)路徑僅處理 patch tokens 以捕獲全局語(yǔ)義。通過將細(xì)節(jié)精修委托給像素級(jí)路徑，本文可以采用較大的 patch 尺寸 ，這縮短了序列長(zhǎng)度并加速了推理，同時(shí)保留了逐像素的保真度。此外，像素級(jí)路徑在縮減的隱藏維度 （例如 ）下運(yùn)行，確保密集的逐像素計(jì)算保持高效。

Pixel Transformer 模塊

每個(gè) PiT 模塊包含兩個(gè)核心組件。首先，逐像素 AdaLN（pixel-wise AdaLN）實(shí)現(xiàn)了單個(gè)像素層面的密集條件化，使逐像素更新與全局上下文對(duì)齊。其次，像素 Token 緊湊化（pixel token compaction）機(jī)制減少了像素 tokens 之間的冗余，使全局注意力能夠在可控的序列長(zhǎng)度上運(yùn)行。

逐像素 AdaLN 調(diào)制。在像素級(jí)路徑中，每個(gè)圖像通過線性層被嵌入為“每像素一個(gè) token”的形式：

為了與 Patch 級(jí)語(yǔ)義 tokens 對(duì)齊，本文將其重塑為  個(gè)序列，每個(gè)序列包含  個(gè)像素 tokens，即 。對(duì)于每個(gè) patch，本文形成一個(gè)總結(jié)全局上下文的語(yǔ)義條件 token 。

如圖 3(B) 所示，一種直接的 patch 級(jí)調(diào)制會(huì)為 patch 內(nèi)的所有  個(gè)像素重復(fù)相同的參數(shù)。然而，這無(wú)法捕獲密集的逐像素變化。相反，本文通過線性投影  將  擴(kuò)展為  組 AdaLN 參數(shù)，從而為每個(gè)像素分配獨(dú)立的調(diào)制：

并將  的最后一個(gè)維度劃分為六組，每組大小為 ，得到 。這些調(diào)制參數(shù)是可學(xué)習(xí)的，并且如圖 3(C) 所示在每個(gè)像素處都是不同的。它們通過逐像素 AdaLN 應(yīng)用于 ，實(shí)現(xiàn)了像素特定的更新；相比之下，Patch 級(jí) AdaLN 將同一組參數(shù)廣播給 patch 內(nèi)的所有像素，因此無(wú)法捕獲此類空間變化。

像素 Token 緊湊化。在像素級(jí)路徑中，直接對(duì)所有  個(gè)像素 tokens 進(jìn)行注意力計(jì)算在計(jì)算上是不可行的。因此，本文在全局注意力之前將每個(gè) patch 內(nèi)的  個(gè)像素 tokens 壓縮為一個(gè)緊湊的 patch token，隨后將注意力后的表示擴(kuò)展回像素。這將注意力序列長(zhǎng)度從  減少到 ，即減少了  倍；當(dāng)  時(shí)，這產(chǎn)生了 256 倍的縮減，同時(shí)通過逐像素 AdaLN 和可學(xué)習(xí)的擴(kuò)展保留了逐像素更新。

本文通過一個(gè)可學(xué)習(xí)的“展平”操作來實(shí)例化緊湊化算子：一個(gè)聯(lián)合混合空間和通道維度的線性映射 ，并配對(duì)一個(gè)擴(kuò)展算子 。這種“壓縮–注意力–擴(kuò)展”流水線保持了全局注意力的高效性。與 VAE 中的有損瓶頸不同，該機(jī)制僅為了注意力操作而瞬間壓縮表示。至關(guān)重要的是，這種緊湊化操作純粹是為了減少自注意力的計(jì)算開銷；它不會(huì)損害細(xì)粒度細(xì)節(jié)，因?yàn)楦哳l信息通過殘差連接和有效地繞過像素 token 瓶頸的可學(xué)習(xí)擴(kuò)展層得以保留。

用于文生圖的 PixelDiT

本文通過多模態(tài) DiT (MM-DiT) 模塊擴(kuò)展了 Patch 級(jí)路徑，該模塊融合了文本和圖像語(yǔ)義，而像素級(jí)路徑保持不變。在每個(gè) MM-DiT 模塊中，圖像和文本 tokens 形成兩個(gè)流，具有獨(dú)立的 QKV 投影。

文本嵌入  由凍結(jié)的 Gemma-2 編碼器生成。遵循文獻(xiàn) [36]，本文在用戶提示前添加簡(jiǎn)潔的系統(tǒng)提示，然后將序列輸入文本編碼器。生成的 token 嵌入被投影到模型寬度，并用作 MM-DiT 中的文本流。

經(jīng)驗(yàn)表明，來自 Patch 級(jí)路徑的語(yǔ)義 tokens 足以將文本意圖傳達(dá)給像素更新。因此，像素級(jí)路徑在架構(gòu)上與類別條件模型相同：它對(duì)像素 tokens 進(jìn)行操作，僅通過語(yǔ)義 tokens 和時(shí)間步進(jìn)行條件化。文本 tokens 不會(huì)被直接路由到像素流。

訓(xùn)練目標(biāo)

本文在像素空間采用 Rectified Flow 公式，并使用其速度匹配（velocity-matching）損失訓(xùn)練模型：

遵循文獻(xiàn) [31]，本文包含一個(gè)對(duì)齊目標(biāo)，鼓勵(lì)中層 Patch 路徑 tokens 與凍結(jié)的 DINOv2 編碼器的特征一致�？傮w目標(biāo)函數(shù)為 。類別條件模型和文本條件模型均使用相同的公式。

實(shí)驗(yàn)

實(shí)驗(yàn)圍繞 PixelDiT 在類條件與文本到圖像兩大任務(wù)上的有效性、可擴(kuò)展性與推斷效率展開。整體思路是先在 ImageNet 256×256 的類條件生成上建立像素空間 DiT 的基線與上限，再擴(kuò)展至 1024² 文本到圖像場(chǎng)景，以驗(yàn)證雙層級(jí)架構(gòu)在高分辨率與復(fù)雜語(yǔ)義條件下的穩(wěn)定性和質(zhì)量。

設(shè)置與規(guī)�；�

本工作實(shí)例化了三種模型規(guī)模（B/L/XL），在 ImageNet-1K 進(jìn)行類條件訓(xùn)練，默認(rèn)采用 PixelDiT-XL。訓(xùn)練細(xì)節(jié)強(qiáng)調(diào)了穩(wěn)定優(yōu)化與收斂速度：bfloat16 混合精度、AdamW、EMA、高梯度裁剪門限的階段性設(shè)定，以及基于 Rectified Flow 的訓(xùn)練范式。文本到圖像方面，采用 Gemma-2 作為凍結(jié)文本編碼器，并在patch級(jí)路徑引入 MM-DiT 融合；預(yù)訓(xùn)練于 512×512，隨后在 1024² 繼續(xù)微調(diào)，數(shù)據(jù)規(guī)模約 2600 萬(wàn)對(duì)，覆蓋多種縱橫比。推斷時(shí)使用 FlowDPM-Solver（Rectified Flow 形式的 DPMSolver++ 變體），類條件默認(rèn) 100 步，文本到圖像默認(rèn) 25 步，以平衡質(zhì)量與時(shí)延。

類條件生成（ImageNet 256×256）

在標(biāo)準(zhǔn) 50K 采樣評(píng)估上，本工作報(bào)告了 gFID、sFID、IS、Precision–Recall。與像素生成系的代表方法（如 PixelFlow-XL、PixNerd-XL、JiT-G 等）相比，PixelDiT-XL 以顯著更低的訓(xùn)練周期開銷達(dá)到更好的或相當(dāng)?shù)?gFID，并在長(zhǎng)訓(xùn)練（320 epoch）時(shí)取得 1.61 的 gFID 與 4.68 的 sFID，IS 為 292.7，Recall 達(dá)到 0.64，顯示出在像素空間中逼近甚至改寫既有上限的潛力。相較于隱空間的 DiT 系列（如 DiT-XL、SiT-XL、MaskDiT、LightningDiT、REPA/RAE 等），本工作雖不依賴 VAE，但質(zhì)量與多樣性指標(biāo)已能與強(qiáng)潛在基線競(jìng)爭(zhēng)，尤其在更長(zhǎng)訓(xùn)練與更優(yōu) CFG 區(qū)間設(shè)置時(shí)進(jìn)一步逼近最佳。

文本到圖像（512×512 與 1024²）

在 GenEval 與 DPG-Bench 上，本工作重點(diǎn)評(píng)測(cè)文本對(duì)齊、計(jì)數(shù)、顏色/位置屬性與組合關(guān)系的穩(wěn)健性。PixelDiT-T2I 在 512×512 達(dá)到 GenEval Overall 0.78，1024² 達(dá)到 0.74；DPG-Bench 上分別為 83.7 與 83.5。與同為像素空間的近期系統(tǒng)相比，PixelDiT-T2I 的綜合得分更高或更均衡；對(duì)比多款主流隱空間擴(kuò)散系統(tǒng)（如 SDXL、Hunyuan-DiT、Playground 等），在 1024² 的綜合表現(xiàn)接近甚至在部分維度上可比，同時(shí)參數(shù)規(guī)模更為緊湊。這表明雙層級(jí) DiT 在高分辨率的語(yǔ)義一致性和構(gòu)圖控制方面具有競(jìng)爭(zhēng)力，縮小了像素空間與隱空間在大模型上的差距。

采樣策略與超參敏感性

采樣器方面，F(xiàn)lowDPM-Solver 在無(wú) CFG 的 100 步對(duì)比中（與 Euler/Heun）綜合了較好的 gFID/sFID 與 IS/精確度/召回率權(quán)衡，因此成為默認(rèn)選擇。步數(shù)上，隨著訓(xùn)練成熟度提升（如 400K、1.6M 步），增加采樣步數(shù)帶來的收益更明顯；在類條件上 100 步可拿到最佳指標(biāo)，而 50 步之后的收益遞減。CFG 的刻度與生效區(qū)間對(duì)質(zhì)量–多樣性平衡影響顯著：較早期（80 epoch）更偏向較強(qiáng)的引導(dǎo)且全程生效（如 3.25，[0.10,1.00]），更長(zhǎng)訓(xùn)練（320 epoch）則偏向較溫和引導(dǎo)且截?cái)鄥^(qū)間（如 2.75，[0.10,0.90]），得到最低 gFID 與較高召回的綜合最優(yōu)。

消融實(shí)驗(yàn)

核心組件的貢獻(xiàn)

表 4 量化了不同模型變體中各像素建模組件的貢獻(xiàn)。注意，表 4 中的標(biāo)簽 A–C 對(duì)應(yīng)圖 3 中的設(shè)計(jì)示意圖。具體而言，本工作使用一個(gè) 30 層、經(jīng)過  patch 化處理且直接在像素空間執(zhí)行去噪的 DiT 作為基線模型（標(biāo)記為“Vanilla DiT/16”）。該基線模型僅在 patch token 上操作，沒有專門的像素級(jí)路徑，將每個(gè)  patch 視為一個(gè)高維向量。其在 80 epoch 時(shí)取得了 9.84 的 gFID。

若引入雙層級(jí)架構(gòu)但不使用像素 token 緊湊化（pixel token compaction），會(huì)導(dǎo)致全局注意力計(jì)算量隨像素?cái)?shù)量呈二次方增長(zhǎng)，從而引發(fā)顯存溢出（OOM）。加入像素 token 緊湊化解決了這一瓶頸，它將全局注意力的序列長(zhǎng)度從  個(gè)像素縮短為  個(gè) patch，在同樣的 80 epoch 預(yù)算下將質(zhì)量顯著提升至 3.50 gFID。

引入逐像素 AdaLN（pixel-wise AdaLN）進(jìn)一步將逐像素更新與 patch 級(jí)路徑產(chǎn)生的語(yǔ)義上下文對(duì)齊，使 gFID 在 80 epoch 時(shí)提升至 2.36，并在 320 epoch 時(shí)達(dá)到 1.61。

模型變體 A、B 和 C 之間的比較證明了每個(gè)提議組件的重要性。更重要的是，完整 PixelDiT 模型 C 與 Vanilla DiT/16 A 的對(duì)比表明，像素級(jí) token 建模在像素生成模型中起著關(guān)鍵作用。若沒有像素建模，即視覺內(nèi)容僅在 patch 級(jí)別進(jìn)行學(xué)習(xí)，模型將難以學(xué)習(xí)精細(xì)細(xì)節(jié)，視覺質(zhì)量也會(huì)顯著下降。

像素 Token 緊湊化分析

Token 緊湊化對(duì)于實(shí)現(xiàn)像素空間訓(xùn)練的可行性至關(guān)重要。對(duì)  個(gè)像素 token 進(jìn)行全局注意力會(huì)產(chǎn)生  的顯存占用和  的 FLOPs，即便在  分辨率下也會(huì)產(chǎn)生數(shù)十億個(gè)注意力條目，如表 5 中該變體報(bào)告的 82,247 GFLOPs 所示。利用像素 token 緊湊化將像素分組為  的 patch，可將序列長(zhǎng)度減少至 ，從而產(chǎn)生  倍的注意力開銷縮減。

為了分析像素級(jí)路徑中注意力的作用，本文包含了一個(gè)“無(wú)像素路徑注意力（No Pixel-Pathway Attention）”的消融實(shí)驗(yàn)，該實(shí)驗(yàn)移除了注意力操作，僅在像素級(jí)保留逐像素 AdaLN 和 MLP。如表 5 所示，盡管該變體減少了 GFLOPs，但在不同的訓(xùn)練迭代次數(shù)下（例如從 80 到 160 epoch），其表現(xiàn)始終遜色于完整的 PixelDiT 模型，gFID 和 IS 均有明顯的性能下降（degradation）。這表明緊湊的全局注意力對(duì)于將局部更新與全局上下文對(duì)齊是必要的。

模型規(guī)模與 Patch 大小的影響

本文研究了 Patch 大小  對(duì)不同規(guī)模模型性能的影響：PixelDiT-B、PixelDiT-L 和 PixelDiT-XL。對(duì)于所有評(píng)估，本文使用相同的 CFG 引導(dǎo)比例 3.25，區(qū)間為 。本文在 ImageNet 256×256 上評(píng)估了 4、8、16 和 32 的 Patch 大��；圖 5(a) 可視化了相應(yīng)的收斂行為。

對(duì)于 Base 模型，將  減小到  和  顯著加速了收斂：在 200K 次迭代時(shí)，gFID 從 48.5 (B/32) 降至 15.1 (B/16) 和 6.7 (B/4)，且 B/4 最終在 500K 次迭代時(shí)達(dá)到 3.4 gFID。更大的模型遵循類似的趨勢(shì)，但極小 patch 帶來的收益隨著模型規(guī)模的增加而遞減。對(duì)于 PixelDiT-L，使用  而非  僅適度改善了 gFID（在 300K 迭代時(shí)從 2.72 降至 2.15），而對(duì)于 PixelDiT-XL，（ 和  之間的差距進(jìn)一步縮小... 注：原文此處截?cái)啵��?/p>

結(jié)論

本文重新審視了像素空間的擴(kuò)散建模，并證明通過適當(dāng)?shù)募軜?gòu)設(shè)計(jì)，像素空間擴(kuò)散Transformer能夠?qū)崿F(xiàn)高保真度和高效率，而無(wú)需依賴預(yù)訓(xùn)練的自編碼器。PixelDiT將像素建模分解為雙層級(jí)Transformer設(shè)計(jì)，引入像素級(jí)AdaLN和像素token壓縮技術(shù)，在保持注意力計(jì)算可承受的同時(shí)，將全局語(yǔ)義與逐像素token學(xué)習(xí)解耦。在類別條件圖像生成和文本到圖像生成任務(wù)上的實(shí)驗(yàn)表明，該設(shè)計(jì)顯著縮小了隱空間方法與像素空間方法之間的性能差距，并在高分辨率下實(shí)現(xiàn)了強(qiáng)勁性能。

盡管由于原始數(shù)據(jù)維度較高，像素空間擴(kuò)散相比隱空間方法需要更高的計(jì)算成本，但本文的工作有效縮小了這一效率差距�？傮w而言，PixelDiT證明了實(shí)用像素空間擴(kuò)散的主要障礙并非表示空間本身，而是缺乏高效的像素建模架構(gòu)。

參考文獻(xiàn)

[1] PixelDiT: Pixel Diffusion Transformers for Image Generation

       原文標(biāo)題 : 再見VAE！英偉達(dá)PixelDiT硬剛SD/FLUX：破局像素生成，端到端效果比肩隱空間模型