論文鏈接：https://arxiv.org/pdf/2502.03444
項目鏈接：https://github.com/Hhhhhhao/continuous_tokenizer

亮點直擊

• 理論與實驗分析：通過實驗和理論分析建立了隱空間結構與擴散模型性能之間的聯系。揭示了具有更少高斯混合模型（GMM）模式的結構化潛空間能夠更有效地訓練和生成擴散模型。
• MAETok：使用mask建模訓練普通自編碼器（AE），并證明了具有更具判別性的潛空間的簡單 AE 能夠加速學習、提高生成質量，并提升擴散模型的吞吐量，從而表明 VAE 的變分正則化并非必要。
• SOTA 生成性能：在 MAETok 訓練的擴散模型（675M 參數，128 token）在 256×256 ImageNet 生成任務上達到了與先前最佳模型相當的性能，并在 512 分辨率上超越了 2B USiT，取得了 1.69 gFID 和 304.2 IS。

結果展示

總結速覽

解決的問題

現有的擴散模型通常使用變分自編碼器（VAE）作為tokenizer，但VAE的變分約束可能限制了重建質量，影響生成效果。此外，現有的自動編碼器（AE）雖然能提供更高保真度的重建，但其隱空間可能過于混亂，不利于擴散模型的訓練和采樣。因此，如何構建一個既能保持高重建質量，又能提供良好結構化隱空間的tokenizer仍然是一個未解決的問題。

提出的方案

本文提出MAETok，一種基于 mask 建模（Masked Autoencoder, MAE）的自編碼器（AE）。它通過隨機屏蔽圖像token并重建缺失部分，以學習更加語義豐富、結構更清晰的隱空間。核心設計包括：

• mask建模：在編碼器端隨機屏蔽一部分圖像token，要求解碼器恢復原始特征，以迫使AE學習更具判別性的特征表示。
• 輔助淺層解碼器：在標準像素解碼器之外，增加輔助解碼器，用于預測latent token的特征（如HOG、DINOv2、CLIP等），增強隱空間表示的語義信息。
• 編碼器凍結策略：在高 mask 比（40%-60%）訓練后，固定編碼器并微調解碼器，以恢復像素級重建質量，同時保留高質量的隱空間表示。

應用的技術

• 擴散模型（Diffusion Models）
• Mask自編碼器（MAE）進行特征學習
• 高效Transformer架構用于tokenizer設計
• 高 mask 比訓練策略
• 隱空間分析（使用高斯混合模型GMM）優化擴散模型的學習

達到的效果

• 更優的隱空間：隱空間特征更加判別清晰，且具有更少的GMM模式，提升擴散模型的生成質量。
• 高效訓練與推理：相比基于VAE的擴散模型，訓練速度提升76倍，推理吞吐量提升31倍（512×512圖像）。
• 優越的生成質量：在ImageNet 256×256 和 512×512 基準上，MAETok 使用僅128個token即可實現最先進（SOTA）的生成質量，gFID 達到 1.69。
• 突破VAE限制：證明高質量的隱空間比變分約束（如VAE）更關鍵，實現無需VAE的擴散模型訓練。

方法

VAE 的變分形式對于擴散模型可能并非必要，簡單的 AE 便足以在 128 token的條件下實現 SOTA 生成性能，只要它們具有判別性的潛空間，即具有更少的 GMM 模式。本文將該方法稱為 MAETok，其詳細信息如下。

結構

在最近的 1D 分詞器 設計基礎上構建 MAETok，該分詞器采用可學習的latent tokens。編碼器 (E) 和解碼器 (D) 都采用 Vision Transformer (ViT) 結構，但經過調整以同時處理 圖像tokens 和 latent tokens，如下圖 3 所示。

Mask 建模

像素解碼器微調

雖然 Mask 建模促使編碼器學習更好的隱空間，但較高的 Mask 比例可能會降低即時重建效果。為了解決這一問題，在使用 Mask 建模訓練自編碼器之后，凍結編碼器，從而保留隱空間表示，并僅對像素解碼器進行少量額外訓練。這一過程使解碼器能夠更緊密地適應凍結的干凈圖像隱空間編碼，恢復在 Mask 訓練過程中丟失的細節。在像素解碼器微調過程中使用與公式 (6) 相同的損失，并在此階段丟棄所有輔助解碼器。

