大道至簡(jiǎn)，何愷明團(tuán)隊(duì)新作pMF開啟像素級(jí)「無潛、單步」生成范式

機(jī)器之心編輯部

何愷明團(tuán)隊(duì)新論文，再次「大道至簡(jiǎn)」。

此次研究直指當(dāng)前以 DiT 為代表的主流擴(kuò)散模型與流匹配模型存在的通病，并提出了一種用于單步、無潛空間（Latent-free）的圖像生成新框架

• 論文標(biāo)題：One-step Latent-free Image Generation with Pixel Mean Flows
• arXiv 地址：https://arxiv.org/pdf/2601.22158v1

在生成式 AI 領(lǐng)域，追求更高效、更直接的生成范式一直是學(xué)界的核心目標(biāo)。

當(dāng)前，以 DiT 為代表的主流擴(kuò)散模型與流匹配模型主要依賴兩大支柱來降低生成難度，一是通過多步采樣將復(fù)雜的分布轉(zhuǎn)換分解為微小的步進(jìn)，二是在預(yù)訓(xùn)練 VAE（變分自編碼器）的潛空間中運(yùn)行以降低計(jì)算維度。

盡管這些設(shè)計(jì)在圖像質(zhì)量上取得了巨大成功，但從深度學(xué)習(xí)「端到端」的精神來看，這種對(duì)多步迭代和預(yù)置編碼器的依賴，無疑增加了系統(tǒng)的復(fù)雜性和推理開銷。

面對(duì)這些挑戰(zhàn)，何愷明團(tuán)隊(duì)提出了用于單步、無潛空間圖像生成的 pixel MeanFlow（pMF）框架。該框架繼承了改進(jìn)均值流（improved MeanFlow，MF）的思路，通過在瞬時(shí)速度（即 v）空間內(nèi)定義損失函數(shù)，來學(xué)習(xí)平均速度場(chǎng)（即 u）。

與此同時(shí)，受 Just image Transformers（JiT）的啟發(fā)，pMF 直接對(duì)類似于去噪圖像的物理量（即 x-prediction 值）進(jìn)行參數(shù)化，并預(yù)期該物理量位于低維流形上。

為了兼容這兩種設(shè)計(jì)，團(tuán)隊(duì)引入了一種轉(zhuǎn)換機(jī)制，將 v、u 和 x 三個(gè)場(chǎng)聯(lián)系起來。實(shí)驗(yàn)證明，這種設(shè)計(jì)更符合流形假設(shè)，并且產(chǎn)生了一個(gè)更易于學(xué)習(xí)的目標(biāo)（見下圖 1）。

概括來說，pMF 訓(xùn)練了一個(gè)能將噪聲輸入直接映射為圖像像素的網(wǎng)絡(luò)。它具備「所見即所得」的特性，而這在多步采樣或基于潛空間的方法中是不存在的。這一特性使得感知損失能夠自然地集成到 pMF 中，從而進(jìn)一步提升生成質(zhì)量。

實(shí)驗(yàn)結(jié)果顯示，pMF 在單步、無潛空間生成方面表現(xiàn)強(qiáng)勁，在 ImageNet 數(shù)據(jù)集上，256x256 分辨率下的 FID 達(dá)到 2.22，512x512 分辨率下達(dá)到 2.48。團(tuán)隊(duì)進(jìn)一步證明，選擇合適的預(yù)測(cè)目標(biāo)至關(guān)重要：在像素空間直接預(yù)測(cè)速度場(chǎng)會(huì)導(dǎo)致性能崩潰。

本文驗(yàn)證了：單步、無潛空間生成正變得既可行又具競(jìng)爭(zhēng)力，這標(biāo)志著向構(gòu)建單一、端到端神經(jīng)網(wǎng)絡(luò)形式的直接生成建模邁出了堅(jiān)實(shí)的一步

框架方法

