CONTAINER: 基于信息聚合，統一視角深度剖析視覺網絡 | 論文解讀

關于機器學習和計算神經科學的全球人工智能頂會NeurIPS（神經信息處理系統大會）于12月6日—14日正式召開。據大會官方統計，NeurIPS 2021共收到9122篇投稿，其中2334篇論文被接收，接收率為26%。上海人工智能實驗室多篇論文入選，本期與讀者分享論文解讀：《Container: Context Aggregation Network》。該論文對Transformer、深度卷積以及MLP-Mixer提供了一個更廣義的統一視角：它們均是更廣義視角下通過神經網絡聚合空間信息的特例。由此，文中提出了一種新穎的模塊CONTAINER(Context Aggregation Network)，它通過可學習參數和響應的架構混合使用了靜態與動態親和度矩陣(Affinity Matrix)，并在圖像分類任務中表現出了很強的結果。此外，在CONTAINER的基礎上，文中還提出了一種更加高效的擴展CONTAINER-LIGHT，該方案在目標檢測與實例分割方面取得了顯著的性能提升。

論文鏈接：https://arxiv.org/pdf/2106.01401.pdf

代碼鏈接：https://github.com/gaopengcuhk/Container

1. Introduction

在計算機視覺領域，Transformer結構已成為新范式，用于在時間（例如視頻理解^[1]）和空間上(例如目標檢測^[2])的長輸入信息聚合。最近有研究表明，精心設計的多層感知器網絡(MLP-mixers)^[3]在圖像分類方面也同樣有效。目前，從表面來看，CNNs、Vision Transformers (ViTs^[4])和MLP-mixers都是不同的結構，但是進一步分析，這些網絡內在呈現出相似的設計。本文旨在提供一種統一視角，表明常見的CNNs、ViTs和MLP-mixers實際上是在神經網絡模塊聚合空間信息的通用方法的特例。

文中提出目前主流網絡都是由神經網絡模塊級聯而成，一個神經網絡模塊則可以分為聚合模塊（Aggregation Module）與混合模塊（Fusion Module）。聚合模塊在模塊輸入的預定義內容窗口進行信息共享與累計（比如Transformer中的自注意力機制），而混合模塊結合位置特征并產生模塊的輸出（比如ResNet^[5]中的前向層）。本文歸納了網絡結構最主要的區別來源于信息聚合模塊的不同，這些差異實際上可以表征為聚合器內的親和度矩陣（Affinity Matrix）的變體，該矩陣用于確定查詢向量與其內容之間的信息傳播。舉例來說，ViTs中的親和度矩陣由key和query動態生成，而Xception^[6]網絡中的depthwise卷積的親和度矩陣則是靜態的，與輸入位置無關并在輸入中保持一致。

基于以上，文中用統一視角提出一種利用多頭信息聚合機制的普適網絡模塊，命名為CONTAINER容器網絡（Context AggregatIon Network）。一個CONTAINER模塊同時包含了基于靜態與動態的親和度的聚合，利用可學習的混合系數將靜態與動態信息混合。CONTAINER模塊十分簡單而且有效，可以方便地插入如今的復雜網絡設計中。在此基礎上，論文中還簡化了CONTAINER網絡，提出一種更加高效的輕型網絡，命名為CONTAINER-LIGHT。

2. CONTAINER

1）視覺信息聚合

文中首先提供當前神經網絡中常用的領域信息聚合模塊的統一視角。

考慮一個輸入圖片為   ,   和  表示輸入圖片的通道和空間維度。圖片可以展開成：

常見的視覺網絡都可以看作一些網絡組件利用殘差進行堆疊，表示為：

其中  表示信息是如何從X聚合來計算某個位置的特征。定義一個親疏度矩陣   來表示代表信息聚合的鄰域，因此，上式可重寫為：

   和   都是可學習的參數。這種信息聚合模塊的建模能力可以通過引入多個親疏矩陣來增加，允許網絡有多個途徑來獲取  的信息。用  表示的切片，是親疏關系矩陣的個數，也可以表示為頭的個數，因此上式可以表示為多頭表達式的形式：  

   表示每個頭的親疏關系矩陣。與單頭機制相比，不同的  可以潛在地捕獲特征空間內的不同關系，從而提高聚合相關信息的能力。需要注意的是，在聚合信息時只有空間信息才會傳播，不發生跨通道信息交換，也不使用非線性激活函數。

