摘要

我們提出了 PixelFlow，這是一系列直接在原始像素空間中運行的圖像生成模型，與主流的隱空間模型形成對比。這種方法通過消除對預訓練變分自編碼器（VAE）的需求，簡化了圖像生成過程，并使整個模型能夠端到端訓練。通過高效的級聯流建模，PixelFlow 在像素空間中實現了可承受的計算成本。在 256×256 ImageNet 類條件圖像生成基準上，它實現了 1.98 的 FID。定性的文本到圖像結果表明，PixelFlow 在圖像質量、藝術性和語義控制方面表現出色。我們希望這種新范式能為下一代視覺生成模型帶來啟發并開辟新的機會。代碼和模型可在??https://github.com/ShoufaChen/PixelFlow??獲取。

1. 引言

“如無必要，勿增實體。”

—— 奧卡姆的威廉

在 Stable Diffusion（SD）模型系列 [17,46,47,50] 成功的推動下，隱擴散模型（LDMs）[50] 已成為跨多種模態生成建模的事實上的標準，涵蓋圖像 [17,35,45]、視頻 [7,8,23,66,69]、音頻 [18,39] 和 3D [57,67] 等領域。如圖 1（a）所示，LDMs 使用預訓練的變分自編碼器（VAEs）將原始數據壓縮到緊湊的隱空間中。這種壓縮降低了計算需求，并促進了高效的擴散去噪。盡管取得了廣泛的成功，但 LDMs 將 VAE 和擴散組件分離，阻礙了聯合優化并使整體診斷復雜化。

另一種方法是在原始像素空間中實現擴散模型。雖然直觀，但由于處理每個像素相關性需要大量資源，對于高分辨率圖像來說，這在計算上是不可承受的。考慮到這一點，先前的研究 [20,22,44,51,52] 通常采用級聯方法：首先生成低分辨率圖像，然后使用額外的上采樣器生成高質量輸出，低分辨率圖像作為條件輸入，如圖 1（b）所示。然而，這些級聯方法也為不同階段引入了獨立的網絡，仍然限制了端到端設計的優勢。

在這項工作中，我們引入 PixelFlow，這是一個簡單而有效的端到端框架，用于在原始像素空間中直接生成圖像，無需像 VAE 或上采樣器這樣的獨立網絡。如圖 1（c）所示，PixelFlow 使用統一的參數集通過流匹配 [38,40] 對級聯分辨率的多尺度樣本進行建模。在去噪的早期階段，當噪聲水平較高時，PixelFlow 在較低分辨率的樣本上運行。隨著去噪的進行，分辨率逐漸增加，直到在最后階段達到目標分辨率。這種漸進策略避免了在全分辨率下執行所有去噪步驟，從而顯著降低了生成過程的總體計算成本。

在訓練期間，不同時間步的跨尺度樣本通過以下方式構建：（1）將圖像調整為連續的尺度，并向每個縮放后的圖像添加高斯噪聲；（2）在相鄰尺度的噪聲圖像之間進行插值作為模型輸入并進行速度預測。整個模型使用從所有階段均勻采樣的訓練示例進行端到端訓練。在推理期間，過程從最低分辨率的純高斯噪聲開始。然后模型逐步對圖像進行去噪和上采樣，直到達到目標分辨率。

圖片

圖1. 基于隱變量的擴散模型（LDMs）、基于像素的擴散模型（PDMs）與PixelFlow的設計范式對比：(a) LDMs將訓練分為兩個獨立階段——先獨立訓練現成的VAE模型，再在預訓練VAE提取的token上訓練擴散模型；(b) 傳統PDMs通常需訓練兩個獨立模型：針對低分辨率圖像的擴散模型和用于高分辨率合成的上采樣器；(c) 相比之下，PixelFlow提供了端到端的像素級生成方案，兼具高效性與強大生成性能。

我們在類條件和文本到圖像生成任務上對 PixelFlow 進行了評估。與已建立的隱空間擴散模型 [42,45,50] 相比，PixelFlow 提供了有競爭力的性能。例如，在 256×256 ImageNet 類條件生成基準上，PixelFlow 實現了 1.98 的 FID。對于文本到圖像生成，PixelFlow 在廣泛使用的基準上進行了評估，在 GenEval [19] 上得分為 0.64，在 DPG-Bench [26] 上得分為 77.93。此外，圖 5 和圖 6 的定性結果表明，PixelFlow 具有很強的視覺保真度和文本 - 圖像對齊，突出了像素空間生成在未來研究中的潛力。

PixelFlow 的貢獻總結為以下三點：

