Dynamic Grained Encoder for Vision Transformers

1 Introduction . Following the evolution of network architectures in natural language processing ( NLP ) , Vision Transformers [ 1–5 ] have recently attracted increasing research attention and demonstrated promising results on several vision tasks , such as image classification , object detection , and other pixel-level tasks . Vision transformers are notable for modeling long-range dependencies and introducing less inductive bias , considered to be a solid alternative to CNNs for vision tasks . One of the eminent obstacles for vision transformers is the high computational cost . Vision tasks typically require high-resolution image features to obtain detail and structure representation , which is critical for pixel-level tasks [ 6–10 ] . However , since the encoders in vision transformers need to establish pairwise relationships , high-resolution features could impose unacceptable computational and memory costs . Therefore , similar to the efficient transformers [ 11–13 ] in NLP , many variants [ 2– 4 ] of vision transformers are proposed to perform sparse self-attentions with dense queries and sparse key-value pairs based on fixed pattern or heuristic rules . In this paper , we notice that different from natural language , natural images involve much spatial redundancy , especially in flat or low-texture regions [ 14–18 ] . This could enable the image features to have a low resolution in some regions while maintaining similar representational capabilities . To verify the spatial redundancy in vision transformers , we give an empirical analysis for DeiT [ 19 ] on ImageNet [ 20 ] classification dataset ( the details refer to Sec . 3.1 ) . It demonstrates the existence ∗Equal contribution . This work was done in Megvii Research . †Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . of spatial redundancy in queries , and the complexity can be dramatically reduced by downsampling some highly redundant regions while maintaining comparable performance . These properties allow the queries to use mixed granularity to achieve a balance between effectiveness and efficiency , i.e. , more tokens in more discriminative regions while fewer tokens in less informative regions . However , the distribution of spatial redundancy varies greatly among different input images , making it difficult for a static method to handle complex and variable features . We thus attempt to explore a new perspective : introducing dynamic network mechanism into vision transformers to reduce the spatial redundancy of image features . As shown in Fig . 1 , we propose a Dynamic Grained Encoder ( DGE ) to replace the vanilla encoder in vision transformers . It could assign a suitable number of queries for each region by using a dynamic grained router , e.g. , the foreground regions of the cat head in Fig . 1 are assigned more queries than the background regions . Concretely , a reshaped 2D feature is first divided into regions using a fixed window . For each region , the number of patches is decided by a data-dependent routing process , and each patch is average pooled to obtain a 1D token . All the tokens are then concatenated into a sequence as the queries . Since our method focuses on the sparsity of queries , it is compatible with many efficient transformer encoders [ 2 , 3 , 11–13 ] , making our approach available as a generic plugin in most vision transformers [ 1–3 , 19 , 21 ] . Furthermore , the output of the encoder is restored to the input resolution by an un-pooling operation and compensates for detailed information with the input feature . To demonstrate the effectiveness , we conduct extensive experiments on three typical vision transformers , i.e. , DeiT [ 19 ] , PVT [ 3 ] and DPVT , where DPVT is a new framework based on the deformable attention [ 2 ] . In the image classification task , our dynamic grained encoder allows these models to reduce computational complexity by 40 % -60 % while maintaining comparable performance . On the other hand , with lower computational complexity , the accuracy can be improved by up to 4.4 % on ImageNet val set . In addition , the experiments on object detection and segmentation show the strong robustness and generalization of our method . 2 Related Work . 2.1 Vision Transformer . Recently , Vision Transformers , inspired by the significant success of transformer [ 22 ] achieved in the NLP field , have received more attention in the vision community . ViT [ 1 ] , which converts the image into a sequence and applies the transformer encoder structure directly on it for image classification , has pioneered this direction in visual recognition . To tackle the issue of training efficiency and data efficiency , DeiT [ 19 ] introduces several training strategies to enable learning the vision transformer on ImageNet . PVT [ 3 ] further develops a feature pyramid based on the transformer structure and makes it applicable for the various downstream vision tasks . Swin [ 21 ] introduces the local window idea to improve the efficiency of the transformer structure . Our work mainly focuses on reducing the spatial redundancy and improving the model efficiency in a data-dependent manner , which is rarely explored in previous works and complementary with various vision transformer structures . 2.2 Efficient Transformer . To improve the efficiency of transformers , prior works mainly concentrate on reducing the quadratic computation of self-attention . These works can be roughly summarized as three types : learnable/fixed pattern based methods , low-rank/kernel based methods and memory based methods . Some recent approaches [ 12 , 23–26 ] try to reduce the complexity of the self-attention mechanism by using a heuristic method to generate fixed or learnable patterns . Other efforts [ 11 , 13 , 27 , 28 ] focus on utilizing the low-rank property of the attention matrix or introducing kernels to avoid computing the attention matrix explicitly . Moreover , some works [ 29–31 ] also explore the memory mechanism to improve efficiency . However , previous attempts mainly concentrate on the NLP tasks . Different from the language sequence , which has a highly abstract representation of information , natural images typically have much spatial redundancy . It makes the vision transformers require expensive costs for downstream vision tasks , especially the dense-prediction tasks , e.g. , object detection , segmentation . Our work tries to utilize this intrinsic property of natural images to achieve redundancy reduction in a data-dependent manner . 2.3 Dynamic Network . Dynamic networks [ 32 ] are proposed to adaptively change the network architecture and parameters according to input , which have been widely explored in computer vision and natural language processing tasks . Most of the dynamic networks focus on coarse-grained strategy by dropping blocks [ 33–36 ] , pruning channels [ 37 , 38 ] or adjusting layer-level scales [ 39 , 40 ] . For instance , MSDNet [ 34 ] proposes an early existing mechanism to achieve efficient inference for image classification . Switch Transformer [ 41 ] uses the Mixture of Experts ( MoE ) model [ 42 ] to select different parameters for each input sample . DRNet [ 39 ] attempts to perform adaptive scale transformation in a feature pyramid network for semantic segmentation . The closest works to ours are probably the Dynamic Convolution [ 43 ] and the Dynamic Head [ 10 ] , which use a learnable mask to skip specific spatial locations . However , they are only applicable to the CNN-based networks , and the skipping-location strategy could result in significant performance degradation for vision transformers ( refer to Sec . 4.1.2 ) . Different from them , our method adapts the region-level granularity to the input feature for the vision transformers , which is more general and flexible . 3 Method . 3.1 Empirical Analyses on Spatial Redundancy . To investigate the spatial redundancy of vision transformer on image data , we conduct a series of experiments on the ImageNet [ 20 ] val set with a pre-trained DeiT-S [ 19 ] model . Our main purpose is to explore the relationship among the granularity of queries , computational complexity , and classification performance . Specifically , for each encoder layer in DeiT-S , we reshape its input queries ( excluding the extra embedding ) as a 2D feature map and split it into 2 × 2 non-overlap patches . For each patch , we calculate its average token , and measure the similarity of each token in the patch with the average token by using the Pearson Correlation Coefficient ( PCC ) metric . Then we have three valuable observations . ( 1 ) Queries share similar patterns in a local region . From the correlation coefficient histogram plotted in Fig . 2 ( a ) , most of the correlation coefficients are greater than 0.8 , which indicates the queries typically have a strong correlation in a local region . ( 2 ) Large potential of reducing spatial redundancy . Furthermore , in each patch , we replace the tokens with the average token when their correlation coefficient is above a given threshold . As shown in Fig . 2 ( b ) , we illustrate the accuracy/complexity curve varying correlation thresholds . When the threshold is 0.9 , the complexity decreases by 27 % , but the top-1 accuracy decreases by only 0.3 % . This evidence demonstrates the potential of reducing the spatial redundancy on vision transformers . ( 3 ) Static strategy is sub-optimal . As shown in Fig . 2 ( c ) , some encoders have large variance of correlation coefficients among different images . Thus , using data-independent methods to reduce spatial redundancy is sub-optimal , which may lead to considerable performance degradation . These observations motivate us to explore a data-dependent manner to reduce spatial redundancy . 3.2 Dynamic Grained Encoder . 3.2.1 Overall Architecture . In this paper , we propose a new encoder block for vision transformers , called Dynamic Grained Encoder ( DGE ) . As shown in Fig . 1 , the proposed encoder consists of two main modules , i.e. , dynamic grained router and vanilla encoder block . Specifically , the dynamic grained router adaptively generates mixed-grained patches for a 2D feature . The vanilla encoder block can be a standard encoder block [ 22 ] or other efficient variants [ 2 , 9 , 11–13 , 44 ] , which is made up of a multi-head attention and a feed-forward network . If there are extra tokens in the input sequence , such as class embedding in ViT [ 1 ] , we handle them separately with the vanilla encoder . For ease of presentation , the rest of this section only considers the input sequence without extra tokens . Given an input sequence x ∈ R ( H×W ) ×C for the dynamic grained encoder , ( H , W ) denotes the resolution of the feature , C is the number of channels . To compatible with most vanilla encoders , we only generate sparse queries q ∈ RN×C by the dynamic grained router , where N indicates the number of queries . Then the sparse queries as well as dense keys k and values v are transformed by a vanilla encoder . It is worth mentioning that keys and values can be sparse in the vanilla encoder to improve efficiency further . The output sequence of the vanilla encoder is restored to a 2D feature with the original resolution by using an un-pooling operation . Furthermore , to enhance the details of the output feature and alleviate the vanishing gradient problem , we add a residual connection [ 45 ] to fuse the input sequence . 3.2.2 Dynamic Grained Router . To achieve dynamic grained patches in space , we first partition the 2D feature , denoting as z , into multiple regions , which can perform in regular or irregular ways . Although the irregular ways , e.g. , superpixels [ 46 ] and segmentation [ 47 ] , may lead to better performance , it is very unfriendly to memory access and inducing inefficiency . Therefore , as shown in Fig . 3 , we adopt a S × S non-overlap window3 to split image features into multiple regular regions . Furthermore , we define a set of candidate granularities Φ = { φ1 , φ2 , ... , φK } to represent the optional patch size in a region , where K is the number of candidate granularities . The granularity denotes the side length of a patch , e.g. , φ = 8 corresponds to an 8× 8 patch . Since each patch is pooled into one query in the encoder , larger granularity indicates fewer queries and less computation . For convenience , we set the region size with the maximum granularity , i.e. , S = max ( Φ ) , in the experiments . Inference . For a region i ∈ { 1 , 2 , ... , dHS e·d W S e } , we use a gating network to select a granularity from the set of candidate granularities . Concretely , we reduce the region feature zi into a representative token by using the average pooling operation and linearly project it to the gating logits : h ( zi ) = 1 S2 S2∑ j=1 zi , jW + b , ( 1 ) 3Bottom-right padding is adopted on the feature if needed . where W ∈ RC×K and b ∈ R1×K indicate the weight and bias , respectively . The gating logits is used to decide the granularity for the region by calculating the gating indices : θi = arg max k ( h ( zi ) k ) ∈ { 1 , 2 , ... , K } . ( 2 ) As shown in Fig . 3 , we split the region feature into multiple groups of patches1 with K granularities . We then choose a group of specific granularity according to the gating indices . We denote the selected group as z′i ∈ RNi×φ 2 θi ×C , where Ni = d Sφθi e · d S φθi e is the number of patches in the group . As shown in Fig . 1 , to construct a sequence as queries , we use the spatial mean vector of each patch as the representative token by a pooling operation and concatenate all the tokens for the vanilla encoder : ŷi = VanillaEncoder ( qi , k , v ) , where qi = 1 φ2θi φ2θi∑ j=1 z′i , j ∈ RNi×C . ( 3 ) Compared with the previous encoders [ 2 , 11–13 ] , the number of queries is reduced to 1/φ2θi of the original , the efficiency of the encoder can be improved , and the acceleration is more significant when selected granularity θi is larger . Training . To enable the end-to-end training for the gating network , motivated by [ 43 , 48–50 ] , we replace the determined decisions in Eq . 2 with a stochastic sampling process during the training phase . Specifically , given a categorical distribution with unnormalized log probabilities , a discrete gating index can be yielded with noise samples gj drawn from a standard Gumbel distribution : θi = arg max k ( h ( zi ) k + gk ) , where gk ∼ Gumbel ( 0 , 1 ) . ( 4 ) Furthermore , since the Eq . 4 is a hard decision process , it is not straightforward to train the gating logits . To enable the back-propagation , we adopt the Gumbel-Softmax technique [ 51 ] to give a continuous and differentiable approximation by replacing the argmax with a softmax operation . The soft gating score for a region is then selected by the gating index : pi = exp ( ( h ( xi ) θi + gθi ) /τ ) ∑K k exp ( ( h ( xi ) k + gk ) /τ ) ∈ [ 0 , 1 ] , ( 5 ) where a fixed temperature τ = 1 is used in our experiments for convenience . Similar with [ 43 , 52 ] , we further use a straight-through estimator for the gradients of gating logits , which are obtained through the soft gating score pi during the backward pass : y′i = { ŷ forward pi · ŷ backward ( 6 ) The above stochastic process is only adopted in the training phase . Our method requires no random sampling and exponential functions during inference , guaranteeing high efficiency in practice .

Summary: The paper proposes a dynamic approach to split an patch into different sizes to save computations for vision transformers. In particular, the authors first split an image into different regions. And for each region, the authors propose a dynamic approach to select one of three predefined patch sizes. This is achieved with a gumbel-softmax function, which enables deriving categorical decisisons in a differentiable way.  The authors experimented with DeiT , PVT  and DPVT.