作者：Yiyang Lu等

解讀：AI生成未來

亮點直擊

Pixel MeanFlow (pMF) ，這是一種針對一步生成（one-step generation）的創(chuàng)新圖像生成模型。pMF 的核心突破在于成功地在無隱空間（latent-free）的像素級建模中實現(xiàn)了高效的一步生成。

pMF不僅擺脫了對預訓練潛在編碼器（如 VQ-GAN 或 VAE）的依賴，直接在原始像素空間操作，而且在生成質量上達到了與最先進的多步隱空間擴散模型相媲美的水平。

解決的問題

現(xiàn)代生成模型通常在兩個核心維度上進行權衡：

采樣效率：多步采樣雖然質量高但推理慢。

空間選擇：隱空間（Latent Space）通過壓縮降低了維度，但引入了復雜的編碼器/解碼器，且丟失了像素級的直接控制；像素空間（Pixel Space）雖然直觀（"所見即所得"），但高維數(shù)據(jù)建模難度極大。

將“一步生成”與“像素空間建模”結合是一個極具挑戰(zhàn)性的任務，因為單一神經(jīng)網(wǎng)絡需要同時承擔極其復雜的軌跡建模（trajectory modeling）和圖像壓縮/抽象（manifold learning）任務�，F(xiàn)有的方法難以兼顧這兩者。

提出的方案

pMF 的核心思想是將網(wǎng)絡的預測目標與損失函數(shù)的計算空間解耦：

預測目標 (Prediction Target) ：網(wǎng)絡直接預測去噪后的“干凈”圖像 （即 -prediction）�；�于流形假設，干凈圖像位于低維流形上，更易于神經(jīng)網(wǎng)絡擬合。

損失空間 (Loss Space) ：損失函數(shù)定義在速度場（velocity space）中，遵循 MeanFlow 的公式，通過最小化瞬時速度誤差來學習平均速度場 。

轉換機制：引入了一個簡單的轉換公式，在圖像流形  和平均速度場  之間建立聯(lián)系：。這一轉換使得模型能夠利用像素空間的流形結構，同時在速度空間進行有效的軌跡匹配。

應用的技術

Pixel-space Prediction：直接在像素空間參數(shù)化去噪圖像 ，利用低維流形假設降低學習難度，避免直接預測高頻噪聲或速度場帶來的困難。

MeanFlow Formulation：利用 Improved MeanFlow (iMF) 框架，通過瞬時速度  的損失來學習平均速度場 。

Flow Matching：基于流匹配理論，建立從噪聲分布到數(shù)據(jù)分布的概率流。

Perceptual Loss：由于模型直接輸出像素，天然適合引入感知損失（LPIPS 和 ConvNeXt 特征），進一步提升生成圖像的視覺質量，彌補了像素級 MSE 損失的不足。

達到的效果

pMF 在 ImageNet 數(shù)據(jù)集上展現(xiàn)了強大的性能，證明了一步無潛在生成的可行性：

ImageNet 256×256: FID 分數(shù)達到 2.22，超越了許多多步隱空間模型。

ImageNet 512×512: FID 分數(shù)達到 2.48。

這表明一步像素級生成模型已經(jīng)具備了極強的競爭力，且不需要額外的解碼器開銷（解碼器本身在隱空間模型中占據(jù)顯著計算量）。

背景

本工作的 pMF 建立在 Flow Matching、MeanFlow 以及 JiT的基礎之上。

Flow Matching. Flow Matching (FM) 學習一個速度場 ，將先驗分布  映射到數(shù)據(jù)分布 。本文考慮標準的線性插值調度：

其中數(shù)據(jù) ，噪聲 （例如高斯分布），時間 。在  時，有：。該插值產(chǎn)生一個條件速度 ：

FM 通過最小化 -空間中的損失函數(shù)（即“-loss”）來優(yōu)化由  參數(shù)化的網(wǎng)絡 ：

已有研究表明 (Lipman et al., 2023)， 的潛在目標是邊緣速度 。

在推理階段，通過求解常微分方程 (ODE)： 從  到  生成樣本，其中 。這可以通過 Euler 或基于 Heun 的數(shù)值求解器來實現(xiàn)。