?通過消除對預訓練 VAE 的需求，我們直接在原始像素空間中建立了端到端可訓練的圖像生成模型。

?通過從低分辨率到高分辨率的級聯流建模，我們的模型在訓練和推理中都實現了可承受的計算成本。

?PixelFlow 在視覺質量方面取得了有競爭力的性能，包括在 256×256 ImageNet 類條件生成基準上的 1.98 FID，以及在文本到圖像生成上的吸引人的特性。

2. 相關工作

2.1 隱空間擴散 / 流模型

變分自編碼器（VAEs）已成為許多最新生成模型 [16,17,35,47,48,50,59,66] 的核心組件，能夠將視覺數據從像素空間映射到低維的感知等效隱空間。這種緊湊的表示便于更高效的訓練和推理。然而，VAEs 常常犧牲高頻細節 [47]，導致生成輸出中不可避免的低級偽影。出于對算法簡單性和完全端到端優化的渴望，我們放棄了 VAE，直接在像素空間中操作。

2.2 像素空間擴散 / 流模型

早期的擴散模型 [2,21,56] 主要直接在像素空間中運行，旨在單階段捕獲圖像的分布。然而，這種方法對于高分辨率圖像生成既具有挑戰性又效率低下，導致了級聯模型 [20,22,30,52] 的發展，這些模型通過一系列階段生成圖像。這些級聯模型通常從生成低分辨率圖像開始，然后通過超分辨率模型進行上采樣以獲得更高的分辨率。然而，基于擴散的超分辨率過程通常需要從純噪聲開始，以低分辨率輸出為條件，導致耗時且低效的生成過程。此外，在孤立階段訓練這些模型阻礙了端到端優化，并需要精心設計的策略來確保超分辨率階段。

此外，像素空間生成的最新進展引入了創新的架構。Simple Diffusion [24,25] 提出了一種用于高分辨率圖像合成的簡化擴散框架，通過調整模型架構和噪聲調度在 ImageNet 上取得了強大性能。FractalGen [37] 通過遞歸調用原子生成模塊構建分形生成模型，形成自相似架構，在逐像素圖像生成中表現出強大性能。TarFlow [68] 提出了一種基于 Transformer 的歸一化流架構，能夠直接建模和生成像素。

3. PixelFlow

3.1 預備知識：流匹配

流匹配算法 [1,38,40] 逐步將樣本從先驗分布（通常是標準正態分布）轉換為目標數據分布。這是通過定義一個前向過程來實現的，該過程由一系列直接連接先驗分布樣本和目標分布對應樣本的線性路徑組成。在訓練期間，通過首先采樣目標樣本x_1，從標準正態分布中抽取噪聲x_0～N(0,1)，并選擇時間步t∈[0,1]來構建訓練示例。然后通過線性插值定義訓練示例：

模型被訓練來近似由常微分方程（ODE）定義的速度v_t = \frac{dx_t}{dt}，使其能夠有效地引導中間樣本x_t到真實數據樣本x_1的轉換。

流匹配的一個顯著優點是它能夠在兩個任意分布之間進行插值，而不僅限于使用標準高斯作為源域。因此，在圖像生成任務中，流匹配超越了噪聲到圖像的場景，可以有效地應用于圖像到圖像翻譯等各種應用。

3.2 像素空間中的多尺度生成

PixelFlow 通過多階段去噪過程逐步增加圖像的分辨率來生成圖像。為此，我們通過在每個尺度上以因子 2 遞歸下采樣目標圖像x_1來構建其多尺度表示。如圖 2 所示，PixelFlow 將圖像生成過程分為 S 個階段。每個階段s\in\{0,1,\ldots,S-1\}在由起始和結束狀態(x_{t_0^s}, x_{t_1^s})定義的時間間隔內運行。在S=1的退化情況下，PixelFlow 簡化為用于圖像生成的標準單階段流匹配方法，類似于最近的工作 [17,42]，但關鍵是在像素空間而非隱空間中操作。

圖2. 基于像素空間的級聯圖像生成框架PixelFlow。我們將完整生成過程劃分為多級分辨率階段：每個分辨率階段起始時，對前階段含較多噪聲的結果進行上采樣，并將其作為當前階段的起始點。通過這種漸進式優化機制，隨著分辨率提升可獲得更精細的生成樣本。

對于每個階段 s，我們定義起始和結束狀態如下：

其中，\text{Down()}和\text{Up()}分別表示下采樣和上采樣操作。除非另有說明，我們采用雙線性插值進行下采樣，最近鄰插值進行上采樣。

