內(nèi)容來自：機(jī)器之心

當(dāng)我們看到一張貓咪照片時(shí)，大腦自然就能識別「這是一只貓」。但對計(jì)算機(jī)來說，它看到的是一個(gè)巨大的數(shù)字矩陣 —— 假設(shè)是一張 1000×1000 像素的彩色圖片，實(shí)際上是一個(gè)包含 300 萬個(gè)數(shù)字的數(shù)據(jù)集（1000×1000×3 個(gè)顏色通道）。每個(gè)數(shù)字代表一個(gè)像素點(diǎn)的顏色深淺，從 0 到 255。

為了更加高效地從成千上萬張圖像中學(xué)習(xí)，AI 模型需要對圖片進(jìn)行壓縮。比如當(dāng)前最先進(jìn)的圖像生成模型，第一步就是一個(gè)名叫 tokenization 的操作，用于執(zhí)行此操作的組件叫 tokenizer。tokenizer 的主要目標(biāo)是將原始圖像壓縮到一個(gè)更小、更易處理的潛在空間，使得生成模型能夠更高效地學(xué)習(xí)和生成。因此，如何得到更好的 tokenizer 是該領(lǐng)域的研究者非常關(guān)心的問題。

在一篇新論文中，來自斯坦福大學(xué)李飛飛、吳佳俊團(tuán)隊(duì)的研究者提出了一種名叫「FlowMo」的改進(jìn)方案（論文一作是斯坦福大學(xué)計(jì)算機(jī)科學(xué)博士生 Kyle Sargent）。

FlowMo 的訓(xùn)練分為兩個(gè)階段：第一階段先學(xué)習(xí)如何全面捕捉圖像的多種可能重建結(jié)果，第二階段則學(xué)習(xí)如何從這些可能中選擇最接近原圖的重建方案。這種方法既保證了圖像重建的多樣性，又確保了重建質(zhì)量，使得 FlowMo 在 ImageNet-1K 數(shù)據(jù)集上展現(xiàn)出了領(lǐng)先的重建性能。

論文標(biāo)題： Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization

論文地址： https://arxiv.org/pdf/2503.11056v1

項(xiàng)目主頁： https://kylesargent.github.io/flowmo

研究背景

自從 VQGAN 和潛在擴(kuò)散模型等視覺生成框架問世以來，最先進(jìn)的圖像生成系統(tǒng)通常采用兩階段設(shè)計(jì)：先將視覺數(shù)據(jù)壓縮到低維潛在空間進(jìn)行 tokenization，再學(xué)習(xí)生成模型。

Tokenizer 訓(xùn)練一般遵循標(biāo)準(zhǔn)流程，即在均方誤差（MSE）、感知損失和對抗損失的組合約束下壓縮并重建圖像。擴(kuò)散自編碼器曾被提出作為學(xué)習(xí)端到端感知導(dǎo)向圖像壓縮的方法，但在 ImageNet-1K 重建這一競爭性任務(wù)上尚未達(dá)到最先進(jìn)水平。

李飛飛團(tuán)隊(duì)提出了 FlowMo，一種基于 Transformer 的擴(kuò)散自編碼器，它在多種壓縮率下實(shí)現(xiàn)了圖像 tokenization 的新性能標(biāo)準(zhǔn)，且無需使用卷積、對抗損失、空間對齊的二維潛在編碼，或從其他 tokenizer 中蒸餾知識（這與傳統(tǒng)的基于 GAN 的 tokenizer，如 VQGAN，非常不同）。

研究的關(guān)鍵發(fā)現(xiàn)是 FlowMo 訓(xùn)練應(yīng)分為模式匹配預(yù)訓(xùn)練階段和模式尋求后訓(xùn)練階段。此外，研究者進(jìn)行了廣泛分析，并探索了基于 FlowMo tokenizer 的生成模型訓(xùn)練。

圖 1：無論是在低比特率訓(xùn)練（FlowMo-Lo）還是高比特率訓(xùn)練（FlowMo-Hi）下，F(xiàn)lowMo 模型都實(shí)現(xiàn)了最先進(jìn)的 image tokenization 性能。

作者強(qiáng)調(diào)，盡管基于 GAN 的 tokenizer 在圖像 tokenization 任務(wù)上已經(jīng)取得了很好的性能，但 FlowMo 提供了一種簡單且不同的方法。

