又一爆款！港大&字節(jié)開源GigaTok首次破解圖像重建與生成矛盾，30億參數(shù)登頂SOTA！
文章鏈接：https://arxiv.org/pdf/2504.08736
項目鏈接：https://silentview.github.io/GigaTok/
代碼鏈接：https://github.com/SilentView/GigaTok
亮點直擊
- 分詞器規(guī)模擴展中存在的重建與生成困境源于更大規(guī)模分詞器隱空間復雜度的提升。本文提出語義正則化方法，有效緩解該困境從而實現(xiàn)分詞器的規(guī)模擴展。
- 探索了分詞器規(guī)模擴展的最佳實踐方案，包括：采用混合CNN-Transformer架構的一維分詞器、非對稱編碼器-解碼器擴展策略，以及針對十億級分詞器的熵損失優(yōu)化。
- GigaTok成為首個規(guī)模擴展至30億參數(shù)的分詞器，在ImageNet數(shù)據(jù)集上同步實現(xiàn)了重建質量、下游自回歸生成和下游自回歸表征三項指標的突破性性能。
2.9 B GigaTok 在 ImageNet 256 × 256 分辨率上以 1.4 B AR 模型實現(xiàn)了 SOTA 自回歸圖像生成
總結速覽
解決的問題
- 視覺分詞器（visual tokenizer）縮放困境：現(xiàn)有方法在擴大視覺分詞器規(guī)模時，雖然提升了圖像重建質量，但導致下游自回歸生成質量下降（即“重建 vs. 生成”矛盾）。
- 隱空間復雜性增長：分詞器規(guī)模擴大時，隱空間復雜度增加，使得下游自回歸模型難以有效學習離散token的分布。
- 大規(guī)模分詞器訓練穩(wěn)定性：億級參數(shù)分詞器的訓練面臨收斂困難問題。
提出的方案
- 語義正則化（Semantic Regularization）：
a.利用預訓練視覺編碼器（如DINOv2）的特征對齊分詞器特征，約束隱空間復雜度，避免生成過于復雜的token依賴關系。
- 分詞器縮放關鍵策略：
a.1D分詞器：相比2D分詞器，1D結構更具可擴展性。
b.非對稱模型縮放：優(yōu)先擴展解碼器而非編碼器。
c.熵損失（Entropy Loss）：用于穩(wěn)定億級參數(shù)分詞器的訓練。
應用的技術
- 混合架構：基于CNN-Transformer的向量量化（VQ）分詞器，支持1D/2D token生成。
- 自回歸探測（AR Probing）：通過輕量級下游AR模型監(jiān)控分詞器訓練過程，分析token分布學習難度。
- 預訓練模型特征對齊：使用DINOv2等模型的語義特征作為正則化目標。
達到的效果
- 性能提升：
a.重建質量：達到SOTA圖像重建效果。
b.生成質量：下游自回歸生成任務性能顯著提升。
c.表征學習：下游AR模型的表征能力增強。
- 規(guī)模化能力：成功將GigaTok擴展至30億參數(shù)，突破“重建 vs. 生成”矛盾。
- 訓練穩(wěn)定性：通過熵損失實現(xiàn)億級參數(shù)分詞器的穩(wěn)定訓練。
先導研究
本文首先引入自回歸探測(AR Probing)作為有效監(jiān)控分詞器對下游生成效果的代理方法，隨后通過先導實驗研究當簡單擴展視覺分詞器時遇到的重建與生成挑戰(zhàn)。
用于分詞器評估的自回歸探測
在自回歸視覺生成中，分詞器和下游AR模型的訓練分兩個獨立階段進行。第一階段訓練視覺分詞器將圖像壓縮為離散token，以重建目標進行優(yōu)化。第二階段基于預訓練分詞器產生的離散token訓練下游生成模型。然而，在第一階段重建效果良好的分詞器，并不一定能帶來下游生成模型的更好表現(xiàn)。因此，在評估訓練好的分詞器時，除了重建質量外，評估其對下游生成的有效性同樣至關重要。
盡管很重要，但評估分詞器如何影響下游生成模型的計算成本可能很高。例如，充分訓練一個3.43億參數(shù)的下游AR生成器需要在64塊V100 GPU上耗時170小時。為解決這一挑戰(zhàn)，受表征學習中線性探測(Linear Probing)的啟發(fā)，引入了自回歸探測方法。其核心思想是使用小型AR模型的性能作為反映大規(guī)模AR模型性能趨勢的代理指標。
本文使用分詞器訓練一個小型Llama風格模型(1.11億參數(shù))50個周期，并評估其gFID、驗證損失和線性探測準確率，以便在不同分詞器之間進行公平比較。訓練所提出的自回歸探測模型來評估分詞器的效率比訓練原始3.43億下游AR模型高10倍。下圖6實驗表明，自回歸探測觀察到的趨勢與充分訓練后大規(guī)模AR模型的性能表現(xiàn)一致。
gFID。自回歸探測的生成FID表明兩階段框架的整體圖像生成性能。它既反映了分詞器的重建保真度，也反映了下游自回歸探測模型學習視覺token依賴關系的能力(即token分布的可學習性)。
驗證損失。本文使用自回歸探測模型的驗證損失作為潛在token可學習性的解耦衡量指標。驗證損失計算為ImageNet 50k驗證集上基于下一token預測范式中各token交叉熵損失的平均值。在相同詞匯量、相同視覺token數(shù)量和結構、相同自回歸探測模型的情況下，更大的驗證損失表明隱空間對AR模型來說更難學習。因此，我們用驗證損失來反映AR模型的隱空間復雜度和可學習性。
線性探測準確率。除了視覺生成質量外，本文還研究擴展分詞器是否會促使AR模型學習到更好的視覺表征，這可能為未來使用AR模型進行統(tǒng)一多模態(tài)理解和生成的研究提供啟發(fā)。為評估AR模型的表征質量，本文采用標準做法，使用自回歸探測模型中間Transformer層的特征評估線性探測準確率。
簡單擴展分詞器不可行
為研究簡單擴展視覺分詞器時面臨的挑戰(zhàn)，在ImageNet 256×256分辨率下訓練了三個規(guī)模遞增的向量量化分詞器。如下圖3所示，隨著分詞器規(guī)模增大，雖然重建質量(rFID)持續(xù)提升，但AR生成性能(gFID)顯著下降。這凸顯了分詞器擴展中的重建與生成困境。此外，我們觀察到自回歸探測的驗證損失隨分詞器規(guī)模擴大持續(xù)增加，表明更大分詞器會導致更復雜的token依賴關系，使AR模型更難學習。這一發(fā)現(xiàn)促使我們設計語義正則化方法來約束分詞器的隱空間復雜度，從而解決重建與生成困境。
GigaTok
本節(jié)介紹可擴展視覺分詞器GigaTok的模型結構和訓練策略。首先提出支持1D和2D token結構的分詞器主干，并討論編碼器-解碼器的非對稱擴展策略。然后引入語義正則化，通過用預訓練視覺表征約束隱空間復雜度來解決重建與生成困境。最后展示熵損失如何促進十億級規(guī)模分詞器的收斂。
架構
當前研究中，CNN架構因其捕捉細粒度局部細節(jié)的有效性成為圖像分詞器的主流選擇。而Transformer則是更具擴展性且歸納偏置更少的架構。因此，本文設計了結合CNN和Transformer的混合架構向量量化(VQ)分詞器主干(下圖4)。編碼器包含系列CNN塊(將輸入圖像逐步下采樣p倍)，后接Transformer層和向量量化器生成離散潛在編碼。解碼器包含多個Transformer層，后接CNN解碼器(將特征上采樣獲得重建圖像)。通過采用下文介紹的兩種Transformer設計，我們的分詞器架構可適配1D和2D分詞器。
2D分詞器(ViT實現(xiàn)) ：對2D分詞器，編解碼器的Transformer均采用ViT架構實現(xiàn)，全程保持潛在特征和token的2D結構。
1D分詞器(Q-Former實現(xiàn)) ：對1D分詞器，編解碼器的Transformer模塊均實現(xiàn)為Q-Former。編碼器的Q-Former使用1D查詢，將2D輸入特征轉為1D潛在token；解碼器的Q-Former使用2D查詢，將1D潛在token轉回2D特征后輸入CNN解碼器重建圖像。1D分詞器消除了2D歸納偏置，實驗表明其擴展性優(yōu)于2D分詞器。
語義正則化
在先導研究中，發(fā)現(xiàn)隨著分詞器規(guī)模的擴大，隱空間復雜度顯著增加，這可能導致更大規(guī)模的分詞器在下游AR生成任務中表現(xiàn)更差。假設，更大的分詞器傾向于捕捉過多細粒度的低級細節(jié)以獲得更好的重建效果，從而導致潛在token分布過于復雜，使得AR模型更難有效學習token依賴關系。
為解決這一問題，本文引入了語義正則化方法，引導分詞器編碼出語義更一致的隱空間——這種空間復雜度更低，更有利于下游生成建模。在分詞器訓練目標中加入了一個簡單的語義正則化項。該正則化將分詞器解碼器的中間特征與預訓練凍結DINOv2模型提取的特征表示對齊。
十億級分詞器的熵損失
在訓練 29 億參數(shù)的分詞器時，發(fā)現(xiàn)使用與 6.22 億參數(shù)分詞器相同的訓練方案會導致感知損失和重建損失無法收斂，并持續(xù)出現(xiàn)碼本使用率低下的問題。我們假設低碼本使用率是造成收斂困難的原因。為解決這個問題，本文引入熵懲罰來提升碼本利用率：
實驗
實驗設置
為擴展視覺分詞器，本文采用下表1總結的GigaTok分詞器Transformer架構配置。從三個維度評估分詞器：重建質量、下游AR生成質量和下游AR表征質量。使用rFID和LPIPS評估重建保真度，gFID評估生成性能，線性探測評估下游AR模型的表征質量。我們的下游AR模型采用具有1D絕對位置編碼的LlamaGen。擴展實驗和消融研究使用自回歸探測的驗證損失、gFID和線性探測分別反映token可學習性、生成性能和表征質量。系統(tǒng)級對比則訓練14億參數(shù)AR模型與之前工作進行對比。
自回歸探測的有效性。如前面圖6所示，自回歸探測的gFID和線性探測準確率等指標與更大規(guī)模的LlamaGen-XL模型結果一致。因此除系統(tǒng)級對比外，后續(xù)實驗均采用自回歸探測。
語義正則化的擴展效果
本文證明所提出的語義正則化解決了分詞器擴展中的重建與生成困境。
帶語義正則化的模型擴展。結果如下圖7所示：(1)語義正則化通過降低rFID改善了重建保真度；(2)更重要的是，無語義正則化時，自回歸探測驗證損失和gFID隨分詞器增大而惡化，顯示重建與生成困境。語義正則化通過相對受限的驗證損失和持續(xù)下降的gFID解決了該困境；(3)線性探測結果表明，語義正則化幫助AR模型在分詞器擴展時學習到更好的表征。
分詞器特征空間可視化。本文對一組圖像的首個Transformer解碼層特征進行前三主成分可視化。如下圖8所示，普通分詞器編碼的隱空間語義一致性有限，可能損害下游AR模型的可學習性；而GigaTok展現(xiàn)出語義一致的模式(圖8)，表明其隱空間具有意義和一致性。
非對稱1D分詞器更具可擴展性
解碼器需要更多參數(shù)。為確定擴展時應優(yōu)先擴展編碼器還是解碼器，在下表2中比較了S-B4和B-S分詞器（兩者在相同設置下訓練100個epoch）。結果表明，相比編碼器，擴展解碼器能帶來更大的重建和下游生成性能提升，說明應優(yōu)先擴展解碼器。
編碼器擴展同樣重要。雖然優(yōu)先擴展解碼器收益顯著，但發(fā)現(xiàn)擴展編碼器也能進一步提升下游模型性能。上表2顯示，B-L分詞器相比S-L分詞器獲得顯著提升。因此建議同時擴展編碼器和解碼器，但保持解碼器規(guī)模大于編碼器以獲得最優(yōu)性能。
1D分詞器比2D分詞器更具擴展性。在相同設置下訓練了S-S、S-B和B-L規(guī)模的1D/2D分詞器（均使用語義正則化）。如下圖9所示，1D分詞器始終獲得更好的rFID和自回歸探測線性分類準確率。在自回歸探測gFID指標上，1D分詞器展現(xiàn)出更陡峭的擴展趨勢，最終在模型規(guī)模擴大時超越2D分詞器。我們將1D分詞器的優(yōu)越擴展性歸因于其更低的歸納偏置。
系統(tǒng)級對比
實驗設置：使用GigaTok進行token化，在256×256 ImageNet訓練集上對LlamaGen AR模型進行300個epoch的訓練，與其他方法進行對比。未采用AdaLN方法，因其專為類別條件生成設計。為公平比較rFID，提供了使用DINO判別器訓練的B-L分詞器結果。但實際發(fā)現(xiàn)DINO判別器對LPIPS提升有限，且可能影響十億級分詞器的訓練穩(wěn)定性，故未將其納入主要設計。
實驗結果：如下表3所示，29億參數(shù)GigaTok在所有離散分詞器中實現(xiàn)了最優(yōu)的重建性能(rIFD)。此外，配合29億參數(shù)分詞器，下游14億參數(shù)AR模型在LLM風格的自回歸下一token預測模型中取得了最佳的圖像生成性能(gFID)。VAR采用下一尺度預測而非下一token預測，與語言模型兼容性較差。本文的模型僅使用簡單的LLM風格下游AR生成器（不含VAR等視覺專用設計）就達到了與VAR相當?shù)膅FID。更重要的是，基于29億分詞器訓練的14億AR模型通過視覺生成式預訓練獲得了最優(yōu)的線性探測準確率，表明GigaTok能幫助下游生成模型學習更好的表征。這種通過生成式預訓練獲得的高質量表征，可能有助于未來原生多模態(tài)模型統(tǒng)一生成與理解能力。
討論與消融研究
生成成本討論
在圖像生成過程中，AR模型需多次前向傳播預測token，而分詞器僅需單次前向傳播。因此，相較于AR模型，分詞器將token解碼為圖像的時間占比極小。表4記錄了不同分詞器/AR模型中分詞器解碼的時間占比：對于14億參數(shù)AR模型，我們最大的29億參數(shù)分詞器僅占總推理時間的約10%。
語義正則化最佳層搜索
探索了Transformer解碼器中用于計算語義正則化（公式1）的中間
預訓練語義編碼器選擇
本文比較了CLIP-B (DFN)、SigLIP-400和DINOv2-B作為S-B分詞器的語義正則化源。下表6顯示：
- 采用DINOv2-B作為語義編碼器時，分詞器在重建、下游類別條件生成和表征質量上均表現(xiàn)最佳
語義正則化權重探索
結論
本文研究并解決了視覺分詞器擴展過程中的重建與生成困境。發(fā)現(xiàn)該困境源于更大規(guī)模分詞器中不斷增加的隱空間復雜度。通過注入預訓練表征并與分詞器特征對齊，提出語義正則化方法來有效約束分詞器隱空間。這種語義正則化結合我們探索的若干關鍵實踐，最終實現(xiàn)了首個30億參數(shù)的分詞器GigaTok，在重建質量、下游AR生成性能和下游AR表征質量上均達到SOTA水平。
本文轉自AI生成未來 ，作者：AI生成未來
