A Transformer-in-Transformer Network Utilizing Knowledge Distillation for Image Recognition

TL;DR

This work addresses the limitations of pure vision transformers in terms of data, compute, and the ability to capture both local and global image structure. The authors introduce Transformer-in-Transformer (TITN), an inner-outer transformer framework that leverages knowledge distillation from a larger teacher, including a learnable distillation token and a hybrid loss combining Cross-Entropy, CutMix, and distillation terms. TITN achieves state-of-the-art top-1/top-5 accuracy on CIFAR-100 (74.71%/92.28%), CIFAR-10 (92.03%/99.80%), and MNIST (99.56% top-1), while aiming for improved training efficiency. The approach demonstrates a promising pathway to efficient, high-accuracy image recognition by jointly modeling local and global features and transferring knowledge from larger models.

Abstract

This paper presents a novel knowledge distillation neural architecture leveraging efficient transformer networks for effective image classification. Natural images display intricate arrangements encompassing numerous extraneous elements. Vision transformers utilize localized patches to compute attention. However, exclusive dependence on patch segmentation proves inadequate in sufficiently encompassing the comprehensive nature of the image. To address this issue, we have proposed an inner-outer transformer-based architecture, which gives attention to the global and local aspects of the image. Moreover, The training of transformer models poses significant challenges due to their demanding resource, time, and data requirements. To tackle this, we integrate knowledge distillation into the architecture, enabling efficient learning. Leveraging insights from a larger teacher model, our approach enhances learning efficiency and effectiveness. Significantly, the transformer-in-transformer network acquires lightweight characteristics by means of distillation conducted within the feature extraction layer. Our featured network's robustness is established through substantial experimentation on the MNIST, CIFAR10, and CIFAR100 datasets, demonstrating commendable top-1 and top-5 accuracy. The conducted ablative analysis comprehensively validates the effectiveness of the chosen parameters and settings, showcasing their superiority against contemporary methodologies. Remarkably, the proposed Transformer-in-Transformer Network (TITN) model achieves impressive performance milestones across various datasets: securing the highest top-1 accuracy of 74.71% and a top-5 accuracy of 92.28% for the CIFAR100 dataset, attaining an unparalleled top-1 accuracy of 92.03% and top-5 accuracy of 99.80% for the CIFAR-10 dataset, and registering an exceptional top-1 accuracy of 99.56% for the MNIST dataset.

Figures (8)

• Figure 1: Inner & outer transformer block of the proposed transformer-in-transformer network. Here, both blocks perform sequentially to generate class and distillation tokens. The inner transformer block rearranges the layer to be fed for the architecture.
• Figure 2: Neural architecture of the proposed TITN. The converted image patches have been fed into the inner transformer block, then the outer transformer block, and finally processed into class and distillation tokens.
• Figure 3: Proposed loss function leveraging both distillation loss and cross-entropy loss. The loss function utilizes both cutmix and cross-entropy criteria for student and teacher models respectively.
• Figure 4: Performance comparison on CIFAR-100 dataset against various pre-trained teacher models.
• Figure 5: Performance evaluation on CIFAR-100 using VGG16-BN as teacher model.