HRFormer: High-Resolution Vision Transformer for Dense Predict

1 Introduction . Vision Transformer ( ViT ) [ 13 ] shows promising performance on ImageNet classification tasks . Many follow-up works boost the classification accuracy through knowledge distillation [ 42 ] , adopting deeper architecture [ 43 ] , directly introducing convolution operations [ 16 , 48 ] , redesigning input image tokens [ 54 ] , and etc . Besides , some studies attempt to extend the transformer to address broader vision tasks such as object detection [ 4 ] , semantic segmentation [ 63 , 37 ] , pose estimation [ 51 , 23 ] , video understanding [ 61 , 2 , 30 ] , and so on . This work focuses on the transformer for dense prediction tasks , including pose estimation and semantic segmentation . Vision Transformer splits an image into a sequence of image patches of size 16× 16 , and extracts the feature representation of each image patch . Thus , the output representations of Vision Transformer lose the fine-grained spatial details that are essential for accurate dense predictions . The Vision Transformer only outputs a single-scale feature representation , and thus lacks the capability to handle multi-scale variation . To mitigate the loss of feature granularity and model the multi-scale variation , we present High-Resolution Transformer ( HRFormer ) that contains richer spatial information and constructs multi-resolution representations for dense predictions . The High-Resolution Transformer is built by following the multi-resolution parallel design that is adopted in HRNet [ 46 ] . First , HRFormer adopts convolution in both the stem and the first stage as several concurrent studies [ 11 , 50 ] also suggest that convolution performs better in the early stages . Second , HRFormer maintains a high-resolution stream through the entire process with parallel medium- and low-resolution streams helping boost high-resolution representations . With feature maps of different resolutions , thus HRFormer is capable to model the multi-scale variation . Third , ∗Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . HRFormer mixes the short-range and long-range attention via exchanging multi-resolution feature information with the multi-scale fusion module . At each resolution , the local-window self-attention mechanism is adopted to reduce the memory and computation complexity . We partition the representation maps into a set of non-overlapping small image windows and perform self-attention in each image window separately . This reduces the memory and computation complexity from quadratic to linear with respect to spatial size . We further introduce a 3× 3 depth-wise convolution into the feed-forward network ( FFN ) that follows the local-window self-attention , to exchange information between the image windows which are disconnected in the local-window self-attention process . This helps to expand the receptive field and is essential for dense prediction tasks . Figure 1 shows the details of an HRFormer block . We conduct experiments on image classification , pose estimation , and semantic segmentation tasks , and achieve competitive performance on various benchmarks . For example , HRFormer-B gains +1.0 % top-1 accuracy on ImageNet classification over DeiT-B [ 42 ] with 40 % fewer parameters and 20 % fewer FLOPs . HRFormer-B gains 0.9 % AP over HRNet-W48 [ 41 ] on COCO val set with with 32 % fewer parameters and 19 % fewer FLOPs . HRFormer-B + OCR gains +1.2 % and +2.0 % mIoU over HRNet-W48 + OCR [ 55 ] with 25 % fewer parameters and slightly more FLOPs on PASCAL-Context test and COCO-Stuff test , respectively . 2 Related work . Vision Transformers . With the success of Vision Transformer ( ViT ) [ 13 ] and the data-efficient image transformer ( DeiT ) [ 42 ] , various techniques are proposed to improve the ImageNet classification accuracy of Vision Transformer [ 12 , 43 , 48 , 16 , 54 , 17 , 5 , 27 , 22 , 40 ] . Among the very recent advancements , the community has verified several effective improvements such as multi-scale feature hierarchies and incorporating convolutions . For example , the concurrent works MViT [ 14 ] , PVT [ 47 ] , and Swin [ 27 ] introduce the multi-scale feature hierarchies into transformer following the spatial configuration of a typical convolutional architecture such as ResNet-50 . Different from them , our HRFormer incorporates the multi-scale feature hierarchies through exploiting the multi-resolution parallel design inspired by HRNet . CvT [ 48 ] , CeiT [ 53 ] , and LocalViT [ 25 ] propose to enhance the locality of transformer via inserting depth-wise convolutions into either the self-attention or the FFN . The purpose of the inserted convolution within our HRFormer is different , apart from enhancing the locality , it also ensures information exchange across the non-overlapping windows . Several previous studies [ 36 , 19 ] have proposed similar local self-attention schemes for image classification . They construct the overlapped local windows following the strided convolution , resulting in heavy computation cost . Similar to [ 21 , 44 , 27 ] , we propose to apply the local-window self-attention scheme to divide the input feature map into non-overlapping windows . Then we apply the self-attention within each window independently so as to improve the efficiency significantly . There are several concurrently-developed works [ 63 , 37 ] use the Vision Transformer to address the dense predict tasks such as semantic segmentation . They have shown that increasing the spatial resolution of the representations output by the Vision Transformer is important for semantic segmen- tation . Our HRFormer provides a different path to address the low-resolution problem of the Vision Transformer via exploiting the multi-resolution parallel transformer scheme . High-Resolution CNN for Dense Prediction . The high-resolution convolutional schemes have achieved great success on both pose estimation and semantic segmentation tasks . In the development of high-resolution convolutional neural networks , the community has developed three main paths including : ( i ) applying dilated convolutions to remove some down-sample layers [ 6 , 52 ] , ( ii ) recovering high-resolution representations from low-resolution representations with decoders [ 38 , 1 , 31 , 32 ] , and ( iii ) maintaining high-resolution representations throughout the network [ 46 , 15 , 39 , 64 , 45 , 59 , 20 ] . Our HRFormer belongs to the third path , and retains the advantages of both vision transformer and HRNet [ 46 ] . 3 High-Resolution Transformer . Multi-resolution parallel transformer . We follow the HRNet [ 46 ] design and start from a highresolution convolution stem as the first stage , gradually adding high-to-low resolution streams one by one as new stages . The multi-resolution streams are connected in parallel . The main body consists of a sequence of stages . In each stage , the feature representation of each resolution stream is updated with multiple transformer blocks independently and the information across resolutions is exchanged repeatedly with the convolutional multi-scale fusion modules . Figure 2 illustrates the overall HRFormer architecture . The design of convolutional multi-scale fusion modules exactly follows HRNet . We illustrate the details of the transformer block in the following discussion and more details are presented in Figure 1 . Local-window self-attention . We divide the feature maps X ∈ RN×D into a set of non-overlapping small windows : X → { X1 , X2 , · · · , XP } , where each window is of size K × K. We perform multi-head self-attention ( MHSA ) within each window independently . The formulation of multi-head self-attention on the p-th window is given as : MultiHead ( Xp ) = Concat [ head ( Xp ) 1 , · · · , head ( Xp ) H ] ∈ RK 2×D , ( 1 ) head ( Xp ) h = Softmax [ ( XpW h q ) ( XpW h k ) T√ D/H ] XpW h v ∈ RK 2×DH , ( 2 ) X̂p = Xp +MultiHead ( Xp ) Wo ∈ RK 2×DH , ( 3 ) where Wo ∈ RD×D , Whq ∈ R D H×D , Whk ∈ R D H×D , and Whv ∈ R D H×D for h ∈ { 1 , · · · , H } . H represents the number of heads , D represents the number of channels , N represents the input resolutions , and X̂p represents the output representation of MHSA . We also apply the relative position embedding scheme introduced in the T5 model [ 35 ] to incorporate the relative position information into the local-window self-attention . With MHSA aggregates information within each window , we merge them to compute the output XMHSA : { X̂1 , X̂2 , · · · , X̂P } Merge−−−−→ XMHSA . ( 4 ) The left part of Figure 1 illustrates how local-window self-attention updates the 2D input representations , where the multi-head self-attention operates within each window independently . FFN with depth-wise convolution . Local-window self-attention performs self-attention over the non-overlapping windows separately . There is no information exchange across the windows . To handle this issue , we add a 3× 3 depth-wise convolution between the two point-wise MLPs that form the FFN in Vision transformer : MLP ( DW-Conv. ( MLP ( ) ) ) . The right part of Figure 1 shows an example of how FFN with 3× 3 depth-wise convolution updates the 2D input representations . Representation head designs . As shown in Figure 2 , the output of HRFormer consists of four feature maps of different resolutions . We illustrate the details of the representation head designs for different tasks as following : ( i ) ImageNet classification , we send the four-resolution feature maps into a bottleneck and the output channels are changed to 128 , 256 , 512 , and 1024 respectively . Then , we apply the strided convolutions to fuse them and output a feature map of the lowest resolution with 2048 channels . Last , we apply a global average pooling operation followed by the final classifier . ( ii ) pose estimation , we only apply the regression head over the highest resolution feature map . ( iii ) semantic segmentation , we apply the semantic segmentation head over the concatenated representations , which are computed by first upsampling all the low-resolution representations to the highest resolution and then concatenate them together . depth-wise convolution . We have not included the convolutional multi-scale fusion modules in Table 1 for simplicity . In our implementation , we set the size of the windows on four resolution streams as ( 7 , 7 , 7 , 7 ) by default . Table 2 illustrates the configuration details of three different HRFormer instances with increasing complexities , where the MLP expansion ratios ( R1 , R2 , R3 , R4 ) are set as ( 4 , 4 , 4 , 4 ) for all models and are not shown . Analysis . The benefits of 3× 3 depth-wise convolution are twofold : one is enhancing the locality and the other one is enabling the interactions across windows . We illustrate how the FFN with depth-wise convolution is capable to expand the interactions beyond the non-overlapping local windows and model the relations between them in Figure 3 . Therefore , based on the combination of the localwindow self-attention and the FFN with 3× 3 depth-wise convolution , we can build the HRFormer block that improves the memory and computation efficiency significantly . 4 Experiments . 4.1 Human Pose Estimation . Training setting . We study the performance of HRFormer on the COCO [ 26 ] human pose estimation benchmark , which contains more than 200K images and 250K person instances labeled with 17 keypoints . We train our model on COCO train 2017 dataset , including 57K images and 150K person instances . We evaluate our approach on the val 2017 set and test-dev 2017 , containing 5K images and 20K images , respectively . We follow most of the default training and evaluation settings of mmpose [ 8 ] 2 , and change the optimizer from Adam to AdamW . For the training batch size , we choose 256 for HRFormer-T and HRFormer-S and 128 for HRFormer-B due to limited GPU memory . Each HRFormer experiment on COCO pose estimation task takes 8× 32G-V100 GPUs . Results . Table 3 reports the comparisons on COCO val set . We compare HRFormer to the representative convolutional method such as HRNet [ 41 ] and several recent transformer methods , including PRTR [ 23 ] , TransPose-H-A6 [ 51 ] , and TokenPose-L/D24 [ 24 ] . HRFormer-B gains 0.9 % with 32 % fewer parameters and 19 % fewer FLOPs when compared to HRNet-W48 with an input size of 384 × 288 . Therefore , our HRFormer-B already achieves 77.2 % w/o using any advanced techniques such as UDP [ 20 ] and DARK [ 59 ] . We believe that our HRFormer-B could achieve better results by exploiting either UDP or DARK scheme . We also report the comparisons on COCO test-dev set in Table 4 . Our HRFormer-B outperforms HRNet-W48 by around 0.7 % with fewer parameters and FLOPs . Figure 4 shows some example results of human pose estimation on COCO val set . 2https : //github.com/open-mmlab/mmpose , Apache License 2.0

Summary: This paper proposes a way to replace the main convolution layers in the HRNet with self-attention layers. A feedforward network with 3x3 depth-wise convolution is also proposed to replace the MLP in transformer. The proposed High-Resolution Transformer achieves performance improvements for dense prediction tasks such as pose estimation and semantic segmentation while saving memory and computation.