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