[Paper Review] PyramidTNT: Improved Transformer-in-Transformer Baselines with Pyramid Architecture

Kai Han, Jianyuan Guo+2 more|arXiv (Cornell University)|Jan 4, 2022

Advanced Memory and Neural ComputingEngineering18 citations

TL;DR

PyramidTNT introduces a hierarchical pyramid architecture and convolutional stem to improve the Transformer-in-Transformer (TNT) backbone, achieving state-of-the-art performance on image classification and object detection. It achieves 82.0% top-1 accuracy on ImageNet-1K with only 3.3B FLOPs, outperforming Swin-T and original TNT while maintaining lower computational cost.

ABSTRACT

Transformer networks have achieved great progress for computer vision tasks. Transformer-in-Transformer (TNT) architecture utilizes inner transformer and outer transformer to extract both local and global representations. In this work, we present new TNT baselines by introducing two advanced designs: 1) pyramid architecture, and 2) convolutional stem. The new "PyramidTNT" significantly improves the original TNT by establishing hierarchical representations. PyramidTNT achieves better performances than the previous state-of-the-art vision transformers such as Swin Transformer. We hope this new baseline will be helpful to the further research and application of vision transformer. Code will be available at https://github.com/huawei-noah/CV-Backbones/tree/master/tnt_pytorch.

Motivation & Objective

• To improve the performance of the Transformer-in-Transformer (TNT) architecture for vision tasks by introducing hierarchical feature learning.
• To enhance training stability and representation learning through a convolutional stem replacing the standard patchify stem.
• To establish a new, stronger baseline for vision transformers that outperforms existing models like Swin Transformer and original TNT.
• To enable efficient multi-scale feature extraction via a pyramid structure with gradually reduced spatial resolution across stages.

Proposed method

• Introduces a pyramid architecture with four stages, each reducing spatial resolution by a factor of 2, enabling hierarchical feature learning.
• Employs a convolutional stem with five 3x3 convolutions to improve patch embedding and training stability.
• Uses Patch Merging layers for downsampling between stages, preserving feature hierarchy and reducing sequence length.
• Applies a dual-branch transformer design: outer transformers process global representations, inner transformers model local patch relationships.
• Incorporates Learnable Spatial Reduction Attention (LSRA) with learnable reduction ratios to compress local tokens.
• Uses standard training techniques including weight decay, AdamW optimizer, and multi-scale data augmentation for robustness.

Experimental results

Research questions

• RQ1 Can a pyramid architecture improve the representational capacity of the TNT model in vision tasks?
• RQ2 Does replacing the standard patchify stem with a convolutional stem enhance training stability and performance?
• RQ3 Can PyramidTNT achieve better accuracy than existing state-of-the-art vision transformers like Swin Transformer with reduced FLOPs?
• RQ4 How does the hierarchical feature learning in PyramidTNT affect object detection performance on COCO?

Key findings

• PyramidTNT-S achieves 82.0% top-1 accuracy on ImageNet-1K with only 3.3B FLOPs, outperforming both original TNT-S and Swin-T.
• On COCO object detection, PyramidTNT-S achieves 42.0 mAP with RetinaNet under 1x schedule, surpassing Swin-T by 0.5 mAP.
• With multi-scale training and 3x schedule, PyramidTNT-S achieves 51.0 mAP on Mask R-CNN, exceeding Hire-MLP-S by 0.9 mAP while using fewer FLOPs.
• The model maintains strong performance across all backbone sizes (Ti, S, M, B), with FLOPs ranging from 0.6B to 16.0B.
• The combination of pyramid structure and convolutional stem enables better feature representation, particularly for large objects in detection tasks.