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

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

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

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

引言

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

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

前情回顧

ConvNeXt

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

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

然后便是 Swin Transformer received ICCV 2021 best paper award (Marr Prize)，一時(shí)間轟動(dòng)華語(yǔ)樂(lè)壇視覺(jué)領(lǐng)域：

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

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

共享權(quán)重；

分層特征；

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

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

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

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

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

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

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

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

MAE

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

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

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

訓(xùn)練速度提升3倍，因?yàn)槟Ｐ蛢H需處理更少(unmasked)的圖像塊；

強(qiáng)迫模型學(xué)習(xí)出泛化能力，提高特征提取的能力；

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

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

本文方法

動(dòng)機(jī)

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

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

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

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

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

FCMAE

FCMAE Framework

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

Masking

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

Encoder design

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

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

Decoder design

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

Reconstruction target

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

Global Response Normalization

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

Feature collapse

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

Feature cosine distance analysis

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

Approach

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

全局特征聚合;

特征歸一化；

特征校準(zhǔn)；

下面給出偽代碼示例：

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

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

Impact of GRN

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

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

Relation to feature normalization methods

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

LRN 缺乏全局上下文，因?yàn)樗粚?duì)比附近領(lǐng)域的通道；

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

LN 則通過(guò)全局均值和方差標(biāo)準(zhǔn)化隱含地鼓勵(lì)特征競(jìng)爭(zhēng)，但效果顯然不如 GRN。

Relation to feature gating methods

另一種增強(qiáng)神經(jīng)元間競(jìng)爭(zhēng)的方法是使用動(dòng)態(tài)特征門(mén)控方法，例如 SE 和 CBAM，一個(gè)注重通道，一種注重空間。這兩個(gè)模塊都可以增加單個(gè)通道的對(duì)比度，類(lèi)似于 GRN 所做的。不過(guò) GRN 顯然更簡(jiǎn)單、更高效，因?yàn)樗恍枰~外的參數(shù)層（例如 FC 層）。

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

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

實(shí)驗(yàn)

參數(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è)計(jì)

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

與 SOTA 方法對(duì)比

總結(jié)

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

編輯：黃飛

?