李飛飛、吳佳俊團隊新作：FlowMo如何以零卷積、零對抗損失實現ImageNet重構新巔峰

一、摘要

二、引言

三、相關工作

四、方法

基于擴散先前的離散標記化器利用廣告

架構

階段 1A：模式匹配預訓練

階段 1B：模式搜索后訓練

采樣

第二階段：潛在生成建模

五、Coovally AI模型訓練與應用平臺

六、實驗

主要結果

分析

七、結論

圖片1.png

論文題目：Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization

論文鏈接：

https://arxiv.org/pdf/2503.11056v1

一、摘要

自VQGAN和潛在擴散模型等流行的視覺生成框架出現以來，最先進的圖像生成系統一般都是兩階段系統，首先將視覺數據標記化或壓縮到低維潛在空間，然后再學習生成模型。標記化訓練通常采用標準方法，即根據MSE、實際損失和對抗損失的組合對圖像進行壓縮和重建。擴散自動編碼器在之前的工作中已被提出，作為一種學習端到端感知導向圖像壓縮的方法，但在ImageNet-1K重構這一競爭性任務中尚未顯示出最先進的性能。我們提出的FlowMo是一種基于變換器的擴散自動編碼器，它能以多種壓縮率實現最新的圖像標記化，而無需使用卷積、對抗損失、空間對齊的二維拉編碼或從其他標記化器中提煉。我們的主要見解是，FlowMo訓練應分為模式匹配前訓練階段和模式搜索后訓練階段。此外，我們還進行了廣泛的分析，并探索了在FlowMo標記器上訓練生成模型的方法。

二、引言

擴散模型和離散自回歸模型等生成模型在圖像和視頻內容創建方面的應用引人注目。雖然級聯系統和單級系統等不太常見的系統也被探索過，但最先進的視覺生成系統一般包括兩個階段：首先，學習一個“標記器”，將像素數據壓縮到一個更小、更容易處理的離散或連續潛空間；其次，在這個壓縮的潛空間上訓練一個生成模型。由于這種兩階段模式的優勢，標記化本身已成為一個活躍的研究領域，許多重要的研究工作都非常關注標記化的訓練和設計。在這項工作中，我們重點研究具有離散潛空間的標記化器。

自VQGAN以來，離散圖像標記符的主流架構和訓練方案已經出現。最先進的圖像標記器通常是經過訓練的卷積自動編碼器，先將視覺數據下采樣為二維、空間對齊的潛碼，然后再將潛碼上采樣為經過精算和感知損失正則化的重構。有人提出了與這一設置不同的方法：TiTok使用了基于變壓器的架構和一維潛碼，但在初始蒸餾階段依賴于傳統的基于CNN的標記化器；ViT-VQGAN使用了基于變壓器的編碼器和解碼器。但總的來說，圖像標記化的基本設置仍然占主導地位。

在這項工作中，我們提出了FlowMo（Flow to the Mode，流向模式），它在技術上大大突破了當前圖像標記化的技術水平。首先，在處理圖像標記化任務時，我們建議使用解碼器的整流目標來模擬多模態再構造分布。其次，我們使用完全基于變壓器的架構，對一維潛碼進行編碼和解碼。第三，我們對FlowMo進行了端到端優化，而不是從已有的二維標記符中進行提煉，或在二維標記符的潛在空間之上進行編碼。

最重要的是，FlowMo通過一個關鍵的洞察力實現了最先進的性能。FlowMo是一種擴散編碼器，這一系統已在之前的工作中進行過探索。然而，在最直觀的感知重建基準ImageNet-1K方面，最先進的技術仍然被更傳統的基于CNN或GAN的標記化器所主導。FlowMo通過以下關鍵見解實現了最先進的標記化：對于感知重建任務來說，對重建分布中實際上接近原始圖像的模式進行采樣比試圖匹配所有模式更好。因此，我們將FlowMo的訓練分為模式匹配前訓練階段和模式搜索后訓練階段，前者是對系統進行端到端訓練，解碼器上有擴散損失，就像典型的擴散自動編碼器一樣；后者是對系統進行訓練，有選擇性地放棄重建分布中不匹配的模式。

