Title: ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders

單位:?KAIST, FAIR(陳鑫磊, 劉壯等人), 紐約大學(謝賽寧)

Paper: https://arxiv.org/abs/2301.00808

Github: https://github.com/facebookresearch/ConvNeXt-V2

引言

受 MAE 的啟發(fā)，本文在 ConvNeXt 的架構(gòu)基礎(chǔ)上延伸出了一個完全卷積掩碼自編碼器框架——ConvNeXt V2，同時作者設(shè)計了一個全新的全局響應(yīng)歸一化(Global Response Normalization, GRN)層以增強原始 ConvNeXt 模塊通道間的特征競爭，從而捕獲更具有判別力的通道特征。

ConvNeXt V2 最終在各種識別基準上的性能，包括 ImageNet 分類、COCO 檢測和 ADE20K 分割任務(wù)上均取得了極具競爭力的結(jié)果，其中最小的模型僅 3.7M 參數(shù)可獲取 76.7% top-1 準確率，而最大的模型約 650M 參數(shù)則能達到 88.9% 準確率。

前情回顧

ConvNeXt

先來看下經(jīng)典卷積神經(jīng)網(wǎng)絡(luò)的發(fā)展史：

緊接著 ViT 的提出帶火了一眾基于 Transformer 的模型：

然后便是 Swin Transformer received ICCV 2021 best paper award (Marr Prize)，一時間轟動華語樂壇視覺領(lǐng)域：

可以看出，Swin 其實是將傳統(tǒng) CNNs 三個重要的先驗信息融入到了 Vision Transformer 的設(shè)計：

帶局部窗口的注意力設(shè)計；

共享權(quán)重；

分層特征；

另一方面，我們從各大CV頂會接收的論文情況也能看出一些端倪：

可以發(fā)現(xiàn)，基于各種 **Former 的模型遍地開花，而基于傳統(tǒng)的卷積神經(jīng)網(wǎng)絡(luò)幾乎沒有增長。當然，除了 idea 新穎外，基于 Transformer 的模型效果也確實驚艷，例如：

那么，這兩者的本質(zhì)區(qū)別在哪里呢？

首先，Transformer 模型由于缺乏固有的歸納偏置屬性，這也意味著在小數(shù)據(jù)量的情況下它是很難“硬train一發(fā)”，因此大部分 Vision Transformer 模型訓練的時候都會加入一堆 DataAug 和 Training Tricks，如果我們將這些 Tricks 應(yīng)用到原始的 ResNet 上也能獲得不錯的性能增益：

好了，讓我們看下 ConvNeXt 的設(shè)計思路是怎么樣的：

整體設(shè)計遵循由粗粒度到細粒度的原則，可以看到，通過最終一步步的實驗論證，ConvNeXt 的精度也隨即提升，最終獲得了完全可以媲美 Vision Transformer 的結(jié)果：

下面展示的是不同模塊之間的結(jié)構(gòu)對比圖：

總的來說，ConvNeXt 是一個基于純 ConvNets 的網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計，其通過借鑒 Vision Transformer 的一些優(yōu)化技巧，在 ResNet 的基礎(chǔ)上一步步扎實的進行實驗論證，最終獲得了一個極具實際應(yīng)用價值的網(wǎng)絡(luò)結(jié)構(gòu)模型，最終實現(xiàn)了優(yōu)于 Swin Transformers 的性能，同時保持了標準 ConvNets 的簡單性和效率，非常值得大家伙借鑒。

MAE

MAE 想必大家伙都認識了，畢竟CV界盛傳的一句經(jīng)典諺語——“凱明出品，必屬精品”！自編碼器一開始是應(yīng)用在 NLP 領(lǐng)域，例如 BERT 模型便是被提出應(yīng)用在句子的不同部分屏蔽單詞，并嘗試通過預(yù)測要填充到空白處的單詞來重建完整的句子。因此自然而然的一個想法便是是否可以將這一想法嫁接到計算機視覺領(lǐng)域。

作為CV界的頭號種子選手，老何當然是不負眾望，憑借著多年來異常敏銳的嗅覺，不費吹灰之力便提出了一種使用自編碼器對計算機視覺模型進行自監(jiān)督預(yù)訓練的新方法：

在 MAE 中，作者提倡使用非常高的遮擋率，例如 75%，這能帶來兩個好處：

訓練速度提升3倍，因為模型僅需處理更少(unmasked)的圖像塊；

強迫模型學習出泛化能力，提高特征提取的能力；

此處，這是一個非對稱的編解碼器，這意味著其應(yīng)用了一個相當輕量的解碼器，而編碼器的輸出便是輸入圖像塊的潛在向量表示。而 MAE 的關(guān)鍵步驟之一便是利用解碼器重建初始圖像，其中每個掩碼標記都是一個共享的學習向量，表示存在缺失的補丁。解碼器則接收潛在表示以及掩碼標記作為輸入，并輸出每個補丁的像素值，包括掩碼。根據(jù)此信息，可以將原始圖像拼湊在一起，以根據(jù)用作輸入的掩碼圖像形成完整圖像的預(yù)測版本。