為了訓練模型，我們通過在起始和結束狀態之間進行線性插值來采樣中間表示：

其中，\tau = \frac{t - t_0^s}{t_1^s - t_0^s}是第 s 階段內的重新縮放時間步 [29,65]。

然后，我們的目標是訓練一個模型\mu_\theta(\cdot)來預測速度\mu_\theta(x_{t_\tau^s}, \tau)，目標為v_t = x_{t_1^s} - x_{t_0^s}。我們使用均方誤差（MSE）損失，形式上表示為：

\mathbb{E}_{s,t,(x_{t_1^s}, x_{t_1^s})} \left\| \mu_\theta(x_{t_\tau^s}, \tau) - v_t \right\|^2

3.3 模型架構

我們使用基于 Transformer 的架構 [62] 實例化\mu_\theta(\cdot)，選擇它是因為其簡單性、可擴展性和在生成建模中的有效性。具體來說，我們的實現基于標準擴散 Transformer（DiT）[45]，在所有實驗中采用 XL 規模配置。為了更好地與 PixelFlow 框架保持一致，我們引入了以下幾個修改：

3.3.1 補丁化（Patchify）

遵循視覺 Transformer（ViT）設計 [15,45]，PixelFlow 的第一層是補丁嵌入層，通過線性投影將輸入圖像的空間表示轉換為 1D 令牌序列。與在 VAE 編碼的隱空間上操作的先前隱 Transformer [17,42,45] 不同，PixelFlow 直接對原始像素輸入進行令牌化。為了支持批處理中多個分辨率的高效注意力，我們應用序列打包策略 [11]，沿著序列維度連接對應不同分辨率的扁平令牌序列。

3.3.2 旋轉位置編碼（RoPE）

補丁化后，我們將原始的正弦位置編碼 [45] 替換為 RoPE [58]，以更好地處理不同的圖像分辨率。RoPE 在支持長度外推方面表現出強大性能，尤其是在大型語言模型中。為了將其適應 2D 圖像數據，我們通過對高度和寬度維度獨立應用 1D-RoPE 來應用 2D-RoPE，每個維度占據隱藏狀態的一半。

3.3.3 分辨率嵌入（Resolution Embedding）

由于 PixelFlow 使用共享的模型參數在多個分辨率上運行，我們引入了額外的分辨率嵌入來區分不同的分辨率。具體來說，我們將補丁嵌入后特征圖的絕對分辨率作為條件信號。該信號使用正弦位置嵌入 [62] 進行編碼，并在傳遞到模型之前添加到時間步嵌入中。

3.3.4 文本到圖像生成

雖然類條件圖像生成通常通過自適應層歸一化（adaLN）[45] 整合條件信息，但我們通過在每個 Transformer 塊內的每個自注意力層之后引入交叉注意力層 [6,7]，將 PixelFlow 擴展為支持文本到圖像生成。這種設計使模型能夠在生成過程的每個階段有效地將視覺特征與文本輸入對齊。遵循最近的工作 [8,59]，我們采用 Flan-T5-XL 語言模型 [10] 來提取豐富的文本嵌入，作為整個網絡的條件信號。

3.4 訓練和推理

為了促進高效訓練，我們使用公式（4）中定義的插值方案從所有分辨率階段均勻采樣訓練示例。此外，我們采用序列打包技術 [11]，允許在單個小批量中聯合訓練尺度可變的示例，提高了效率和可擴展性。

在推理期間，生成過程從最低分辨率的純高斯噪聲開始，通過多個階段逐步過渡到更高的分辨率。在每個分辨率階段內，我們應用標準的基于流的采樣，根據速度和準確性的權衡，使用歐拉離散采樣器 [17] 或 Dopri5 求解器。為了確保跨尺度的平滑和連貫過渡，我們采用重新去噪策略 [29,60]，有效緩解了多尺度生成管道中常見的跳躍點問題 [4]。

4. 實驗

4.1 實驗設置

我們在 ImageNet-1K [12] 數據集上評估 PixelFlow 的類條件圖像生成能力。除非另有說明，我們在 256×256 分辨率下訓練 PixelFlow。所有模型使用 AdamW 優化器 [32,41] 進行訓練，學習率為1?—10^{-4}。性能主要通過 Fréchet Inception Distance（FID）使用標準評估工具包 1 進行測量。我們還報告了 Inception Score（IS）[53]、sFID [43] 和 Precision/Recall [33]。

對于文本條件圖像生成，我們逐步將 PixelFlow 從 256×256 訓練到 1024×1024 分辨率。我們與當前最先進的生成模型進行定性比較，并在流行的基準上進行定量評估，如 T2ICompBench [27]、GenEval [19] 和 DPG-Bench [26]。

