Source-linked AI summary
Transformer in Transformer
Kai Han, An Xiao, Enhua Wu, Jianyuan Guo, Chunjing Xu, Yunhe Wang
TL;DR
Visual transformers’ patch-level representations may be too coarse for the detail and diversity of natural images. TNT introduces nested transformers that model both patches as visual sentences and sub-patches as visual words, achieving 81.5% ImageNet top-1 accuracy, 1.7% above DeiT-S.
Problem
Patch-based visual transformers may not capture the fine-grained details and internal relationships needed for complex natural images.
Method
TNT applies an outer transformer to visual-sentence embeddings and an inner transformer to visual-word embeddings within each patch, then aggregates their information.
Results
81.5% top-1 accuracy on ImageNet is achieved by TNT-S, 1.7% higher than DeiT-S.
Takeaways & Limitations
TNT provides fine-grained visual information while preserving patch-level representations for visual recognition.
Abstract
from arXiv · showhide
Transformer is a new kind of neural architecture which encodes the input data as powerful features via the attention mechanism. Basically, the visual transformers first divide the input images into several local patches and then calculate both representations and their relationship. Since natural images are of high complexity with abundant detail and color information, the granularity of the patch dividing is not fine enough for excavating features of objects in different scales and locations. In this paper, we point out that the attention inside these local patches are also essential for building visual transformers with high performance and we explore a new architecture, namely, Transformer iN Transformer (TNT). Specifically, we regard the local patches (e.g., 16$\times$16) as "visual sentences" and present to further divide them into smaller patches (e.g., 4$\times$4) as "visual words". The attention of each word will be calculated with other words in the given visual sentence with negligible computational costs. Features of both words and sentences will be aggregated to enhance the representation ability. Experiments on several benchmarks demonstrate the effectiveness of the proposed TNT architecture, e.g., we achieve an 81.5% top-1 accuracy on the ImageNet, which is about 1.7% higher than that of the state-of-the-art visual transformer with similar computational cost. The PyTorch code is available at https://github.com/huawei-noah/CV-Backbones, and the MindSpore code is available at https://gitee.com/mindspore/models/tree/master/research/cv/TNT.
1 Introduction
Visual transformers extend attention-based architectures to images, but patch-only representations can miss important sub-patch structure. TNT addresses this by modeling both visual sentences and finer visual words with nested transformers.
- Transformer architectures: Transformer architectures use self-attention to model relationships between features and have been extended from NLP to visual tasks.Prior work applied transformers to long-range dependencies, object detection, and image recognition.
- Visual transformers: ViT represents an image as a sequence of local patches and computes attention between patches for recognition.This patch-based formulation enabled transformer models to transfer successfully to visual tasks.
- Motivation: Natural-image diversity and detail motivate examining similar sub-patches within local patches rather than modeling only relationships between patches.Patch representations can contain internal regions with high similarity that warrant finer-grained modeling.
- Proposed approach: TNT divides images into visual sentences and further divides each sentence into visual words processed by an inner transformer.Word features and within-sentence word attentions are computed with a shared network, adding negligible parameters and FLOPs.
2 Approach
TNT extends patch-based visual transformers with nested word- and sentence-level processing, combining local feature extraction with relationships among image patches. Its additional local modeling raises computation modestly while preserving practical model variants and efficiency.
- Transformer in Transformer: TNT divides images into patches as visual sentences and further splits each patch into sub-patches as visual words.The word embeddings are formed by vectorizing and linearly projecting each sub-patch.
- Transformer in Transformer: The inner transformer models relationships among visual words, while the outer transformer models relationships among sentence embeddings.The two data flows separately process words within each sentence and sentences across the image.
- Transformer in Transformer: Word features are projected into sentence-embedding space and added to augment sentence representations before outer-transformer processing.A fully connected layer matches dimensions for the addition.
- Position encoding: Sentence and word position encodings preserve global spatial information and local relative position, respectively.Both embedding levels receive learnable 1D positional encodings.
- Complexity analysis: 1.14× FLOPs and 1.08× parameters are the TNT-to-standard-transformer ratios in the DeiT-S configuration.The corresponding standard-transformer and TNT-block costs are 376M and 429M FLOPs, respectively.
- Network architecture: TNT-Ti, TNT-S, and TNT-B contain 6.1M, 23.8M, and 65.6M parameters, with 1.4B, 5.2B, and 14.1B FLOPs at 224×224 resolution.The default configuration uses 16×16 patches and m = 16 sub-patches.
3 Experiments
Experiments evaluate TNT on ImageNet, downstream classification, detection, and segmentation, alongside efficiency, ablation, and visualization studies. TNT improves visual-transformer performance while modeling both patch-level and within-patch relationships.
- 3.2 TNT on ImageNet: TNT-S achieves 81.5% top-1 accuracy on ImageNet, 1.7% higher than DeiT-S under the same training settings.The comparison includes representative transformer and CNN backbones.
- 3.2 TNT on ImageNet: TNT models consistently outperform other transformer-based models across accuracy-parameters and accuracy-FLOPs comparisons.These comparisons are visualized in Figure 2.
- 3.2 TNT on ImageNet: TNT is more efficient than DeiT and PVT in GPU throughput experiments using 224×224 inputs on an NVIDIA V100.The architecture can reduce TNT blocks and replace some with vanilla transformer blocks.
- 3.3 Ablation Studies: Including both sentence and word position encodings gives TNT-S its best result, while removing them causes accuracy drops of 0.8% and 0.7%, respectively.Sentence encodings preserve global spatial information; word encodings preserve local relative position.
- 3.3 Ablation Studies: A proper inner-transformer head count, such as 2 or 4, performs best, while changing the number of visual words has only slight performance influence.The default configuration uses m = 16 visual words for efficiency.
- 3.5 Transfer Learning: On downstream tasks, TNT outperforms DeiT on most image-classification datasets, improves DETR over PVT-Small by 3.5 AP, and reaches 43.6% mIoU on ADE20K.The ADE20K result is 1.0% higher than PVT-small and 2.8% higher than DeiT-S with similar parameters.
4 Conclusion
TNT models both global and local image information by nesting an inner word-level transformer within an outer sentence-level transformer. The architecture preserves and models local patch information, with effectiveness demonstrated on ImageNet and downstream tasks.
- TNT processes visual sentences with an outer transformer and visual words with an inner transformer to model global and local information.Word embeddings are projected and added to sentence embeddings before stacking TNT blocks.
- TNT is designed to better preserve and model local patch structure than conventional vision transformers such as ViT.
- TNT achieves effective performance on ImageNet and downstream visual recognition tasks.
A.1 Visualization of Attention Maps
Attention visualizations compare how DeiT-S and TNT-S distribute attention across patches and how class-token attention relates to recognized objects. TNT focuses on meaningful patches more clearly in a deep block, while output features primarily attend to object-related patches.
- In deeper layers, both DeiT-S and TNT-S relate more patches as information is communicated across the image.
- In Block-12, TNT focuses attention on meaningful panda patches, whereas DeiT still attends to an unrelated tree.
- Class-token attention maps show that the output feature mainly focuses on patches related to the object being recognized.
- Figure 6 visualizes attention maps between all patches in the outer transformer block.
- Figure 7 shows example attention maps from the output token to the input space.
A.2 Exploring SE module in TNT
The TNT architecture is extended with squeeze-and-excitation channel-wise attention. This adds dimension-wise feature enhancement with few extra parameters and further improves accuracy.
- The SE module averages sentence or word embeddings and uses a two-layer MLP to calculate attention values applied to all embeddings.
- SE adds only a few parameters while providing dimension-wise attention for feature enhancement.
- Table 12 reports the accuracy effects of adding the SE module to TNT.
A.3 Object Detection with Faster RCNN
TNT is evaluated as a multi-scale backbone for Faster RCNN with FPN on COCO object detection. Features from multiple TNT layers are combined, and TNT outperforms the compared ResNet and DeiT backbones in the reported results.
- TNT supplies multi-scale features to Faster RCNN by selecting representations from its 3rd, 6th, 9th, and 12th layers.FPN receives four feature levels, with deconvolution or convolution layers used to match feature shapes.
- TNT-S and DeiT-S are evaluated with FPN for Faster RCNN object detection on the COCO2017 validation set.
- TNT achieves much better performance than ResNet and DeiT backbones in the reported COCO object-detection results.
- Table 13 reports Faster RCNN object-detection results on the COCO minival set with ImageNet pre-training.