李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer

2025年03月20日18:53:08 科技 1754

機器之心報道

機器之心編輯部

當我們看到一張貓咪照片時,大腦自然就能識別「這是一隻貓」。但對計算機來說,它看到的是一個巨大的數字矩陣 —— 假設是一張 1000×1000 像素的彩色圖片,實際上是一個包含 300 萬個數字的數據集(1000×1000×3 個顏色通道)。每個數字代表一個像素點的顏色深淺,從 0 到 255。

為了更加高效地從成千上萬張圖像中學習,AI 模型需要對圖片進行壓縮。比如當前最先進的圖像生成模型,第一步就是一個名叫 tokenization 的操作,用於執行此操作的組件叫 tokenizer。tokenizer 的主要目標是將原始圖像壓縮到一個更小、更易處理的潛在空間,使得生成模型能夠更高效地學習和生成。因此,如何得到更好的 tokenizer 是該領域的研究者非常關心的問題。

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

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

FlowMo 的訓練分為兩個階段:第一階段先學習如何全面捕捉圖像的多種可能重建結果,第二階段則學習如何從這些可能中選擇最接近原圖的重建方案。這種方法既保證了圖像重建的多樣性,又確保了重建質量,使得 FlowMo 在 ImageNet-1K 數據集上展現出了領先的重建性能。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

  • 論文標題:Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization
  • 論文地址:https://arxiv.org/pdf/2503.11056v1
  • 項目主頁:https://kylesargent.github.io/flowmo

研究背景

自從 VQGAN 和潛在擴散模型等視覺生成框架問世以來,最先進的圖像生成系統通常採用兩階段設計:先將視覺數據壓縮到低維潛在空間進行 tokenization,再學習生成模型。

Tokenizer 訓練一般遵循標準流程,即在均方誤差(MSE)、感知損失和對抗損失的組合約束下壓縮並重建圖像。擴散自編碼器曾被提出作為學習端到端感知導向圖像壓縮的方法,但在 ImageNet-1K 重建這一競爭性任務上尚未達到最先進水平。

李飛飛團隊提出了 FlowMo,一種基於 Transformer 的擴散自編碼器,它在多種壓縮率下實現了圖像 tokenization 的新性能標準,且無需使用卷積、對抗損失、空間對齊的二維潛在編碼,或從其他 tokenizer 中蒸餾知識(這與傳統的基於 GAN 的 tokenizer,如 VQGAN,非常不同)。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

研究的關鍵發現是 FlowMo 訓練應分為模式匹配預訓練階段和模式尋求後訓練階段。此外,研究者進行了廣泛分析,並探索了基於 FlowMo tokenizer 的生成模型訓練。

圖 1:無論是在低比特率訓練(FlowMo-Lo)還是高比特率訓練(FlowMo-Hi)下,FlowMo 模型都實現了最先進的 image tokenization 性能。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

作者強調,儘管基於 GAN 的 tokenizer 在圖像 tokenization 任務上已經取得了很好的性能,但 FlowMo 提供了一種簡單且不同的方法。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

FlowMo 方法

眾所周知,基於 Transformer 的擴散自編碼器包含編解碼結構,因此 FlowMo 也是由編碼器 e_θ 和解碼器 d_θ 組成,其核心架構遵循了 MMDiT,在 Transformer 的架構中學習一維潛在空間。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

圖 2:FlowMo 架構概覽

首先,編碼器將輸入圖像轉換為潛空間向量 c,然後解碼器則根據潛空間 c 學習重建圖像的條件分布,與旨在產生單一確定性輸出的傳統方法不同,FlowMo 的解碼器會生成可能的重建分布,從而更好地捕捉圖像重建中固有的模糊性。

FlowMo 架構的主要包括以下四點:

  • 基於 Transformer 的設計:編碼器和解碼器都使用 Transformer 架構,從而能夠更有效地處理圖像數據。
  • 一維潛空間表示:FlowMo 產生緊湊的潛在表示,使其適用於下游生成建模任務。
  • 量化層:編碼器的輸出被量化以創建離散 token,從而實現更高效的壓縮。
  • 擴散過程:解碼器使用擴散過程逐漸將隨機輸入去噪為高質量重建。

