自何愷明MAE橫空出世以來，MIM（Masked Image Modeling）這一自監(jiān)督預(yù)訓(xùn)練表征越來越引發(fā)關(guān)注。

但與此同時(shí)， 研究人員也不得不思考它的局限性。

MAE論文中只嘗試了使用原版ViT架構(gòu)作為編碼器，而表現(xiàn)更好的分層設(shè)計(jì)結(jié)構(gòu)（以Swin Transformer為代表），并不能直接用上MAE方法。

于是，一場(chǎng)整合的范式就此在研究團(tuán)隊(duì)中上演。

代表工作之一是來自清華、微軟亞研院以及西安交大提出SimMIM，它探索了Swin Transformer在MIM中的應(yīng)用。

但與MAE相比，它在可見和掩碼圖塊均有操作，且計(jì)算量過大。有研究人員發(fā)現(xiàn)，即便是SimMIM的基本尺寸模型，也無法在一臺(tái)配置8個(gè)32GB GPU的機(jī)器上完成訓(xùn)練。

基于這樣的背景，東京大學(xué)&商湯&悉尼大學(xué)的研究員，提供一個(gè)新思路。

不光將Swin Transformer整合到了MAE框架上，既有與SimMIM相當(dāng)?shù)娜蝿?wù)表現(xiàn)，還保證了計(jì)算效率和性能——

將分層ViT的訓(xùn)練速度提高2.7倍，GPU內(nèi)存使用量減少70%。

來康康這是一項(xiàng)什么研究？

當(dāng)分層設(shè)計(jì)引入MAE

這篇論文提出了一種面向MIM的綠色分層視覺Transformer。

即允許分層ViT丟棄掩碼圖塊，只對(duì)可見圖塊進(jìn)行操作。

具體實(shí)現(xiàn)，由兩個(gè)關(guān)鍵部分組成。

首先，設(shè)計(jì)了一種基于分治策略的群體窗口注意力方案。

將具有不同數(shù)量可見圖塊的局部窗口聚集成幾個(gè)大小相等的組，然后在每組內(nèi)進(jìn)行掩碼自注意力。

其次，把上述分組任務(wù)視為有約束動(dòng)態(tài)規(guī)劃問題，受貪心算法的啟發(fā)提出了一種分組算法。

它可以自適應(yīng)選擇最佳分組大小，并將局部窗口分成最少的一組，從而使分組圖塊上的注意力整體計(jì)算成本最小。

表現(xiàn)相當(dāng)，訓(xùn)練時(shí)間大大減少

結(jié)果顯示，在ImageNet-1K和MS-COCO數(shù)據(jù)集上實(shí)驗(yàn)評(píng)估表明，與基線SimMIM性能相當(dāng)?shù)耐瑫r(shí)，效率提升2倍以上。

而跟SimMIM相比，這一方法在所需訓(xùn)練時(shí)間大大減少，消耗GPU內(nèi)存也小得多。具體而言，在相同的訓(xùn)練次數(shù)下，在Swin-B上提高2倍的速度和減少60%的內(nèi)存。

值得一提的是，該研究團(tuán)隊(duì)在有8個(gè)32GB V100 GPU的單機(jī)上進(jìn)行評(píng)估的，而SimMIM是在2或4臺(tái)機(jī)器上進(jìn)行評(píng)估。

研究人員還發(fā)現(xiàn)，效率的提高隨著Swin-L的增大而變大，例如，與SimMIM192相比，速度提高了2.7倍。

實(shí)驗(yàn)的最后，提到了算法的局限性。其中之一就是需要分層次掩碼來達(dá)到最佳的效率，限制了更廣泛的應(yīng)用。這一點(diǎn)就交給未來的研究。

而談到這一研究的影響性，研究人員表示，主要就是減輕了MIM的計(jì)算負(fù)擔(dān)，提高了MIM的效率和有效性。

審核編輯 ：李倩