論文鏈接：https://arxiv.org/pdf/2502.03444
項(xiàng)目鏈接：https://github.com/Hhhhhhao/continuous_tokenizer

亮點(diǎn)直擊

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

結(jié)果展示

總結(jié)速覽

解決的問(wèn)題

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

提出的方案

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

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

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

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

達(dá)到的效果

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

方法

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

結(jié)構(gòu)

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

Mask 建模

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

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

實(shí)驗(yàn)

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

實(shí)驗(yàn)設(shè)置

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

評(píng)估  

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

MAETok的設(shè)計(jì)選擇

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

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

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

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

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

潛在空間分析

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

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

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

主要結(jié)果

生成  

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

重建  

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

討論

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

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

無(wú)條件生成  

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

結(jié)論

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

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

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