何愷明MAE局限性被打破,與Swin Transformer結(jié)合,訓(xùn)練速度提升
本文經(jīng)AI新媒體量子位(公眾號ID:QbitAI)授權(quán)轉(zhuǎn)載,轉(zhuǎn)載請聯(lián)系出處。
自何愷明MAE橫空出世以來,MIM(Masked Image Modeling)這一自監(jiān)督預(yù)訓(xùn)練表征越來越引發(fā)關(guān)注。
但與此同時, 研究人員也不得不思考它的局限性。
MAE論文中只嘗試了使用原版ViT架構(gòu)作為編碼器,而表現(xiàn)更好的分層設(shè)計結(jié)構(gòu)(以Swin Transformer為代表),并不能直接用上MAE方法。
于是,一場整合的范式就此在研究團(tuán)隊中上演。
代表工作之一是來自清華、微軟亞研院以及西安交大提出SimMIM,它探索了Swin Transformer在MIM中的應(yīng)用。
但與MAE相比,它在可見和掩碼圖塊均有操作,且計算量過大。有研究人員發(fā)現(xiàn),即便是SimMIM的基本尺寸模型,也無法在一臺配置8個32GB GPU的機(jī)器上完成訓(xùn)練。
基于這樣的背景,東京大學(xué)&商湯&悉尼大學(xué)的研究員,提供一個新思路。
不光將Swin Transformer整合到了MAE框架上,既有與SimMIM相當(dāng)?shù)娜蝿?wù)表現(xiàn),還保證了計算效率和性能——
將分層ViT的訓(xùn)練速度提高2.7倍,GPU內(nèi)存使用量減少70%。
來康康這是一項什么研究?
當(dāng)分層設(shè)計引入MAE
這篇論文提出了一種面向MIM的綠色分層視覺Transformer。
即允許分層ViT丟棄掩碼圖塊,只對可見圖塊進(jìn)行操作。
具體實現(xiàn),由兩個關(guān)鍵部分組成。
首先,設(shè)計了一種基于分治策略的群體窗口注意力方案。
將具有不同數(shù)量可見圖塊的局部窗口聚集成幾個大小相等的組,然后在每組內(nèi)進(jìn)行掩碼自注意力。
其次,把上述分組任務(wù)視為有約束動態(tài)規(guī)劃問題,受貪心算法的啟發(fā)提出了一種分組算法。
它可以自適應(yīng)選擇最佳分組大小,并將局部窗口分成最少的一組,從而使分組圖塊上的注意力整體計算成本最小。
表現(xiàn)相當(dāng),訓(xùn)練時間大大減少
結(jié)果顯示,在ImageNet-1K和MS-COCO數(shù)據(jù)集上實驗評估表明,與基線SimMIM性能相當(dāng)?shù)耐瑫r,效率提升2倍以上。
而跟SimMIM相比,這一方法在所需訓(xùn)練時間大大減少,消耗GPU內(nèi)存也小得多。具體而言,在相同的訓(xùn)練次數(shù)下,在Swin-B上提高2倍的速度和減少60%的內(nèi)存。
值得一提的是,該研究團(tuán)隊在有8個32GB V100 GPU的單機(jī)上進(jìn)行評估的,而SimMIM是在2或4臺機(jī)器上進(jìn)行評估。
研究人員還發(fā)現(xiàn),效率的提高隨著Swin-L的增大而變大,例如,與SimMIM192相比,速度提高了2.7倍。
實驗的最后,提到了算法的局限性。其中之一就是需要分層次掩碼來達(dá)到最佳的效率,限制了更廣泛的應(yīng)用。這一點就交給未來的研究。
而談到這一研究的影響性,研究人員表示,主要就是減輕了MIM的計算負(fù)擔(dān),提高了MIM的效率和有效性。
感興趣的旁友,可戳下方鏈接了解更多~
論文鏈接:
https://arxiv.org/abs/2205.13515
GitHub鏈接:
https://github.com/LayneH/GreenMIM
SimMIM論文鏈接:
https://arxiv.org/abs/2111.09886