2）Transformer、Depthwise卷積和MLP-Mixer

親疏度矩陣是由輸入影響，動態產生，從而可以更好地捕捉到物體層面的信息。但是由于復雜度是次方級的，所以在高分辨率時需要很大的計算資源。

如圖1所示，本節介紹如何使用上述介紹的信息聚合機制來統一這三種網絡結構。

Transformer: 

在Transformers中的自注意力機制，親和度矩陣是由投影的查詢向量與鍵值對之間的相似性建模，在M個頭的情況下，矩陣可寫為： 

親疏度矩陣是由輸入影響，動態產生，從而可以更好地捕捉到物體層面的信息。但是由于復雜度是次方級的，所以在高分辨率時需要很大的計算資源。

圖1 三種信息聚合機制

Depthwise卷積：

Depthwise卷積是一種組卷積的特例。假設信息聚合模塊的頭的個數等于depthwise卷積中的通道數量，可以得到1d的卷積化的親疏度關系矩陣  : 

與自注意力機制的相比，這里的親疏度關系矩陣是靜態的，與輸入特征無關，只是關系到局部聯系并且在全圖共享。

MLP-Mixer: 

MLP-Mixer的核心在于轉置的多層線性層操作，可以寫成： 

關系矩陣可以定義為： 

這個簡單的等式表明，轉置的MLP算子是具有密集關系矩陣的單個特征組上的信息聚合算子。與自注意力和depthwise卷積相比，轉置MLP矩陣是靜態的、密集的并且沒有參數共享。

圖2 CONTAINER動態與靜態信息融合

3）CONTAINER模塊

根據上面的介紹，以往的操作都是只利用動態或者靜態產生的親疏度關系，動態和靜態的矩陣都有各自的優點。本文中的CONTAINER模塊綜合動態與靜態的優點，如圖2所示，公式可表示為: 

  是表示關系矩陣  從輸入  動態產生，而單個的表示是靜態矩陣，對于動態和靜態不同的組合，可以得到不同的網絡模塊。

當  : 是自注意力機制。

當  ：是depthwise卷積模塊。當   時，可以看作是一個多頭depthwise卷積（MH-DW）。

當  ：是MLP-Mixer模塊，當  時，命名為多頭MLP（MH-MLP）。MH-MLP將通道劃分成M組，在每組上分別進行轉置MLP來抓取靜態編碼信息。

當  ：這種模塊混合動態和靜態信息，靜態信息由MLP-Mixer呈現，命名為CONTAINER-PAM(Pay Attention to MLP)。

當  ：這種模塊混合了動態信息和depthwise卷積的靜態信息。這種靜態關系矩陣包含了旋轉不變的局部約束，讓這種結構更適合視覺任務。

CONTAINER模塊非常容易實現，可以輕松地嵌入到現有的神經網絡中。

4）CONTAINER網絡結構

文中的基本網絡結構包含4個階段，每個階段都會不斷的降采樣圖像的分辨率。在每個階段，都包含了模塊的級聯，每個模塊都由空間信息聚合和通道信息混合構成。文中固定使用兩層MLP網絡來做通道信息的混合。4個階段分別使用了2、3、8、2個模塊，每個階段都使用了patch embedding的方式來將的空間塊混合成一個向量，對于每個階段。每個階段的特征維度保持一直，分別設置為128、256、320、512。此外，還提出了一種更加輕量級的網絡，命名為CONTAINER-LIGHT。CONTAINER-LIGHT使用與CONTAINER相似的網絡結構，但是在前三個階段不使用，從而更節省計算資源并且在下游任務上表現更好。CONTAINER-LIGHT可表示為： 

3. Performance

