谷歌、MIT提出統一框架MAGE：表征學習超MAE，無監督圖像生成超越 Latent Diffusion
識別和生成是人工智能領域中的兩大核心任務，如果能將二者合并到一個統一的系統中，這兩個任務應該能實現互補。事實上，在自然語言處理中，像 BERT [1] 這樣的模型不僅能夠生成高質量的文本，還能夠提取文本中的特征。
然而，在計算機視覺領域，目前的圖像生成模型和識別模型大多是分開進行訓練，沒有充分利用這兩個任務的協同作用。這主要是由于圖像生成和圖像識別的模型通常具有本質上的結構差異：圖像生成的輸入是低維度的特征或噪聲，而輸出是高維度的原始圖像；與之相反，圖像識別的輸入是高維度的原始圖像，而輸出是低維度的特征。
最近，來自 MIT 和 Google Research 的研究人員提出了一種基于圖像語義符掩碼的表征學習方法，首次在一個統一的框架中實現了圖像生成和表征學習，并在多個數據集上取得了 SOTA 表現。研究論文已被 CVPR 2023 接收，相關代碼與預訓練模型已開源。
- 論文地址：https://arxiv.org/abs/2211.09117
- 代碼地址：https://github.com/LTH14/mage
在 CVPR 2022 上，MAE [2] 提出了一種基于圖像掩碼（MIM）的表征學習方法，并在多個子任務上取得了非常好的效果。在高達 75% 的掩碼率下，MAE 可以重構出與原圖語義十分貼合的圖像，并借此讓網絡能夠自監督地學習圖像中的特征。然而，如圖 1 所示， MAE 重建的圖像雖然具有與原始圖像相似的語義信息，但會出現嚴重的模糊與失真問題。類似的問題也出現在所有基于 MIM 的表征學習方法中。同時，目前的生成模型，不管是擴散模型還是 GAN，都缺乏提取高質量圖像特征的能力。
圖 1：MAE 與 MAGE 重構對比
方法概述
針對上述問題，本文作者提出了 MAGE（Masked Generative Encoder），首次實現了統一的圖像生成和特征提取模型。與MIM直接作用于圖像的掩碼方法不同，MAGE 提出了基于圖像語義符的 masked image token modeling 方法。如圖所示，MAGE 首先使用 VQGAN [3] 編碼器將原始圖像轉換為離散的語義符。之后，MAGE 對其進行隨機掩碼，并使用基于 transformer 的 encoder-decoder 結構對掩碼進行重構，重構后的語義符可以通過 VQGAN 解碼器生成原始圖像。通過在訓練中使用不同的掩碼率，MAGE 可以同時進行生成模型（接近 100% 掩碼率）和表征學習（50%-80% 掩碼率）的訓練。如圖 1 所示，MAGE 重建出的圖像不僅具有與原始圖像一致的語義信息，還能夠同時保證生成圖像的多樣性與真實性。
圖 2：MAGE 結構圖
實驗結果
MAGE 在多個圖像生成與圖像識別任務上都達到或超過了 SOTA。
在 ImageNet 的無監督圖像生成任務中，MAGE 的 FID 從之前的 > 20 降至 7.04，甚至達到了有監督圖像生成的水準（有監督 Latent Diffusion 在 ImageNet 上的 FID 為 3.60）：
圖3：MAGE 無監督圖像生成樣例
MAGE 還能夠進行各類圖像編輯工作，包括 image inpainting、outpainting、uncropping：
圖 4：MAGE 圖像編輯樣例
在表征學習方面，MAGE 在 ImageNet linear probing、少樣本學習、遷移學習等任務中，相較于目前的 MIM 方法有了大幅提升，并且可以達到或超過目前最優的自監督學習方法的水平。
結語
本文旨在將圖像生成與表征學習統一起來。為此，本文作者提出了 MAGE，一種基于圖像語義符掩碼的自監督學習框架。該框架簡潔、高效，并首次在圖像生成和表征學習上都達到或超越了 SOTA 的表現。感興趣的讀者可以查看論文原文，以了解更多研究細節。
