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

2025-11-27 15:29

AI生成未來

關(guān)注

作者：Yongsheng Yu等

解讀：AI生成未來

亮點(diǎn)直擊

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

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

引入了像素級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)致：

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

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

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

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

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

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

提出的方案

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

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

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

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

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

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

像素token壓縮機(jī)制：在執(zhí)行全局注意力之前壓縮像素token，之后再進(jìn)行解壓縮。這使得像素級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)勢體現(xiàn)：避免了VAE偽影：由于直接在像素空間操作，在圖像編輯任務(wù)中能更好地保留內(nèi)容細(xì)節(jié)。高效與細(xì)節(jié)并存：雙層級設(shè)計(jì)結(jié)合高效像素建模，實(shí)現(xiàn)了高訓(xùn)練效率和快速收斂，同時(shí)保留了精細(xì)細(xì)節(jié)。

方法

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

雙層級 DiT 架構(gòu)

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

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

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

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

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

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

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

Pixel Transformer 模塊

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

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

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

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

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

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

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

用于文生圖的 PixelDiT

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

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

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

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

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

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

實(shí)驗(yàn)

實(shí)驗(yàn)圍繞 PixelDiT 在類條件與文本到圖像兩大任務(wù)上的有效性、可擴(kuò)展性與推斷效率展開。整體思路是先在 ImageNet 256×256 的類條件生成上建立像素空間 DiT 的基線與上限，再擴(kuò)展至 1024² 文本到圖像場景，以驗(yàn)證雙層級架構(gòu)在高分辨率與復(fù)雜語義條件下的穩(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級路徑引入 MM-DiT 融合；預(yù)訓(xùn)練于 512×512，隨后在 1024² 繼續(xù)微調(diào)，數(shù)據(jù)規(guī)模約 2600 萬對，覆蓋多種縱橫比。推斷時(shí)使用 FlowDPM-Solver（Rectified Flow 形式的 DPMSolver++ 變體），類條件默認(rèn) 100 步，文本到圖像默認(rèn) 25 步，以平衡質(zhì)量與時(shí)延。

類條件生成（ImageNet 256×256）

在標(biāo)準(zhǔn) 50K 采樣評估上，本工作報(bào)告了 gFID、sFID、IS、Precision–Recall。與像素生成系的代表方法（如 PixelFlow-XL、PixNerd-XL、JiT-G 等）相比，PixelDiT-XL 以顯著更低的訓(xùn)練周期開銷達(dá)到更好的或相當(dāng)?shù)?gFID，并在長訓(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)潛在基線競爭，尤其在更長訓(xùn)練與更優(yōu) CFG 區(qū)間設(shè)置時(shí)進(jìn)一步逼近最佳。

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

在 GenEval 與 DPG-Bench 上，本工作重點(diǎn)評測文本對齊、計(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 的綜合得分更高或更均衡；對比多款主流隱空間擴(kuò)散系統(tǒng)（如 SDXL、Hunyuan-DiT、Playground 等），在 1024² 的綜合表現(xiàn)接近甚至在部分維度上可比，同時(shí)參數(shù)規(guī)模更為緊湊。這表明雙層級 DiT 在高分辨率的語義一致性和構(gòu)圖控制方面具有競爭力，縮小了像素空間與隱空間在大模型上的差距。

采樣策略與超參敏感性

采樣器方面，F(xiàn)lowDPM-Solver 在無 CFG 的 100 步對比中（與 Euler/Heun）綜合了較好的 gFID/sFID 與 IS/精確度/召回率權(quán)衡，因此成為默認(rèn)選擇。步數(shù)上，隨著訓(xùn)練成熟度提升（如 400K、1.6M 步），增加采樣步數(shù)帶來的收益更明顯；在類條件上 100 步可拿到最佳指標(biāo)，而 50 步之后的收益遞減。CFG 的刻度與生效區(qū)間對質(zhì)量–多樣性平衡影響顯著：較早期（80 epoch）更偏向較強(qiáng)的引導(dǎo)且全程生效（如 3.25，[0.10,1.00]），更長訓(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 對應(yīng)圖 3 中的設(shè)計(jì)示意圖。具體而言，本工作使用一個(gè) 30 層、經(jīng)過 patch 化處理且直接在像素空間執(zhí)行去噪的 DiT 作為基線模型（標(biāo)記為“Vanilla DiT/16”）。該基線模型僅在 patch token 上操作，沒有專門的像素級路徑，將每個(gè) patch 視為一個(gè)高維向量。其在 80 epoch 時(shí)取得了 9.84 的 gFID。

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

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

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

像素 Token 緊湊化分析

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

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

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

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

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

結(jié)論

本文重新審視了像素空間的擴(kuò)散建模，并證明通過適當(dāng)?shù)募軜?gòu)設(shè)計(jì)，像素空間擴(kuò)散Transformer能夠?qū)崿F(xiàn)高保真度和高效率，而無需依賴預(yù)訓(xùn)練的自編碼器。PixelDiT將像素建模分解為雙層級Transformer設(shè)計(jì)，引入像素級AdaLN和像素token壓縮技術(shù)，在保持注意力計(jì)算可承受的同時(shí)，將全局語義與逐像素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：破局像素生成，端到端效果比肩隱空間模型

像素模型空間圖像