FlowMo 方法

眾所周知，基于 Transformer 的擴(kuò)散自編碼器包含編解碼結(jié)構(gòu)，因此 FlowMo 也是由編碼器 e_θ 和解碼器 d_θ 組成，其核心架構(gòu)遵循了 MMDiT，在 Transformer 的架構(gòu)中學(xué)習(xí)一維潛在空間。

圖 2：FlowMo 架構(gòu)概覽

首先，編碼器將輸入圖像轉(zhuǎn)換為潛空間向量 c，然后解碼器則根據(jù)潛空間 c 學(xué)習(xí)重建圖像的條件分布，與旨在產(chǎn)生單一確定性輸出的傳統(tǒng)方法不同，F(xiàn)lowMo 的解碼器會(huì)生成可能的重建分布，從而更好地捕捉圖像重建中固有的模糊性。

FlowMo 架構(gòu)的主要包括以下四點(diǎn)：

• 基于 Transformer 的設(shè)計(jì)：編碼器和解碼器都使用 Transformer 架構(gòu)，從而能夠更有效地處理圖像數(shù)據(jù)。

• 一維潛空間表示：FlowMo 產(chǎn)生緊湊的潛在表示，使其適用于下游生成建模任務(wù)。

• 量化層：編碼器的輸出被量化以創(chuàng)建離散 token，從而實(shí)現(xiàn)更高效的壓縮。

• 擴(kuò)散過程：解碼器使用擴(kuò)散過程逐漸將隨機(jī)輸入去噪為高質(zhì)量重建。

在 FlowMo 架構(gòu)中的一個(gè)核心創(chuàng)新點(diǎn)是其兩階段訓(xùn)練策略，這一策略使得重建分布偏向于原始圖像具有高度感知相似性的模式。

階段 1A：模式匹配預(yù)訓(xùn)練

圖 3：FlowMo 的訓(xùn)練過程結(jié)合了基于流的損失和感知損失，以引導(dǎo)模型實(shí)現(xiàn)高質(zhì)量的重建。

在階段 1A 中，F(xiàn)lowMo 通過聯(lián)合訓(xùn)練編碼器與解碼器，以實(shí)現(xiàn)兩個(gè)核心目標(biāo)：最大化潛在編碼的信息量，并使其重建分布與真實(shí)分布相匹配。這一訓(xùn)練過程巧妙地結(jié)合了多種損失函數(shù)，展現(xiàn)出其獨(dú)特的技術(shù)優(yōu)勢：

• 修正流損失（Rectified flow loss）：引導(dǎo)擴(kuò)散過程向目標(biāo)圖像分布靠攏，確保生成結(jié)果的準(zhǔn)確性；

• 感知損失（Perceptual loss）：保證了重建圖像在視覺上與原始圖像高度相似；

• 熵?fù)p失（Entropy loss）：鼓勵(lì)生成多樣化的潛在編碼，避免模式單一化；

• 承諾損失（Commitment loss）：使得編碼器輸出與量化表示盡可能接近，進(jìn)一步優(yōu)化了模型的穩(wěn)定性與效率。

具體而言，F(xiàn)lowMo 作為擴(kuò)散自動(dòng)編碼器進(jìn)行端到端訓(xùn)練，以優(yōu)化解碼器輸出上的修正流損失 L_flow，在過程中使用了 L_perc 來監(jiān)督圖像生成中的去噪預(yù)測，同時(shí)在潛空間 c 上，作者還結(jié)合了 LFQ 的熵?fù)p失和承諾損失來進(jìn)行訓(xùn)練。其中損失函數(shù)的數(shù)學(xué)表達(dá)式如下所示：

結(jié)合這些損失函數(shù)，并最終得到了第一階段的損失表達(dá)式：

階段 1B：模式探索后訓(xùn)練

在第二階段中，F(xiàn)lowMo 的核心目標(biāo)是優(yōu)化解碼器分布 pθ(x∣c)，以尋找那些在感知上與原始圖像高度相似的模式。為實(shí)現(xiàn)這一目標(biāo)，F(xiàn)lowMo 采用了創(chuàng)新的訓(xùn)練策略：首先凍結(jié)編碼器，隨后在 Lflow 的基礎(chǔ)上，聯(lián)合訓(xùn)練解碼器，并引入受擴(kuò)散模型訓(xùn)練后的 x_0 來生成目標(biāo) Lsample。這一過程通過以下步驟實(shí)現(xiàn)：

