Vision Transformers for Dense Prediction

1 Vision Transformers for Dense PredictionRen e RanftlAlexey BochkovskiyIntel KoltunAbstractWe introduce Dense Prediction Transformers , an archi-tecture that leverages Vision Transformers in place of con-volutional networks as a backbone for Dense predictiontasks. We assemble tokens from various stages of the vi-sion transformer into image-like representations at vari-ous resolutions and progressively combine them into full-resolution predictions using a convolutional decoder. Thetransformer backbone processes representations at a con-stant and relatively high resolution and has a global re-ceptive field at every stage. These properties allow thedense Prediction transformer to provide finer-grained andmore globally coherent predictions when compared to fully-convolutional networks. Our experiments show that thisarchitecture yields substantial improvements on Dense pre-diction tasks, especially when a large amount of train-ing data is available.

2 For monocular depth estimation,we observe an improvement of up to 28% in relativeperformance when compared to a state-of-the-art fully-convolutional network. When applied to semantic segmen-tation, Dense Prediction Transformers set a new state ofthe art on ADE20K with mIoU. We further showthat the architecture can be fine-tuned on smaller datasetssuch as NYUv2, KITTI, and Pascal Context where it alsosets the new state of the art. Our models are available IntroductionVirtually all existing architectures for Dense predictionare based on convolutional networks [6, 33, 36, 44, 51,52, 55]. The design of Dense Prediction architectures com-monly follows a pattern that logically separates the networkinto an encoder and a decoder. The encoder is frequentlybased on an image classification network, also called thebackbone, that is pretrained on a large corpus such as Im-ageNet [9].

3 The decoder aggregates features from the en-coder and converts them to the final Dense predictions. Ar-chitectural research on Dense Prediction frequently focuseson the decoder and its aggregation strategy [6, 7, 52, 55].However, it is widely recognized that the choice of back-bone architecture has a large influence on the capabilitiesof the overall model, as any information that is lost in theencoder is impossible to recover in the backbones progressively downsample theinput image to extract features at multiple scales. Down-sampling enables a progressive increase of the receptivefield, the grouping of low-level features into abstract high-level features, and simultaneously ensures that memoryand computational requirements of the network remaintractable. However, downsampling has distinct drawbacksthat are particularly salient in Dense Prediction tasks: fea-ture resolution and granularity are lost in the deeper stagesof the model and can thus be hard to recover in the feature resolution and granularity may not matter forsome tasks, such as image classification, they are criticalfor Dense Prediction , where the architecture should ideallybe able to resolve features at or close to the resolution of theinput techniques to mitigate the loss of feature gran-ularity have been include training athigher input resolution (if the computational budget per-mits), dilated convolutions [51] to rapidly increase the re-ceptive field without downsampling, appropriately-placedskip connections from multiple stages of the encoder tothe decoder [33], or, more recently, by connecting multi-resolution representations in parallel throughout the net-work [44].

4 While these techniques can significantly im-prove Prediction quality, the networks are still bottleneckedby their fundamental building block: the convolution. Con-volutions together with non-linearities form the fundamen-tal computational unit of image analysis networks. Convo-lutions, by definition, are linear operators that have a lim-ited receptive field. The limited receptive field and the lim-ited expressivity of an individual convolution necessitate se-quential stacking into very deep architectures to acquire suf-ficiently broad context and sufficiently high representationalpower. This, however, requires the production of many in-termediate representations that require a large amount ofmemory. Downsampling the intermediate representations12179is necessary to keep memory consumption at levels that arefeasible with existing computer this work, we introduce the Dense Prediction trans-former (DPT).

5 DPT is a Dense Prediction architecture that isbased on an encoder-decoder design that leverages a trans-former as the basic computational building block of the en-coder. Specifically, we use the recently proposed visiontransformer (ViT) [11] as a backbone architecture. We re-assemble the bag-of-words representation that is providedby ViT into image-like feature representations at variousresolutions and progressively combine the feature repre-sentations into the final Dense Prediction using a convolu-tional decoder. Unlike fully-convolutional networks, the vi-sion transformer backbone foregoes explicit downsamplingoperations after an initial image embedding has been com-puted and maintains a representation with constant dimen-sionality throughout all processing stages. It furthermorehas a global receptive field at every stage. We show thatthese properties are especially advantageous for Dense pre-diction tasks as they naturally lead to fine-grained and glob-ally coherent conduct experiments on monocular depth estimationand semantic segmentation.