一旦重建了目標圖像，我們只需測量它與原始輸入圖像的差異并將其用作損失執(zhí)行反向傳播即可。最終，當模型訓練完成后，解碼器將會被丟棄掉，只保留編碼器部分，以供下游任務(wù)使用。

本文方法

動機

上面我們提到過，ConvNeXt 架構(gòu)是對傳統(tǒng)的 ConvNet 進行的一次現(xiàn)代化改造，作者證明了純卷積模型也可以成為像 Vision Transformer 一樣的可擴展架構(gòu)。然而，探索神經(jīng)網(wǎng)絡(luò)架構(gòu)設(shè)計空間的最常用方法仍然是通過在 ImageNet 上對監(jiān)督學習性能進行基準測試。

近年來，自監(jiān)督學習也是非常熱門，相比之下，自監(jiān)督學習不需要任何人工創(chuàng)建的標簽。顧名思義，模型學會自我監(jiān)督。在計算機視覺中，對這種自監(jiān)督進行建模的最常見方法是對圖像進行不同裁剪或?qū)ζ鋺?yīng)用不同的增強，并將修改后的輸入傳遞給模型。即使圖像包含相同的視覺信息但看起來不一樣，我們讓模型學習這些圖像仍然包含相同的視覺信息，即相同的對象。這導致模型為相同的對象學習相似的潛在表示（輸出向量）。特別地，MAE 的提出成功地將掩碼語言建模成功帶到了視覺領(lǐng)域，并迅速成為視覺表示學習的流行方法。然而，自監(jiān)督學習中的一種常見做法是使用為監(jiān)督學習設(shè)計的預(yù)定架構(gòu)，并假設(shè)模型架構(gòu)是固定的。例如，MAE 是基于 ViT 開發(fā)的。

既然 ConvNeXt 和 MAE 有這么大的優(yōu)勢，那么簡單將 MAE 的思想套在 ConvNeXt 上會有什么問題呢？首先，正如上述所說，MAE 有一個特定的編解碼器設(shè)計，是專門針對 Transformer 的序列處理能力進行了優(yōu)化，這使得計算量大的編碼器可以專注于可見的補丁，從而降低預(yù)訓練成本。此設(shè)計可能與使用密集滑動窗口的標準 ConvNet 不兼容。此外，如果不考慮體系結(jié)構(gòu)和訓練目標之間的關(guān)系，則很大可能獲取不到最佳性能。事實上，相關(guān)研究表明，使用基于掩碼的自監(jiān)督學習訓練 ConvNets 是比較困難的，并且 Transformers 和 ConvNets 本質(zhì)上提取到的特征本身就沒有良好的兼容性。

因此，本文的重點便是探討如何在同一框架下共同設(shè)計網(wǎng)絡(luò)架構(gòu)和掩碼自編碼器，目的是使基于掩碼的自監(jiān)督學習對 ConvNeXt 模型有效，并獲得與使用 Transformer 相當?shù)男阅堋?/p>

具體地，在設(shè)計 MAE 時，作者將 mask 輸入視為一組稀疏補丁，同時使用稀疏卷積僅處理可見部分，這點跟原始的 MAE 保持一致。這個想法的靈感來自于在處理大規(guī)模 3D 點云時所采用的稀疏卷積一致。在實踐中，我們可以用稀疏卷積實現(xiàn) ConvNeXt，在微調(diào)時，權(quán)重被轉(zhuǎn)換回標準的密集層，不需要特殊處理。為了進一步提高預(yù)訓練效率，作者應(yīng)用了單個 ConvNeXt 塊替換了原始的基于 Transformer 的解碼器，使整個設(shè)計完全卷積化。最后，作者對 ConvNeXt 的不同訓練配置進行特征空間分析。當直接在掩碼輸入上訓練 ConvNeXt 時，會發(fā)現(xiàn) MLP 層的潛在特征崩潰問題。為了解決這個問題，本文額外設(shè)計了一個全新的 Global Response Normalization 層來鼓勵通道間的特征競爭。

FCMAE

FCMAE Framework

上圖左半部分展示了一種全卷積掩碼自編碼器 (FCMAE)。它是由一個基于稀疏卷積的 ConvNeXt 編碼器和一個輕量級的 ConvNeXt 塊解碼器組成。整體來看，這是一個非對稱的編解碼器。其中編碼器僅處理可見像素，解碼器則使用編碼像素和掩碼標記重建圖像，損失僅在 mask 區(qū)域上計算，與 MAE 保持一致。大家可以明顯的看出本文方法在概念上很簡單，并且以純卷積的方式運行。學習信號是通過以高屏蔽率隨機屏蔽原始輸入視覺效果并讓模型在給定剩余上下文的情況下預(yù)測缺失部分來生成的。下面讓我們拆解下整個網(wǎng)絡(luò)結(jié)構(gòu)具體看看作者作出了什么改進和優(yōu)化。