在這一階段，系統會接受訓練，選擇性地放棄重建分布中不接近原始圖像的模式。我們將在第3節解釋這兩個階段。

盡管FlowMo在方法上與之前的研究有很大的不同，但與現有最強的多BPP標記化器相比，FlowMo仍是最先進的。我們的主要貢獻如下：

我們基于擴散自動編碼器和多模態擴散圖像轉換器 (MMDiT)，提出了一種簡單但新穎的圖像標記化架構。
我們為擴散自編碼器提出了一種新穎的兩階段訓練方案，包括模式匹配前訓練和模式搜索后訓練。
在每像素0.07比特和每像素0.22比特兩種情況下，我們在感知圖像標記化的rFID、PSNR和其他指標方面都達到了最新水平。我們還表明，使用FlowMo標記化器訓練的生成模型可以與使用傳統標記化器訓練的生成模型相媲美（盡管不能超越）。
我們對FlowMo中的系統設計選擇進行了廣泛的分析，概述了噪聲調度、采樣器設計、模型設計、量化策略和后期訓練中幾個微妙而關鍵的決定。

圖片2.png

三、相關工作

圖像標記化。最先進的視覺生成系統一般由兩個階段組成（少數例外情況是嘗試直接在像素空間中學習）。第一階段是“標記化”，通過投射到連續或離散的潛在空間，重新生成像素數據的空間或時空維度。這一范例的主導地位促使許多作品將標記化器本身作為重要組件進行研究。在這項工作中，我們研究的是具有離散潛變量的標記化器。與之前的工作不同，FlowMo是第一個基于擴散自動編碼器的架構，在ImageNet-1K重建中達到了最先進的性能。擴散自動編碼器擴散模型在視覺生成中非常流行，而簡化框架（如整流模型）已被進一步采用。利用擴散解碼器端對端學習自動編碼器的想法最早是在中提出的。隨后又有許多研究跟進了這一想法，研究了用于表征學習的擴散自動編碼器，特別是用于感知導向的圖像壓縮。擴散后訓練。各種研究都試圖通過專門的后訓練策略來提高擴散模型樣本的質量，其目的通常是灌輸所需的屬性，如美學質量。DDPO和DPOK探索的是強化學習目標，而DRAFT和AlignProp探索的是通過采樣鏈進行反向傳播。FlowMo采用基于反向傳播的后訓練策略，以適應整流的連續噪聲時間表和圖像標記化的獨特設置。

并行工作DiTo在與我們同時進行的工作中研究了用于連續圖像標記化任務的擴散自動編碼器。與DiTo不同，我們專注于離散圖像標記化，并與現有最強大的離散標記化器進行比較。同時，Flex-Tok提出了一種系統，在該系統中，擴散標記化器是在傳統連續變異自動編碼器（VAE）的潛空間之上學習的，該自動編碼器經過感知損失、對抗損失和重建損失的訓練。FlowMo不依賴輔助VAE。

四、方法

現有的先進標記化器在視覺生成方面取得了相當大的進步。不過，它們也有很多缺點。首先，它們需要對抗損失，這可能不穩定且難以調整。它們幾乎都需要在至少一個訓練階段使用CNN，因此很難利用變壓器的硬件效率和廣為人知的擴展行為。最后，它們通常利用以前訓練過的標記化器進行提煉，以達到最先進的性能。

FlowMo有幾個關鍵目標。我們提出的標記符是：

基于擴散先前的離散標記化器利用廣告

基于擴散。先前的離散標記化器利用廣告

這就需要自適應梯度尺度計算、LeCam正則化，或對損失權重進行仔細調整以確保穩定性。相反，我們將使用擴散解碼器，因為擴散模型已被證明簡單、非常適合多模態建模，而且在大規模應用中非常可靠。

純變換器，一維潛碼。

幾乎所有最先進的圖像標記化框架都采用了基于CNN的架構，或預測了一個局部對齊的二維潛碼，但TiTok除外，它仍然依賴于從一個基于CNN的預訓練標記化器中提煉出二維潛碼。雖然這些選擇可能會提供有用的歸納偏差，但我們認為基于變換器的一維潛碼標記器最終可能會在大型數據和模型規模上提供更高的效率和靈活性。FlowMo在基于轉換器的架構中學習一維潛碼，這與MMDiT類似。

