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

01

總  述

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

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

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

02 背景&動(dòng)機(jī)

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

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

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

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

03 新框架

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

MULTI-LEVEL REVERSIBLE UNIT

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

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

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

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

REVERSIBLE COLUMN ARCHITECTURE

• 宏觀設(shè)計(jì)
  

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

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

對(duì)于其他下游任務(wù)，如目標(biāo)檢測(cè)和語(yǔ)義分割，在最后一列中使用所有四個(gè)級(jí)別的特征圖，因?yàn)樗鼈儼图?jí)和語(yǔ)義信息。

• 微觀設(shè)計(jì)

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

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

04

實(shí)驗(yàn)

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