為了實(shí)現(xiàn)單步、無潛空間的生成，團(tuán)隊(duì)引入了 pMF（pixel MeanFlow），它的核心設(shè)計(jì)在于建立 u、 v 和 x 這三個(gè)不同場(chǎng)之間的關(guān)聯(lián)。團(tuán)隊(duì)希望網(wǎng)絡(luò)能像 JiT 那樣直接輸出 x，而單步建模則像均值流 (MeanFlow) 一樣在 u 和 v 空間內(nèi)進(jìn)行。

去噪圖像場(chǎng)

iMF 和 JiT 都可以被視為在最小化 v-loss，不同之處在于 iMF 執(zhí)行的是 u-prediction，而 JiT 執(zhí)行的是 x-prediction。團(tuán)隊(duì)在 u 與廣義形式的 x 之間引入了一種聯(lián)系。

原論文等式 (5) 中定義的平均速度場(chǎng) u 代表了一個(gè)潛在的基準(zhǔn)真值（ground-truth），它取決于 p_data、p_prior 以及時(shí)間調(diào)度，但與網(wǎng)絡(luò)無關(guān)（因此不依賴于參數(shù) θ）。團(tuán)隊(duì)引出了一個(gè)定義為 x (z_t, r, t) 的新場(chǎng)：

可泛化的流形假設(shè)

上圖 1 通過模擬從預(yù)訓(xùn)練流匹配（FM）模型中獲得的一條 ODE 軌跡，可視化了 u 場(chǎng)和 x 場(chǎng)。u 包含噪聲圖像，這是因?yàn)樽鳛樗俣葓?chǎng)，u 同時(shí)包含了噪聲和數(shù)據(jù)成分。相比之下，x 場(chǎng)具有去噪圖像的外觀：它們或是近乎清晰的圖像，或是因過度去噪而顯得模糊的圖像。接下來，團(tuán)隊(duì)討論了如何將流形假設(shè)泛化到一物理量 x 上。

請(qǐng)注意，MeanFlow 中的時(shí)間步 r 滿足：。團(tuán)隊(duì)首先展示了 r=t 和 r=0 這兩種邊界情況可以近似滿足流形假設(shè)；隨后討論了 0＜r＜t 的情況。

算法

上文公式 (8) 中導(dǎo)出的 x 場(chǎng)為 MeanFlow 網(wǎng)絡(luò)提供了一種重參數(shù)化方法。具體而言，團(tuán)隊(duì)讓網(wǎng)絡(luò) net_θ 直接輸出 x，并根據(jù)公式 (8) 計(jì)算出相應(yīng)的速度場(chǎng) u：

接著將公式 (11) 中的 u_θ 納入 iMF 表述中，即結(jié)合 v-loss 使用原論文公式 (7)。具體的優(yōu)化目標(biāo)如下：

從概念上講，這是基于 x-prediction 的 v-loss，其中 x 通過 x→u→v 的關(guān)系轉(zhuǎn)換為 v 空間，從而對(duì) v 進(jìn)行回歸。相應(yīng)的偽代碼見算法 1。遵循 iMF 的思路，該算法可以擴(kuò)展以支持無分類器引導(dǎo)（CFG）。

帶有感知損失的像素均值

網(wǎng)絡(luò) x_θ(z_t,r,t) 直接將噪聲輸入 z_t 映射為去噪圖像，這使得模型在訓(xùn)練時(shí)具備了「所見即所得」的特性。因此團(tuán)隊(duì)進(jìn)一步引入了感知損失，基于潛空間的方法在 tokenizer 重構(gòu)訓(xùn)練中獲益于感知損失，而基于像素的方法此前尚未能輕易利用這一優(yōu)勢(shì)。

實(shí)驗(yàn)結(jié)果

玩具（Toy）實(shí)驗(yàn)

團(tuán)隊(duì)首先通過一個(gè) 2D 玩具實(shí)驗(yàn)表明，「當(dāng)?shù)讓訑?shù)據(jù)位于低維流形上時(shí)，在 MeanFlow 中使用 x-prediction 更加理想。」

圖 2 顯示，x-prediction 的表現(xiàn)相當(dāng)出色，而隨著維度 D 的增加，u-prediction 的性能迅速退化。團(tuán)隊(duì)觀察到，這種性能差距反映在訓(xùn)練損失的差異上：x-prediction 的訓(xùn)練損失低于對(duì)應(yīng)的 u-prediction。這表明，對(duì)于容量有限的網(wǎng)絡(luò)而言，預(yù)測(cè) x 更加容易。

