何愷明MAE局限性被打破，與Swin Transformer結(jié)合，訓(xùn)練速度大大提升 | 東大&商湯&悉大
GPU內(nèi)存使用量減少70%
楊凈 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI
自何愷明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方法。
于是，一場整合的范式就此在研究團隊中上演。
代表工作之一是來自清華、微軟亞研院以及西安交大提出SimMIM，它探索了Swin Transformer在MIM中的應(yīng)用。
但與MAE相比，它在可見和掩碼圖塊均有操作，且計算量過大。有研究人員發(fā)現(xiàn)，即便是SimMIM的基本尺寸模型，也無法在一臺配置8個32GB GPU的機器上完成訓(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)存。
值得一提的是，該研究團隊在有8個32GB V100 GPU的單機上進(jìn)行評估的，而SimMIM是在2或4臺機器上進(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
- 英偉達(dá)最新芯片B30A曝光2025-08-20
- AI應(yīng)用如何落地政企？首先不要卷通用大模型2025-08-12
- 本科必學(xué)Dijkstra算法被超越！清華段然團隊打破圖靈獎得主證明的普遍最優(yōu)性2025-08-09
- 智能體邁入L4 時代！納米AI多智能體蜂群，可創(chuàng)作最長10分鐘AI視頻2025-08-06
相關(guān)閱讀