實驗

本文進行全面的實驗來驗證 MAETok 的設計選擇，分析其隱空間，并基準測試生成性能，展示其優越性。

實驗設置

擴散模型的實現細節
在訓練 MAETok 之后，使用 SiT和 LightningDiT來執行基于擴散的圖像生成任務。將 Patch 大小設置為 1，并使用 1D 位置 embedding，其他參數遵循它們的原始訓練設置。使用 458M 參數的 SiT-L 進行分析和消融研究。對于主要結果，訓練了 675M 參數的 SiT-XL，經過 4M 步，和 LightningDiT 訓練 400K 步，在分辨率為 256 和 512 的 ImageNet 上進行訓練。

評估  

對于 Tokenizer 評估，報告了 ImageNet 和 MS-COCO驗證集上的重建 Fréchet Inception Distance（rFID）、峰值信噪比（PSNR）和結構相似性指數（SSIM）。對于 Tokenizer 的潛在空間評估，我們對平坦的潛在表示進行線性探測（LP），并報告其準確性。為了評估生成任務的性能，報告了生成 FID（gFID）、Inception Score（IS）（Salimans et al., 2016）、精度和召回率（Precision and Recall），有無無分類器引導（CFG），使用 250 次推理步驟。

MAETok的設計選擇

首先進行了一項廣泛的消融研究，以了解Mask建模和不同設計如何影響tokenizer的重建效果，更重要的是，如何影響擴散模型的生成效果。從一個AE開始，逐步添加不同的組件，研究AE的rFID和SiT-L的gFID。

Mask建模
在下表1a中，我們比較了AE和VAE的Mask建模效果，并研究了像素解碼器微調的提議。對于AE，Mask建模顯著提高了gFID，并略微降低了rFID，這可以通過解碼器微調階段恢復，而不會犧牲生成性能。相反，Mask建模對VAE的gFID僅有輕微改善，因為施加的KL約束可能會妨礙潛在空間的學習。

重建目標
在上表1b中，我們研究了不同的重建目標如何影響Mask建模中的潛在空間學習。我們展示了使用低級重建特征，如原始像素（僅使用像素解碼器）和HOG特征，已經能夠學習到更好的潛在空間，從而實現更低的gFID。采用語義教師，如DINO-v2和CLIP，則可以顯著提高gFID。結合不同的重建目標可以在重建保真度和生成質量之間實現平衡。

Mask比例
在上表1c中，我們展示了合適的Mask比例對于使用HOG目標學習潛在空間的重要性，正如前人工作中所強調的。較低的Mask比例會阻止AE學習到更具區分性的潛在空間。較高的Mask比例則在重建保真度和潛在空間質量之間存在權衡，從而影響生成性能。

輔助解碼器深度
在上表1d中，我們研究了具有多重重建目標的輔助解碼器的深度。我們發現，解碼器過淺或過深都會損害重建保真度和生成質量。當解碼器過淺時，組合目標特征可能會混淆潛在空間中的高級語義和低級細節，從而導致更差的重建保真度。然而，過深的輔助解碼器可能會憑借其強大的能力學習到AE的潛在空間中較少的區分特征，從而也導致更差的生成性能。

潛在空間分析

本文進一步分析了AE變體的潛在空間與SiT-L生成性能之間的關系。

潛在空間可視化
提供了UMAP可視化如下圖4所示，以直觀地比較不同AE變體所學習到的潛在空間。值得注意的是，AE和VAE展示了更為糾纏的潛在 embedding，其中不同類別的樣本往往有較大的重疊。MAETok展示了明顯分離的聚類，并且類別之間的邊界較為清晰，這表明MAETok學習到了更具區分性的潛在表示。與下圖2中的分析一致，MAETok的潛在表示更加區分且分離，導致了更少的GMM模式，并提升了生成性能。

