自VQGAN和Latent Diffusion Models等視覺生成框架問世以來,先進的圖像生成系統(tǒng)通常采用兩階段架構(gòu):首先將視覺數(shù)據(jù)Token化或壓縮至低維潛在空間,隨后學(xué)習(xí)生成模型。傳統(tǒng)Token化器訓(xùn)練遵循標準范式,通過MSE、感知損失和對抗性損失的組合約束來實現(xiàn)圖像壓縮與重建。雖然擴散自編碼器曾被提出作為端到端感知導(dǎo)向圖像壓縮的學(xué)習(xí)方法,但在ImageNet1K重建等競爭性任務(wù)上尚未展現(xiàn)出優(yōu)越性能。

這個研究提出了FlowMo,一種基于Transformer的擴散自編碼器,在多種比特率條件下實現(xiàn)了圖像Token化的最新技術(shù)水平。其顯著特點在于無需依賴卷積網(wǎng)絡(luò)、對抗性損失、空間對齊的二維潛在編碼或從其他Token化器中提取表征。FlowMo的關(guān)鍵技術(shù)創(chuàng)新在于其訓(xùn)練流程被劃分為模式匹配預(yù)訓(xùn)練階段和模式搜索后訓(xùn)練階段。通過廣泛的分析與消融實驗,驗證了該方法的有效性,并在FlowMo Token化器基礎(chǔ)上訓(xùn)練了生成模型,進一步確認了其性能優(yōu)勢。

打開網(wǎng)易新聞 查看精彩圖片

FlowMo在低比特率(FlowMo-Lo)和高比特率(FlowMo-Hi)條件下均達到了圖像Token化的領(lǐng)先性能。作為一種基于Transformer的擴散自編碼器,F(xiàn)lowMo摒棄了卷積層、對抗性損失以及來自輔助Token化器的代理目標,顯示出其獨特的技術(shù)優(yōu)勢。

圖像Token化技術(shù)概述

在深入探討FlowMo的技術(shù)創(chuàng)新前,有必要了解當(dāng)前圖像生成領(lǐng)域的主流范式?,F(xiàn)代高性能圖像生成系統(tǒng)通常采用二階段方法:

  1. Token化階段:將高維像素數(shù)據(jù)壓縮至低維潛在空間,實現(xiàn)數(shù)據(jù)降維和特征提取
  2. 生成階段:在壓縮表征基礎(chǔ)上訓(xùn)練生成模型,實現(xiàn)圖像合成

自VQGAN提出以來,Token化器通常構(gòu)建為卷積自編碼器,將視覺數(shù)據(jù)下采樣為空間對齊的二維潛在編碼。這類模型通常綜合應(yīng)用重建損失、感知損失和對抗性損失進行訓(xùn)練,以確保壓縮表征保持高視覺保真度。然而這種方法存在若干技術(shù)局限:對抗性損失的訓(xùn)練不穩(wěn)定性和調(diào)優(yōu)難度、對卷積神經(jīng)網(wǎng)絡(luò)的過度依賴、對空間對齊二維潛在編碼的要求,以及可能依賴于預(yù)訓(xùn)練Token化器的知識提取。

FlowMo:圖像Token化的技術(shù)突破

FlowMo(Flow towards Modes)代表了圖像Token化領(lǐng)域的重大技術(shù)突破,其方法論在多個關(guān)鍵維度上突破了傳統(tǒng)Token化器的限制:

基于擴散的解碼機制:FlowMo摒棄對抗性損失,轉(zhuǎn)而采用基于擴散的解碼器,該解碼器能自然地對多模態(tài)分布進行精確建模。

純Transformer架構(gòu):模型的編碼和解碼組件均采用完全基于Transformer的架構(gòu),完全避免了對CNN的依賴。

一維潛在表征:與大多數(shù)采用空間對齊二維潛在編碼的Token化器不同,F(xiàn)lowMo將圖像編碼為一維潛在序列。

端到端優(yōu)化范式:FlowMo采用直接優(yōu)化策略,無需從預(yù)訓(xùn)練Token化器中提取知識。

架構(gòu)設(shè)計與實現(xiàn)

打開網(wǎng)易新聞 查看精彩圖片

FlowMo的架構(gòu)設(shè)計圖:作為擴散自編碼器,F(xiàn)lowMo將輸入圖像x編碼為潛在表征?,隨后量化為c。模型進一步解碼以c為條件的修正流速度v,同時結(jié)合噪聲水平t和噪聲圖像xt。

打開網(wǎng)易新聞 查看精彩圖片

