Transformer in Transformer

1 Introduction . In the past decade , the mainstream deep neural architectures used in the computer vision ( CV ) are mainly established on convolutional neural networks ( CNNs ) [ 17 , 12 , 11 ] . Differently , transformer is a type of neural network mainly based on self-attention mechanism [ 35 ] , which can provide the relationships between different features . Transformer is widely used in the field of natural language processing ( NLP ) , e.g. , the famous BERT [ 8 ] and GPT-3 [ 2 ] models . The power of these transformer models inspires the whole community to investigate the use of transformer for visual tasks . To utilize the transformer architectures for conducting visual tasks , a number of researchers have explored for representing the sequence information from different data . For example , Wang et al . explore self-attention mechanism in non-local networks [ 37 ] for capturing long-range dependencies in video and image recognition . Carion et al . present DETR [ 3 ] , which treats object detection as a direct set prediction problem and solve it using a transformer encoder-decoder architecture . Chen et al . propose the iGPT [ 5 ] , which is the pioneering work applying pure transformer model ( i.e. , without convolution ) on image recognition by self-supervised pre-training . ∗Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Different from the data in NLP tasks , there exists a semantic gap between input images and the ground-truth labels in CV tasks . To this end , Dosovitskiy et al . develop the ViT [ 9 ] , which paves the way for transferring the success of transformer based NLP models . Concretely , ViT divides the given image into several local patches as a visual sequence . Then , the attention can be naturally calculated between any two image patches for generating effective feature representations for the recognition task . Subsequently , Touvron et al . explore the data-efficient training and distillation to enhance the performance of ViT on the ImageNet benchmark and obtain an about 81.8 % ImageNet top-1 accuracy , which is comparable to that of the state-of-the-art convolutional networks . Chen et al . further treat the image processing tasks ( e.g. , denosing and super-resolution ) as a series of translations and develop the IPT model for handling multiple low-level computer vision problems [ 4 ] . Nowadays , transformer architectures have been used in a growing number of computer vision tasks [ 10 ] such as image recognition [ 6 , 40 , 29 ] , object detection [ 46 ] , and segmentation [ 43 , 38 ] . Although the aforementioned visual transformers have made great efforts to boost the models ’ performances , most of existing works follow the conventional representation scheme used in ViT , i.e. , dividing the input images into patches . Such a exquisite paradigm can effectively capture the visual sequential information and estimate the attention between different image patches . However , the diversity of natural images in modern benchmarks is very high , e.g. , there are over 120 M images with 1000 different categories in the ImageNet dataset [ 26 ] . As shown in Figure 1 , representing the given image into local patches can help us to find the relationship and similarity between them . However , there are also some sub-patches inside them with high similarity . Therefore , we are motivated to explore a more exquisite visual image dividing method for generating visual sequences and improve the performance . In this paper , we propose a novel Transformer-iN-Transformer ( TNT ) architecture for visual recognition as shown in Figure 1 . To enhance the feature representation ability of visual transformers , we first divide the input images into several patches as “ visual sentences ” and then further divide them into sub-patches as “ visual words ” . Besides the conventional transformer blocks for extracting features and attentions of visual sentences , we further embed a sub-transformer into the architecture for excavating the features and details of smaller visual words . Specifically , features and attentions between visual words in each visual sentence are calculated independently using a shared network so that the increased amount of parameters and FLOPs ( floating-point operations ) is negligible . Then , features of words will be aggregated into the corresponding visual sentence . The class token is also used for the subsequent visual recognition task via a fully-connected head . Through the proposed TNT model , we can extract visual information with fine granularity and provide features with more details . We then conduct a series of experiments on the ImageNet benchmark and downstream tasks to demonstrate its superiority and thoroughly analyze the impact of the size for dividing visual words . The results show that our TNT can achieve better accuracy and FLOPs trade-off over the state-of-the-art transformer networks . 2 Approach . In this section , we describe the proposed transformer-in-transformer architecture and analyze the computation and parameter complexity in details . 2.1 Preliminaries . We first briefly describe the basic components in transformer [ 35 ] , including MSA ( Multi-head Self-Attention ) , MLP ( Multi-Layer Perceptron ) and LN ( Layer Normalization ) . MSA . In the self-attention module , the inputs X ∈ Rn×d are linearly transformed to three parts , i.e. , queries Q ∈ Rn×dk , keys K ∈ Rn×dk and values V ∈ Rn×dv where n is the sequence length , d , dk , dv are the dimensions of inputs , queries ( keys ) and values , respectively . The scaled dot-product attention is applied on Q , K , V : Attention ( Q , K , V ) = softmax ( QKT√ dk ) V. ( 1 ) Finally , a linear layer is used to produce the output . Multi-head self-attention splits the queries , keys and values to h parts and perform the attention function in parallel , and then the output values of each head are concatenated and linearly projected to form the final output . MLP . The MLP is applied between self-attention layers for feature transformation and non-linearity : MLP ( X ) = FC ( σ ( FC ( X ) ) ) , FC ( X ) = XW + b , ( 2 ) where W and b are the weight and bias term of fully-connected layer respectively , and σ ( · ) is the activation function such as GELU [ 13 ] . LN . Layer normalization [ 1 ] is a key part in transformer for stable training and faster convergence . LN is applied over each sample x ∈ Rd as follows : LN ( x ) = x− µ δ ◦ γ + β ( 3 ) where µ ∈ R , δ ∈ R are the mean and standard deviation of the feature respectively , ◦ is the element-wise dot , and γ ∈ Rd , β ∈ Rd are learnable affine transform parameters . 2.2 Transformer in Transformer . Given a 2D image , we uniformly split it into n patches X = [ X1 , X2 , · · · , Xn ] ∈ Rn×p×p×3 , where ( p , p ) is the resolution of each image patch . ViT [ 9 ] just utilizes a standard transformer to process the sequence of patches which corrupts the local structure of a patch , as shown in Fig . 1 ( a ) . Instead , we propose Transformer-iN-Transformer ( TNT ) architecture to learn both global and local information in an image . In TNT , we view the patches as visual sentences that represent the image . Each patch is further divided into m sub-patches , i.e. , a visual sentence is composed of a sequence of visual words : Xi → [ xi,1 , xi,2 , · · · , xi , m ] , ( 4 ) where xi , j ∈ Rs×s×3 is the j-th visual word of the i-th visual sentence , ( s , s ) is the spatial size of sub-patches , and j = 1 , 2 , · · · , m. With a linear projection , we transform the visual words into a sequence of word embeddings : Y i = [ yi,1 , yi,2 , · · · , yi , m ] , yi , j = FC ( Vec ( xi , j ) ) , ( 5 ) where yi , j ∈ Rc is the j-th word embedding , c is the dimension of word embedding , and Vec ( · ) is the vectorization operation . In TNT , we have two data flows in which one flow operates across the visual sentences and the other processes the visual words inside each sentence . For the word embeddings , we utilize a transformer block to explore the relation between visual words : Y ′ i l = Y i l−1 + MSA ( LN ( Y i l−1 ) ) , ( 6 ) Y il = Y ′i l + MLP ( LN ( Y ′i l ) ) . ( 7 ) where l = 1 , 2 , · · · , L is the index of the l-th block , and L is the total number of stacked blocks . The input of the first block Y i0 is just Y i in Eq . 5 . All word embeddings in the image after transformation are Yl = [ Y 1l , Y 2l , · · · , Y nl ] . This can be viewed as an inner transformer block , denoted as Tin . This process builds the relationships among visual words by computing interactions between any two visual words . For example , in a patch of human face , a word corresponding to the eye is more related to other words of eyes while interacts less with forehead part . For the sentence level , we create the sentence embedding memories to store the sequence of sentencelevel representations : Z0 = [ Zclass , Z10 , Z20 , · · · , Zn0 ] ∈ R ( n+1 ) ×d where Zclass is the class token similar to ViT [ 9 ] , and all of them are initialized as zero . In each layer , the sequence of word embeddings are transformed into the domain of sentence embedding by linear projection and added into the sentence embedding : Zil−1 = Z i l−1 + FC ( Vec ( Y i l ) ) , ( 8 ) where Zil−1 ∈ Rd and the fully-connected layer FC makes the dimension match for addition . With the above addition operation , the representation of sentence embedding is augmented by the word-level features . We use the standard transformer block for transforming the sentence embeddings : Z ′l = Zl−1 + MSA ( LN ( Zl−1 ) ) , ( 9 ) Zl = Z ′l + MLP ( LN ( Z ′l ) ) . ( 10 ) This outer transformer block Tout is used for modeling relationships among sentence embeddings . In summary , the inputs and outputs of the TNT block include the visual word embeddings and sentence embeddings as shown in Fig . 1 ( b ) , so the TNT can be formulated as Yl , Zl = TNT ( Yl−1 , Zl−1 ) . ( 11 ) In our TNT block , the inner transformer block is used to model the relationship between visual words for local feature extraction , and the outer transformer block captures the intrinsic information from the sequence of sentences . By stacking the TNT blocks for L times , we build the transformerin-transformer network . Finally , the classification token serves as the image representation and a fully-connected layer is applied for classification . Position encoding . Spatial information is an important factor in image recognition . For sentence embeddings and word embeddings , we both add the corresponding position encodings to retain spatial information as shown in Fig . 1 . The standard learnable 1D position encodings are utilized here . Specifically , each sentence is assigned with a position encodings : Z0 ← Z0 + Esentence , ( 12 ) where Esentence ∈ R ( n+1 ) ×d are the sentence position encodings . As for the visual words in a sentence , a word position encoding is added to each word embedding : Y i0 ← Y i0 + Eword , i = 1 , 2 , · · · , n ( 13 ) where Eword ∈ Rm×c are the word position encodings which are shared across sentences . In this way , sentence position encoding can maintain the global spatial information , while word position encoding is used for preserving the local relative position .

Summary: The work proposes a nested vision Transformer variant (TNT), which adds an additional inner Transformer to capture within-patch information. By setting the hidden dimension of the inner Transformer much smaller, the FLOPs and #Params of the proposed TNT won't be overly large. Empirically, authors show that the proposed model performs better than ViT/DeiT on ImageNet classification, COCO object detection and ADE20K semantic segmentation.