在 FlowMo 架構中的一個核心創新點是其兩階段訓練策略,這一策略使得重建分布偏向於原始圖像具有高度感知相似性的模式。

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

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

圖 3:FlowMo 的訓練過程結合了基於流的損失和感知損失,以引導模型實現高質量的重建。

在階段 1A 中,FlowMo 通過聯合訓練編碼器與解碼器,以實現兩個核心目標:最大化潛在編碼的信息量,並使其重建分布與真實分布相匹配。這一訓練過程巧妙地結合了多種損失函數,展現出其獨特的技術優勢:

  • 修正流損失(Rectified flow loss):引導擴散過程向目標圖像分布靠攏,確保生成結果的準確性;
  • 感知損失(Perceptual loss):保證了重建圖像在視覺上與原始圖像高度相似;
  • 熵損失(Entropy loss):鼓勵生成多樣化的潛在編碼,避免模式單一化;
  • 承諾損失(Commitment loss):使得編碼器輸出與量化表示儘可能接近,進一步優化了模型的穩定性與效率。

具體而言,FlowMo 作為擴散自動編碼器進行端到端訓練,以優化解碼器輸出上的修正流損失 L_flow,在過程中使用了 L_perc 來監督圖像生成中的去噪預測,同時在潛空間 c 上,作者還結合了 LFQ 的熵損失和承諾損失來進行訓練。其中損失函數的數學表達式如下所示:

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

結合這些損失函數,並最終得到了第一階段的損失表達式:

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

階段 1B:模式探索後訓練

在第二階段中,FlowMo 的核心目標是優化解碼器分布 pθ(x∣c),以尋找那些在感知上與原始圖像高度相似的模式。為實現這一目標,FlowMo 採用了創新的訓練策略:首先凍結編碼器,隨後在 Lflow 的基礎上,聯合訓練解碼器,並引入受擴散模型訓練後的 x_0 來生成目標 Lsample。這一過程通過以下步驟實現:

  • 概率流 ODE:通過少量步驟的概率流常微分方程(ODE)集成;
  • 感知損失計算:在生成樣本後,模型會計算其與原始圖像之間的感知損失,確保重建結果在視覺上與原始圖像保持一致;
  • 解碼器參數更新:基於感知損失,FlowMo 對解碼器參數進行優化。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

圖 4:模式搜索訓練過程,編碼器處於凍結狀態(雪花表示),而解碼器則進行微調以提高感知質量。

如上圖所示,其中 FlowMo 通過凍結編碼器,集中精力優化解碼器,使其在重建圖像時更加註重感知相似性,從而進一步提升生成圖像的質量與真實感。對概率流 ODE 進行積分的 n 步樣本感知損失 Lsample 如下所示:

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

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

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

採樣過程

為了生成重構圖像,FlowMo 通過求解概率流 ODE,對給定一維潛空間 c 的重建圖像的多模態分布進行採樣

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

FlowMo 採樣方法的一項關鍵創新是使用「移位」採樣器。FlowMo 不使用統一的時間步長間隔,而是採用可調的移位超參數,將採樣步驟集中在擴散過程的某些區域,從而提高感知質量。

採樣過程需要多次前向通過解碼器模型,這在計算上很昂貴,但可以產生高質量的結果。

實驗結果分析

主要結果

FlowMo 在多個比特率設置下(0.07 BPP 和 0.22 BPP)與當前最先進的 tokenizer 進行了比較,在重建 FID(rFID)、PSNR 和 SSIM 指標上均取得了最佳結果。在 0.07 BPP 設置下,FlowMo-Lo 的 rFID 為 0.95,相比 OpenMagViT-V2 的 1.17 有顯著提升;在 0.22 BPP 設置下,FlowMo-Hi 的 rFID 為 0.56,略優於 LlamaGen-32 的 0.59。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

表 1. tokenization 結果。

消融實驗分析

研究團隊進行了大量消融實驗,分析了 FlowMo 設計中的關鍵決策:雜訊調度、量化策略、模型架構和後訓練策略等。結果表明,thick-tailed logit-normal 雜訊分布、shifted sampler 和後訓練階段對模型性能至關重要。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

圖 5:雜訊調度導致失真的可視化案例。