Flow Matching with x-prediction. 等式 (2) 中的量  是一個帶噪聲的圖像。為了便于使用在像素上操作的 Transformer，JiT 選擇通過神經(jīng)網(wǎng)絡參數(shù)化數(shù)據(jù) ，并通過以下方式將其轉換為速度 ：

其中  是 Vision Transformer (ViT) 的直接輸出。這種公式被稱為 -prediction，而在訓練中使用等式 (2) 中的 -loss。表 1 列出了這種關系。

Mean Flows. MeanFlow (MF) 框架學習一個平均速度場  用于少步/一步生成。將 FM 的  視為瞬時速度，MF 定義平均速度  為：

其中  和  是兩個時間步：。該定義引出了 MeanFlow 恒等式：

該恒等式提供了一種通過網(wǎng)絡  定義預測函數(shù)的方法：

這里，大寫  對應于等式 (6) 的左側，而在右側，JVP 表示用于計算  的 Jacobian-vector product，“sg”表示停止梯度（stop-gradient）。本文遵循 iMF的 JVP 計算和實現(xiàn)，這不是本文的重點。根據(jù)等式 (7) 的定義，iMF 像等式 (3) 一樣最小化 -loss，即 。這種公式可以被視為帶有 -loss 的 -prediction（參見表 1）。

Pixel MeanFlow

為了實現(xiàn)一步、無潛在生成，本文提出了 Pixel MeanFlow (pMF)。pMF 的核心設計是在 、 和  的不同場之間建立聯(lián)系。本文希望網(wǎng)絡像 JiT一樣直接輸出 ，而一步建模則像 MeanFlow一樣在  和  空間上進行。

去噪圖像場

如前所述，iMF 和 JiT均可視為在最小化瞬時速度  的損失（-loss），區(qū)別在于 iMF 執(zhí)行的是平均速度預測（-prediction），而 JiT 執(zhí)行的是原始數(shù)據(jù)預測（-prediction）�；�于這一觀察，本工作在平均速度  與一種廣義形式的  之間建立了一種映射聯(lián)系。

考慮等式 (5) 中定義的平均速度場 ：該場代表了一個由數(shù)據(jù)分布 、先驗分布  以及時間調度決定的底層真實量，它與具體的網(wǎng)絡參數(shù)  無關。由此，本文推導出一個誘導場（induced field），定義如下：

如下文詳述，該場  扮演了類似于“去噪圖像”的角色。需要注意的是，本工作定義的  與以往文獻中提及的  不同，它是一個受兩個時間戳  索引的二元變量：對于給定的觀測值 ，本文的  是一個隨  變化的二維場，而非僅受  索引的一維軌跡。

廣義流形假設

圖 1 通過模擬從預訓練 FM 模型獲得的一條 ODE 軌跡，可視化了  場和  場。如圖所示， 由含噪圖像組成，因為作為速度場， 包含噪聲和數(shù)據(jù)成分。相比之下， 場具有去噪圖像的外觀：它們是接近干凈的圖像，或者是因過度去噪而顯得模糊的圖像。接下來，本文討論流形假設如何推廣到這個量 。

注意 MF 中的時間步  滿足：。本文首先展示在  和  處的邊界情況可以近似滿足流形假設；然后討論  的情況。

邊界情況 I: . 當  時，平均速度  退化為瞬時速度 ，即 。在這種情況下，等式 (8) 變?yōu)椋?/p>

這本質上是 JiT 中使用的 -prediction 目標。直觀地說，這個  是 JiT 要預測的去噪圖像。如果噪聲水平很高，這個去噪圖像可能是模糊的。正如經(jīng)典圖像去噪研究中廣泛觀察到的那樣，可以假設這些去噪圖像近似位于低維（或較低維）流形上。

邊界情況 II: . 等式 (5) 中  的定義給出：。將其代入等式 (8) 得到：

即，它是 ODE 軌跡的終點。對于真實的 ODE 軌跡，有 ，即它應遵循圖像分布。因此，本文可以假設  近似位于圖像流形上。

