如何擴(kuò)大卷積來消除與Transformer的性能差距

本文分析了是否可以通過策略性地?cái)U(kuò)大卷積來消除與Transformer的性能差距。

首先，先讓我 brainstorm 一下。當(dāng)你看到 neural network scaling 這個(gè)詞的時(shí)候你能想到什么？先不要看下文，把你想到的東西記下來。說不定這個(gè)簡(jiǎn)單的 brainstorm 能讓你找到絕妙的 idea。

我想大多數(shù)人想到的應(yīng)該是模型大小（寬度 + 深度），數(shù)據(jù)大小，或者圖片像素等等。有沒有哪位小科學(xué)家曾經(jīng)想過去 scale convolutional kernels？scale 卷積核同樣能增大模型的參數(shù)，但能帶來像寬度和深度一樣的增益嗎？我這篇文章從這個(gè)角度出發(fā)深入探究了超大卷積核對(duì)模型表現(xiàn)的影響。我發(fā)現(xiàn)現(xiàn)有的大卷積核訓(xùn)練方法的瓶頸：現(xiàn)有的方法都無法無損的將卷積核 scale 到 31x31 以上，更別說用更大的卷積來進(jìn)一步獲得收益。

我這篇文章的貢獻(xiàn)可以總結(jié)為以下幾點(diǎn)：

（1）現(xiàn)有的方法可以將卷積核增大到 31x31。但在更大的卷積上，例如 51x51 和 61x61，開始出現(xiàn)明顯掉點(diǎn)的現(xiàn)象。（2）經(jīng)典的 CNN 網(wǎng)絡(luò)，如 ResNet 和 ConvNeXt，stem cell 都采用了 4× 降采樣。所以對(duì)典型的 224×224 ImageNet 來說，51×51 的極端核已經(jīng)大致等于全局卷積。和全局注意力機(jī)制一樣，我推測(cè)全局卷積也存在著捕捉局部低級(jí)特征的能力不足的問題。（3）基于此觀察，我提出了一套訓(xùn)練極端卷積核的 recipe，能夠絲滑的將卷積核增大到 61x61，并進(jìn)一步提高模型的表現(xiàn)。我的方法論主要是基于人類視覺系統(tǒng)中普遍存在的稀疏性提出來的。在微觀層面上，我將 1 個(gè)方形大卷積核分解為 2 個(gè)具有動(dòng)態(tài)稀疏結(jié)構(gòu)的，平行的長方形卷積核，用來提高大卷積的可擴(kuò)展性；在宏觀層面上，我構(gòu)建了一個(gè)純粹的稀疏網(wǎng)絡(luò)，能夠在提升網(wǎng)絡(luò)容量的情況下保持著和稠密網(wǎng)絡(luò)一樣的參數(shù)和 FLOPs。（4）根據(jù)這個(gè) recipe，我構(gòu)造了一個(gè)新型網(wǎng)絡(luò)結(jié)構(gòu) Sparse Large Kernel Network，簡(jiǎn)稱 SLaK。SLaK 搭載著有史以來最大的 51x51 卷積核，能夠在相似的參數(shù)量和 FLOPs 的條件下，獲得比最新先進(jìn)的 ConvNeXt，Swin Transformer 和 RepLKNet 更好的性能。（5）最后，作者認(rèn)為本文最重要的貢獻(xiàn)是 sparsity，通常作為模型壓縮的“老伙計(jì)”，can be a promising tool to boost neural network scaling。

Pytorch 開源代碼：https://github.com/VITA-Group/SLaK

一、引言

隨著 vison transformer 在各個(gè)領(lǐng)域的大放異彩，CNN 和 Vision Transformer 的競(jìng)爭(zhēng)也愈演愈烈。在愈發(fā)強(qiáng)大的各類 attention 變種的推進(jìn)下，ViT 取代 CNN 這個(gè)視覺老大哥的野心已經(jīng)路人皆知。而 CNN 在全局和局部注意力的啟發(fā)下也帶著大卷積乘風(fēng)破浪回來。前浪有 ConvNeXts 配備著 7x7 卷積核和 swin transformer 精巧的模型設(shè)計(jì)，成功的超越了后者的表現(xiàn)。

后浪中 RepLKNet 用結(jié)構(gòu)再參數(shù)化成功的克服了大卷積在訓(xùn)練上的困難，一度將卷積增大到了 31x31，并達(dá)到了和 Swin Transformer 相當(dāng)?shù)谋憩F(xiàn)。在下游任務(wù)上更是超過了后者。但與 Swin Transformer 等高級(jí) ViT 的擴(kuò)展趨勢(shì)相比，隨著卷積核的持續(xù)擴(kuò)大，大卷積核有著明顯的疲軟趨勢(shì)。

二、超越 31x31 超大卷積核的訓(xùn)練 recipe

