RevCol：大模型架構(gòu)設(shè)計新范式，給神經(jīng)網(wǎng)絡(luò)架構(gòu)增加了一個維度！

01

總  述

這樣的架構(gòu)方案使RevCol的行為與傳統(tǒng)網(wǎng)絡(luò)大不相同：在前向傳播過程中，RevCol中的特征在通過每一列時都會逐漸解開，其總信息會被保留，而不是像其他網(wǎng)絡(luò)那樣被壓縮或丟棄。

實驗表明，CNN風(fēng)格的RevCol模型可以在圖像分類、目標檢測和語義分割等多項計算機視覺任務(wù)上取得非常有競爭力的性能，尤其是在參數(shù)預(yù)算大、數(shù)據(jù)集大的情況下。例如，在ImageNet-22K預(yù)訓(xùn)練后，RevCol-XL在ImageNet-1K上獲得88.2%的準確度。給定更多的預(yù)訓(xùn)練數(shù)據(jù)，最大的模型RevCol-H在ImageNet-1K上達到90.0%，在COCO檢測最小值集上達到63.8%的APbox，在ADE20k分割上達到61.0%的mIoU。

據(jù)所知，這是純（靜態(tài)）CNN模型中最好的COCO檢測和ADE20k分割結(jié)果。此外，作為一種通用的宏架構(gòu)方式，RevCol還可以引入到Transformer或其他神經(jīng)網(wǎng)絡(luò)中，這被證明可以提高計算機視覺和NLP任務(wù)中的性能。

02 背景&動機

Information Bottleneck principle（IB）統(tǒng)治著深度學(xué)習(xí)世界?？紤]一個典型的監(jiān)督學(xué)習(xí)網(wǎng)絡(luò)，如下圖a所示：靠近輸入的層包含更多的低級信息，而靠近輸出的特征具有豐富的語義。

換句話說，在逐層傳播期間，與目標無關(guān)的信息被逐漸壓縮。盡管這種學(xué)習(xí)范式在許多實際應(yīng)用中取得了巨大成功，但從特征學(xué)習(xí)的角度來看，它可能不是最佳選擇——如果學(xué)習(xí)到的特征被過度壓縮，或者學(xué)習(xí)到的語義信息與目標任務(wù)無關(guān)，則下游任務(wù)的性能可能較差，特別是在源任務(wù)和目標任務(wù)之間存在顯著的領(lǐng)域差距的情況下。研究人員付出了巨大的努力，使學(xué)習(xí)到的特征更加普遍適用，例如通過自監(jiān)督的預(yù)訓(xùn)練或多任務(wù)學(xué)習(xí)。

在今天分享中，研究者主要關(guān)注另一種方法：建立一個網(wǎng)絡(luò)來學(xué)習(xí)解耦表征。與IB學(xué)習(xí)不同，去解耦特征學(xué)習(xí)不打算提取最相關(guān)的信息，而丟棄不太相關(guān)的信息；相反，它旨在將與任務(wù)相關(guān)的概念或語義詞分別嵌入到幾個解耦的維度中。同時，整個特征向量大致保持與輸入一樣多的信息。這與生物細胞中的機制非常相似，每個細胞共享整個基因組的相同拷貝，但具有不同的表達強度。因此，在計算機視覺任務(wù)中，學(xué)習(xí)解耦特征也是合理的：例如，在ImageNet預(yù)訓(xùn)練期間調(diào)整高級語義表示，同時，在目標檢測等下游任務(wù)的需求下，還應(yīng)在其他特征維度上保持低級信息（如邊緣的位置）。