FlowMo采用簡潔而高效的架構(gòu)設(shè)計,主要包含以下核心組件:

  1. 編碼器(eθ):一個專用Transformer模型,負責(zé)將分塊化(Patchified)圖像編碼為潛在Token序列
  2. 量化層:采用無查找表量化(Lookup-Free Quantization, LFQ)技術(shù)將連續(xù)潛在特征二值化
  3. 解碼器(dθ):基于Transformer的擴散模型,用于從量化潛在變量重建原始圖像

編碼器和解碼器均基于多模態(tài)擴散圖像Transformer(MMDiT)架構(gòu)構(gòu)建,其中解碼器在參數(shù)規(guī)模和網(wǎng)絡(luò)深度上均大于編碼器。該模型將圖像和潛在編碼作為獨立數(shù)據(jù)流處理,這些數(shù)據(jù)流通過自注意力機制實現(xiàn)高效交互。完整的處理流程如下:

x (輸入圖像) → Patchify → eθ (編碼器) → ? → 量化 → c → dθ (解碼器) → 重建圖像

模式搜索:FlowMo的核心技術(shù)創(chuàng)新

該研究的核心發(fā)現(xiàn)是:對于高質(zhì)量感知重建,有選擇地采樣感知上接近原始圖像的重建分布模式,優(yōu)于嘗試匹配所有可能的分布模式。基于這一關(guān)鍵見解,F(xiàn)lowMo采用創(chuàng)新性的兩階段訓(xùn)練方法:

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

打開網(wǎng)易新聞 查看精彩圖片

階段1A:編碼器和解碼器通過輸出損失函數(shù)Lperc、Lflow和潛在損失函數(shù)Lcommit、Lent進行端到端聯(lián)合訓(xùn)練。

打開網(wǎng)易新聞 查看精彩圖片

在此階段,編碼器和解碼器進行端到端訓(xùn)練,優(yōu)化多個技術(shù)目標:

  1. 流匹配目標(Lflow):訓(xùn)練解碼器精確建模修正流常微分方程的速度場
  2. 感知損失(Lperc):保證原始圖像與重建圖像間的高感知相似性
  3. 熵與承諾損失:對量化潛在空間進行正則化處理

通過這一階段訓(xùn)練,模型能夠準確匹配給定潛在編碼中有限信息所對應(yīng)的多模態(tài)重建分布。

階段1B:模式搜索后訓(xùn)練

打開網(wǎng)易新聞 查看精彩圖片

階段1B:凍結(jié)的編碼器eθ將輸入圖像編碼為c,用于調(diào)節(jié)解碼器dθ,解碼器通過整個采樣鏈的反向傳播進行訓(xùn)練優(yōu)化。同時維持Lflow損失的共同訓(xùn)練。

打開網(wǎng)易新聞 查看精彩圖片

這項關(guān)鍵技術(shù)創(chuàng)新使FlowMo超越了傳統(tǒng)擴散自編碼器。在此階段:

  1. 編碼器參數(shù)被凍結(jié)固定
  2. 解碼器進行微調(diào),以識別并偏好感知上與原始圖像高度相似的分布模式
  3. 訓(xùn)練采用創(chuàng)新的目標函數(shù)(Lsample),該函數(shù)通過整個采樣鏈實現(xiàn)端到端反向傳播

該階段有效地指導(dǎo)模型將其重建分布偏向于能夠保留原始圖像關(guān)鍵感知細節(jié)的高質(zhì)量模式。

高效采樣策略設(shè)計

FlowMo引入了優(yōu)化的采樣策略,用于在推理階段高效求解概率流常微分方程。其關(guān)鍵技術(shù)創(chuàng)新在于引入"shift"超參數(shù)(ρ),該參數(shù)將采樣步驟集中于較低噪聲水平,從而同時提升重建質(zhì)量和峰值信噪比(PSNR)。這一方法顯著優(yōu)于傳統(tǒng)的在速率-失真-感知曲線上簡單權(quán)衡各指標的方法。通過智能分配計算資源,F(xiàn)lowMo在不犧牲數(shù)據(jù)保真度的前提下,實現(xiàn)了感知質(zhì)量的顯著提升。

實驗結(jié)果與技術(shù)評估

實驗數(shù)據(jù)清晰展示了FlowMo的技術(shù)優(yōu)勢。在多個壓縮率下,F(xiàn)lowMo均實現(xiàn)了圖像Token化的最先進性能:

低每像素比特數(shù)條件(BPP = 0.07)

  • FlowMo-Lo: 0.95 rFID, 22.07 PSNR, 0.649 SSIM
  • OpenMagViT-V2: 1.17 rFID, 21.63 PSNR, 0.640 SSIM**

高每像素比特數(shù)條件(BPP = 0.22)

  • FlowMo-Hi: 0.56 rFID, 24.93 PSNR, 0.785 SSIM
  • LlamaGen-32: 0.59 rFID, 24.44 PSNR, 0.768 SSIM