最新技術。

在擴散自動編碼器方面已經有相當多的研究，但FlowMo是第一個在ImageNet-1K上實現最先進的感知重建的技術。

FlowMo實現了這些目標，我們現在將通過對系統的闡述以及隨后（在第 4 節中）通過大量的實驗、消融和分析加以說明。作為一種擴散自動編碼器，鑒于潛在代碼c中的信息有限，重建圖像x的分布（表示為p(x|c)）必然是多模態的。我們的主要見解是，為了利用擴散自動編碼器實現最先進的標記化，應采取多種措施使 p(x|c) 偏向于與原始圖像具有高感知相似性的模式。我們通過兩個關鍵要素來實現這一目標： (1) 3.3節中解釋的模式搜索后訓練，以及 (2) 3.4節中解釋的移位采樣器。

現在我們將詳細介紹FlowMo。這是一種基于MMDiT的模擬變壓器架構。然后，我們將解釋兩個訓練階段：模式匹配預訓練（階段1A）和模式搜索微調（階段1B）。最后，我們將討論在FlowMo潛在空間上的生成建模（階段 2）和采樣。

架構

FlowMo的結構圖如圖3所示。FlowMo是一個擴散自動編碼器。解碼器是一個條件擴散模型，用于學習重建圖像的條件分布p(x|c)。

圖片3.png

FlowMo由編碼器eθ和解碼器dθ組成。兩者都是基于多模態擴散圖像變換器（MMDiT）的變換器。給定一個經過修補的圖像x∈Rn和一個初始潛碼c0∈Rd（一個全為零的向量），編碼器產生一個潛標記序列

圖片4.png

c?然后通過無查找量化（LFQ）的量化操作q進行元素二值化，得到

圖片5.png

按照整流，解碼器被訓練成從噪聲到數據的速度場v建模，定義為

圖片6.png

解碼器處理xt和c的方式與編碼器處理x和c0的方式相同，但通過AdaLN調制]，解碼器還接受時間（或噪聲水平）參數t來調節每個MMDiT塊。編碼器eθ和解碼器dθ在結構上是對稱的，但大小不同，解碼器更大更深。我們對所有模型都使用了μP參數化，以簡化探索性配置和放大配置之間的超參數轉移。

階段 1A：模式匹配預訓練

在第一個訓練階段，我們的目標是聯合訓練編碼器eθ和解碼器dθ，使量化后的c對x具有最大的信息量，從而使pθ(x|c)與真實分布p(x|c)相匹配，由于量化后的c只包含有限的信息，因此真實分布必然是多模態的。圖 4 給出了階段1A的完整示意圖。關于理論依據，我們參考了之前在擴散自動編碼器方面的研究，其中指出端到端擴散目標對應于非高斯解碼器 p(x|c) 的修正變分下限。

圖片7.png

FlowMo被訓練成端到端擴散自動編碼器。具體來說，FlowMo編碼器和解碼器是端對端訓練，以優化解碼器輸出上的整流損失Lflow。給定噪聲z～N (0, I )、數據 x～px 和時間（或噪聲水平）t～pt,t∈[0,1]，我們定義

圖片8.png

并優化流量匹配目標

圖片9.png

我們還使用學習到的感知距離dperc，通過以下方式對網絡的1步去噪預測進行優化

圖片10.png

最后，在潛碼c上，我們按照的方法使用LFQ的熵損失和承諾損失，即

圖片11.png

為簡單起見，這兩個損失可以用有限標量量化（FSQ）代替，但我們發現LFQ 的性能稍好一些。我們在1A階段的訓練損耗為

圖片12.png

損失權重和更多細節見補充材料。圖像被重新調整為位于[-1, 1]范圍內。噪聲水平t是按照 Stable Diffusion 3 ，從厚尾logit-normal分布中采樣的。

階段 1B：模式搜索后訓練

