何愷明最新工作：簡單實(shí)用的自監(jiān)督學(xué)習(xí)方案MAE，ImageNet-1K 87.8%！

論文鏈接：https://arxiv.org/pdf/2111.06377.pdf

愷明出品，必屬精品！這篇文章延續(xù)了其一貫的風(fēng)格：簡單且實(shí)用。這篇文章仍屬于愷明最近兩年的研究領(lǐng)域：自監(jiān)督領(lǐng)域(自監(jiān)督學(xué)習(xí)就是被他帶火的吧)。本文的出發(fā)點(diǎn)則是BERT的掩碼自編碼機(jī)制：移除一部分?jǐn)?shù)據(jù)并對移除的內(nèi)容進(jìn)行學(xué)習(xí)。掩碼自編碼源于CV但盛于NLP，愷明對此提出了疑問：是什么導(dǎo)致了掩碼自編碼在視覺與語言之間的差異？嘗試從不同角度進(jìn)行解釋并由此引申出了本文的MAE。

Abstract

愷明提出一種用于計算機(jī)視覺的可擴(kuò)展自監(jiān)督學(xué)習(xí)方案Masked AutoEncoders(MAE)。所提MAE極為簡單：對輸入圖像進(jìn)行塊隨機(jī)mask并對遺失像素進(jìn)行重建。它基于以下兩個核心設(shè)計：

• 我們設(shè)計了一種非對稱編解碼架構(gòu)，其中編碼器僅作用于可見塊(無需mask信息)，而解碼器則通過隱表達(dá)與mask信息進(jìn)行原始圖像重建；

• 我們發(fā)現(xiàn)對輸入圖像進(jìn)行高比例mask(比如75%)可以產(chǎn)生一項重要且有意義的自監(jiān)督任務(wù)。

上述兩種設(shè)計促使我們可以更高效的訓(xùn)練大模型：我們加速訓(xùn)練達(dá)3x甚至更多，同時提升模型精度。所提方案使得所得高精度模型具有很好的泛化性能：僅需ImageNet-1K，ViT-Huge取得了87.8%的top1精度 。下游任務(wù)的遷移取得了優(yōu)于監(jiān)督訓(xùn)練的性能，證實(shí)了所提方案的可擴(kuò)展能力。

極致精簡版

用下面幾句話來簡單說明下這篇文章：

• 愷明出品，必屬精品！MAE延續(xù)了其一貫的研究風(fēng)格：簡單且實(shí)用；

• MAE興起于去噪自編碼，但興盛于NLP的BERT。那么是什么導(dǎo)致了MAE在CV與NLP中表現(xiàn)的差異呢？這是本文的出發(fā)點(diǎn)。

• 角度一：CV與NLP的架構(gòu)不同。CV中常采用卷積這種具有”規(guī)則性“的操作，直到近期ViT才打破了架構(gòu)差異；

• 角度二：信息密度不同。語言是人發(fā)明的，具有高語義與信息稠密性；而圖像則是自然信號具有重度空間冗余：遺失塊可以通過近鄰塊重建且無需任何全局性理解。為克服這種差異，我們采用了一種簡單的策略：高比例隨機(jī)塊掩碼，大幅降低冗余。

• 角度三：自編碼器的解碼器在重建方面的作用不同。在視覺任務(wù)方面，解碼器進(jìn)行像素重建，具有更低語義信息；而在NLP中，解碼器預(yù)測遺失的詞，包含豐富的語義信息。

• 基于上述三點(diǎn)分析，作者提出了一種非常簡單的用于視覺表達(dá)學(xué)習(xí)的掩碼自編碼器MAE。

• MAE采用了非對稱的編解碼器架構(gòu)，編碼器僅作用于可見圖像塊(即輸入圖像塊中一定比例進(jìn)行丟棄，丟棄比例高達(dá)75%)并生成隱式表達(dá)，解碼器則以掩碼token以及隱式表達(dá)作為輸入并對遺失塊進(jìn)行重建。

• 搭配MAE的ViT-H取得了ImageNet-1K數(shù)據(jù)集上的新記錄：87.8%；同時，經(jīng)由MAE預(yù)訓(xùn)練的模型具有非常好的泛化性能。

Method

