Container: Context Aggregation Networks

1 Introduction . Convolutional neural networks ( CNNs ) have become the de facto standard for extracting visual representations , and have proven remarkably effective at numerous downstream tasks such as object detection [ 37 ] , instance segmentation [ 22 ] and image captioning [ 1 ] . Similarly , in natural language processing , Transformers rule the roost [ 13 , 43 , 42 , 4 ] . Their effectiveness at capturing short and long range information have led to state-of-the-art results across tasks such as question answering [ 45 ] and language understanding [ 58 ] . In computer vision , Transformers were initially employed as long range information aggregators across space ( e.g. , in object detection [ 5 ] ) and time ( e.g. , in video understanding [ 61 ] ) , but these methods continued to use CNNs [ 34 ] to obtain raw visual representations . More recently however , CNN-free visual backbones employing Transformer modules [ 54 , 14 ] have shown impressive performance on image classification benchmarks such as ImageNet [ 33 ] . The race to dethrone CNNs has now begun to expand beyond Transformers – a recent unexpected result shows that a multi-layer perceptron ( MLP ) exclusive network [ 52 ] can be just as effective at image classification . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . On the surface , CNNs [ 34 , 8 , 63 , 23 ] , Vision Transformers ( ViTs ) [ 14 , 54 ] and MLP-mixers [ 52 ] are typically presented as disparate architectures . However , taking a step back and analyzing these methods reveals that their core designs are quite similar . Many of these methods adopt a cascade of neural network blocks . Each block typically consists of aggregation modules and fusion modules . Aggregation modules share and accumulate information across a predefined context window over the module inputs ( e.g. , the self attention operation in a Transformer encoder ) , while fusion modules combine position-wise features and produce module outputs ( e.g. , feed forward layers in ResNet ) . In this paper , we show that the primary differences in many popular architectures result from variations in their aggregation modules . These differences can in fact be characterized as variants of an affinity matrix within the aggregator that is used to determine information propagation between a query vector and its context . For instance , in ViTs [ 14 , 54 ] , this affinity matrix is dynamically generated using key and query computations ; but in the Xception architecture [ 8 ] ( that employs depthwise convolutions ) , the affinity matrix is static – the affinity weights are the same regardless of position , and they remain the same across all input images regardless of size . And finally the MLP-Mixer [ 52 ] also uses a static affinity matrix which changes across the landscape of the input . Along this unified view , we present CONTAINER ( CONText AggregatIon NEtwoRk ) , a general purpose building block for multi-head context aggregation . A CONTAINER block contains both static affinity as well as dynamic affinity based aggregation , which are combined using learnable mixing coefficients . This enables the CONTAINER block to process long range information while still exploiting the inductive bias of the local convolution operation . CONTAINER blocks are easy to implement , can easily be substituted into many present day neural architectures and lead to highly performant networks whilst also converging faster and being data efficient . Our proposed CONTAINER architecture obtains 82.7 % Top-1 accuracy on ImageNet using 22M parameters , improving +2.8 points over DeiT-S [ 54 ] with a comparable number of parameters . It also converges faster , hitting DeiT-S ’ s accuracy of 79.9 % in just 200 epochs compared to 300 . We also propose a more efficient model , named CONTAINER-LIGHT that employs only static affinity matrices early on but uses the learnable mixture of static and dynamic affinity matrices in the latter stages of computation . In contrast to ViTs that are inefficient at processing large inputs , CONTAINER-LIGHT can scale to downstream tasks such as detection and instance segmentation that require high resolution input images . Using a CONTAINER-LIGHT backbone and 12 epochs of training , RetinaNet [ 37 ] is able to achieve 43.8 mAP , while Mask-RCNN [ 22 ] is able to achieve 45.1 mAP on box and 41.3 mAP on instance mask prediction , improvements of +7.3 , +6.9 and +6.6 respectively , compared to a ResNet-50 backbone . The more recent DETR and its variants SMCA-DETR and Deformable DETR [ 5 , 19 , 75 ] also benefit from CONTAINER-LIGHT and achieve 38.9 , 43.0 and 44.2 mAP , improving significantly over their ResNet-50 backbone baselines . CONTAINER-LIGHT is data efficient . Our experiments show that it can obtain an ImageNet Top-1 accuracy of 61.8 using just 10 % of training data , significantly better than the 39.3 accuracy obtained by DeiT . CONTAINER-LIGHT also convergences faster and achieves better kNN accuracy ( 71.5 ) compared to DeiT ( 69.6 ) under DINO self-supervised training framework [ 6 ] . The CONTAINER unification and framework enable us to easily reproduce several past models and even extend them with just a few code and parameter changes . We extend multiple past models and show improved performance – for instance , we produce a Hierarchical DeiT model , a multi-head MLP-Mixer and add a static affinity matrix to the DeiT architecture . Our code base and models will be released publicly . Finally , we analyse a CONTAINER model containing both static and dynamic affinities and show the emergence of convolution-like local affinities in the early layers of the network . In summary , our contributions include : ( 1 ) A unified view of popular architectures for visual inputs – CNN , Transformer and MLP-mixer , ( 2 ) A novel network block – CONTAINER , which uses a mix of static and dynamic affinity matrices via learnable parameters and the corresponding architecture with strong results in image classification and ( 3 ) An efficient and effective extension – CONTAINER-LIGHT with strong results in detection and segmentation . Importantly , we see that a number of concurrent works are aiming to fuse the CNN and Transformer architectures [ 36 , 64 , 40 , 24 , 55 , 69 , 64 , 47 ] , validating our approach . We hope that our unified view helps place these different concurrent proposals in context and leads to a better understanding of the landscape of these methods . 2 Related Work . Visual Backbones . Since AlexNet [ 33 ] revolutionized computer vision , a host of CNN based architectures have provided further improvements in terms of accuracy including VGG [ 46 ] , ResNet [ 23 ] , Inception Net [ 48 ] , SENet [ 28 ] , ResNeXt [ 63 ] and Xception [ 8 ] and efficiency including Mobile-net v1 [ 26 ] , Mobile-net v2 [ 26 ] and Efficient-net v2 [ 50 ] . With the success of Transformers [ 56 ] in NLP such as BERT [ 13 ] and GPT [ 43 ] , researchers have begun to apply them towards solving the long range information aggregation problem in computer vision . ViT [ 14 ] /DeiT [ 54 ] are transformers that achieve better performance on ImageNet than CNN counterparts . Recently , several concurrent works explore integrating convolutions with transformers and achieve promising results . ConViT [ 11 ] explores soft convolutional inductive bias for enhancing DeiT . CeiT [ 66 ] directly incorporates CNNs into the Feedforward module of transformers to enhance the learned features . PVT [ 60 ] proposes a pyramid vision transformer for efficient transfer to downstream tasks . Pure Transformer models such as ViT/DeiT however , require huge GPU memory and computation for detection [ 60 ] and segmentation [ 73 ] tasks , which need high resolution input . MLP-Mixer [ 52 ] shows that simply performing transposed MLP followed by MLP can achieve near state-of-the-art performance . We propose CONTAINER , a new visual backbone that provides a unified view of these different architectures and performs well across several vision tasks including ones that require a high resolution input . Transformer Variants . Vanilla Transformers are unable to scale to long sequences or high-resolution images due to the quadratic computation in self-attention . Several methods have been proposed to make Transformer computations more efficient for high resolution input . Reformer [ 32 ] , Clusterform [ 57 ] , Adaptive Clustering Transformer [ 73 ] and Asymmetric Clustering [ 10 ] propose to use Locality Sensitivity Hashing to cluster keys or queries and reduce quadratic computation into linear computation . Lightweight convolution [ 62 ] explore convolution architectures for replacing Transformers but only explore applications in NLP . RNN Transformer [ 31 ] builds a connection between RNN and Transformer and results in attention with linear computation . Linformer [ 59 ] changes the multiplication order of key , query , value into query , value , key by deleting the softmax normalization layer and achieve linear complexity . Performer [ 9 ] uses Orthogonal Random Feature to approximate full rank softmax attention . MLIN [ 18 ] performs interaction between latent encoded nodes , and its complexity is linear with respect to input length . Bigbird [ 3 ] breaks the full rank attention into local , randomly selected and global attention . Thus the computation complexity becomes linear . Longformer [ 68 ] uses local Transformers to tackle the problem of massive GPU memory requirements for long sequences . MLP-Mixer [ 52 ] is a pure MLP architecture for image recognition . In the unified formulation we provide , MLP-Mixer can be considered as a single-head Transformer with static affinity matrix weight . MLP-Mixer can provide more efficient computation than vanilla transformer due to no need to calculate affinity matrix using key query multiplication . Efficient Transformers mostly use approximate message passing which results in performance deterioration across tasks . Lightweight Convolution [ 62 ] , Involution [ 35 ] , Synthesizer [ 51 ] , and MUSE [ 71 ] explored the relationship between Depthwise Convolution and Transformer . Our CONTAINER unification performs global and local information exchange simultaneously using a mixture affinity matrix , while CONTAINER-LIGHT switches off the dynamic affinity matrix for high resolution feature maps to reduce computation . Although switching off the dynamic affinity matrix slightly hinders classification performance , CONTAINER-LIGHT still provides effective and efficient generalization to downstream tasks compared with popular backbones such as ViT and ResNet . Transformers for Vision . Transformers enable high degrees of parallelism and are able to capture long-range dependencies in the input . Thus Transformers have gradually surpassed other architectures such as CNN [ 34 ] and RNN [ 25 ] on image [ 14 , 5 , 70 ] , audio [ 2 ] , multi-modality [ 17 , 21 , 20 ] , and language understanding [ 13 ] . In computer vision , Non-local Neural Network [ 61 ] has been proposed to capture long range interactions to compensate for the local information captured by CNNs and used for object detection [ 27 ] and semantic segmentation [ 16 , 29 , 76 , 67 ] . However , these methods use Transformers as a refinement module instead of treating the transformer as a first-class citizen . ViT [ 14 ] introduces the first pure Transformer model into computer vision and surpasses CNNs with large scale pretraining on the non publicly available JFT dataset . DeiT [ 54 ] trains ViT from scratch on ImageNet-1k and achieve better performance than CNN counterparts . DETR [ 5 ] uses Transformer as an encoder and decoder architecture for designing the first end-to-end object detection system . Taming Transformer [ 15 ] use Vector Quantization [ 41 ] GAN and GPT [ 43 ] for high quality highresolution image generation . Motivated by the success of DETR on object detection , Transformers have been applied widely on tasks such as semantic segmentation [ 74 ] , pose estimation [ 65 ] , trajectory estimation [ 39 ] , 3D representation learning and self-supervised learning with MOCO v3 [ 7 ] and DINO [ 6 ] . ProTo [ 72 ] verify the effective of transformer on reasoning tasks .

Summary: This paper first provides a unified view on three existing network architectures: depthwise  CNNs, Transformers and MLP-Mixers. It then proposes Container and a light-weight version Container-Light, which combine CNNs and self-attention to aggregate spatial information. Experiments on ImageNet classification and COCO object detection verify the high modeling ability of Container.