XCiT: Cross-Covariance Image Transformers

1 Introduction . Transformers architectures [ 68 ] have provided quantitative and qualitative breakthroughs in speech and natural language processing ( NLP ) . After a few attempts to incorporate wide-range self-attention in vision architectures [ 71 , 82 ] , Dosovitskiy et al . [ 21 ] established transformers as a viable architecture for learning visual representations , reporting competitive results for image classification while relying on large-scale pre-training . Touvron et al . [ 64 ] have shown on par or better accuracy/throughput compared to strong convolutional baselines such as EfficientNets [ 58 ] when training transformers on ImageNet-1k using extensive data augmentation and improved training schemes . Promising results have been obtained for other vision tasks , including image retrieval [ 22 ] , object detection and semantic segmentation [ 44 , 70 , 81 , 83 ] , as well as video understanding [ 2 , 7 , 23 ] . One major drawback of transformers is the time and memory complexity of the core self-attention operation , that increases quadratically with the number of input tokens , or similarly number of patches in computer vision . For w⇥h images , this translates to a complexity of O ( w2h2 ) , which is prohibitive for most tasks involving high-resolution images , such as object detection and segmentation . Various strategies have been proposed to alleviate this complexity , for instance using approximate forms of self-attention [ 44 , 81 ] , or pyramidal architectures which progressively downsample the feature maps [ 70 ] . However , none of the existing solutions are fully satisfactory , as they either trade complexity for accuracy , or their complexity remains excessive for processing very large images . We replace the self-attention , as originally introduced by Vaswani et al . [ 68 ] , with a “ transposed ” attention that we denote as “ cross-covariance attention ” ( XCA ) . Cross-covariance attention substi- Code : https : //github.com/facebookresearch/xcit 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . tutes the explicit full pairwise interaction between tokens by self-attention among features , where the attention map is derived from the cross-covariance matrix computed over the key and query projections of the token features . Importantly , XCA has a linear complexity in the number of patches . To construct our Cross-Covariance Image Transformers ( XCiT ) , we combine XCA with local patch interaction modules that rely on efficient depth-wise convolutions and point-wise feedforward networks commonly used in transformers , see Figure 1 . XCA can be regarded as a form of a dynamic 1⇥1 convolution , which multiplies all tokens with the same data-dependent weight matrix . We find that the performance of our XCA layer can be further improved by applying it on blocks of channels , rather than directly mixing all channels together . This “ block-diagonal ” shape of XCA further reduces the computational complexity with a factor linear in the number of blocks . Given its linear complexity in the number of tokens , XCiT can efficiently process images with more than thousand pixels in each dimension . Notably , our experiments show that XCiT does not compromise the accuracy and achieves similar results to DeiT [ 64 ] and CaiT [ 67 ] in comparable settings . Moreover , for dense prediction tasks such as object detection and image segmentation , our models outperform popular ResNet [ 28 ] backbones as well as the recent transformer-based models [ 44 , 70 , 81 ] . Finally , we also successfully apply XCiT to the self-supervised feature learning using DINO [ 12 ] , and demonstrate improved performance compared to a DeiT-based backbone [ 64 ] . Overall , we summarize our contributions as follows : • We introduce cross-covariance attention ( XCA ) , which provides a “ transposed ” alternative to conventional self-attention , attending over channels instead of tokens . Its complexity is linear in the number of tokens , allowing for efficient processing of high-resolution images , see Figure 2 . • XCA attends to a fixed number of channels , irrespective of the number of tokens . As a result , our models are significantly more robust to changes in image resolution at test time , and are therefore more amenable to process variable-size images . • For image classification , we demonstrate that our models are on par with state-of-the-art vision transformers for multiple model sizes using a simple columnar architecture , i.e. , in which we keep the resolution constant across layers . In particular , our XCiT-L24 model achieves 86.0 % top-1 accuracy on ImageNet , outperforming its CaiT-M24 [ 67 ] and NFNet-F2 [ 10 ] counterparts with comparable numbers of parameters . • For dense prediction tasks with high-resolution images , our models outperform ResNet and multiple transformer-based backbones . On the COCO benchmark , we achieve a strong performance of 48.5 % and 43.7 % mAP for object detection and instance segmentation respectively . Moreover , we report 48.4 % mIoU for semantic segmentation on the ADE20k benchmark , outperforming the state-of-the-art Swin Transformer [ 44 ] backbones across all comparable model sizes . • Finally , our XCiT model is highly effective in self-supervised learning setups , achieving 80.9 % top-1 accuracy on ImageNet-1k using DINO [ 12 ] . 2 Related work . Deep vision transformers . Training deep vision transformers can be challenging due to instabilities and optimization issues . Touvron et al . [ 67 ] successfully train models with up to 48 layers using LayerScale , which weighs contributions of residual blocks across layers and improves optimization . Additionally , the authors introduce class attention layers which decouple the learning of patch features and the feature aggregation stage for classification . Spatial structure in vision transformers . Yuan et al . [ 79 ] propose applying a soft split for patch projection with overlapping patches which is applied repeatedly across model layers , reducing the number of patches progressively . Han et al . [ 27 ] introduce a transformer module for intra-patch structure , exploiting pixel-level information and integrating with an inter-patch transformer to attain higher representation power . d ’ Ascoli et al . [ 18 ] consider the initialization of self-attention blocks as a convolutional operator , and demonstrate that such initialization improves the performance of vision transformers in low-data regimes . Graham et al . [ 26 ] introduce LeViT , which adopts a multistage architecture with progressively reduced feature resolution similar to popular convolutional architectures , allowing for models with high inference speed while retaining a strong performance . Moreover , the authors adopt a convolution-based module for extracting patch descriptors . Yuan et al . [ 78 ] improve both the performance and the convergence speed of vision transformers by replacing the linear patch projection with convolutional layers and max-pooling , as well as modifying the feed-forward networks in each transformer layer to incorporate depth-wise convolutions . Efficient attention . Numerous methods for efficient self-attention have been proposed in the literature to address the quadratic complexity of self-attention in the number of input tokens . These include restricting the span of the self-attention to local windows [ 48 , 50 ] , strided patterns [ 14 ] , axial patterns [ 30 ] , or an adaptive computation across layers [ 57 ] . Other methods provide an approximation of the self-attention matrix which can be achieved by a projection across the token dimension [ 69 ] , or through a factorization of the softmax-attention kernel [ 15 , 37 , 56 , 77 ] , which avoids explicit computation of the attention matrix . While conceptually different , our XCA performs similar computations without being sensitive to the choice of the kernel . Similarly , Lee-Thorp et al . [ 41 ] achieve faster training by substituting self-attention with unparametrized Fourier Transform . Other efficient attention methods rely on local attention and adding a small number of global tokens , thus allowing interaction among all tokens only by hopping through the global tokens [ 1 , 5 , 34 , 80 ] . Similarly , Goyal et al . [ 25 ] use a global workspace though which items interact , albeit one that is shared across layers . Transformers for high-resolution images . Several works adopt visual transformers to highresolution image tasks beyond image classification , such as object detection and image segmentation . Wang et al . [ 70 ] design a model with a pyramidal architecture and address complexity by gradually reducing the spatial resolution of keys and values . Similarly , for video recognition Fan et al . [ 23 ] utilize pooling to reduce the resolution across the spatial and temporal dimensions to allow for an efficient computation of the attention matrix . Zhang et al . [ 81 ] adopt global tokens and local attention to reduce the model complexity , while Liu et al . [ 44 ] provide an efficient method for local attention with shifted windows . In addition , Zheng et al . [ 83 ] and Ranftl et al . [ 54 ] study problems like semantic segmentation and monocular depth estimation with the quadratic self-attention operation . Data-dependent layers . Our XCiT layer can be regarded as a “ dynamic ” 1⇥1 convolution , which multiplies all token features with the same data-dependent weight matrix , derived from the key and query cross-covariance matrix . In the context of convolutional networks , Dynamic Filter Networks [ 9 ] explore a related idea , using a filter generating subnetwork to produce convolutional filters based on features in previous layers . Squeeze-and-Excitation networks [ 32 ] use data dependent 1⇥1 convolutions in convolutional architectures . Spatially average-pooled features are fed to a 2-layer MLP which produces per channel scaling parameters . Closer in spirit to our work , Lambda layers propose a way to ensure global interaction in ResNet models [ 4 ] . Their “ content-based lambda function ” is computing a similar term as our cross-covariance attention , but differing in how the softmax and ` 2 normalizations are applied . Moreover , Lambda layers also include specific positionbased lambda functions , and LambdaNetworks are based on ResNets while XCiT follows the ViT architecture . Recently data-independent analogues of self-attention have also been found to be an effective alternative to convolutional and self-attention layers for vision tasks [ 20 , 46 , 62 , 66 ] . These methods treat entries in the attention map as learnable parameters , rather than deriving the attention map dynamically from queries and keys , but their complexity remains quadratic in the number of tokens . Zhao et al . [ 82 ] consider alternative attention forms in computer vision . 3 Method . In this section , we first recall the self-attention mechanism , and the connection between the Gram and covariance matrices , which motivated our work . We then propose our cross-covariance attention operation ( XCA ) – which operates along the feature dimension instead of token dimension in conventional transformers – and combine it with local patch interaction and feedforward layers to construct our Cross-Covariance Image Transformer ( XCiT ) . See Figure 1 for an overview . 3.1 Background . Token self-attention . Self-attention , as introduced by Vaswani et al . [ 68 ] , operates on an input matrix X 2 RN⇥d , where N is the number of tokens , each of dimensionality d. The input X is linearly projected to queries , keys and values , using the weight matrices Wq 2 Rd⇥dq , Wk 2 Rd⇥dk and Wv 2 Rd⇥dv , such that Q=XWq , K=XWk and V =XWv , where dq = dk . Keys and values are used to compute an attention map A ( K , Q ) = Softmax ( QK > / p dk ) , and the output of the self-attention operation is defined as the weighted sum of N token features in V with the weights corresponding to the attention map : Attention ( Q , K , V ) = A ( K , Q ) V . The computational complexity of self-attention scales quadratically in N , due to pairwise interactions between all N elements . Relationship between Gram and covariance matrices . To motivate our cross-covariance attention operation , we recall the relation between Gram and covariance matrices . The unnormalised d⇥d covariance matrix is obtained as C=X > X . The N⇥N Gram matrix contains all pairwise innerproducts : G=XX > . The non-zero part of the eigenspectrum of the Gram and covariance matrix are equivalent , and the eigenvectors of C and G can be computed in terms of each other . If V are the eigenvectors of G , then the eigenvectors of C are given by U=XV . To minimise the computational cost , the eigendecomposition of either the Gram or covariance matrix can be obtained in terms of the decomposition of the other , depending on which of the two matrices is the smallest.1 We draw upon this strong connection between the Gram and covariance matrices to consider whether it is possible to avoid the quadratic cost to compute the N⇥N attention matrix , which is computed from the analogue of the N ⇥N Gram matrix QK > =XWqW > k X > . Below we consider how we can use the dk⇥dq cross-covariance matrix , K > Q=W > k X > XWq , which can be computed in linear time in the number of elements N , to define an attention mechanism .

Summary: This paper presents a modification of vision transformers by introducing Cross-Covariance Attention (XCA) and Local Patch Interaction (LPI) which substantially reduce the computational overhead and enhances the stability of training. The authors also adopt the DeiT training strategy with distillation using a convolutional teacher to improve the performance. Experiments on Imagenet classification, self-supervised training with DINO and many downstream tasks (COCO object detection and instance segmentation, ADE20K semantic segmentation) show the superior performance.