特別是,沒有模式尋求後訓練階段,FlowMo-Lo 的 rFID 會從 0.95 下降到 1.10,FlowMo-Hi 的 rFID 會從 0.56 下降到 0.73。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

表 2:後訓練消融實驗結果。

生成任務驗證

在生成任務中,基於 FlowMo 訓練的 MaskGiT 在某些指標上表現優於基於 OpenMagViT-V2 訓練的模型,但在 FID 上略遜一籌(4.30 vs 3.73)。這表明 tokenizer 質量與下游生成模型質量之間存在複雜關係,需要進一步研究。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

表 3:生成模型指標對比。

李飛飛、吳佳俊團隊新作:不需要卷積和GAN,更好的圖像tokenizer - 天天要聞

圖 6:生成圖像對比。

科技分類資訊推薦

從陪跑個體到企業培訓,我的IP陪跑之路,越走越寬了 - 天天要聞

從陪跑個體到企業培訓,我的IP陪跑之路,越走越寬了

大家好,我是Tina。來繼續通過文章,分享我的自媒體創業生涯。來說說最近在乾的事兒。一今天給江南布衣的全國經銷商做了小紅書的業務輔導培訓。很難想像6年的時間,我從一名職場人,慢慢成長為一個自媒體人,然後成為超級個體,到最後一步步做到可以給企
vivo Y300 GT續航超耐用:內置7620mAh電池 還有直驅供電 - 天天要聞

vivo Y300 GT續航超耐用:內置7620mAh電池 還有直驅供電

【TechWeb】去年底以來,vivo推出了vivo Y300系列的多款機型,其中vivo Y300 Pro首發搭載了年度最大的6500mAh超薄藍海電池,打破了藍海電池容量紀錄,也刷新了vivo電池容量新高。而在近期,該系列的又一款新機——vivo Y300 GT也得到官宣並開啟預約,將在5月9日也就是今天10:00正式開售。現在有最新消息,近日官方進一步
人才需求達百萬級!人工智慧如何催生職業新賽道 - 天天要聞

人才需求達百萬級!人工智慧如何催生職業新賽道

人工智慧的應用場景,正像星火燎原般迅速鋪展延伸到各行各業,一系列充滿科技張力和未來想像的新興職業也應運而生,成為年輕一代心馳神往的職業新選擇。在內蒙古包頭市包鋼白雲鄂博鐵礦,電鏟式挖掘機的巨型鏟斗,一次可以鏟起10立方米的礦石。然而這個「巨無霸」的駕駛室里,卻空無一人。操作這輛挖掘機的工人正坐在距離采...
即時零售,美團、京東、淘寶們的新戰場 - 天天要聞

即時零售,美團、京東、淘寶們的新戰場

當京東殺入到外賣市場的時候,很多人僅僅只是關注的是「京東開始做外賣」這個點,但卻並未真正了解「京東開始做外賣」背後延伸開來的新邏輯。 隨著淘寶的加入,特別是隨著越來越多的解讀開始出現,人們開始發現,京東做外賣,並不僅僅只是做外賣,而是為了在即時零售的新戰場上佔據一席之地。 於是,「即時零售」的概念,被...
潮聲丨今天的機器人賽場,或許正是未來社會的預演 - 天天要聞

潮聲丨今天的機器人賽場,或許正是未來社會的預演

潮新聞 執筆 薛文春北京「半馬」才結束,世界人形機器人運動會又來了!5月7日,北京宣布全球首個為人形機器人組織的綜合性競技賽事——世界人形機器人運動會,將於8月15日在國家體育場(鳥巢)和國家速滑館(冰絲帶)舉辦。不得不說,最近這段時間,人
聯想拯救者十年硬核進化,Y9000P 2025成六邊形性能戰神 - 天天要聞

聯想拯救者十年硬核進化,Y9000P 2025成六邊形性能戰神

5月8日,以「熱AI自發光」為主題的聯想天禧AI生態春季新品超能之夜在上海璀璨啟幕。發布會以一場充滿科技感與年輕活力的「熱AI時尚秀」,演繹了AI熱力和青春活力的雙向奔赴。聯想重磅發布了天禧和想幫幫兩款智能體,面向年輕客戶發布全場景AI終端,更為大學生群體發布了量身定製的「青春有AI」教育特惠計劃。聯想集團高級副...