4.2 模型設計

4.2.1 起始序列長度

原則上，PixelFlow 可以訓練為從非常低的分辨率（例如 1×1）逐步增加到目標分辨率。然而，這種方法在實踐中效率低下，因為極低分辨率的令牌傳達的有意義信息有限。此外，為非常短的序列分配過多的時間步會未充分利用現代 GPU 的計算能力，導致模型 FLOPS 利用率下降。因此，我們探索圖像生成開始時的分辨率（我們稱為起始圖像分辨率）如何影響整體性能。

對于我們的 Transformer 骨干網絡，注意力操作中涉及的令牌數量由原始圖像分辨率和補丁大小決定。在這個實驗中，我們保持 2×2 的補丁大小 [45]，使起始序列長度直接取決于起始圖像分辨率。具體來說，我們評估了三種起始序列長度 ——2×2、8×8 和 32×32—— 同時將目標分辨率固定為 64×64。值得注意的是，32×32 設置代表了沒有跨分辨率級聯的香草像素基方法。

如表 1 所示，在這些配置中，8×8 起始序列長度與 32×32 基線相比，實現了相當甚至稍好的 FID。這表明，從適當較小的分辨率開始生成并逐步擴展，可以在保持生成質量的同時，通過為最大分辨率階段分配更少的計算來提高計算效率。相反，將起始序列長度進一步減少到 2×2 會導致性能下降，可能是因為極低分辨率的令牌提供的有用信息有限，無法為后續生成步驟提供足夠的指導。考慮到生成質量和計算效率，我們因此采用 8×8 作為默認起始序列長度。

4.2.2 補丁大小

接下來，我們在保持起始序列長度為 2×2 的同時，研究補丁大小對模型性能的影響。首先，我們在目標分辨率為 64×64 的情況下進行實驗，比較兩種補丁大小 ——2×2 和 4×4—— 結果如表 2 的上部所示。我們觀察到 PixelFlow 在這兩種設置下表現非常相似，4×4 補丁在五個評估指標中的四個上略優于 2×2 補丁。此外，使用 4×4 補丁大小消除了 2×2 補丁大小配置所需的最高分辨率階段，從而提高了效率。

當擴展到更大的目標分辨率（即 256×256）時，由于資源需求巨大，使用 2×2 補丁大小在計算上變得不可行，限制我們只能進行 100K 次訓練迭代（表 2 的中間部分）。這一限制促使我們采用更大的補丁大小。盡管將補丁大小進一步增加到 8×8 顯著提高了計算效率，但它導致性能質量明顯下降。此外，即使經過擴展訓練（1600K 次迭代），這種性能差距仍然存在，如表 2 的底部所示。考慮到生成質量和計算成本，我們因此選擇 4×4 作為默認補丁大小。

4.3 推理調度

在表 3 中，我們詳細分析了推理配置空間，包括每個分辨率階段的推理步驟數、ODE 求解器的選擇以及無分類器引導（CFG）的調度。

4.3.1 采樣步驟數

在表 3a 中，我們評估了每個分辨率階段的推理步驟數對生成質量的影響。隨著步驟數的增加，我們觀察到 FID、sFID 和 IS 的持續改進，在 30 步時達到最佳整體性能。超過這個點后，增益飽和甚至略有下降，表明收益遞減。

PixelFlow 的一個顯著優點是它在推理期間為每個分辨率階段分配不同數量采樣步驟的靈活性。這種自適應配置允許對采樣過程進行細粒度控制，實現性能 - 效率權衡。超越統一設置并探索更細粒度的階段特定步驟分配有可能帶來進一步的性能提升。

4.3.2 ODE 求解器

我們進一步研究 ODE 求解器類型對生成質量的影響。如表 3b 所示，我們將一階歐拉求解器與自適應高階 Dormand-Prince（Dopri5）求解器 [14] 進行比較。結果表明，Dopri5 在大多數評估指標上始終優于歐拉求解器，實現了更低的 FID 和 sFID 分數、更高的 Inception Score 以及稍好的精度，同時保持相似的召回率。這表明，更準確和自適應的求解器（如 Dopri5）可以更好地捕捉生成動態，從而生成更高質量的樣本 —— 盡管通常伴隨著更高的計算成本。

4.3.3 CFG 調度