潛在分布與生成性能
我們通過研究潛在空間上的線性探測（LP）準確度，作為潛在代碼中語義信息保留程度的代理，并與生成性能的gFID進行關聯，來評估潛在空間的質量。在圖5a中，我們觀察到，潛在分布更具區分性的tokenizer，表現在更高的LP準確度，相應地也能實現更低的gFID。這個發現表明，當特征在潛在空間中良好聚類時，生成器能夠更容易學習生成高保真樣本。我們進一步通過跟蹤訓練過程中gFID的變化來驗證這一直覺，如下圖5b所示，MAETok使得訓練更快收斂，gFID迅速下降，且低于AE或VAE基準。高質量的潛在分布被證明是實現強大的最終生成指標并加速訓練的關鍵因素。

主要結果

生成  

在下表2和表3中，分別基于256×256和512×512的ImageNet基準，比較了SiT-XL和LightningDiT的MAETok變體，并與其他SOTA生成模型進行了比較。值得注意的是，使用僅128個token和普通AE架構訓練的SiT-XL，在沒有使用CFG的情況下，始終能取得更好的gFID和IS：它在256分辨率下比REPA提高了3.59的gFID，并在512分辨率下達到了2.79的SOTA可比gFID。當使用CFG時，SiT-XL與基于VAEs的競爭性自回歸和擴散基準在256分辨率下的表現相當。它擊敗了使用256個token的2B USiT，并且在512分辨率下實現了新的SOTA，gFID為1.69，IS為304.2。使用更多先進技巧訓練的LightningDiT表現更好，它在沒有CFG的情況下超越了1B參數的MAR-H和2B參數的USiT，gFID為2.56，IS為224.5，并且在使用CFG時，gFID為1.72。這些結果表明，潛在空間的結構（見上圖4），而非tokenizer的變分形式，對于擴散模型的有效學習至關重要。

重建  

MAETok在ImageNet和MS-COCO上的重建能力也非常強，如下表4所示。與之前的連續tokenizer（包括SD-VAE、DC-AE、VA-VAE、SoftVQ-VAE和TexTok）相比，MAETok在重建質量和潛在空間大小之間實現了一個良好的平衡。在256×256的ImageNet上，使用128個token，MAETok的rFID為0.48，SSIM為0.763，超越了SoftVQ等方法，在保真度和感知相似度方面均表現優越，同時使用了TexTok一半的token數。在MS-COCO上，盡管tokenizer沒有直接訓練，MAETok仍能提供強大的重建能力。在512分辨率下，MAETok通過平衡壓縮比和重建質量，保持了其優勢。

討論

高效的訓練與生成  

1D tokenizer設計的一個顯著優勢是，它能夠支持任意數量的潛在token。通常256×256和512×512圖像被編碼為256和1024個token，而MAETok在兩種情況下都使用128個token。這大大提高了擴散模型的訓練和推理效率。例如，在使用512×512圖像的1024個token時，SiT-XL的Gflops和推理吞吐量分別為373.3和每秒0.1張圖像。而MAETok將Gflops降至48.5，將吞吐量提高到每秒3.12張圖像。通過改進的收斂性，MAETok使得訓練速度比REPA快了76倍，表現相當。

無條件生成  

從我們的結果中得到的一個有趣觀察是，與以前的方法相比，使用MAETok訓練的擴散模型在沒有CFG的情況下通常展現出顯著更好的生成性能，同時在使用CFG時性能差距較小。我們推測原因是，無條件類別也學習了潛在空間中的語義，如表5中的無條件生成性能所示。隨著潛在空間變得更加區分性，無條件生成性能也顯著提高。這意味著CFG線性組合方案可能變得不那么有效，這與我們在附錄C.2中提供的CFG調優結果一致。

結論

本文對擴散模型的隱空間特性進行了理論和實證分析，證明了隱空間分布中較少的模式有助于更有效的學習和更好的生成質量。基于這些見解，開發了MAETok，它通過mask建模實現了最先進的性能，而無需使用變分約束。僅使用128個token，本文的方法顯著提高了ImageNet上的計算效率和生成質量。研究結果表明，具有更強區分性的隱空間，而非變分約束，是有效的擴散模型的關鍵，為大規模高效生成建模開辟了新的方向。

本文轉自AI生成未來 ，作者：AI生成未來

原文鏈接:??https://mp.weixin.qq.com/s/ZWhWR_Sc275SDl1BRVCH2g??