Fast-iTPN: Integrally Pre-Trained Transformer Pyramid Network with Token Migration

TL;DR

This work introduces iTPN, an integrally pre-trained transformer pyramid network that jointly optimizes backbone and neck by embedding a feature pyramid into the pre-training stage and applying masked feature modeling (MFM) for multi-stage supervision. Fast-iTPN further accelerates inference with token migration and token gathering to reduce global attention overhead while preserving core advantages from integral pre-training. On ImageNet-1K, iTPN achieves state-of-the-art top-1 accuracies for base and large models, and Fast-iTPN further improves speed with competitive or superior downstream performance on COCO and ADE20K. The combination of integral pre-training, MFM, and efficiency strategies demonstrates strong potential for a versatile, scalable backbone for image classification, detection, and segmentation.

Abstract

We propose integrally pre-trained transformer pyramid network (iTPN), towards jointly optimizing the network backbone and the neck, so that transfer gap between representation models and downstream tasks is minimal. iTPN is born with two elaborated designs: 1) The first pre-trained feature pyramid upon vision transformer (ViT). 2) Multi-stage supervision to the feature pyramid using masked feature modeling (MFM). iTPN is updated to Fast-iTPN, reducing computational memory overhead and accelerating inference through two flexible designs. 1) Token migration: dropping redundant tokens of the backbone while replenishing them in the feature pyramid without attention operations. 2) Token gathering: reducing computation cost caused by global attention by introducing few gathering tokens. The base/large-level Fast-iTPN achieve 88.75%/89.5% top-1 accuracy on ImageNet-1K. With 1x training schedule using DINO, the base/large-level Fast-iTPN achieves 58.4%/58.8% box AP on COCO object detection, and a 57.5%/58.7% mIoU on ADE20K semantic segmentation using MaskDINO. Fast-iTPN can accelerate the inference procedure by up to 70%, with negligible performance loss, demonstrating the potential to be a powerful backbone for downstream vision tasks. The code is available at: github.com/sunsmarterjie/iTPN.

Figures (7)

• Figure 1: Comparison of ImageNet-1K classification accuracy. iTPN/Fast-iTPN shows significant advantages over prior models, either using pixel supervision (upper) or knowledge distillation from a teacher (lower).
• Figure 2: The comparison between a conventional pre-training (left) and the proposed integral pre-training framework (right). We use a feature pyramid as the unified neck module and apply masked feature modeling to pre-train the feature pyramid. The blocks in green and red respectively indicate that the network weights are pre-trained and un-trained (i.e., randomly initialized for fine-tuning).
• Figure 3: Fast-iTPN architecture during fine-tuning. Illustration of (a) token dropping from the backbone, (b) token gathering on the backbone, and (c) token replenishing to the feature pyramid.
• Figure 4: Token gathering. Top: global attention within the respective window for each gathering token. Bottom: integrating information from all gathering tokens to all tokens.
• Figure 5: Speed-up of Fast-iTPN over iTPN. The experiments is done using Mask-RCNN for object detection on COCO by fine-tuning the pre-trained Fast-iTPN-B model for 12 epochs. Fast-iTPN achieves a speedup of approximately 20% with minimal performance loss ($\sim$ 0.2) and a speedup of 70% with a marginal decrease in mAP of 0.5%.