所提MAE是一種非常簡單的自編碼器方案：基于給定部分觀測信息對原始信號進(jìn)行重建 。類似于其他自編碼器，所提MAE包含一個將觀測信號映射為隱式表達(dá)的編碼器，一個用于將隱式表達(dá)重建為原始信號的解碼器。與經(jīng)典自編碼器不同之處在于：我們采用了非對稱設(shè)計，這使得編碼器僅依賴于部分觀測信息(無需掩碼token信息)，而輕量解碼器則接與所得隱式表達(dá)與掩碼token進(jìn)行原始信號重建(可參見下圖)。

Masking ?參考ViT，我們將輸入圖像拆分為非重疊塊，然后采樣一部分塊并移除其余塊(即Mask)。我們的采樣策略非常簡單：服從均勻分布的無重復(fù)隨機(jī)采樣 。我們將該采樣策略稱之為“隨機(jī)采樣”。具有高掩碼比例的隨機(jī)采樣可以極大程度消除冗余，進(jìn)而構(gòu)建一個不會輕易的被近鄰塊推理解決的任務(wù) (可參考下面圖示)。而均勻分布則避免了潛在的中心偏置問題。

MAE Encoder ?MAE中的編碼器是一種ViT，但僅作用于可見的未被Mask的塊。類似于標(biāo)準(zhǔn)ViT，該編碼器通過線性投影于位置嵌入對塊進(jìn)行編碼，然后通過一系列Transformer模塊進(jìn)行處理。然而，由于該編解碼僅在較小子集塊(比如25%)進(jìn)行處理，且未用到掩碼Token信息。這就使得我們可以訓(xùn)練一個非常大的編碼器 。

MAE Decoder ?MAE解碼器的輸入包含：(1) 編碼器的輸出；(2) 掩碼token。正如Figure1所示，每個掩碼Token共享的可學(xué)習(xí)向量，它用于指示待預(yù)測遺失塊。此時，我們對所有token添加位置嵌入信息。解碼器同樣包含一系列Transformer模塊。

注：MAE解碼器僅在預(yù)訓(xùn)練階段用于圖像重建，編碼器則用來生成用于識別的圖像表達(dá) 。因此，解碼器的設(shè)計可以獨(dú)立于編碼設(shè)計，具有高度的靈活性。在實(shí)驗(yàn)過程中，我們采用了窄而淺的極小解碼器，比如默認(rèn)解碼器中每個token的計算量小于編碼器的10% 。通過這種非對稱設(shè)計，token的全集僅被輕量解碼器處理，大幅減少了預(yù)訓(xùn)練時間。

Reconstruction target ?該MAE通過預(yù)測每個掩碼塊的像素值進(jìn)行原始信息重建 。解碼器的最后一層為線性投影，其輸出通道數(shù)等于每個塊的像素數(shù)量。編碼器的輸出將通過reshape構(gòu)建重建圖像。損失函數(shù)則采用了MSE，注：類似于BERT僅在掩碼塊計算損失。

我們同時還研究了一個變種：其重建目標(biāo)為每個掩碼塊的規(guī)范化像素值 。具體來說，我們計算每個塊的均值與標(biāo)準(zhǔn)差并用于對該塊進(jìn)行歸一化，最后采用歸一化的像素作為重建目標(biāo)提升表達(dá)能力。

Simple implementation ?MAE預(yù)訓(xùn)練極為高效，更重要的是：它不需要任何特定的稀疏操作。實(shí)現(xiàn)過程可描述如下：

• 首先，我們通過線性投影與位置嵌入對每個輸入塊生成token；

• 然后，我們隨機(jī)置換(random shuffle)token序列并根據(jù)掩碼比例移除最后一部分token；

• 其次，完成編碼后，我們在編碼塊中插入掩碼token并反置換(unshuffle)得到全序列token以便于與target進(jìn)行對齊；

• 最后，我們將解碼器作用于上述全序列token。

正如上所述：MAE無需稀疏操作。此外，shuffle與unshuffle操作非?？欤氲挠嬎懔靠梢院雎?。

Experiments

我們在ImageNet-1K數(shù)據(jù)集上進(jìn)行自監(jiān)督預(yù)訓(xùn)練，然后再通過監(jiān)督訓(xùn)練評估預(yù)訓(xùn)練模型的表達(dá)能力。

Main Properties

