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

2026-02-02 15:51

AI生成未來

關注

作者：Yiyang Lu等

解讀：AI生成未來

亮點直擊

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

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

解決的問題

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

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

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

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

提出的方案

pMF 的核心思想是將網絡的預測目標與損失函數的計算空間解耦：

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

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

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

應用的技術

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

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

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

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

達到的效果

pMF 在 ImageNet 數據集上展現了強大的性能，證明了一步無潛在生成的可行性：

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

ImageNet 512×512: FID 分數達到 2.48。

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

背景

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

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

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

FM 通過最小化 -空間中的損失函數（即“-loss”）來優化由參數化的網絡：

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

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

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

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

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

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

該恒等式提供了一種通過網絡定義預測函數的方法：

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

Pixel MeanFlow

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

去噪圖像場

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

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

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

廣義流形假設

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

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

邊界情況 I: . 當時，平均速度退化為瞬時速度，即。在這種情況下，等式 (8) 變為：

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

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

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

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

算法

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

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

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

其中。

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

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

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

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

與前人工作的關系

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

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

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

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

實驗

本文通過 2D 玩具實驗（圖 2）證明，當底層數據位于低維流形上時，在 MeanFlow 中使用 -prediction 是更可取的。實驗設置遵循。

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

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

ImageNet 實驗

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

網絡的預測目標

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

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

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

消融研究

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

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

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

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

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

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

系統級比較

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