Masking

首先，本文使用 mask ratio 為 0.6 的隨機 mask 策略，相對之下 MAE 是建議使用 0.75。由于卷積模型具有分層設(shè)計，其特征會在不同階段被下采樣，因此掩碼在最后階段生成并遞歸上采樣到最精細的分辨率。為了在實踐中實現(xiàn)這一點，作者從原始輸入圖像中隨機刪除 60% 的 32 × 32 patch，同時使用最少的數(shù)據(jù)增強，僅包括隨機調(diào)整大小的裁剪策略。

Encoder design

ConvNeXt V2 顧名思義是使用 ConvNeXt 模型作為編碼器。使 mask 圖像建模有效的一個挑戰(zhàn)是防止模型學習允許它從 mask 區(qū)域復(fù)制和粘貼信息的快捷方式。這在基于 Transformer 的模型中相對容易防止，它可以將可見補丁作為編碼器的唯一輸入。然而，使用 ConvNets 其實是比較難實現(xiàn)這一點，因為必須保留 2D 圖像結(jié)構(gòu)。一種最簡單的解決方案是在輸入端引入可學習的掩碼標記，如 BEiT 和 Simmim，但這些方法降低了預(yù)訓練的效率并導致訓練和測試時間不一致，因為在測試時沒有掩碼標記。特別是當遮擋率過高時問題便會尤其凸顯。

為了解決這個問題，本文借鑒在 3D 任務(wù)中學習稀疏點云的道理，將 mask 圖表示為一個二維的稀疏像素陣列。因此，一種自然而然的想法便是引入稀疏卷積，以促進 MAE 的預(yù)訓練。在具體的代碼實現(xiàn)中，我們可以將標準卷積層轉(zhuǎn)換為稀疏卷積，這使得模型可以僅對可見數(shù)據(jù)點進行操作，而在 fine-tune 階段，完全可以轉(zhuǎn)換為標準卷積，而無需額外處理。作為替代方案，也可以在密集卷積運算之前和之后應(yīng)用 binary masking operation。此操作在數(shù)值上與稀疏卷積等價，但理論上計算量更大些，不過在 TPU 等 AI 加速器上更友好。

Decoder design

同 MAE 一致，ConvNeXt V2 也采用輕量級的解碼器設(shè)計。其次，作者還考慮了更復(fù)雜的解碼器，例如分層解碼器如 FPN 和 U-Net 或 ViT 和 MAE，不過最終的實驗表明更簡單的單個 ConvNeXt 塊解碼器效果其實就很不錯了，而且還可以顯著減少預(yù)訓練時間，本文將將解碼器的維度設(shè)置為 512。

Reconstruction target

這里目標重構(gòu)的方式同 MAE 類似，也是采用 MSE 來衡量損失，目標是原始輸入的 patch-wise 歸一化圖像，并且損失僅應(yīng)用于 mask 過后的 patch。在具體的訓練步驟中，作者基于 ImageNet-1K 數(shù)據(jù)集進行了 800 輪 epochs 的預(yù)訓練并額外微調(diào)了 100 輪 epochs。

Global Response Normalization

如上所述，當直接在 mask 輸入上訓練 ConvNeXt 時，會導致特征崩潰的問題。為此，作者引入了一種全新的全局響應(yīng)歸一化層，以結(jié)合 ConvNeXt 架構(gòu)使 FCMAE 預(yù)訓練更加有效，下面讓我們介紹下具體細節(jié)。

Feature collapse

為了更深入地了解學習行為，我們首先在特征空間中進行定性分析。具體的，可以將 FCMAE 預(yù)訓練的 ConvNeXt-Base 模型的激活可視化，由此可以觀察到一個有趣的“特征崩潰”現(xiàn)象：即存在許多飽和的特征圖，并且激活在通道之間變得很多余。如下圖所示，這些特征圖可以通過可視化 ConvNeXt 塊中的維度擴展 MLP 層中觀察到：

Feature cosine distance analysis

上圖展示了定量分析結(jié)果。為了執(zhí)行此分析，本文隨機選擇 ImageNet-1K 驗證集中不同類別的 1,000 張圖像，并從不同模型的每一層提取高維特征，包括 FCMAE 模型、ConvNeXt 監(jiān)督模型和 MAE 預(yù)訓練 ViT 模型。隨后作者計算了每個圖像的每層距離，并對所有圖像像素值求平均。從給出的結(jié)果圖可以看出，F(xiàn)CMAE 預(yù)訓練的 ConvNeXt 模型表現(xiàn)出明顯的特征崩潰趨勢，這與我們從之前的激活可視化中觀察到的結(jié)果一致。這促使作者進一步考慮在學習過程中使特征多樣化并防止特征崩潰的方法。