• 概率流 ODE：通過少量步驟的概率流常微分方程（ODE）集成；

• 感知損失計(jì)算：在生成樣本后，模型會(huì)計(jì)算其與原始圖像之間的感知損失，確保重建結(jié)果在視覺上與原始圖像保持一致；

• 解碼器參數(shù)更新：基于感知損失，F(xiàn)lowMo 對解碼器參數(shù)進(jìn)行優(yōu)化。

圖 4：模式搜索訓(xùn)練過程，編碼器處于凍結(jié)狀態(tài)（雪花表示），而解碼器則進(jìn)行微調(diào)以提高感知質(zhì)量。

如上圖所示，其中 FlowMo 通過凍結(jié)編碼器，集中精力優(yōu)化解碼器，使其在重建圖像時(shí)更加注重感知相似性，從而進(jìn)一步提升生成圖像的質(zhì)量與真實(shí)感。對概率流 ODE 進(jìn)行積分的 n 步樣本感知損失 Lsample 如下所示：

第二階段模式探索損失如下所示：

采樣過程

為了生成重構(gòu)圖像，F(xiàn)lowMo 通過求解概率流 ODE，對給定一維潛空間 c 的重建圖像的多模態(tài)分布進(jìn)行采樣

FlowMo 采樣方法的一項(xiàng)關(guān)鍵創(chuàng)新是使用「移位」采樣器。FlowMo 不使用統(tǒng)一的時(shí)間步長間隔，而是采用可調(diào)的移位超參數(shù)，將采樣步驟集中在擴(kuò)散過程的某些區(qū)域，從而提高感知質(zhì)量。

采樣過程需要多次前向通過解碼器模型，這在計(jì)算上很昂貴，但可以產(chǎn)生高質(zhì)量的結(jié)果。

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

主要結(jié)果

FlowMo 在多個(gè)比特率設(shè)置下（0.07 BPP 和 0.22 BPP）與當(dāng)前最先進(jìn)的 tokenizer 進(jìn)行了比較，在重建 FID（rFID）、PSNR 和 SSIM 指標(biāo)上均取得了最佳結(jié)果。在 0.07 BPP 設(shè)置下，F(xiàn)lowMo-Lo 的 rFID 為 0.95，相比 OpenMagViT-V2 的 1.17 有顯著提升；在 0.22 BPP 設(shè)置下，F(xiàn)lowMo-Hi 的 rFID 為 0.56，略優(yōu)于 LlamaGen-32 的 0.59。

表 1. tokenization 結(jié)果。

消融實(shí)驗(yàn)分析

研究團(tuán)隊(duì)進(jìn)行了大量消融實(shí)驗(yàn)，分析了 FlowMo 設(shè)計(jì)中的關(guān)鍵決策：噪聲調(diào)度、量化策略、模型架構(gòu)和后訓(xùn)練策略等。結(jié)果表明，thick-tailed logit-normal 噪聲分布、shifted sampler 和后訓(xùn)練階段對模型性能至關(guān)重要。

圖 5：噪聲調(diào)度導(dǎo)致失真的可視化案例。

特別是，沒有模式尋求后訓(xùn)練階段，F(xiàn)lowMo-Lo 的 rFID 會(huì)從 0.95 下降到 1.10，F(xiàn)lowMo-Hi 的 rFID 會(huì)從 0.56 下降到 0.73。

表 2：后訓(xùn)練消融實(shí)驗(yàn)結(jié)果。

生成任務(wù)驗(yàn)證

在生成任務(wù)中，基于 FlowMo 訓(xùn)練的 MaskGiT 在某些指標(biāo)上表現(xiàn)優(yōu)于基于 OpenMagViT-V2 訓(xùn)練的模型，但在 FID 上略遜一籌（4.30 vs 3.73）。這表明 tokenizer 質(zhì)量與下游生成模型質(zhì)量之間存在復(fù)雜關(guān)系，需要進(jìn)一步研究。

表 3：生成模型指標(biāo)對比。

圖 6：生成圖像對比。

更多詳情請見論文。