在這一階段，我們的目標是優化解碼器分布pθ(x|c)，以尋求在感知上與原始圖像高度相似的模式。為了實現這一目標，我們凍結了編碼器，并在Lflow和后訓練目標Lsample的基礎上對解碼器進行聯合訓練，后訓練目標Lsample的靈感來源于之前在擴散模型后訓練方面的工作。

Lsample只是對概率流ODE的n步采樣的感知損失，我們通過它進行區分。我們隨機采樣積分的時間步t1、......、tn，以便在測試時使用不同的采樣計劃進行實驗，并使用梯度檢查點來反向傳播采樣鏈。圖 5 給出了階段1B的完整示意圖。

圖片13.png

讓dti (xt)表示流量采樣更新函數，即?

圖片14.png

然后，我們定義

圖片15.png

第 1B 階段的全部損失為

圖片16.png

我們發現本階段的λsample值尤為重要，因此使用λsample=0.01。Lsample的權重太小會導致rFID效果不佳，因為網絡會遺忘在第1A階段獲得的實際特征；而權重太大則會導致所謂的“獎勵劫持”（即解碼器dθ與dperc過度擬合）或訓練發散。我們在第1A階段使用LPIPS-VGG網絡作為感知網絡，但在第 1B 階段，我們使用ResNet作為感知網絡。我們在補充材料中解釋了這一設計選擇和其他設計選擇，并通過消融法證明，僅僅在Lperc中的1步去噪預測中使用ResNet網絡是無效的；因此有必要引入后訓練階段。這一階段的計算成本很高，因為它需要通過整個采樣鏈進行反向傳播。我們使用n=8，并發現它在采樣時能很好地適用于其他 n。

采樣

給定一個量化的潛碼c，通過求解概率流ODE來采樣給定c的重建圖像的多模態分布（表示為 p(x|c)）。

圖片17.png

給定初始 x1～N (0, I)。另外，FlowMo 也可用于計算樣本對數似然，方法是按照流匹配反向求解流 ODE。

在推理時，我們用可調整的移動超參數給出的時間步距對整流ODE進行積分。圖5.第1B階段。凍結編碼器eθ將輸入圖像編碼為c，作為解碼器 dθ 的條件，解碼器dθ通過整個采樣鏈的反向傳播進行訓練。我們還使用 Lflow 進行協同訓練，這與圖4中的方法相同。

ρ 表示采樣時間步向低噪聲水平集中。對于時間步距：

圖片18.png

設置ρ=1相當于通常的線性間隔整流ODE采樣器。在極端情況下，讓ρ→∞意味著在t=1時邁出一大步，相當于給定c對x進行回歸。我們使用ρ=4，這相當于在接近 t=1時邁出一大步，并將采樣集中到 p(x|c) 的平均值上，同時在低噪聲水平上仍然花費了大量的采樣FLOP，而之前的研究表明，這一選擇對rFID至關重要。我們的采樣器大大提高了rFID和PSNR。

第二階段：潛在生成建模

與其他使用離散空間訓練自動編碼器的工作一樣，我們驗證了我們的標記器可以用來訓練高質量的第二階段生成模型。我們使用MaskGiT，這一階段的設置主要來自MaskGiT和TiTok。

五、Coovally AI模型訓練與應用平臺

如果你也想要進行模型訓練或模型改進，Coovally平臺滿足你的要求！

Coovally平臺整合了國內外開源社區1000+模型算法和各類公開識別數據集，無論是YOLO系列模型還是Transformer系列視覺模型算法，平臺全部包含，均可一鍵下載助力實驗研究與產業應用。

而且在該平臺上，無需配置環境、修改配置文件等繁瑣操作，一鍵上傳數據集，使用模型進行訓練與結果預測，全程高速零代碼！

具體操作步驟可參考：YOLO11全解析：從原理到實戰，全流程體驗下一代目標檢測

平臺鏈接：https://www.coovally.com

如果你想要另外的模型算法和數據集，歡迎后臺或評論區留言，我們找到后會第一時間與您分享！

六、實驗

主要結果

標記化。對于標記化這一主要任務，所有標記化器都采用圖像輸入，將其編碼為量化潛像，然后重建圖像。所有標記化器都是在ImageNet-1K上訓練的。重建質量用rFID 、PSNR、SSIM和LPIPS等指標來衡量。我們在分辨率為256×256的ImageNet-1K驗證集上進行評估。

