作者：Yongsheng Yu等

解讀：AI生成未來

亮點直擊

PixelDiT，一種單階段、完全基于Transformer的像素空間擴散模型，無需單獨的自編碼器即可進行端到端訓練。

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

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

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

總結(jié)速覽

解決的問題

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

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

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

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

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

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

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

提出的方案

PixelDiT模型：一個單階段、端到端的完全基于Transformer的擴散模型。

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

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

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

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

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

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

達到的效果

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

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

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

方法

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

雙層級 DiT 架構(gòu)

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

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

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

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

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

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

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

Pixel Transformer 模塊

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

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

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

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

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

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

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

用于文生圖的 PixelDiT

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

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

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

訓練目標

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

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

實驗

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

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

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

類條件生成（ImageNet 256×256）

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

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

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

采樣策略與超參敏感性

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

消融實驗

核心組件的貢獻

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

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

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

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

像素 Token 緊湊化分析

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

為了分析像素級路徑中注意力的作用，本文包含了一個“無像素路徑注意力（No Pixel-Pathway Attention）”的消融實驗，該實驗移除了注意力操作，僅在像素級保留逐像素 AdaLN 和 MLP。如表 5 所示，盡管該變體減少了 GFLOPs，但在不同的訓練迭代次數(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 大��；圖 5(a) 可視化了相應(yīng)的收斂行為。

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

結(jié)論

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

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

參考文獻

[1] PixelDiT: Pixel Diffusion Transformers for Image Generation

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