Baseline：ViT-Large 。我們采用ViT-Large作為消融實(shí)驗(yàn)的骨干，上表為從頭開始訓(xùn)練與MAE微調(diào)的性能對比?？梢钥吹剑簭念^開始訓(xùn)練(200epoch)，ViT-L的性能為82.5%且無需強(qiáng)正則技術(shù)；而MAE(注：僅微調(diào)50epoch)則取得了大幅性能提升。

上表則從不同角度進(jìn)行了消融實(shí)驗(yàn)對比，一一道來。

Decoder Design ?從Table1a與Table1b可以看到：解碼器的設(shè)計可以非常靈活 ?？偠灾?，默認(rèn)解碼器非常輕量，僅有8個模塊，維度為512，每個token的計算量僅為編碼的9%。

Mask Token ?MAE的的重要設(shè)計：在編碼階段跳過掩碼token，在解碼階段對其進(jìn)行處理。Table1c給出了性能對比，可以看到：編碼器使用掩碼token則會導(dǎo)致性能下降 。

Recontruction target ?Table1d比較了不同重建目標(biāo)的性能，可以看到：引入塊歸一化可以進(jìn)一步提升模型精度 。

Data Augmentation ?Table1e比較了不同數(shù)據(jù)增廣的影響，可以看到：MAE僅需crop即可表現(xiàn)非常好，添加ColorJitter反而會影響性能 。另外，令人驚訝的是：當(dāng)不使用數(shù)據(jù)增廣時，MAE性能也非常優(yōu)秀 。

Mask Sampling ?Table1f比較了不同搞得掩碼采樣策略，可以看到：不同的采樣策略均具有比較好的性能，而隨機(jī)采樣則具有最佳性能 。

Masking ?ratio 下圖給出了掩碼比例的影響，可以看到：最優(yōu)比例驚人的高 。掩碼比例為75%有益于兩種監(jiān)督訓(xùn)練方式(端到端微調(diào)與linear probing)。這與BERT的行為截然相反，其掩碼比例為15%。

與此同時，從上圖可以看到：端到端微調(diào)與linear probing兩種方式存在不同的趨勢：

• 對于linear probing而言，模型性能隨掩碼比例非常穩(wěn)定的提升直到達(dá)到最高點(diǎn)，精度差約為20%；

• 對于微調(diào)來說，模型性能再很大范圍的掩碼比例內(nèi)均極度不敏感，而且所有微調(diào)結(jié)果均優(yōu)于linear probing方式。

Training Schedule 下圖給出了不同訓(xùn)練機(jī)制的性能對比(此時采用了800epoch預(yù)訓(xùn)練)，可以看到：更長的訓(xùn)練可以帶來更定的精度提升 。作者還提到：哪怕1600epoch訓(xùn)練也并未發(fā)現(xiàn)linear probing方式的性能飽和。這與MoCoV3中的300epoch訓(xùn)練飽和截然相反 ：在每個epoch，MAE僅能看到25%的圖像塊；而MoCoV3則可以看到200%，甚至更多的圖像塊。

Comparisons with Previous Results

上表給出了所提MAE與其他自監(jiān)督方案的性能對比，從中可以看到：

• 對于ViT-B來說，不同方案的性能非常接近；對于ViT-L來說，不同方案的性能差異則變大。這意味著：對更大模型降低過擬合更具挑戰(zhàn)性 。

• MAE可以輕易的擴(kuò)展到更大模型并具有穩(wěn)定的性能提升。比如：ViT-H取得了86.9%的精度，在448尺寸微調(diào)后，性能達(dá)到了87.8% ，超越了此前VOLO的最佳87.1%(尺寸為512)。注：該結(jié)果僅使用了ViT，更優(yōu)秀的網(wǎng)絡(luò)表達(dá)可能會更好。

Transfer Learning Experiments

上表給出了COCO檢測與分割任務(wù)上的遷移性能對比，可以看到：相比監(jiān)督預(yù)訓(xùn)練，MAE取得了全配置最佳 。當(dāng)骨干為ViT-B時，MAE可以取得2.4AP指標(biāo)提升；當(dāng)骨干為ViT-L時，性能提升達(dá)4.0AP。

上表給出了ADE20K語義分割任務(wù)上的遷移性能對比，可以看到：MAE可以大幅改善ViT-L的性能，比監(jiān)督訓(xùn)練高3.7。

全文到此結(jié)束，更多實(shí)驗(yàn)結(jié)果與分析建議查看原文。

如果覺得有用，就請分享到朋友圈吧！