一般情況: . 與邊界情況不同，量  不保證對應于來自數(shù)據(jù)流形的（可能模糊的）圖像樣本。然而，根據(jù)經(jīng)驗，本文的模擬（圖 1 右）表明  看起來像去噪圖像。這與速度空間量（圖 1 中的 ）形成鮮明對比，后者噪聲明顯更多。這種比較表明，通過神經(jīng)網(wǎng)絡對  進行建�？赡鼙葘Ω�嘈雜的  進行建模更容易。實驗表明，對于像素空間模型，-prediction 表現(xiàn)有效，而 -prediction 則嚴重退化。

算法

等式 (8) 中的誘導場  提供了 MeanFlow 網(wǎng)絡的一種重參數(shù)化。具體來說，本文讓網(wǎng)絡  直接輸出 ，并通過等式 (8) 計算相應的速度場 ：

這里， 是網(wǎng)絡的直接輸出，遵循 JiT。這個公式是等式 (4) 的自然擴展。

本文將 (11) 中的  納入 iMF 公式，即使用帶有 -loss 的等式 (7)。具體來說，本文的優(yōu)化目標是：

其中 。

從概念上講，這是帶有 -prediction 的 -loss，其中  通過  的關系轉換為  空間以回歸 。表 1 總結了這種關系。相應的偽代碼在 Alg. 1 中。

帶有感知損失的像素平均流

網(wǎng)絡  直接將含噪輸入  映射到去噪圖像。這使得在訓練時能夠實現(xiàn)“所見即所得”的行為。因此，除了  損失外，本文還可以進一步結合感知損失�；�于潛在的方法在 tokenizer 重建訓練期間受益于感知損失，而基于像素的方法尚未能利用這一優(yōu)勢。

形式上，由于  是像素中的去噪圖像，本文直接對其應用感知損失（例如 LPIPS）。本文的總體訓練目標是 ，其中  表示  和真實干凈圖像  之間的感知損失， 是權重超參數(shù)。在實踐中，僅當添加的噪聲低于某個閾值（即 ）時才應用感知損失，以使去噪圖像不會太模糊。本文研究了基于 VGG 分類器的標準 LPIPS 損失和基于 ConvNeXt-V2 的變體。

與前人工作的關系

本文的 pMF 與幾種先前的少步/一步方法密切相關，討論如下。

Consistency Models (CM): 學習從含噪樣本  直接到生成圖像的映射。在本文的符號中，這對應于固定終點 。此外，CM 通常采用預處理器 (Pre-conditioner)，其形式為 。除非  為零，否則網(wǎng)絡不執(zhí)行純粹的 -prediction。

Consistency Trajectory Models (CTM): 制定了一個雙時間量。與基于導數(shù)公式的 MeanFlow 不同，CTM 依賴于在訓練期間對 ODE 進行積分。

Flow Map Matching (FMM): 也是基于雙時間量。在本文符號中，F(xiàn)low Map 扮演位移的角色，即 。該量通常不位于低維流形上（例如  是含噪圖像）。

實驗

本文通過 2D 玩具實驗（圖 2）證明，當?shù)讓訑?shù)據(jù)位于低維流形上時，在 MeanFlow 中使用 -prediction 是更可取的。實驗設置遵循。

形式上，本文考慮定義在 2D 空間上的底層數(shù)據(jù)分布（此處為 Swiss roll）。數(shù)據(jù)使用  列正交矩陣投影到  維觀測空間。本文在  維觀測空間上訓練 MeanFlow 模型，其中 。本文比較了-prediction 與本文的 -prediction。

圖 2 顯示，-prediction 表現(xiàn)相當不錯，而當  增加時，-prediction 迅速退化。本文觀察到這種性能差距反映在訓練損失的差異上：-prediction 產(chǎn)生的訓練損失低于 -prediction 對應物。這表明對于容量有限的網(wǎng)絡，預測  更容易。

ImageNet 實驗

本文默認在分辨率 256×256 的 ImageNet 上進行消融實驗。報告基于 50,000 個生成樣本的 FID。所有模型均通過單次函數(shù)評估 (1-NFE) 生成原始像素圖像。

網(wǎng)絡的預測目標

本文的方法基于流形假設，即  位于低維流形上且更容易預測。本文在表 2 中驗證了這一假設。