Approach

在這項工作中，全局響應(yīng)歸一化的引入主要旨在增加通道的對比度和選擇性。具體地，給定一個輸入特征 ，所提出的 GRN 單元包括三個步驟：

全局特征聚合;

特征歸一化；

特征校準；

下面給出偽代碼示例：

通過將 GRN 層整合到原始的 ConvNeXt 塊中，我們可以得到下面的結(jié)構(gòu)：

此外，作者發(fā)現(xiàn)，當應(yīng)用 GRN 時，LayerScale 其實是沒有必要的，因此在此版本中將其刪掉。使用這種新的模塊設(shè)計，我們可以創(chuàng)建具有不同效率和容量的各種模型，分別適用于不同的應(yīng)用場景。

Impact of GRN

從圖 3 中的可視化和圖 4 中的余弦距離分析，我們可以觀察到 ConvNeXt V2 有效地緩解了特征崩潰問題。其中余弦距離值一直很高，這表明跨層機制保持了特征的多樣性。這種行為類似于 MAE 預(yù)訓練的 ViT 模型?？傮w而言，這表明在類似的掩碼圖像預(yù)訓練框架下，ConvNeXt V2 的學習行為可以類似于 ViT。

此外，當配備 GRN 時，F(xiàn)CMAE 預(yù)訓練模型可以顯著優(yōu)于 300 個 epoch 的監(jiān)督模型。GRN 通過增強特征多樣性來提高表示質(zhì)量，這在原始的 ConvNeXt 模型中是不存在的，但已被證明對于基于掩碼的預(yù)訓練至關(guān)重要。值得一提的是，這種改進是在不增加額外參數(shù)開銷或增加 FLOPS 的情況下實現(xiàn)的。

Relation to feature normalization methods

如上表所示，作者將 GRN 與三個廣泛使用的歸一化層進行了比較，即局部響應(yīng)歸一化(LRN)、批量歸一化(BN)和層歸一化(LN)。從結(jié)果可以觀察到只有 GRN 可以顯著優(yōu)于 baseline，這是由于：

LRN 缺乏全局上下文，因為它只對比附近領(lǐng)域的通道；

BN 沿批處理軸在空間上歸一化，這不適用于 mask 輸入；

LN 則通過全局均值和方差標準化隱含地鼓勵特征競爭，但效果顯然不如 GRN。

Relation to feature gating methods

另一種增強神經(jīng)元間競爭的方法是使用動態(tài)特征門控方法，例如 SE 和 CBAM，一個注重通道，一種注重空間。這兩個模塊都可以增加單個通道的對比度，類似于 GRN 所做的。不過 GRN 顯然更簡單、更高效，因為它不需要額外的參數(shù)層（例如 FC 層）。

The role of GRN in pre-training/fine-tuning

最后我們可以看下 GRN 在預(yù)訓練和微調(diào)中的重要性。從表格最后一欄可以看到，無論我們是從微調(diào)中刪除 GRN，還是僅在微調(diào)時添加新初始化的 GRN，都可以觀察到模型性能顯著下降，這表明在預(yù)訓練和微調(diào)中使用 GRN 層是非常有必要的。

實驗

參數(shù)設(shè)置

Pre-training setting

End-to-end IN-1K fine-tuning setting for Atto (A), Femto (F), Pico (P) and Nano (N) models

End-to-end IN-1K fine-tuning setting for Tiny model

End-to-end IN-1K fine-tuning setting for Base (B), Large (L), and Huge (H) models

End-to-end IN-22K intermediate fine-tuning settings

End-to-end IN-1K fine-tuning settings (after IN-22K intermediate fine-tuning)

協(xié)同設(shè)計

可以發(fā)現(xiàn)，在不修改模型架構(gòu)的情況下使用 FCMAE 框架對表示學習質(zhì)量的影響有限。同樣，新的 GRN 層對監(jiān)督設(shè)置下的性能影響很小。然而，如果我們將兩者結(jié)合起來使用可以令微調(diào)性能顯著提高。

與 SOTA 方法對比

總結(jié)

在本文中，我們介紹了一個名為 ConvNeXt V2 的新 ConvNet 模型系列，一種更適合自監(jiān)督學習而設(shè)計的新型網(wǎng)絡(luò)架構(gòu)。通過使用全卷積掩碼自編碼器預(yù)訓練，可以顯著提高純 ConvNets 在各種下游任務(wù)中的性能，包括 ImageNet 分類、COCO 目標檢測和 ADE20K 分割。

編輯：黃飛

?