本文主要想探究的問題是：是否可以通過采用極致大卷積核（超過 31x31）來進(jìn)一步提高 CNNs 的表現(xiàn)？為了回答這個(gè)問題，我在最近大火的 ConvNeXt 上對(duì)大卷積進(jìn)行了系統(tǒng)的研究。我采用了和 ConvNeXt 一模一樣的訓(xùn)練設(shè)定和超參并將卷積核放大到 31x31，51x51 和 61x61。受限于計(jì)算資源，我這里將模型訓(xùn)練到 120 個(gè) epoch，僅僅用來觀察卷積增大的趨勢(shì)，得到了如下 3 個(gè)主要結(jié)論。

結(jié)論 1：現(xiàn)有的技術(shù)無法無損的將卷積核擴(kuò)展到 31x31 以上

現(xiàn)有的大卷積核技術(shù)主要有兩個(gè)，一是直接暴力 scale up 卷積核的 ConvNeXt，二是增加一個(gè)額外的小卷積層來輔助大卷積的訓(xùn)練，訓(xùn)練完成之后再用結(jié)構(gòu)化再參數(shù)將小卷積核融入大卷積核里，即 RepLKNet。我分別測(cè)試了這兩種方法在極致大卷積上的表現(xiàn)，如下表所示：

原始 ConvNeXt 采用的是 7x7 卷積核，ImageNet 上能達(dá)到 81.0% 的 top1 精度。但是當(dāng)卷積逐漸增大的時(shí)候，ConvNeXt 出現(xiàn)了明顯的掉點(diǎn)。相比之下，RepLKNet 成功的把卷積核增大到 31x31 并帶來了超過 0.5 個(gè)點(diǎn)的可觀提升。但是當(dāng)卷積核增大到 51x51 甚至是 61x61 的時(shí)候，RepLKNet 也逐漸乏力。尤其是在 61x61 上，RepLKNet 的 FLOPs 增加了兩倍，精度卻反而降低了 0.2%。

如果仔細(xì)分析 ConvNeXt 模型的特點(diǎn)，51x51 和 61x61 卷積核帶來的精度下降是可以理解的。如下圖所示，現(xiàn)階段最先進(jìn)的模型的 stem cell 都不約而同的采用了 stride=4 的結(jié)構(gòu)將輸入圖片的分辨率縮減到了原來的 1/4。那么對(duì)經(jīng)典的 224x224 iamgenet 來說，通過 stem cell 之后，feature 的大小就只有 56x56 了。所有 51x51 和 61x61 規(guī)模的卷積核就已經(jīng)是全局水平的卷積核。一種合某些理想的特性，比如有效的局部特性。同樣的現(xiàn)象我在 ViTs 的類似機(jī)制中也觀察到過，即局部注意力通常優(yōu)于全局注意力。在此基礎(chǔ)上，我想到了通過引入局部性來解決這個(gè)問題的機(jī)會(huì)。

結(jié)論 2：用兩個(gè)平行的，長方形卷積來代替方形大卷積可以絲滑的將卷積核擴(kuò)展到 61x61

這里我采用的方法是將一個(gè)常用的 MxM 方形卷積核分解為兩個(gè)平行的 MxN+NxM 長方形卷積核。如下圖最右邊所示。這里經(jīng)驗(yàn)性的設(shè)置 N=5。

這種分解不僅繼承了大卷積捕獲遠(yuǎn)程依賴關(guān)系的能力，而且可以利用短邊來提取局部上下文特征。我選用了兩個(gè) parallel 卷積核相加而不是以往的兩個(gè) sequential 卷積核相疊是因?yàn)橄冗M(jìn)行 MxN 卷積再進(jìn)行 NxM 卷積可能會(huì)因?yàn)?N 過小而丟失一部分長距離的信息（有待于去驗(yàn)證）。果然與預(yù)想的一樣，這樣分解可以讓我逆轉(zhuǎn)大卷積帶來精度下降的趨勢(shì)。由于該分解減少了 FLOPs，相比不分解（即 RepLKNet）在 31x31 卷積上會(huì)犧牲掉少量的精度 （0.2%）。但是，隨著卷積大小增加到全局卷積，它可以驚人地將 Kernel-size 擴(kuò)展到 61x61 并帶來更好的性能指標(biāo)。

更重要的是，隨著卷積核的增大，現(xiàn)有的大卷積訓(xùn)練技術(shù)的內(nèi)存和計(jì)算開銷會(huì)呈現(xiàn)二次方的增長。這種分解方式保持了線性增長的趨勢(shì)并且可以極大的減少大卷積核帶來的開銷。如下圖所示。在中等卷積 31x31 上，參數(shù)量和計(jì)算量基本豆差不多。但是繼續(xù)增大卷積核的時(shí)候，RepLKNet 的計(jì)算量二次方的增長，而我的方法能基本保持不變。不要小看這一點(diǎn)，因?yàn)楝F(xiàn)在已經(jīng)有很多工作指明了一個(gè)明顯的趨勢(shì)：高分辨率訓(xùn)練（Swin Transfermor V2 使用了高達(dá) 1536x1536 的像素）能夠帶來明顯的增益。這種極大分辨率上，51x51 分辨率明顯已經(jīng)不足以去獲得足夠大的感受野。我很可能需要 100 + 的卷積核去獲得足夠大的感受野。

