Sparse DETR: Efficient End-to-End Object Detection with Learnable Sparsity

1 INTRODUCTION . In recent years , we have witnessed the dramatic advancement and the success of object detection in deep learning . Diverse object detection methods have been proposed , but the existing algorithms that perform positive matching with the ground truth as a heuristic way require non-maximum suppression ( NMS ) post-processing of near-duplicate predictions . Recently , Carion et al . ( 2020 ) has introduced a fully end-to-end detector DETR by eliminating the need for NMS post-processing through a set-based objective . The training objective is designed by employing the Hungarian algorithm that considers both classification and regression costs , and achieves highly competitive performance . However , DETR is unable to use multi-scale features such as feature pyramid networks ( Lin et al. , 2017 ) , which are commonly used in object detection to improve the detection of small objects . The main reason is increased memory usage and computation by adding Transformer ( Vaswani et al. , 2017 ) architecture . As a result , its ability to detect small objects is relatively poor . To address this problem , Zhu et al . ( 2021 ) has proposed a deformable-attention inspired by the deformable convolution ( Dai et al. , 2017 ) and reduced the quadratic complexity to linear complexity through key sparsification in the attention module . By using deformable attention , deformable DETR addresses the slow convergence and high complexity issue of DETR , which enables the encoder to use multi-scale features as an input and significantly improves performance on detecting small objects . However , using the multi-scale features as an encoder input increases the number of tokens to be processed by about 20 times . Eventually , despite efficient computation for the same token length , the overall complexity increases back again , making the model inference slower even than vanilla DETR . In general , natural images often contain large background regions irrelevant to the objects of interest , and accordingly , in end-to-end detectors , the tokens corresponding to the background also occupy a significant portion . In addition , the importance of each regional feature is not identical , which has been proven by the two-stage detectors successfully do their job by focusing only on the foreground . It suggests that there exists considerable regional redundancy that can be reduced in the detection tasks and seeking to devise an efficient detector focusing on the salient regions is a necessary and natural direction . In our preliminary experiments , we observe the following : ( a ) during inference of a fully-converged Deformable DETR model on the COCO validation dataset , the encoder tokens referenced by the decoder account for only about 45 % of the total , and ( b ) retraining a new detector while updating only the encoder tokens preferred by the decoder from another fully-trained detector , barely suffers performance loss ( 0.1 AP degradation ) . See Appendix A.9 for the details . Inspired by this observation , we propose a learnable decoder cross-attention map predictor to sparsify encoder tokens . In the existing methods ( Carion et al. , 2020 ; Zhu et al. , 2021 ) , the encoder takes all the tokens , i.e . the backbone features combined with corresponding positional embeddings , as input without discrimination . Meanwhile , our approach distinguishes encoder tokens to be referenced later in the decoder and considers only those tokens in self-attention . Therefore , this can significantly reduce the number of encoder tokens involved in the computation and reduce the total computational cost for detection . We further propose the encoder auxiliary loss only for selected encoder tokens to improve detection performance while minimizing computational overhead . The proposed auxiliary loss not only improves performance , but also allows training a larger number of encoder layers . Extensive experiments on the COCO 2017 benchmark ( Lin et al. , 2014 ) demonstrate that Sparse DETR effectively reduces computational cost while achieving better detection performance . Without bells and whistles , Sparse DETR using Swin-T ( Liu et al. , 2021 ) backbone achieves 48.2 AP with 38 % reduction of the entire computational cost compared to the 48.0 AP baseline and 49.2 AP with 23 % reduction . In the case of the experiment that achieves 48.2 AP using only 10 % of encoder tokens , the computational cost of the transformer encoder block is reduced by approximately 82 % . We summarize our contributions as follows : • We propose encoder token sparsification method for an efficient end-to-end object detector , by which we lighten the attention complexity in the encoder . This efficiency enables stacking more encoder layers than Deformable DETR , leading to performance improvement within the same computational budget . • We propose two novel sparsification criteria to sample the informative subset from the entire token set : Objectness Score ( OS ) and Decoder cross-Attention Map ( DAM ) . Based on the decoder cross-attention map criterion , the sparsified model preserves detection performance even when using only 10 % of the whole tokens . • We adopt an encoder auxiliary loss only for the selected tokens . This additional loss not only stabilizes the learning process , but also greatly improves performance , with only marginally increased training time . 2 RELATED WORK . Efficient computation in vision transformers . It is a well-known problem that the attention computation in Transformers incurs the high time and memory complexity . The vision transformers need to digest even bigger token sets as input so that a large body of works ( Parmar et al. , 2018 ; Child et al. , 2019a ; Ho et al. , 2019 ; Wang et al. , 2020 ; Katharopoulos et al. , 2020 ; Choromanski et al. , 2021 ; Kitaev et al. , 2020 ) has been proposed lightweight attention mechanisms for them . Most of those works shed light on the complexity that resides only in a single-scale attention module , which hinders direct extension to the multi-scale features generally required in object detection . One of the promising approaches for the lighter transformer attention is input-dependent token sparsification . DynamicViT ( Rao et al. , 2021 ) and IA-RED2 ( Pan et al. , 2021 ) , similar to our work , both propose jointly-learned token selectors generating the sparsity patterns to be overlaid on the input tokens . Those approaches mainly focus on sparsifying a backbone network evaluated on the classification tasks , while our interest lies in a sparse encoder of the end-to-end object detectors . On the other hand , there has been a line of works sharing the spirit with ours in that they aim at sparse transformers in the DETR-based framework . Deformable DETR ( Zhu et al. , 2021 ) conducts sparse attention computation by sampling only a fraction of the entire key set with learnable 2-d off- sets , which enables to use multi-scale feature maps with a reasonable computational cost . It can be viewed as a key sparsification method but with dense queries , while our approach further reduces the query set pursuing even more sparsity . PnP-DETR ( Wang et al. , 2021 ) shortens the token length of the transformer encoder by introducing the Polling and Pull ( PnP ) module to sample the foreground tokens and condense the background tokens into a smaller set . However , their method can not naively be integrated with Deformable DETR , since their sparsification breaks the 2d spatial structure of the token set assumed in the deformable attention . On the contrary , Sparse DETR preserves the 2d sample space of the set and can be seamlessly combined with the deformable attention , which facilitates handling the multi-scale features . Thus , our approach gets benefits from both the deformable key sampling and the proposed query sparsification . Most of all , we propose explicit objectives for the token selection network , whereas the aforementioned works have no explicit objective implying their beliefs in a good selection strategy , merely relying on the final detection objective . Auxiliary Loss . Auxiliary loss ( Lee et al. , 2015 ; Szegedy et al. , 2015 ) is widely adopted to deliver gradients to the early layers of deep networks . DETR variants employ auxiliary Hungarian matching objectives at the end of every decoder layer with extra FFN heads so that each decoder layer directly learns to detect the correct number of objects out of the decoder ’ s outputs . Unlike the decoder ’ s object queries whose number is relatively small ( e.g . 300 ) , the number of encoder ’ s tokens has much larger scales when using multi-scale features . Thus , extending the layerwise auxiliary loss to the multi-scale encoder increases the training time cost by feeding too many tokens to the attached FFN heads . In Sparse DETR , thanks to the sparsity already induced in the encoder , we can instantly economize that cost while enjoying the auxiliary gradients in a wider range of intermediate layers . 3 APPROACH . In this section , we present our main contributions : ( a ) formulating a generalized saliency-based token sparsification scheme for the encoder , ( b ) proposing the effective saliency criteria with which that scheme can practically work , and ( c ) employing the encoder auxiliary losses and the top-k decoder query selection to improve the performance . Before describing the details , we revisit the key components of DETR ( Carion et al. , 2020 ) and Deformable DETR ( Zhu et al. , 2021 ) . 3.1 PRELIMINARY . DETR . DETR takes the flattened spatial feature map xfeat ∈ RN×D from a backbone network into the transformer encoder , where N denotes the number of tokens ( i.e . features ) and D denotes token dimension . The encoder iteratively updates xfeat by several vanilla self-attention modules . Then , the transformer decoder takes both the refined encoder tokens ( i.e . encoder output ) and M learnable object queries { qi } i=1···M as inputs and predicts a tuple of a class score c ∈ [ 0 , 1 ] C and a bounding box b ∈ [ 0 , 1 ] 4 for each object query qi , denoted as { ŷi } = { ( ci , bi ) } , where C denotes the number of classes . All components including the backbone network are jointly trained by performing the bipartite matching between the ground truth { yi } and predictions { ŷi } . Deformable DETR . Deformable DETR replaces the vanilla dense attention , which is the main computational bottleneck in DETR , with a deformable attention module . This significantly reduces the computational cost and improves the convergence . Suppose that we have the same size of a set of queries ( denoted as Ωq ) and a set of keys ( denoted as Ωk ) , which means |Ωq| = |Ωk| = N . The conventional dense attention computes the attention weight Aqk for every pair { ( q , k ) : q ∈ Ωq , k ∈ Ωk } , resulting in quadratic complexity with respect to N . Deformable attention reduces this quadratic complexity into the linear one by only considering relevant keys for each query . Specifically , deformable attention computes attention weight Aqk for all queries and a small set of keys : { ( q , k ) : q ∈ Ωq , k ∈ Ωqk } , where Ωqk ⊂ Ωk and |Ωqk| = K N . Due to this key sparsification , Deformable DETR is able to use the multi-scale features of the backbone network , improving the detection performance of small objects significantly . Paradoxically , using the multi-scale feature increases the number of tokens in the transformer encoder by about 20× compared to DETR , making that the encoder becomes the computational bottleneck of deformable DETR . This motivates us to develop a sparsification method to reduce the number of tokens in the encoder aggressively , which is described in the next sections .

Summary: A modification to the recently proposed Transformer-based object detector, DETR, is proposed. The authors are motivated by low performance of DETR on small objects, and high computational cost of self-attention on large feature maps. To improve both, they propose to sparsify input feature maps by learning a classifier whether to include an input feature or not. Supervision for this classifier comes from Transformer decoder attention weights. The authors apply their approach to Deformable DETR and evaluate on the COCO bounding box detection benchmark, and show that it allows to trade detection performance for computational efficiency by tuning the classifier threshold.