這篇AI論文介紹了MAETok：一種基於遮罩自編碼器的標記器，用於高效擴散模型

擴散模型透過逐步改善噪音來生成有結構的圖像。然而，這些模型的計算成本仍然是一個主要挑戰，尤其是在直接處理高維像素數據時。研究人員一直在尋找優化潛在空間表示的方法，以提高效率，同時不影響圖像質量。

擴散模型中的一個關鍵問題是潛在空間的質量和結構。傳統的方法，如變分自編碼器 (Variational Autoencoders, VAE)，被用作標記器來調節潛在空間，確保學習到的表示是平滑且有結構的。然而，VAE在達到高像素級保真度方面經常遇到困難，這是由於正則化所施加的限制。自編碼器 (Autoencoders, AE) 不使用變分約束，能夠以更高的保真度重建圖像，但通常會導致潛在空間的糾纏，妨礙擴散模型的訓練和性能。解決這些挑戰需要一個標記器，能夠提供結構化的潛在空間，同時保持高重建準確性。

之前的研究努力嘗試使用各種技術來解決這些問題。VAE施加了Kullback-Leibler (KL) 約束，以促進平滑的潛在分佈，而表示對齊的VAE則精煉潛在結構，以提高生成質量。一些方法利用高斯混合模型 (Gaussian Mixture Models, GMM) 來結構化潛在空間，或將潛在表示與預訓練模型對齊，以增強性能。儘管這些進展存在，但現有方法仍然面臨計算開銷和可擴展性限制，迫切需要更有效的標記策略。

來自卡內基梅隆大學 (Carnegie Mellon University)、香港大學 (The University of Hong Kong)、北京大學 (Peking University) 和AMD的研究團隊提出了一種新型標記器，稱為遮罩自編碼器標記器 (Masked Autoencoder Tokenizer, MAETok)，以應對這些挑戰。MAETok在自編碼器框架內使用遮罩建模，開發出更結構化的潛在空間，同時確保高重建保真度。研究人員設計MAETok以利用遮罩自編碼器 (Masked Autoencoders, MAE) 的原則，優化生成質量和計算效率之間的平衡。

MAETok的工作方法涉及訓練一個基於視覺變壓器 (Vision Transformer, ViT) 的自編碼器架構，包含編碼器和解碼器。編碼器接收一張被劃分為小塊的輸入圖像，並與一組可學習的潛在標記一起處理。在訓練過程中，部分輸入標記會隨機被遮罩，迫使模型從剩餘可見區域推斷缺失數據。這種機制增強了模型學習區分性和語義豐富表示的能力。此外，輔助淺層解碼器預測被遮罩的特徵，進一步提高潛在空間的質量。與傳統的VAE不同，MAETok消除了變分約束的需要，簡化了訓練，同時提高了效率。

進行了廣泛的實驗評估，以評估MAETok的有效性。該模型在ImageNet生成基準上展示了最先進的性能，同時顯著降低了計算需求。具體而言，MAETok僅使用128個潛在標記，卻在512×512分辨率圖像上達到了1.69的生成Frechet Inception Distance (gFID)。訓練速度比傳統方法快76倍，推理吞吐量提高了31倍。結果顯示，具有較少高斯混合模式的潛在空間產生了較低的擴散損失，從而提高了生成性能。該模型在675M參數的SiT-XL上進行訓練，並超越了之前的最先進模型，包括那些使用VAE訓練的模型。

這項研究強調了在擴散模型中有效結構化潛在空間的重要性。通過整合遮罩建模，研究人員實現了重建保真度和表示質量之間的最佳平衡，顯示潛在空間的結構是生成性能的關鍵因素。這些發現為擴散基圖像合成的進一步發展提供了堅實的基礎，提供了一種在不犧牲輸出質量的情況下增強可擴展性和效率的方法。

查看論文和GitHub頁面。所有對這項研究的讚譽都歸功於這個項目的研究人員。此外，別忘了在Twitter上關注我們，加入我們的Telegram頻道和LinkedIn小組。也別忘了加入我們的75k+ ML SubReddit。

🚨 加入我們在Twitter/X上的機器學習社群