首先，在分類任務上，對比在ImageNet-1K^[7]數據集上的結果如表1所示：

表1 分類任務的對比

其中藍色表示參數比CONTAINER少或者多10%內的模型。

可以看到CONTAINER和CONTAINER-LIGHT在公平比較的基礎上達到很好的表現并且有著更低的FLOPs。

此外，如圖3所示，在更好的表現的基礎上，CONTAINER有更快的收斂速度。

圖3 分類任務收斂速度的對比

此外，本文還在各種下游任務上驗證了方法的有效性。在傳統檢測任務上，本文與其他骨干網絡對比了RetinaNet^[8]和Mask R-CNN^[9]的方法，在分割任務上，與其他骨干網絡在Mask R-CNN方法上進行了對比，結果如表2所示。可以看到在RetinaNet上，與ResNet-50相對比，有7.3的提升。在Mask R-CNN上，與耗費更多資源的大模型ViL-B^[10]具有相當的表現。

表2 在檢測和分割任務上的對比

在端到端的檢測中，在DETR和SMCADETR^[11]上，文中替換ResNet-50為COTAINER-LIGHT，并且都訓練50個epochs，從表3可以看出，CONTAINER-LIGHT作為骨干網絡可以在端到端的檢測上有明顯提升。

表3 端到端的檢測對比

最后，論文中還在自監督任務上進行了對比，使用DINO^[12]結構在無監督的視覺表征上訓練DeiT^[13]和CONTAINER-LIGHT。表4比較了兩者在訓練的不同epochs階段的top-10kNN的精度。COTAINER-LIGHT對比DeiT提升非常顯著。

表4 在DINO自監督任務上的對比

Reference：

[1] Xiaolong Wang, Ross Girshick, AbhinavGupta, and Kaiming He. Non-local neural networks. In CVPR, 2018. 1, 3

[2] Nicolas Carion, Francisco Massa, GabrielSynnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-endobject detection with transformers. In ECCV, 2020. 1, 2, 3,

[3] Ilya Tolstikhin, Neil Houlsby, AlexanderKolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, DanielKeysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mlp-mixer: Anall-mlp architecture for vision. arXiv, 2021. 1, 2, 3, 4, 5, 7

[4] Alexey Dosovitskiy, Lucas Beyer,Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner,Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. Animage is worth 16x16 words: Transformers for image recognition at scale. InICLR, 2021. 1, 2, 3, 6, 7

[5] Kaiming He, Xiangyu Zhang, Shaoqing Ren,and Jian Sun. Deep residual learning for image recognition. In CVPR, 2016. 2,3, 4, 6, 7, 8, 9

[6] Fran?ois Chollet. Xception: Deeplearning with depthwise separable convolutions. In CVPR, 2017. 2, 3, 4, 7

[7] Deng, Jia, et al. "Imagenet: Alarge-scale hierarchical image database." 2009 IEEE conference oncomputer vision and pattern recognition. Ieee, 2009.

[8] Lin, Tsung-Yi, et al. "Focal lossfor dense object detection." Proceedings of the IEEE internationalconference on computer vision. 2017.

[9] He, Kaiming, et al. "Maskr-cnn." Proceedings of the IEEE international conference on computervision. 2017.

[10] Pengchuan Zhang, Xiyang Dai, JianweiYang, Bin Xiao, Lu Yuan, Lei Zhang, and Jianfeng Gao. Multiscale visionlongformer: A new vision transformer for high-resolution image encoding. arXiv,2021. 2, 7, 8, 9

[11] Peng Gao, Minghang Zheng, XiaogangWang, Jifeng Dai, and Hongsheng Li. Fast convergence of detr with spatially modulatedco-attention. arXiv, 2021. 2, 9

[12] Mathilde Caron, Hugo Touvron, IshanMisra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin.Emerging properties in self-supervised vision transformers. arXiv, 2021. 2, 4,9

[13] Hugo Touvron, Matthieu Cord, MatthijsDouze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Trainingdata-efficient image transformers & distillation through attention. arXiv,2020. 1, 2, 3, 6, 7, 9