何愷明一作！極簡(jiǎn)自監(jiān)督學(xué)習(xí)方案MAE | 指導(dǎo)意義極強(qiáng)！

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

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

Abstract

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

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

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

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

極致精簡(jiǎn)版

用下面幾句話(huà)來(lái)簡(jiǎn)單說(shuō)明下這篇文章：

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

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

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

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

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

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

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

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

Method

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

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

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

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

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

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

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

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

• 首先，我們通過(guò)線(xiàn)性投影與位置嵌入對(duì)每個(gè)輸入塊生成token；

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

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

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

正如上所述：MAE無(wú)需稀疏操作。此外，shuffle與unshuffle操作非常快，引入的計(jì)算量可以忽略。

Experiments

我們?cè)贗mageNet-1K數(shù)據(jù)集上進(jìn)行自監(jiān)督預(yù)訓(xùn)練，然后再通過(guò)監(jiān)督訓(xùn)練評(píng)估預(yù)訓(xùn)練模型的表達(dá)能力。

Main Properties

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

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

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

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

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

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

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

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

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

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

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

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

Comparisons with Previous Results

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

• 對(duì)于ViT-B來(lái)說(shuō)，不同方案的性能非常接近；對(duì)于ViT-L來(lái)說(shuō)，不同方案的性能差異則變大。這意味著：對(duì)更大模型降低過(guò)擬合更具挑戰(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á)可能會(huì)更好。

Transfer Learning Experiments

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

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

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

交流群

歡迎加入公眾號(hào)讀者群一起和同行交流，目前有美顏、三維視覺(jué)、計(jì)算攝影、檢測(cè)、分割、識(shí)別、醫(yī)學(xué)影像、GAN、算法競(jìng)賽等微信群

個(gè)人微信（如果沒(méi)有備注不拉群！）
請(qǐng)注明：地區(qū)+學(xué)校/企業(yè)+研究方向+昵稱(chēng)



下載1：何愷明頂會(huì)分享

在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù)：何愷明，即可下載。總共有6份PDF，涉及 ResNet、Mask RCNN等經(jīng)典工作的總結(jié)分析

下載2：終身受益的編程指南：Google編程風(fēng)格指南

在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù)：c++，即可下載。歷經(jīng)十年考驗(yàn)，最權(quán)威的編程規(guī)范！


下載3 CVPR2021

在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù)：CVPR，即可下載1467篇CVPR?2020論文 和 CVPR 2021 最新論文