團(tuán)隊(duì)默認(rèn)在分辨率為 256x256 的 ImageNet 數(shù)據(jù)集上進(jìn)行消融實(shí)驗(yàn)。團(tuán)隊(duì)采用了 iMF 架構(gòu)，它是 DiT 設(shè)計(jì)的一個(gè)變體。除非另有說明，團(tuán)隊(duì)將 Patch 大小設(shè)置為 16× 16（表示為 pMF/16）。消融模型從零開始訓(xùn)練了 160 個(gè) Epoch。

關(guān)于網(wǎng)絡(luò)預(yù)測(cè)目標(biāo)，團(tuán)隊(duì)的方法基于流形假設(shè)，即假設(shè) x 處于低維流形中且更易于預(yù)測(cè)。表 2 驗(yàn)證了這一假設(shè)。

首先將 64×64 分辨率作為較簡(jiǎn)單的設(shè)置。當(dāng) Patch 大小為 4×4 時(shí)，Patch 維度為 48（即 4×4×3）。這一維度遠(yuǎn)低于網(wǎng)絡(luò)容量（隱藏層維度為 768）。因此，pMF 在 x-prediction 和 u-prediction 下均表現(xiàn)良好。

接下來考慮 256×256 分辨率。按照慣例，Patch 大小設(shè)為 16×16，Patch 維度達(dá)到 768（即 16×16×3）。這導(dǎo)致了更高維的觀測(cè)空間，增加了神經(jīng)網(wǎng)絡(luò)建模的難度。在這種情況下，只有 x-prediction 表現(xiàn)良好，表明 x 位于更低維的流形上，因此更易于學(xué)習(xí)。

相比之下，u-prediction 性能徹底崩潰：作為一種含噪物理量，u 在高維空間中具有全支撐，建模難度大得多。

關(guān)于高分辨率生成，團(tuán)隊(duì)在表 4 中研究了分辨率在 256、512 和 1024 下的 pMF。在保持序列長(zhǎng)度不變（16^2）的情況下，不同分辨率下大致維持了相同的計(jì)算成本。這樣做會(huì)導(dǎo)致極其激進(jìn)的 Patch 大小（例如 64^2）和 Patch 維度（例如 12288）。

結(jié)果顯示，pMF 可以有效處理這種極具挑戰(zhàn)性的情況。盡管觀測(cè)空間是高維的，但模型始終預(yù)測(cè) x，其底層維度并不會(huì)成比例增長(zhǎng)。

關(guān)于可擴(kuò)展性，團(tuán)隊(duì)在表 5 中報(bào)告了增加模型大小和訓(xùn)練 Epoch 的結(jié)果。正如預(yù)期的那樣，pMF 從這兩個(gè)維度的擴(kuò)展中均有獲益。

最后，團(tuán)隊(duì)在表 6（256×256）和表 7（512×512）中 ，將 pMF 與之前的模型進(jìn)行了對(duì)比。

其中，在256×256 分辨率下，團(tuán)隊(duì)的方法達(dá)到了 2.22 FID（在 360 個(gè) Epoch 時(shí)），如表 6 所示。據(jù)團(tuán)隊(duì)的了解，該類別中（單步、無潛空間擴(kuò)散 / 流模型）唯一的其他方法是最近提出的 EPG，它在自監(jiān)督預(yù)訓(xùn)練下達(dá)到了 8.82 FID。

512×512 分辨率下，pMF 達(dá)到了 2.48 FID，如表 7 所示。這一結(jié)果的計(jì)算成本（參數(shù)量和 Gflops）與 256×256 版本相當(dāng)。事實(shí)上，唯一的額外開銷僅來自通道數(shù)更多的 Patch 嵌入層和預(yù)測(cè)層，所有的 Transformer 模塊都維持了相同的計(jì)算成本。

更多實(shí)驗(yàn)細(xì)節(jié)請(qǐng)參閱原論文。