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

亮點直擊

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

結(jié)果展示

總結(jié)速覽

解決的問題

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

提出的方案

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

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

應(yīng)用的技術(shù)

• 擴散模型（Diffusion Models）
• Mask自編碼器（MAE）進行特征學(xué)習(xí)
• 高效Transformer架構(gòu)用于tokenizer設(shè)計
• 高 mask 比訓(xùn)練策略
• 隱空間分析（使用高斯混合模型GMM）優(yōu)化擴散模型的學(xué)習(xí)

達到的效果

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

方法

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

結(jié)構(gòu)

在最近的 1D 分詞器 設(shè)計基礎(chǔ)上構(gòu)建 MAETok，該分詞器采用可學(xué)習(xí)的latent tokens。編碼器 (E) 和解碼器 (D) 都采用 Vision Transformer (ViT) 結(jié)構(gòu)，但經(jīng)過調(diào)整以同時處理 圖像tokens 和 latent tokens，如下圖 3 所示。

Mask 建模

像素解碼器微調(diào)

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

實驗

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

實驗設(shè)置

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

評估  

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

MAETok的設(shè)計選擇

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

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

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

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

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

潛在空間分析

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

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

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

主要結(jié)果

生成  

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

重建  

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

討論

高效的訓(xùn)練與生成  

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

無條件生成  

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

結(jié)論

本文對擴散模型的隱空間特性進行了理論和實證分析，證明了隱空間分布中較少的模式有助于更有效的學(xué)習(xí)和更好的生成質(zhì)量?；谶@些見解，開發(fā)了MAETok，它通過mask建模實現(xiàn)了最先進的性能，而無需使用變分約束。僅使用128個token，本文的方法顯著提高了ImageNet上的計算效率和生成質(zhì)量。研究結(jié)果表明，具有更強區(qū)分性的隱空間，而非變分約束，是有效的擴散模型的關(guān)鍵，為大規(guī)模高效生成建模開辟了新的方向。

本文轉(zhuǎn)自AI生成未來 ，作者：AI生成未來

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