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