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: This paper proposes a novel vision transformer (TNT) to better model the fine-grained relationship within local image patches.  Specifically, based on the patch splitting scheme in ViT, the authors further split the local patches ("visual sentence") into smaller ones ("visual words"). Inner transformers are adopted to model relationships between words and improve sentence representation. Experiments on ImageNet classification and several downstream tasks validate the effectiveness of the proposed method.