標記化器量化潛碼所包含的信息量以每像素比特（BPP）為單位。

其中 S是潛碼序列長度，V是標記詞詞匯量，標記化器的BPP計算公式為：

圖片19.png

在不同BPP下訓練的標記化器不能進行蘋果對蘋果的比較，因為獲取更多比特會提高性能。因此，我們訓練了兩種模式，以匹配現有最先進的標記符。

FlowMo-Lo以0.0703 BPP進行訓練，以匹配OpenMagViT-V2的BPP。FlowMo-Hi以0.219 BPP進行訓練，以匹配LlamaGen-32的BPP。除了BPP之外，FlowMo-Lo和FlowMo-Hi都是完全相同的FlowMo架構。為了與我們進行比較的tokens生成器的BPP相匹配，我們修改了FlowMo tokens詞匯的大小或tokens數量。就=rFID、PSNR和SSIM而言，我們在兩個BPP上都取得了最先進的結果。FlowMo性能優異的唯一例外是LPIPS指標，FlowMo在該指標上仍然表現不佳。可在補充材料和我們的網站上查看具體的可視化比較。

圖片20.png

生成。在生成任務中，我們對不同標記化器生成的編碼標記序列訓練生成模型MaskGiT。然后，我們通過各自的標記化器對生成的標記序列進行解碼，并用圖像生成指標來衡量圖像質量，即FID、sFID、Inception Score以及Precision和Recall。生成指標在ImageNet-1K基準上進行評估。盡管實現了更強的標記化，但使用FlowMo作為下游生成模型的標記化器，只能改善部分生成指標，而不能改善所有生成指標。標記化質量和生成質量之間存在有趣而復雜的相互作用，我們希望在未來的工作中改進結果。

表2中的模型大小、批量大小和訓練長度等訓練超參數與托克生成器相同。具體來說，我們在批量大小為128、學習率為1×10-4的兩個token- enizers上都訓練了300個epoch。其他超參數設置見補充材料。雖然我們沒有足夠的資源來訓練最先進的生成模型（例如，需要1,000 - 1,080個歷元的訓練，批量大小為1,024 - 2,048），但我們在這里的目標并不是為生成模型設定最先進的水平，而是從生成模型所提供的潛在空間的角度對標記生成器進行公平的比較。

圖片21.png

其他比較。在表3中，我們還對FlowMo和DiTo進行了標記化比較。為了進行比較，我們使用連續潛空間訓練標記化器，將Lay- erNorm作為最終編碼層，并使DiTo中提出的 “噪聲同步 ”增強技術。我們還均衡了總體潛空間大小，使用256個標記，標記維度為16。實驗細節。我們所有的模型都是μP參數化的，因此超參數可以在探索性配置和放大配置之間“μ-轉移”。我們的大部分實驗是在混合硬件（A6000、L40S、A100、H100 或 H200 GPU）上進行的，ViT補丁大小為8，隱藏維度為768；這也是表4中用于消融研究的配置。

圖片22.png

我們的最終實驗（FlowMo-Lo、FlowMo-Hi）是在8×H100節點上進行的，使用的是放大模型，補丁大小為4，僅解碼器的隱藏維度增加到 1152。所有其他超參數都直接進行了μ轉移。我們對FlowMo-Lo進行了100個epoch的預訓練，對FlowMo-Hi進行了80個epoch的預訓練，然后對兩者進行了約 1個epoch的后訓練。在預訓練中，rFID的性能沒有達到飽和。更長時間的訓練會耗費更多資源，但可能會進一步提高性能。

我們對FlowMo-Lo和FlowMo-Hi進行了批量大小為128、學習率為10-4的訓練。我們按照EDM2，每一步對編碼器和解碼器的MLP塊中的權重矩陣進行強制歸一化，以防止激活和權重矩陣爆炸。所有模型都在PyTorch中以bfloat16的精度進行訓練。我們使用Adam優化器，其值為 (β1，β2 )=(0.9,0.95)，因為我們注意到，可能由于bfloat16精度或變壓器序列長度較長，β2越高越不穩定。我們將編碼器的學習率設置為0，訓練步數為200,000 步。我們使用的是指數移動平均值，速率為0.9999。