64×64 分辨率: patch 維度為 48 ()。這個維度遠低于網(wǎng)絡容量。結果顯示 pMF 在 -prediction 和 -prediction 下都表現(xiàn)良好。

256×256 分辨率: patch 維度為 768 ()。這導致高維觀測空間，神經(jīng)網(wǎng)絡更難建模。在這種情況下，只有 -prediction 表現(xiàn)良好（FID 9.56），這表明  位于較低維流形上，因此更適合學習。相比之下，-prediction 發(fā)生災難性失�。‵ID 164.89）：作為一個含噪量， 在高維空間中具有全支撐，更難建模。

消融研究

優(yōu)化器 本文發(fā)現(xiàn)優(yōu)化器的選擇在 pMF 中起著重要作用。在圖 3a 中，本文比較了標準 Adam 優(yōu)化器與最近提出的 Muon。Muon 表現(xiàn)出更快的收斂速度和大幅提升的 FID（從 11.86 提升至 8.71）。在一步生成設置中，更快的收斂優(yōu)勢被進一步放大，因為更好的網(wǎng)絡能提供更準確的停止梯度目標。

感知損失 在圖 3b 中，本文進一步結合感知損失。使用標準 VGG-based LPIPS 將 FID 從 9.56 提升至 5.62；結合 ConvNeXt-V2 變體進一步將 FID 提升至 3.53�？傮w而言，結合感知損失帶來了約 6 個 FID 點的提升。

替代方案：預處理器 本文比較了三種預處理器變體：(i) 線性；(ii) EDM 風格；(iii) sCM 風格。表 3a 顯示，盡管 EDM 和 sCM 風格優(yōu)于樸素線性變體，但在本文考慮的極高維輸入機制中，簡單的 -prediction 更可取且性能更好。這是因為除非 ，否則網(wǎng)絡預測會偏離  空間，可能位于更高維流形上。

替代方案：時間采樣器 本文研究了限制時間采樣的替代設計：僅 （即 Flow Matching），僅 （類似 CM），或兩者的組合。表 3b 顯示這些受限的時間采樣器都不足以解決本文考慮的挑戰(zhàn)性場景。這表明 MeanFlow 方法利用  點之間的關系來學習場，限制時間采樣可能會破壞這種公式。

高分辨率生成 在表4中，本文研究了分辨率 256、512 和 1024 下的 pMF。通過增加 patch size（例如 ）來保持序列長度不變 ()，導致極大的 patch 維度（例如 12288）。結果顯示 pMF 可以有效處理這種極具挑戰(zhàn)性的情況。即使觀測空間是高維的，模型始終預測 ，其潛在維度不會成比例增長。

可擴展性 表 5 顯示，增加模型大小和訓練周期均能提升結果。

系統(tǒng)級比較

ImageNet 256×256. 表 6 顯示本文的方法達到了 2.22 FID。據(jù)本文所知，該類別中（一步、無潛在擴散/流）唯一的方法是最近提出的 EPG，其 FID 為 8.82。與領先的 GAN 相比，pMF 實現(xiàn)了相當?shù)?FID，但計算量大幅降低（例如 StyleGAN-XL 的計算量是 pMF-H/16 的 5.8 倍）。

ImageNet 512×512. 表 7 顯示 pMF 在 512×512 下達到 2.48 FID。值得注意的是，其計算成本（參數(shù)量和 Gflops）與 256×256 對應物相當。唯一的開銷來自 patch embedding 和預測層。

結論

本質上，圖像生成模型是從噪聲到圖像像素的映射。由于生成建模的固有挑戰(zhàn)，該問題通常被分解為更易處理的子問題，涉及多個步驟和階段。雖然有效，但這些設計偏離了深度學習的端到端精神。

本文關于 pMF 的研究表明，神經(jīng)網(wǎng)絡是具有高度表現(xiàn)力的映射，當設計得當時，能夠學習復雜的端到端映射，例如直接從噪聲到像素。除了其實際潛力外，本文希望本工作將鼓勵未來對直接、端到端生成建模的探索。

參考文獻

[1] One-step Latent-free Image Generation with Pixel Mean Flows

       原文標題 : 擴散模型迎來“終極簡化”！何愷明團隊新作：像素級一步生成，速度質量雙巔峰