打開網(wǎng)易新聞 查看精彩圖片

Token化性能對比結(jié)果。橫線分隔不同BPP下訓(xùn)練的各Token化器。與現(xiàn)有最先進Token化器相比,F(xiàn)lowMo在多個比特率條件下均實現(xiàn)了領(lǐng)先性能。?標記結(jié)果來自原始研究論文。

值得特別注意的是,F(xiàn)lowMo在不依賴以下常規(guī)技術(shù)的情況下達到了這些優(yōu)異結(jié)果:

  • 卷積神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)
  • 對抗性損失函數(shù)
  • 空間對齊的二維潛在編碼
  • 預(yù)訓(xùn)練Token化器的知識提取

視覺質(zhì)量評估

FlowMo重建的視覺質(zhì)量同樣令人矚目。在面部和文本等傳統(tǒng)上對圖像壓縮具有高挑戰(zhàn)性的內(nèi)容重建方面,F(xiàn)lowMo較競爭方法表現(xiàn)出更佳的細節(jié)保留能力。這種優(yōu)勢在面部特征精細度、文本清晰度和整體感知質(zhì)量等關(guān)鍵方面尤為明顯。值得注意的是,即使在后訓(xùn)練階段完成后,F(xiàn)lowMo的重建結(jié)果仍保持適度的多模態(tài)特性,但模型智能地將變異分配至感知上不太關(guān)鍵的區(qū)域(如背景),而非關(guān)鍵前景細節(jié)。

打開網(wǎng)易新聞 查看精彩圖片

生成圖像示例:使用不同Token化器訓(xùn)練的MaskGiT模型所生成的樣例。結(jié)果表明FlowMo可用于訓(xùn)練高質(zhì)量的第二階段生成模型。為便于比較,展示樣例采用相同類別索引。

消融研究與技術(shù)洞察

本研究包含廣泛的消融實驗,揭示了多項關(guān)鍵設(shè)計選擇的技術(shù)意義:

  1. 圖像塊大小(Patch Size):較小的圖像塊尺寸(導(dǎo)致序列長度增加)顯著提升了重建質(zhì)量,表明高分辨率表征對于Token化至關(guān)重要。
  2. 端到端訓(xùn)練策略:使用擴散解碼器從頭訓(xùn)練編碼器的方法,優(yōu)于先使用均方誤差(MSE)目標進行預(yù)訓(xùn)練的傳統(tǒng)方法。
  3. 感知損失的應(yīng)用:對單步去噪預(yù)測應(yīng)用感知損失是獲得高感知質(zhì)量的關(guān)鍵因素,提供了比像素級損失更有意義的優(yōu)化信號。
  4. 噪聲調(diào)度設(shè)計:研究提出了改進的粗尾logit-normal噪聲調(diào)度,包括在t=1處的采樣點,有效防止了圖像變色問題。
  5. 后訓(xùn)練階段的重要性:模式搜索后訓(xùn)練階段證明是不可或缺的,簡單替代方案(如增加感知損失權(quán)重)均未能達到相同效果。

技術(shù)局限與未來研究方向

FlowMo的主要技術(shù)局限在于推理計算開銷。由于采用基于擴散的解碼器,圖像重建需要多次模型前向傳遞(研究實驗中為25步)。這與基于GAN的Token化器能在單次前向傳遞中完成重建形成鮮明對比。這一限制可通過模型蒸餾或一致性模型等技術(shù)手段加以改進,這些方法已在擴散模型加速領(lǐng)域得到驗證。

總結(jié)

FlowMo代表了圖像Token化領(lǐng)域的重大技術(shù)進步,挑戰(zhàn)了關(guān)于實現(xiàn)最先進性能所需技術(shù)組件的傳統(tǒng)認知。通過充分利用擴散模型自然建模多模態(tài)分布的能力,并引入創(chuàng)新的模式搜索訓(xùn)練方法,F(xiàn)lowMo在簡化Token化架構(gòu)的同時實現(xiàn)了卓越的重建質(zhì)量。這一成功為多個研究方向提供了啟示:

  • 將模式搜索技術(shù)擴展應(yīng)用于視頻或三維數(shù)據(jù)等其他媒體形式
  • 探索更高效的加速方法以降低推理計算成本
  • 研究與大規(guī)模Transformer模型的整合可能性

FlowMo的研究成果表明,當(dāng)經(jīng)過專門訓(xùn)練以定位高質(zhì)量模式時,擴散自編碼器能夠顯著超越傳統(tǒng)圖像Token化方法,為各類應(yīng)用場景中的高質(zhì)量視覺內(nèi)容生成開辟了新途徑。

參考文獻:https://avoid.overfit.cn/post/baf2f21e78e9457eb3bfe25d8009c012