分析

消融研究。我們對標記化性能進行了消融研究，分析了不同決策對rFID和其他指標的影響。實驗設置與我們主要的ImageNet-1K標記化實驗相同。

我們首先對表4中第 1A 階段的設計決策進行消融研究。我們對這些決策進行分析，并在適用時參考先前的擴散自動編碼器作品。在這些實驗中，我們使用了減小μP寬度的配置，以提高實驗效率。

圖片23.png

雙倍貼片尺寸。在視覺轉換器和去肋架構中，“補丁尺寸”決定了圖像轉換序列的長度，補丁尺寸越小，模型的計算量越大。之前關于像素空間中基于變換器的擴散模型的研究已經指出了小補丁尺寸的重要性。我們證實了這一點：增大補丁尺寸會影響所有指標。

使用MSE訓練的編碼器。之前的擴散自動編碼器研究已經探索了首先用MSE或基于LPIPS的回歸目標來訓練自動編碼器，也有可能用對抗損失來訓練，然后用訓練得到的凍結特征作為擴散解碼器的條件。雖然這種訓練方式可能會提高PSNR，但為了實現最佳的rFID，必須在所有噪聲水平下對整個系統進行端到端訓練，以確保潛碼包含有助于在所有噪聲水平下進行速度估計的特征。

無感知損失。之前在擴散自動編碼器方面的研究已經注意到感知損失對擴散解碼器1步去噪預測的重要性。如果沒有感知損失，性能就會下降。

FSQ量化。與我們的默認量化策略LFQ相比，我們發現FSQ的最終訓練損失和成對度量性能略勝一籌，但卻以最關鍵的rFID為代價。我們還嘗試了其他量化技術，但LFQ的性能最佳。

對數正態噪聲。用于整流訓練的logit-normal噪聲計劃將 t∈{0, 1} 的概率質量分配為0。雖然這種方法適用于t=0（此時整流目標變得奇異），但由于我們的調節信號非常強，因此在t=1（純噪聲）時對v(xt,t)的估計至關重要。使用通常的logit-normal時間表會導致PSNR和rFID下降。取而代之，我們使用厚尾 logit-normal噪聲計劃：我們從 [0, 1] 上的均勻分布中抽取10%的噪聲水平樣本。如果不做這種修改，t=1時的速度估計值就會不準確，從而導致低PSNR和變色。示例見圖 7。

圖片24.png

非移位采樣器。如第3.4節所述，我們使用時間步距t對整流ODE進行積分，時間步距按比例向t=0集中。使用線性間隔的時間步會降低性能。

無引導。與GAN標記化器相比，FlowMo的一個優勢是能夠利用無分類器引導和引導間隔等技術來提高感知質量。如果不在有限的時間間隔內應用分類器引導，rFID就會惡化。

表5是納入階段1B的消融研究。如果沒有這一關鍵階段，所有指標的性能都會下降。有趣的是，盡管只在Lsample上進行了訓練，但這一階段也提高了 PSNR。我們在圖8中說明了FlowMo解碼器在后期訓練后如何保持多模態。

圖片25.png

圖片26.png

局限性。FlowMo的主要局限在于推理時間。FlowMo需要多次模型前向傳遞（我們在工作中使用了n=25步）才能從給定量化潛碼的解碼器中生成一個樣本。相比之下，幾乎所有其他標記化器都是基于GAN的，因此只需一次前向傳遞即可重新構建給定圖像。盡管如此，將擴散模型提煉為單步或幾步系統的工作仍然相當多，這些技術也適用于FlowMo。

七、結論

我們介紹了基于變壓器的擴散自動編碼器FlowMo，它利用擴散自動編碼器的訓練和采樣應著眼于尋求感知高質量模式這一關鍵見解，實現了最先進的圖像標記化。我們在256×256分辨率的競爭性ImageNet-1K基準上展示了最先進的性能，與之前的許多工作不同，我們沒有使用二維空間對齊潛碼、對抗損失、來自輔助標記化器的代理目標，也沒有進行壓縮。