6 For the task of general-purposemonocular depth estimation [32], where large-scale train-ing data is available, DPT provides a performance increaseof more than 28% when compared to the top-performingfully-convolutional network for this task. The architecturecan also be fine-tuned to small monocular depth predictiondatasets, such as NYUv2 [37] and KITTI [15], where it alsosets the new state of the art. We provide further evidenceof the strong performance of DPT using experiments on se-mantics segmentation. For this task, DPT sets a new state ofthe art on the challenging ADE20K [56] and Pascal Context[28] Related WorkFully-convolutional networks [35, 36] are the prototyp-ical architecture for pixel-level Dense Prediction tasks suchas semantic segmentation [7, 25, 55], monocular depth esti-mation [12, 16, 32], and keypoint detection [21, 57]. Manyvariants of this pattern have been proposed over the years,however, all existing architectures adopt convolution andsubsampling as their fundamental elements in order to learnmulti-scale representations that can leverage an appropri-ately large context.

7 Several works propose to progressivelyupsample representations that have been pooled at differ-ent stages [1, 25, 29, 33], while others use dilated convo-lutions [6, 7, 51] or parallel multi-scale feature aggregationat multiple scales [55] to recover fine-grained predictionswhile at the same time ensuring a sufficiently large recent architectures maintain a high-resolution repre-sentation together with multiple lower-resolution represen-tations throughout the network [39, 44].Attention-based models [2] and in particular transform-ers [41] have been the architecture of choice for learningstrong models for natural language processing (NLP) [4,10, 26] in recent years. Transformers are set-to-set mod-els that are based on the self-attention mechanism. Trans-former models have been particularly successful when in-stantiated as high-capacity architectures and trained on verylarge datasets.

8 There have been several works that adapt at-tention mechanisms to image analysis [3, 30, 31, 43, 54]. Inparticular, it has recently been demonstrated that a direct ap-plication of token-based transformer architectures that havebeen successful in NLP can yield competitive performanceon image classification [11]. A key insight of this work wasthat, like transformer models in NLP, Vision transformersneed to be paired with a sufficient amount of training datato realize their ArchitectureThis section introduces the Dense Prediction maintain the overall encoder-decoder structure that hasbeen successful for Dense Prediction in the past. We lever-age Vision Transformers [11] as the backbone, show how therepresentation that is produced by this encoder can be ef-fectively transformed into Dense predictions, and provideintuition for the success of this strategy. An overview of thecomplete architecture is shown in Figure 1 (left).

9 Transformer a high level, the Vision trans-former (ViT) [11] operates on a bag-of-words representa-tion of the image [38]. Image patches that are individuallyembedded into a feature space, or alternatively deep fea-tures extracted from the image, take the role of words .We will refer to embedded words astokensthroughoutthe rest of this work. Transformers transform the set of to-kens using sequential blocks of multi-headed self-attention(MHSA) [41], which relate tokens to each other to trans-form the for our application, a transformer maintainsthe number of tokens throughout all computations. Since to-kens have a one-to-one correspondence with image patches,this means that the ViT encoder maintains the spatial reso-lution of the initial embedding throughout all transformerstages. Additionally, MHSA is an inherently global oper-ation, as every token can attend to and thus influence ev-ery other token.

10 Consequently, the transformer has a globalreceptive field at every stage after the initial is in stark contrast to convolutional networks, whichprogressively increase their receptive field as features passthrough consecutive convolution and downsampling specifically, ViT extracts a patch embedding fromthe image by processing all non-overlapping square patchesof sizep2pixels from the image. The patches are flattenedinto vectors and individually embedded using a linear pro-12180 TransformerTransformer TransformerFusionFusionFusionFusionHeadE mbedTransformerReassemble32 Reassemble16 Reassemble8 Reassemble4 ReassemblesConcatenateReadResamplesProje ctResidual Conv UnitResidual Conv +ProjectFusionFigure : Architecture overview. The input image is transformed into tokens (orange) either by extracting non-overlapping patchesfollowed by a linear projection of their flattened representation (DPT-Base and DPT-Large) or by applying a ResNet-50 feature extractor(DPT-Hybrid).