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

本文分析了是否可以通過策略性地擴大卷積來消除與Transformer的性能差距。

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

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

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

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

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

一、引言

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

接下來，我展示了動態(tài)稀疏性的 high efficiency 可以有效地轉(zhuǎn)化成 high scalability。例如，使用相同的稀疏性（40%），我可以將模型寬度擴展 1.3 倍，但是總體的模型參數(shù)和 FLOPs 卻仍然和稠密網(wǎng)絡(luò)一樣，并顯著漲點。在極端的 51×51 卷積下，性能可以從 80.5% 直接提高到 81.6%。值得注意的是，配備了 61×61 的內(nèi)核之后，我的模型可以超越 RepLKNet 的精度，同時還節(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 檢測數(shù)據(jù)集，都超過了 Swin Transformer，ConvNeXt，和 RepLKNet 的表現(xiàn)。實驗結(jié)果突出了極致卷積核在下游視覺任務(wù)中的關(guān)鍵作用。

四、感受野分析

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