上圖（b）概述了主要想法：RevCol，它在很大程度上受到了GLOM大局的啟發(fā)。網(wǎng)絡(luò)由結(jié)構(gòu)相同（但其權(quán)重不一定相同）的N個子網(wǎng)絡(luò)（命名列）組成，每個子網(wǎng)絡(luò)接收輸入的副本并生成預(yù)測。因此，多級嵌入，即從低級到高度語義表示，存儲在每一列中。此外，引入可逆變換以在沒有信息損失的情況下將多級特征從第i列傳播到第（i+1）列。在傳播過程中，由于復(fù)雜性和非線性增加，預(yù)計所有特征級別的質(zhì)量都會逐漸提高。因此，最后一列（圖中的第N列）1（b））預(yù)測了輸入的最終解耦表示。

03 新框架

接下來，我們將介紹RevCol的設(shè)計細節(jié)。上圖b說明了頂層體系結(jié)構(gòu)。請注意，對于RevCol中的每一列，為了簡單起見，直接重用ConvNeXt等現(xiàn)有結(jié)構(gòu)，因此在以下主要關(guān)注如何構(gòu)建列之間的可逆連接。此外，在每個列的頂部引入了即插即用的中間監(jiān)督，這進一步提高了訓(xùn)練收斂性和特征質(zhì)量

MULTI-LEVEL REVERSIBLE UNIT

在新提出的網(wǎng)絡(luò)中，可逆變換在不損失信息的情況下對特征解耦著關(guān)鍵作用，其見解來自可逆神經(jīng)網(wǎng)絡(luò)。其中，首先回顧了RevNet的一部代表作。如下圖a所示，RevNet首先將輸入x劃分為兩組，x0和x1。

然后，對于后面的塊，例如塊t，它將前面兩個塊的輸出xt?1和xt?2作為輸入，并生成輸出xt。塊t的映射是可逆的，即xt?2可以由兩個后驗塊xt?1和xt重構(gòu)。形式上，正向和反向計算遵循方程?：

為了解決上面提及的問題，將上面等式概括為以下形式：

因此，可以將上面方程重組為多列形式，如下圖b所示。每列由一組內(nèi)的m個特征圖及其母網(wǎng)絡(luò)組成。將其命名為多級可逆單元，這是RevCol的基本組成部分。

REVERSIBLE COLUMN ARCHITECTURE

• 宏觀設(shè)計
  

上圖2c闡述了框架設(shè)計。按照最近模型的常見做法，首先通過補丁嵌入模塊將輸入圖像分割成不重疊的補丁。然后，將補丁饋送到每個子網(wǎng)絡(luò)（列）中。列可以用任何傳統(tǒng)的單列架構(gòu)來實現(xiàn)，例如ViT或ConvNeXt。從每一列中提取四級特征圖，以在列之間傳播信息；例如，如果列是用廣泛使用的層次網(wǎng)絡(luò)實現(xiàn)的，可以簡單地從每個階段的輸出中提取多分辨率特征。

對于分類任務(wù)，只使用最后一列中最后一級（第4級）的特征圖來獲取豐富的語義信息。

對于其他下游任務(wù)，如目標檢測和語義分割，在最后一列中使用所有四個級別的特征圖，因為它們包含低級和語義信息。

• 微觀設(shè)計

每一個層級中，先用一個Fusion單元把不同尺寸的輸入調(diào)整到同一個shape，后面再經(jīng)過一堆ConvNeXt Blocks得到輸出，這些就是公式中的Ft(·)，然后再和Reversible operation的輸入加和得到最終結(jié)果。

值得注意的是，把原本的ConvNeXt block中7x7的kernel size改成了3x3，大kernel的收益在Revcol上有限，但小kernel特別快。

04

實驗

除了2B參數(shù)的模型以外，還收集了168Million的私有數(shù)據(jù)集，weakly-label的標簽用來預(yù)訓(xùn)練。XL模型(800M param)，在22k下能達到88.2%，經(jīng)過Megdata-168M的訓(xùn)練后能夠漲到89.4%。Huge 224 pretrain，640x640 Finetune能夠達到90.0%Top-1 Accuracy。這個模型的訓(xùn)練開銷：預(yù)訓(xùn)練總共1600個ImageNet Epochs，訓(xùn)練一次使用80塊A100，14天。