摘要
本文展示了掩碼自編碼器(MAE)作為計算機視覺中的可擴展自監督學習方法。我們的MAE方法很簡單:我們對輸入圖像進行隨機掩碼,并重建缺失的像素。該方法基于兩個核心設計。首先,我們開發了一種非對稱的編碼器-解碼器架構,編碼器僅處理可見部分的圖像(不含掩碼標記),而輕量級的解碼器則根據潛在表示和掩碼標記重建原始圖像。其次,我們發現掩碼高比例的輸入圖像(例如75%)可以生成一個非平凡且有意義的自監督任務。這兩種設計的結合使得我們能夠高效、有效地訓練大型模型:我們加速了訓練(提高了3倍或更多),并提高了準確率。我們的可擴展方法使得訓練具有高容量的模型成為可能,這些模型能夠很好地泛化:例如,基礎的ViT-Huge模型在只使用ImageNet-1K數據的情況下達到了最佳準確率(87.8%)。在下游任務中的遷移表現優于監督預訓練,并顯示出良好的擴展性。
1. 引言
深度學習見證了架構的爆炸式增長,其能力和容量不斷提升[29, 25, 52]。得益于硬件的快速進步,今天的模型可以輕松地對超過一百萬張圖像進行過擬合,并且開始需求數億張——通常是公開無法訪問的——標注圖像[16]。
這種對數據的需求在自然語言處理(NLP)中通過自監督預訓練得到了成功解決。基于GPT中的自回歸語言建模[42, 43, 4]和BERT中的掩碼自編碼[14]的解決方案概念上很簡單:它們去除一部分數據并學習預測缺失的內容。這些方法現在使得訓練包含超過一百億參數的通用NLP模型成為可能[4]。
掩碼自編碼器的概念——一種更一般的去噪自編碼器形式[53]——在計算機視覺中也同樣自然且適用。事實上,與BERT密切相關的視覺研究[54, 41]早于BERT出現。
然而,盡管在BERT成功之后這一想法受到極大關注,視覺領域中的自編碼方法進展卻落后于NLP。我們問:在視覺和語言中,掩碼自編碼有什么不同?我們從以下幾個角度來嘗試回答這個問題:
(i) 直到最近,架構是不同的。在視覺中,卷積網絡[30]在過去十年中占主導地位[29]。卷積操作通常作用于規則網格,且并不容易將“指示符”如掩碼標記[14]或位置嵌入[52]集成到卷積網絡中。然而,這一架構差距已通過引入視覺變換器(ViT)[16]得到了彌補,不再構成障礙。
(ii) 語言和視覺之間的信息密度不同。語言是由人類生成的信號,語義豐富且信息密集。當訓練模型預測每個句子中缺失的幾個詞時,這個任務看起來會促使模型進行復雜的語言理解。相反,圖像是自然信號,具有較強的空間冗余——例如,缺失的圖像塊可以通過鄰近的圖像塊恢復,而不需要對部分、物體和場景有太多高層次的理解。為了克服這一差異并鼓勵學習有用的特征,我們展示了一種在計算機視覺中效果良好的簡單策略:掩碼掉高比例的隨機圖像塊。這種策略大大減少了冗余,創造了一個具有挑戰性的自監督任務,要求超越低級圖像統計的整體理解。想要獲得我們重建任務的定性感受,請參見圖2到圖4。
(iii) 自編碼器的解碼器在重建文本和圖像時扮演著不同的角色。在視覺中,解碼器重建像素,因此其輸出在語義層次上低于常見的識別任務。這與語言不同,在語言中,解碼器預測缺失的詞,而這些詞包含豐富的語義信息。雖然在BERT中,解碼器可以是簡單的多層感知機(MLP)[14],但我們發現對于圖像,解碼器的設計在決定學習到的潛在表示的語義層次上起著關鍵作用。
基于這一分析,我們提出了一種簡單、有效且可擴展的掩碼自編碼器(MAE)形式,用于視覺表示學習。我們的MAE方法通過掩碼隨機圖像塊并重建缺失的圖像塊來學習。它采用了非對稱的編碼器-解碼器設計。我們的編碼器僅處理可見部分的圖像塊(不含掩碼標記),解碼器則輕量級,利用潛在表示和掩碼標記重建輸入圖像(見圖1)。將掩碼標記移至小型解碼器,在我們的非對稱編碼器-解碼器架構中極大地減少了計算量。在這種設計下,極高的掩碼比例(例如75%)能夠實現雙贏:既優化了準確性,又使得編碼器僅處理較少的圖像塊(例如25%)。這可以將整體預訓練時間縮短3倍或更多,同時減少內存消耗,從而使我們能夠輕松地將MAE擴展到大型模型。
我們的MAE能夠學習出非常高容量且泛化良好的模型。通過MAE預訓練,我們可以在ImageNet-1K數據上訓練需要大量數據的模型,如ViT-Large/-Huge,并取得更好的泛化性能。使用基礎的ViT-Huge模型,當在ImageNet-1K上進行微調時,我們達到了87.8%的準確率。這超越了所有僅使用ImageNet-1K數據的現有結果。我們還評估了在物體檢測、實例分割和語義分割任務上的遷移學習。在這些任務中,我們的預訓練方法優于監督預訓練方法,且更重要的是,隨著模型規模的擴大,我們觀察到了顯著的提升。這些觀察結果與NLP中自監督預訓練的發現相一致[14, 42, 43, 4],我們希望這些結果能夠推動我們領域探索類似的軌跡。
2. 相關工作
掩蔽語言建模及其自回歸對偶方法,例如 BERT [14] 和 GPT [42, 43, 4],是自然語言處理(NLP)中非常成功的預訓練方法。這些方法通過隱藏部分輸入序列,訓練模型去預測缺失的內容。研究表明,這些方法具有出色的擴展性[4],大量證據表明,這些預訓練表示能很好地推廣到各種下游任務。
自編碼是一種經典的表示學習方法。它包括一個將輸入映射到潛在表示的編碼器和一個從潛在表示重建輸入的解碼器。例如,PCA 和 k-means 是自編碼器[26]。去噪自編碼器(DAE)[53]是一類自編碼器,它通過破壞輸入信號并學習重建原始、未破壞的信號。可以將一系列方法看作是在不同破壞方式下的廣義 DAE,例如,掩蔽像素[54, 41, 6]或去除顏色通道[64]。我們的 MAE 是一種去噪自編碼的形式,但與經典的 DAE 在許多方面有所不同。
掩蔽圖像編碼方法從被掩蔽的圖像中學習表示。早期的工作[54]將掩蔽作為 DAE 中的噪聲類型。Context Encoder [41] 使用卷積網絡填補大范圍的缺失區域。受 NLP 成功的啟發,相關的最近方法[6, 16, 2]基于 Transformer [52]。iGPT [6] 處理像素序列并預測未知像素。ViT 論文[16]研究了自監督學習中的掩蔽補丁預測。最近,BEiT [2] 提出了預測離散 token [39, 45]。自監督學習方法在計算機視覺中引起了顯著關注,通常專注于不同的預訓練任務[15, 55, 37, 64, 40, 17]。最近,對比學習[3, 22]變得流行,例如,[56, 38, 23, 7],它建模了圖像之間的相似性和不相似性(或僅相似性[21, 8])。對比學習及相關方法強烈依賴數據增強[7, 21, 8]。自編碼追求一種概念上不同的方向,并且它展現了不同的行為,正如我們接下來所展示的那樣。
3. 方法
我們的掩蔽自編碼器(MAE)是一種簡單的自編碼方法,給定部分觀察,重建原始信號。像所有自編碼器一樣,我們的方法包含一個將觀察到的信號映射到潛在表示的編碼器和一個從潛在表示重建原始信號的解碼器。與經典自編碼器不同,我們采用不對稱設計,允許編碼器僅在部分觀察到的信號(沒有掩蔽 token)上操作,并且有一個輕量級解碼器從潛在表示和掩蔽 token 中重建完整信號。圖 1 展示了這一思想,接下來將介紹。
掩蔽
沿用 ViT [16],我們將圖像劃分為規則的非重疊補丁。然后,我們從中隨機選擇一部分補丁,并掩蔽(即移除)剩余的補丁。我們的采樣策略非常直接:我們隨機采樣補丁,且不放回,遵循均勻分布。我們簡單地稱之為“隨機采樣”。
具有較高掩蔽比例(即移除的補丁比例)的隨機采樣會大大減少冗余,從而創造出一個無法通過從可見鄰近補丁推斷的任務(見圖 2 – 4)。均勻分布防止了潛在的中心偏差(即更多掩蔽的補丁靠近圖像中心)。最后,極其稀疏的輸入為設計高效的編碼器提供了機會,接下來將介紹。
MAE 編碼器
我們的編碼器是 ViT [16],但僅作用于可見的、未掩蔽的補丁。與標準的 ViT 相似,我們的編碼器通過線性投影嵌入補丁,并添加位置嵌入,然后通過一系列 Transformer 塊處理這些結果。但是,我們的編碼器僅在完整集合的小部分(例如,25%)上操作。掩蔽補丁被移除;沒有使用掩蔽 token。這樣我們就可以用很少的計算和內存訓練非常大的編碼器。完整的集合由輕量級解碼器處理,接下來將介紹。
MAE 解碼器
MAE 解碼器的輸入是包含(i)編碼的可見補丁和(ii)掩蔽 token 的完整 token 集合。見圖 1。每個掩蔽 token [14] 是一個共享的、學習到的向量,表示一個缺失補丁需要被預測。我們為這個完整集合中的所有 token 添加位置嵌入;如果沒有這些位置嵌入,掩蔽 token 將不知道它們在圖像中的位置。解碼器有一系列 Transformer 塊。
MAE 解碼器僅在預訓練期間用于執行圖像重建任務(僅使用編碼器生成圖像表示用于識別)。因此,解碼器架構可以靈活設計,獨立于編碼器設計。我們實驗了非常小的解碼器,其寬度和深度都比編碼器小。例如,我們的默認解碼器在每個 token 上的計算量少于編碼器的 10%。通過這種不對稱設計,完整的 token 集合僅由輕量級解碼器處理,這大大減少了預訓練時間。
重建目標
我們的 MAE 通過預測每個掩蔽補丁的像素值來重建輸入。解碼器輸出中的每個元素是一個表示補丁的像素值向量。解碼器的最后一層是一個線性投影,輸出通道數等于補丁中像素值的數量。解碼器的輸出被重塑為重建的圖像。我們的損失函數計算重建圖像和原始圖像之間的均方誤差(MSE),該計算僅對掩蔽補丁進行,類似于 BERT [14]。我們還研究了一種變體,其重建目標是每個掩蔽補丁的歸一化像素值。具體來說,我們計算補丁中所有像素的均值和標準差,并用它們對該補丁進行歸一化。在我們的實驗中,使用歸一化像素作為重建目標能夠提高表示質量。
簡單實現
我們的 MAE 預訓練可以高效實現,且重要的是,它不需要任何特殊的稀疏操作。首先,我們為每個輸入補丁生成一個 token(通過線性投影并添加位置嵌入)。接下來,我們隨機打亂 token 列表,并根據掩蔽比例移除列表的最后一部分。這個過程生成一個用于編碼器的小 subset token 集合,等同于不放回的補丁采樣。編碼后,我們將掩蔽 token 列表附加到編碼補丁列表中,并重新打亂這個完整列表(反轉隨機打亂操作),以使所有 token 與它們的目標對齊。解碼器應用于這個完整的列表(并添加位置嵌入)。如前所述,不需要稀疏操作。這個簡單的實現幾乎不會引入額外的開銷,因為打亂和還原打亂操作非常快速。
4. ImageNet 實驗
我們在 ImageNet-1K (IN1K) [13] 訓練集上進行自監督預訓練。然后,我們通過 (i) 端到端微調或 (ii) 線性探測來評估表示性能。我們報告了單個 224×224 裁剪的 top-1 驗證準確率。詳細信息見附錄 A.1。
基準:ViT-Large。我們在消融實驗中使用 ViT-Large (ViT-L/16) [16] 作為骨干網絡。ViT-L 非常大(比 ResNet-50 [25] 大一個數量級),且容易發生過擬合。以下是 ViT-L 從頭訓練與從我們的基準 MAE 微調后的比較:
我們注意到,從頭開始訓練監督學習的 ViT-L 是一項復雜的任務,需要一個有效的訓練方案并且強大的正則化(82.5%,見附錄 A.2)。即便如此,我們的 MAE 預訓練仍然帶來了顯著的提升。在這里,微調僅進行 50 輪(與從頭訓練的 200 輪相比),這意味著微調準確率在很大程度上依賴于預訓練。
4.1 主要特性
我們使用表1中的默認設置(見說明)對 MAE 進行了消融實驗,并觀察到幾個有趣的特性。
遮罩比率。圖5展示了遮罩比率的影響。最佳的遮罩比率出乎意料地高。75%的比率對線性探測和微調都表現良好。這種行為與 BERT [14] 相對立,后者的典型遮罩比率為 15%。我們的遮罩比率也遠高于計算機視覺中相關工作的比率[6, 16, 2](20%到50%)。該模型推斷缺失的補丁,生成不同但合理的輸出(圖4)。它能夠理解物體和場景的整體感,這些內容無法僅通過延伸線條或紋理來完成。我們假設,這種類推推理行為與有用表示的學習有關。圖5還顯示,線性探測和微調的結果遵循不同的趨勢。對于線性探測,隨著遮罩比率的增加,準確率穩步上升,直到達到最佳點:準確率差距達到約 20%(54.6% vs. 73.5%)。對于微調,結果對比率的敏感性較低,40%到80%之間的廣泛遮罩比率都表現良好。圖5中的所有微調結果均優于從頭訓練(82.5%)。
解碼器設計。我們的 MAE 解碼器可以靈活設計,如表1a和1b中所研究的那樣。
表1a變化了解碼器的深度(Transformer塊的數量)。一個足夠深的解碼器對于線性探測至關重要。這可以通過像素重建任務與識別任務之間的差異來解釋:自編碼器中的最后幾層更加專注于重建,但與識別任務的相關性較低。一個合理深度的解碼器可以考慮到重建的專業化,將潛在表示保留在一個更抽象的層次。這種設計可以在線性探測中帶來最多 8% 的提升(表1a,“lin”)。然而,如果使用微調,編碼器的最后幾層可以調整以適應識別任務。解碼器深度對于提高微調的效果影響較小(表1a,“ft”)。
有趣的是,我們的 MAE 使用單個塊的解碼器也能在微調中表現出色(84.8%)。請注意,單個 Transformer 塊是將信息從可見標記傳播到遮罩標記的最小要求。如此小的解碼器還可以進一步加速訓練。
在表1b中,我們研究了解碼器的寬度(通道數)。默認情況下,我們使用 512-d,它在微調和線性探測下表現良好。一個更窄的解碼器在微調時也能良好工作。
總體而言,我們的默認 MAE 解碼器是輕量級的。它有 8 個塊和 512-d 的寬度(表1中的灰色部分)。每個標記的 FLOP 僅為 ViT-L(24 個塊,1024-d)的 9%。因此,盡管解碼器處理所有標記,它仍然只占總體計算的很小一部分。
遮罩標記。我們MAE的一個重要設計是跳過編碼器中的遮罩標記 [M],并在輕量級解碼器中稍后應用它。表1c 研究了這個設計。
如果編碼器使用遮罩標記,其性能較差:在線性探測中,準確率下降了 14%。在這種情況下,預訓練和部署之間存在差距:在預訓練中,編碼器輸入中有大量遮罩標記,而在未損壞的圖像中并不存在這些遮罩標記。這種差距可能會在部署時降低準確率。通過從編碼器中去除遮罩標記,我們限制了編碼器始終看到真實的補丁,從而提高了準確率。
此外,通過在編碼器中跳過遮罩標記,我們大大減少了訓練計算量。在表1c中,我們將整體訓練 FLOP 數減少了 3.3 倍。這導致了我們實現中的 2.8 倍墻鐘加速(見表2)。對于較小的解碼器(1塊)、更大的編碼器(ViT-H)或兩者兼有,墻鐘加速甚至更大(3.5–4.1 倍)。請注意,對于 75% 的遮罩比率,加速可能超過 4 倍,部分原因是自注意力的復雜度是二次的。此外,內存消耗大幅減少,這可以使得訓練更大的模型成為可能,或通過大批量訓練進一步加速。時間和內存的高效性使得我們的 MAE 非常適合訓練非常大的模型。
重建目標。我們比較了表1d中的不同重建目標。到目前為止,我們的結果基于未進行(每個補丁)歸一化的像素。使用歸一化的像素可以提高準確率。這種每補丁歸一化在局部增強了對比度。在另一種變體中,我們在補丁空間中執行 PCA,并使用最大的 PCA 系數(此處為 96 個)作為目標。這樣做會降低準確率。兩個實驗都表明,在我們的方法中,高頻成分是有用的。
我們還比較了一個 MAE 變體,它預測token,這是 BEiT [2] 中使用的目標。具體而言,對于這個變體,我們使用經過 DALLE 預訓練的 dVAE [45] 作為分詞器,跟隨 [2]。在這里,MAE 解碼器使用交叉熵損失預測token索引。與未歸一化的像素相比,這種token化在微調時提高了 0.4% 的準確率,但與歸一化像素相比沒有優勢。它還會降低線性探測的準確性。在 §5 中,我們進一步展示了在遷移學習中并不需要token化。
我們的基于像素的 MAE 比token化要簡單得多。dVAE 分詞器需要一個額外的預訓練階段,這可能依賴于額外的數據(250M 圖像 [45])。dVAE 編碼器是一個大型卷積網絡(占 ViT-L FLOP 的 40%),并且增加了不容忽視的開銷。使用像素并不會遇到這些問題。
數據增強。表1e研究了數據增強對我們 MAE 預訓練的影響。
我們的 MAE 在僅使用裁剪的增強(無論是固定大小還是隨機大小,都有隨機水平翻轉)時表現良好。添加顏色抖動會降低結果,因此我們在其他實驗中不使用它。
令人驚訝的是,即使不使用數據增強(僅使用中心裁剪,無翻轉),我們的 MAE 也表現得相當不錯。這個特性與對比學習和相關方法 [56, 23, 7, 21] 大不相同,這些方法嚴重依賴于數據增強。已有研究[21]觀察到,僅使用裁剪增強會分別使 BYOL [21] 和 SimCLR [7] 的準確率降低 13% 和 28%。此外,目前沒有證據表明對比學習可以在沒有增強的情況下工作:圖像的兩個視圖是相同的,容易滿足一個簡單的解決方案。
在 MAE 中,數據增強的作用主要通過隨機遮蔽來實現(下一步將進行消融實驗)。每次迭代的遮蔽是不同的,因此它們會生成新的訓練樣本,無論是否進行數據增強。由于遮蔽任務本身具有挑戰性,因此在訓練過程中對數據增強的需求較少,且仍能有效地正則化訓練。
遮蔽采樣策略。表1f中,我們比較了不同的遮蔽采樣策略,見圖6。
基于塊的遮蔽策略,提出于 [2],傾向于去除較大的塊(見圖6中間)。我們的 MAE 在50%遮蔽比例下使用塊狀遮蔽表現得相當不錯,但在75%時效果下降。這個任務比隨機采樣更加困難,因為觀察到較高的訓練損失,并且重建效果也更加模糊。
我們還研究了網格狀采樣策略,定期保留每四個補丁中的一個(見圖6右)。這是一個較容易的任務,具有較低的訓練損失,且重建效果更清晰。然而,表示的質量較低。
簡單的隨機采樣是我們 MAE 最佳的采樣方式。它允許更高的遮蔽比例,這不僅提供了更大的加速效益,同時也保持了較好的準確率。
訓練計劃。到目前為止,我們的消融實驗基于 800 輪的預訓練。圖7顯示了訓練計劃長度的影響。隨著訓練輪次的增加,準確率穩步提高。事實上,即使訓練到 1600 輪,我們也未觀察到線性探測準確率的飽和現象。這種行為與對比學習方法不同,例如,MoCo v3 [9] 對 ViT-L 的訓練在 300 輪時便達到飽和。需要注意的是,MAE 編碼器每個輪次只看到 25% 的補丁,而在對比學習中,編碼器每個輪次看到的是 200%(雙裁剪)甚至更多(多裁剪)補丁。
4.2 與先前結果的比較
與自監督方法的比較。在表3中,我們比較了自監督 ViT 模型的微調結果。對于 ViT-B,所有方法的表現接近。對于 ViT-L,各方法之間的差距更大,這表明對于更大的模型,減少過擬合是一個挑戰。我們的 MAE 模型可以輕松擴展,并且在更大的模型上取得了穩定的改進。我們使用 ViT-H(224 大小)時獲得了 86.9% 的準確率。通過使用 448 大小進行微調,我們實現了 87.8% 的準確率,僅使用 IN1K 數據。之前最好的準確率(所有僅使用 IN1K 數據的方法中)是 87.1%(512 大小)[61],該結果基于先進的網絡。我們在 IN1K 這一高度競爭的基準測試中,憑借顯著的優勢超越了現有的最先進方法(不使用外部數據)。我們的結果基于原始 ViT,并且我們預期,使用更先進的網絡會取得更好的表現。
與 BEiT [2] 比較,我們的 MAE 方法不僅更加準確,同時也更簡單且更快速。我們的 MAE 通過重建像素來進行訓練,而 BEiT 則是通過預測 token:在使用 ViT-B 時,BEiT 在重建像素時報告了 1.8% 的性能下降[2]。我們不需要進行 dVAE 預訓練。此外,正如表 1c所研究的那樣,我們的 MAE 在每個 epoch 的訓練速度比 BEiT 快 3.5 倍。
MAE模型的訓練時間。表3中的 MAE 模型經過了 1600 個 epoch 的預訓練以獲得更好的準確性(見圖 7)。即便如此,我們的總預訓練時間仍然低于其他方法,且在相同硬件上訓練。例如,使用 128 個 TPU-v3 核心訓練 ViT-L時,我們的 MAE 訓練時間為 1600 個 epoch 需要 31 小時,而 MoCo v3 訓練 300 個 epoch 則需要 36 小時[9]。
與監督預訓練的比較。在原始的 ViT 論文中[16],ViT-L 在 IN1K 數據集上訓練時表現下降。我們實現的監督訓練(見 A.2)效果更好,但準確率會飽和,詳見圖 8。我們的 MAE 預訓練,僅使用 IN1K 數據集,可以更好地泛化:對于更高容量的模型,其從頭開始訓練的增益更大。這一趨勢與 [16] 中的 JFT-300M 監督預訓練相似。這個對比表明,我們的 MAE 能夠幫助擴大模型規模。
4.3 部分微調
表 1 顯示了線性探測和微調結果大體上沒有相關性。線性探測近年來非常流行;然而,它忽略了追求強大但非線性特征的機會——而這恰恰是深度學習的強項。作為折衷方案,我們研究了部分微調協議:只微調最后幾層,其它層被凍結。這一協議也在早期的工作中使用過,例如[59, 64, 37]。
圖 9 顯示了結果。值得注意的是,僅微調一個 Transformer 塊就能顯著提升準確率,從 73.5% 提升至 81.0%。此外,如果我們僅微調最后一塊的“半部分”(即其 MLP 子塊),我們可以達到 79.1%,明顯高于線性探測。這一變體本質上是在微調一個 MLP 頭。微調幾個塊(例如 4 或 6 個塊)可以獲得接近全微調的準確率。
在圖 9中,我們還與 MoCo v3 [9] 進行了對比,后者是一個對比學習方法,且提供了 ViT-L 的結果。MoCo v3 在線性探測準確率上更高;然而,它的所有部分微調結果都遜色于 MAE。當微調 4 個塊時,差距為 2.6%。雖然 MAE 的表示不那么線性可分,但它們是更強的非線性特征,在微調非線性頭時表現優異。
這些觀察表明,線性可分性并不是評估表示質量的唯一指標。也有研究表明(例如,[8]),線性探測與遷移學習性能并沒有很好地相關,例如在物體檢測任務中。據我們所知,線性評估在 NLP 中并不常用于預訓練基準測試。
5. 遷移學習實驗
我們使用表 3 中預訓練的模型評估下游任務中的遷移學習。
物體檢測和分割
我們在 COCO 數據集 [33] 上對 Mask R-CNN [24] 進行端到端微調。ViT 主干網絡被調整為與 FPN [32] 配合使用(見 A.3)。我們對表 4 中的所有條目應用了這一方法。我們報告了物體檢測的框 AP 和實例分割的掩碼 AP。與監督預訓練相比,我們的 MAE 在所有配置下表現更好(見表 4)。對于較小的 ViT-B,MAE 比監督預訓練高出 2.4 分(50.3 vs. 47.9, APbox)。更顯著的是,對于更大的 ViT-L,MAE 預訓練在監督預訓練上超出了 4.0 分(53.3 vs. 49.3)。基于像素的 MAE 優于或與基于 token 的 BEiT 不相上下,而 MAE 更加簡單且更快。MAE 和 BEiT 都優于 MoCo v3,且 MoCo v3 與監督預訓練相當。
語義分割
我們在 ADE20K [66] 上使用 UperNet [57] 進行實驗(見 A.4)。表 5 顯示了我們的預訓練顯著提高了與監督預訓練的結果,例如 ViT-L 提高了 3.7 分。我們的基于像素的 MAE 也優于基于 token 的 BEiT。這些觀察與在 COCO 中的結果一致。
分類任務
表 6 研究了 iNaturalist [51] 和 Places [65] 任務上的遷移學習(見 A.5)。在 iNat 上,我們的方法展現出了強大的擴展能力:隨著模型規模的增大,準確率顯著提升。我們的結果大幅超越了之前的最佳結果。在 Places 上,我們的 MAE 超越了之前的最佳結果 [19, 36],這些結果是通過在數十億張圖像上進行預訓練得到的。
像素與 token
表 7 比較了將像素與 token 作為 MAE 重建目標。雖然使用 dVAE token 比使用未歸一化的像素效果更好,但在我們測試的所有情況下,使用歸一化像素與使用 token 的效果統計上相似。這再次表明,token 化對于我們的 MAE 并非必要。
6. 討論與結論
深度學習的核心是能擴展的簡單算法。在 NLP 中,簡單的自監督學習方法(例如,[42, 14, 43, 4])使得模型能夠從大規模擴展中獲益。在計算機視覺中,盡管自監督學習有所進展,但實用的預訓練范式仍然以監督學習為主(例如,[29, 46, 25, 16])。在這項研究中,我們觀察到,在 ImageNet 和遷移學習中,類似于 NLP 中技術的自監督方法——自動編碼器——提供了可擴展的益處。視覺中的自監督學習可能正走在與 NLP 相似的軌跡上。
另一方面,我們也注意到,圖像和語言是不同性質的信號,這一差異必須仔細處理。圖像僅僅是記錄的光線,而沒有像詞匯那樣的語義分解。我們不嘗試移除物體,而是移除隨機的塊,這些塊最有可能不會形成語義部分。同樣,我們的 MAE 重建的是像素,這些像素并不是語義實體。然而,我們觀察到(例如,圖 4)我們的 MAE 能夠推斷出復雜的、整體的重建,表明它學習了眾多視覺概念,即語義。我們假設,這一行為可能通過 MAE 內部的豐富隱藏表示發生。我們希望這種觀點能激發未來的工作。
更廣泛的影響
所提出的方法基于訓練數據集的統計特征來預測內容,因此將反映這些數據中的偏見,包括那些可能帶來負面社會影響的偏見。模型可能會生成不存在的內容。在基于此工作生成圖像時,這些問題值得進一步研究和關注。
A. 實現細節
A.1 ImageNet 實驗
ViT 架構
我們遵循標準的 ViT 架構 [16]。該架構由一系列 Transformer 塊 [52] 組成,每個塊包括一個多頭自注意力塊和一個 MLP 塊,二者都使用了 LayerNorm (LN) [1]。編碼器最后使用 LN。由于 MAE 的編碼器和解碼器寬度不同,我們在編碼器后面加入了一個線性投影層以匹配其寬度。我們的 MAE 在編碼器和解碼器輸入中都添加了位置嵌入 [52](采用正弦余弦版本)。我們的 MAE 不使用相對位置或層級縮放(這些在 [2] 的代碼中使用)。
我們從編碼器的輸出中提取特征進行微調和線性探測。由于 ViT 有一個類 token [16],為了適應這一設計,我們在 MAE 的預訓練中將一個輔助虛擬 token 附加到編碼器輸入中。該 token 在進行線性探測和微調時會作為訓練分類器的類 token。我們的 MAE 在沒有這個 token(使用平均池化)時也能表現得很好。
預訓練
默認設置見表 8。我們不使用顏色抖動、路徑丟棄或梯度裁剪。我們使用 xavier 均勻初始化 [18] 所有 Transformer 塊,遵循 ViT 官方代碼 [16]。我們采用線性學習率縮放規則 [20]:lr = base lr × batchsize / 256
。
端到端微調
我們的微調方法遵循監督 ViT 訓練的常見做法。默認設置見表 9。我們使用層級學習率衰減 [10],參照 [2] 的做法。
線性探測
我們的線性分類器訓練遵循 [9]。見表 10。我們觀察到線性探測與端到端微調需要完全不同的策略。特別地,正則化通常對線性探測有害。按照 [9] 的方法,我們禁用了許多常見的正則化策略:不使用 mixup [63]、cutmix [62]、路徑丟棄 [27] 或顏色抖動,并將權重衰減設置為零。
訓練傳統線性分類器(如 SVM [11])時,常見的做法是對分類器輸入進行歸一化。同樣,當訓練線性探測分類器時,歸一化預訓練特征也是有益的。參照 [15],我們在預訓練的特征上采用一個額外的 BatchNorm 層 [28],且不進行仿射變換(affine=False
)。該層應用于編碼器生成的預訓練特征,并位于線性分類器之前。我們注意到該層不會破壞線性屬性,且在訓練后可以被吸收進線性分類器中:它本質上是一個重新參數化的線性分類器。引入這個層有助于校準我們不同變體中特征的幅度,使它們能夠在沒有進一步學習率搜索的情況下使用相同的設置。
部分微調
我們的 MAE 部分微調 (§4.3) 遵循表 9 中的設置,不同之處在于我們調整了微調的訓練輪數。我們觀察到,微調較少的塊需要更長的訓練周期。我們將微調的輪數設置為 f50, 100, 200g
,并根據每個微調的塊數選擇最優的輪數。
A.2 從頭訓練監督 ViT-L/H
我們發現,從頭訓練監督 ViT-L/H 在 ImageNet-1K 上是非常復雜的。訓練不穩定。盡管對于較小模型已經有強有力的基準,并且有公開的實現 [48],但對于較大的 ViT-L/H 模型,訓練的策略尚未深入研究。直接將之前的訓練策略應用于這些較大的模型并不起作用。在訓練過程中,常常會觀察到 NaN 損失。
我們提供的訓練策略見表 11。我們使用權重衰減 wd=0.3
,大批量訓練(batch size=4096),并采用較長的預熱階段,遵循原始 ViT [16]。我們使用 β2=0.95,參照 [6]。我們使用表 11 中列出的正則化方法,并禁用其他正則化方法,參照 [58]。所有這些選擇都是為了提高訓練穩定性。我們的策略能夠在沒有 NaN 損失的情況下完成訓練。
訓練結果
ViT-L 的準確率為 82.6%(若不使用 EMA,則為 81.5%),ViT-H 的準確率為 83.1%(若不使用 EMA,則為 80.9%)。如果不使用 EMA,ViT-L 和 ViT-H 都顯示出過擬合的趨勢。
作為附帶結果,我們使用的 ViT-B 配置準確率為 82.3%(不使用 EMA 時為 82.1%),相比之下,[48] 中的 ViT-B 準確率為 81.8%。
A.3. COCO 中的目標檢測和分割
我們將 Vanilla ViT 適配為 Mask R-CNN [24] 中使用 FPN 主干網絡 [32]。ViT 具有一堆 Transformer 塊,所有這些塊都在一個單一的尺度(例如,步長為 16)上生成特征圖。我們將這一堆塊均分為 4 個子集,并通過卷積對中間特征圖進行上采樣或下采樣,以產生不同的尺度(步長為 4、8、16 或 32,與標準 ResNet [25] 相同)。FPN 基于這些多尺度特征圖構建。
為了在不同方法之間進行公平比較,我們為表 4 中的每個條目(包括所有對比方法)搜索超參數。我們搜索的超參數包括學習率、權重衰減、drop path 率和微調輪數。我們將發布代碼及具體配置。有關完整的模型和訓練細節以及更多實驗,請參見 [31]。
A.4. ADE20K 中的語義分割
我們使用 UperNet [57],并按照 [2] 中的語義分割代碼進行訓練。我們進行 100 輪的端到端微調,批量大小為 16。我們為表 5 中的每個條目(包括所有對比方法)搜索最優學習率。
[2] 中的語義分割代碼使用了相對位置偏差 [44]。我們的 MAE 預訓練未使用此項。為了公平比較,我們僅在遷移學習時啟用相對位置偏差,且初始化為零。我們注意到,我們的 BEiT 復現使用了相對位置偏差,并在預訓練和微調階段都啟用了該功能,參照他們的代碼。
A.5. 其他分類任務
對于 iNaturalist 和 Places 的微調(見表 6),我們遵循表 9 中的設置。我們針對每個數據集調整了學習率和微調的輪數。
B. 線性探測結果對比
在 §4.3 中,我們展示了線性探測準確率和微調準確率之間基本上沒有相關性,并且它們對線性可分性的關注點不同。我們注意到,現有的基于掩碼的圖像編碼方法通常在線性探測中不如對比學習等方法有競爭力。為了完整起見,表 12 中對比了基于掩碼的方法與線性探測準確率的比較。
我們使用 ViT-L 的 MAE 達到了 75.8% 的線性探測準確率。這顯著優于以往的基于掩碼的方法。另一方面,在這一協議下,它仍然落后于對比學習方法,例如 MoCo v3 [9] 在 ViT-L 上的線性探測準確率為 77.6%(見圖 9)。
論文名稱:
Masked Autoencoders Are Scalable Vision Learners
論文地址:
https://openaccess.thecvf.com/content/CVPR2022/papers/He_Masked_Autoencoders_Are_Scalable_Vision_Learners_CVPR_2022_paper.pdf