結(jié)論 3：擁有動(dòng)態(tài)稀疏性的卷積核極大的提高了模型的容量同時(shí)又不增加模型大小

最近提出的 ConvNeXt 重新訪問了 ResNeXt 中 “use more groups， expand width” 的準(zhǔn)則，使用增加寬度的 depth-wise 卷積來增加 model capacity。在本文中，我用 dynamic sparsity 進(jìn)一步擴(kuò)展的這一原則，即“use sparse groups， expand more”。關(guān)于 dynamic sparsity 的介紹，請(qǐng)移步 https://zhuanlan.zhihu.com/p/376304225 看我之前的分享。

具體來說，我首先用稀疏卷積代替密集卷積，其中每一層的稀疏度是基于 SNIP 的稀疏比率提前決定的。構(gòu)建完成后，我采用了動(dòng)態(tài)稀疏度方法來訓(xùn)練模型。具體來說就是在模型訓(xùn)練一段時(shí)間后我會(huì)采用參數(shù)剪枝的方法去 prune 掉一部分相對(duì)不重要的參數(shù)，緊接著去隨機(jī)的漲同樣數(shù)量的參數(shù)來保證總體訓(xùn)練參數(shù)的固定。這樣做可以動(dòng)態(tài)地適應(yīng)稀疏權(quán)值，從而獲得更好的局部特征。由于在整個(gè)訓(xùn)練過程中模型都是稀疏的，相應(yīng)的參數(shù)計(jì)數(shù)和訓(xùn)練 / 推理 FLOPs 只與模型的稀疏度成比例。為了評(píng)估，這里以 40% 的稀疏度稀疏化分解后的 kernel，并將其性能報(bào)告為 “稀疏分解” 組?？梢栽诒?2 的中間一列中觀察到，動(dòng)態(tài)稀疏性顯著降低了模型的 FLOPs（超過 2.0G），并導(dǎo)致了暫時(shí)的性能下降。

接下來，我展示了動(dòng)態(tài)稀疏性的 high efficiency 可以有效地轉(zhuǎn)化成 high scalability。例如，使用相同的稀疏性（40%），我可以將模型寬度擴(kuò)展 1.3 倍，但是總體的模型參數(shù)和 FLOPs 卻仍然和稠密網(wǎng)絡(luò)一樣，并顯著漲點(diǎn)。在極端的 51×51 卷積下，性能可以從 80.5% 直接提高到 81.6%。值得注意的是，配備了 61×61 的內(nèi)核之后，我的模型可以超越 RepLKNet 的精度，同時(shí)還節(jié)省了 55% 的 FLOPs。

三、Sparse Large Kernel Network - SLaK

我利用上面發(fā)現(xiàn)的 recipe 在 ConvNeXt 上直接加載了 51x51 的卷積核，得到了 SLaK-T/S/B 模型。在不改變?nèi)魏?ConvNeXt 原有的訓(xùn)練設(shè)置和超參的情況下，SLaK 在 ImageNet-1K 分類數(shù)據(jù)集，ADE20K 分割數(shù)據(jù)集、PASCAL VOC 2007 檢測(cè)數(shù)據(jù)集，都超過了 Swin Transformer，ConvNeXt，和 RepLKNet 的表現(xiàn)。實(shí)驗(yàn)結(jié)果突出了極致卷積核在下游視覺任務(wù)中的關(guān)鍵作用。

四、感受野分析

前面我猜測(cè)新方法既能夠保證對(duì)遠(yuǎn)距離相關(guān)性的獲取，又能夠捕捉到近距離重要的特征。接下來通過對(duì)感受野的分析來證明這種想法。我計(jì)算了輸入圖片的像素對(duì)不同模型決策的貢獻(xiàn)度，并把貢獻(xiàn)度加加回到 1024x1024 的圖片上。我可以發(fā)現(xiàn)原始的 ConvNeXt 用 7X7 卷積只用了中間很小一部分的像素來做決策；RepLKNet 用 31x31 的卷積把感受野擴(kuò)大了許多；而 SLaK 進(jìn)一步用 51x51 的卷積核幾乎達(dá)到了全局感受野。值得注意的在大的感受野之上，能明顯的看到一個(gè)小的正方形堆疊著，證明了我能捕捉小范圍的低級(jí)特征的猜想。