受最近工作 [5,34,63] 的啟發，我們提出了一種階段式 CFG 調度，其中不同階段應用不同的 CFG 值，并且從早期階段到后期階段，值從 1 增加到CFG_{max}。在 4 個階段的情況下，我們發現0, 1/6, 2/3和1的(CFG_{max}-1)給出了最佳的 FID 性能。全局常數 CFG 和階段式 CFG 之間的比較如表 3c 所示，其中我們為每種方法搜索最佳 CFG 值。我們提出的階段式 CFG 將 FID 性能從 2.43 提升到 1.98。

4.4 在 ImageNet 基準上的比較

在表 4 中，我們將 PixelFlow 與基于隱空間和像素空間的圖像生成模型在 ImageNet 256×256 基準上進行了比較。PixelFlow 實現了 1.98 的 FID，代表了相對于最先進的隱空間方法的極具競爭力的性能。例如，它優于 LDM [50]（FID 3.60）、DiT [45]（FID 2.27）和 SiT [42]（FID 2.06），同時實現了可比的 IS 和召回分數。這些結果突出了我們設計的有效性，表明 PixelFlow 可以作為高質量視覺生成系統的強大原型。

與最近的像素基模型相比，PixelFlow 實現了卓越的樣本質量。它顯著優于 FractalMAR-H [37]，并且與 ADM-U [13]、SiD2 [25] 和 VDM++[31] 等強基線相比，提供了具有競爭力或更好的結果。

我們在圖 4 中可視化了 256×256 分辨率下 PixelFlow 的類條件圖像生成。我們可以觀察到我們的模型能夠跨廣泛的類別生成高視覺質量的圖像。

4.5 文本到圖像生成

4.5.1 設置

我們為 PixelFlow 的文本到圖像生成采用兩階段訓練策略。首先，使用 ImageNet 預訓練的檢查點在 256×256 分辨率下初始化模型，并在 LAION 數據集 [55] 的子集上以相同分辨率進行訓練。在第二階段，我們在精心挑選的高審美質量圖像集上以更高的 512×512 分辨率對模型進行微調。所有報告的 PixelFlow 結果均基于此最終的 512×512 分辨率模型。

4.5.2 定量結果

如表 5 所示，PixelFlow 在所有基準上均實現了有競爭力的性能，展示了在自由形式文本到圖像生成中的強大組合理解能力。它在 T2I-CompBench 上表現尤為出色，在顏色和紋理綁定方面得分很高，在 GenEval（0.64）和 DPG-Bench（77.93）上取得了堅實的結果，超越了許多已建立的模型。這些結果強調了 PixelFlow 作為基于自然語言條件的像素空間圖像生成的有前途方向 —— 展示了其在開放式文本驅動圖像合成中的潛力。

4.5.3 可視化

我們在圖 3 中可視化了采樣過程中的中間結果，具體顯示了每個分辨率階段的最后一步。隨著分辨率的增加，出現了明顯的去噪趨勢 —— 圖像在每個階段逐漸變得更清晰，噪聲更少。圖 5（512×512）和圖 6（1024×1024）中顯示了額外的生成樣本及其輸入文本提示。PixelFlow 展示了高視覺保真度和強文本 - 圖像對齊，有效地捕捉了復雜提示中的關鍵視覺元素及其關系。值得注意的是，它生成了細粒度的細節 —— 如動物毛發、人類頭發和帽子紋理 —— 突出了其在像素空間中對細節的強烈關注。

圖3. 級聯生成階段的中間結果可視化。我們從四個生成階段分別提取中間結果進行直接可視化觀測，可以清晰觀察到不同分辨率階段逐步去噪的過程。

5. 結論

我們引入了 PixelFlow，這是一種新穎的圖像生成模型，通過直接在原始像素空間中操作，重新思考了基于隱空間模型的主導地位。通過在不同分辨率階段之間直接轉換，我們的模型在簡單性和端到端可訓練性方面展現出引人注目的優勢。在類條件圖像生成和文本到圖像生成基準上，PixelFlow 已被證明與流行的基于隱空間的方法相比，展示了有競爭力的圖像生成能力。我們希望這種新視角將激發未來視覺生成模型的研究。

局限性

盡管具有優勢，PixelFlow 仍然面臨某些局限性。盡管模型避免了所有階段的全分辨率計算，但最后階段需要全分辨率注意力，這約占總推理時間的 80%。此外，我們觀察到隨著序列長度的增加，訓練收斂速度變慢。解決這些挑戰為未來提高效率和可擴展性提供了機會。

本文轉載自公眾號AIRoobt ，作者：Shoufa Chen等

原文鏈接：??https://mp.weixin.qq.com/s/ktz8qQ0C-l_W_nuQvxkKkw??