PaperlayerView exact revision on arXiv

Source-linked AI summary

TinyViT: Fast Pretraining Distillation for Small Vision Transformers

Kan Wu, Jinnian Zhang, Houwen Peng, Mengchen Liu, Bin Xiao, Jianlong Fu, Lu Yuan

arXiv:2207.10666v1cs.CV

TL;DR

Large vision transformers are difficult to deploy under limited computational budgets, while small models can struggle to benefit from massive pretraining data. TinyViT addresses this with fast pretraining distillation and constrained scaling of small transformer architectures, achieving strong ImageNet-1k performance and downstream transfer. Its main reported limitation is the training cost associated with large teacher models.

  • Problem

    Large vision transformers require substantial resources, while small models saturate when pretrained directly on massive datasets and have limited capacity to absorb their benefits.

  • Method

    TinyViT distills knowledge during pretraining using sparsified, precomputed teacher logits and scales down large models under parameter and throughput constraints.

  • Results

    84.8% top-1 accuracy on ImageNet-1k is achieved with 21M parameters, and distillation improves small-model gains from ImageNet-21k pretraining while supporting downstream transfer.

  • Takeaways & Limitations

    Small vision transformers can leverage large-scale pretraining data and transfer knowledge to downstream tasks when guided by large pretrained models.

  • Takeaways & Limitations

    Large teacher models can occupy substantial GPU memory and lengthen training time.

Abstract

from arXiv · show

Vision transformer (ViT) recently has drawn great attention in computer vision due to its remarkable model capability. However, most prevailing ViT models suffer from huge number of parameters, restricting their applicability on devices with limited resources. To alleviate this issue, we propose TinyViT, a new family of tiny and efficient small vision transformers pretrained on large-scale datasets with our proposed fast distillation framework. The central idea is to transfer knowledge from large pretrained models to small ones, while enabling small models to get the dividends of massive pretraining data. More specifically, we apply distillation during pretraining for knowledge transfer. The logits of large teacher models are sparsified and stored in disk in advance to save the memory cost and computation overheads. The tiny student transformers are automatically scaled down from a large pretrained model with computation and parameter constraints. Comprehensive experiments demonstrate the efficacy of TinyViT. It achieves a top-1 accuracy of 84.8% on ImageNet-1k with only 21M parameters, being comparable to Swin-B pretrained on ImageNet-21k while using 4.2 times fewer parameters. Moreover, increasing image resolutions, TinyViT can reach 86.5% accuracy, being slightly better than Swin-L while using only 11% parameters. Last but not the least, we demonstrate a good transfer ability of TinyViT on various downstream tasks. Code and models are available at https://github.com/microsoft/Cream/tree/main/TinyViT.

1 Introduction

TinyViT addresses the resource demands of large vision transformers by distilling knowledge into small models during pretraining and scaling architectures under efficiency constraints. The resulting models leverage large-scale data and transfer effectively to downstream tasks.

  • Large vision transformers achieve strong results but require substantial parameters, data, and computation, limiting use on mobile and IoT edge devices.
  • Small models can saturate during direct pretraining on massive datasets, motivating knowledge transfer from large pretrained teachers.
  • Pretraining distillation unlocks large-scale data for small models and improves their downstream transfer capacity.
  • Sparse teacher logits and stored augmentation information remove teacher forward computation and memory overhead during student training, enabling larger batch sizes.
  • 84.8% top-1 accuracy on ImageNet-1k is achieved by TinyViT with 21M parameters, compared with Swin-B's 85.2% using 88M parameters.
  • TinyViT introduces a family of tiny vision transformers designed to leverage large-scale pretraining data through fast distillation.

2 Related Work

Related work spans large-scale pretraining, efficient vision transformers, and teacher-student distillation. TinyViT differs by targeting pretraining-stage distillation for small transformers and by extending efficient stored-label distillation to broader augmentations.

  • Large-scale pretraining: Large-scale vision foundation models offer strong transfer and zero-shot capabilities but remain unsuitable for applications with limited computational budgets.
  • Small vision transformers: Efficient vision transformer research combines convolution, neural architecture search, and optimization to reduce computation while preserving performance.
  • Knowledge distillation: Vision transformer distillation methods include distillation tokens and techniques for extracting knowledge from transformer teachers.
  • Knowledge distillation: TinyViT uses stored teacher logits like FKD but stores a 4-byte random seed and supports augmentations including Mixup and CutMix.
  • Knowledge distillation: Unlike FKD's finetune-stage distillation for CNNs, TinyViT focuses on pretraining-stage distillation for transformers.

3 TinyViT

TinyViT combines fast pretraining distillation with progressively contracted hierarchical vision transformers to build efficient small models. It stores sparse teacher predictions and compact augmentation information in advance, then reconstructs them during asynchronous student training.

  • 3.1 Fast Pretraining Distillation: Fast pretraining distillation stores teacher predictions and augmentation information before student training, avoiding repeated teacher computation and reducing memory demands.The framework uses independent, asynchronous branches for saving teacher information and training the student.
  • 3.1 Fast Pretraining Distillation: The student optimizes cross entropy between stored teacher predictions and its outputs, making the framework label-free and applicable to unlabeled web data.Only teacher-generated soft labels are required during training; ground-truth labels are not used.
  • 3.1 Fast Pretraining Distillation: Sparse soft labels retain the top-K teacher logits and recover the remaining values for distillation, reducing storage by orders of magnitude when K ≪ C.The recovered labels can achieve performance comparable to dense labels for knowledge distillation.
  • 3.1 Fast Pretraining Distillation: Data augmentation is encoded by a single parameter and decoded during student training, accurately reconstructing transformations such as random cropping and RandAugment.A pseudo-random number generator can serve as the decoder, allowing the stored parameter to reproduce the augmentation sequence.
  • 3.2 Model Architectures: TinyViT architectures are generated by progressively contracting a large model seed while enforcing parameter and throughput constraints.At each step, constrained candidate models are evaluated, and the best validation model is further reduced.
  • 3.2 Model Architectures: TinyViT uses a hierarchical four-stage transformer with gradually reduced resolution, supporting dense prediction tasks that require multi-scale features.The architecture includes convolutional patch embedding and lightweight MBConvs in early stages and downsampling blocks.

4 Analysis and Discussions

The analysis attributes small models’ limited use of large datasets to hard samples and examines how distillation addresses this issue. Teacher predictions provide class-relationship information and filter noisy labels, producing student predictions closer to the teacher.

  • 4 Analysis and Discussions: ImageNet-21k contains hard samples such as wrong labels and visually similar images with different labels, which small models struggle to fit.Approximately 10% of ImageNet images are considered hard samples; TinyViT-21M reaches 53.2% training accuracy versus 57.1% for Swin-L-197M.
  • 4 Analysis and Discussions: 0.7% gains result from original ImageNet-21k pretraining, compared with 1.0%/1.1% after removing hard samples for Swin-T/TinyViT-21M and 2.2%/1.7% with distillation.The comparisons are reported for ImageNet-1k performance after pretraining on ImageNet-21k.
  • 4 Analysis and Discussions: Distillation avoids ground-truth-label defects by replacing polluted labels with teacher soft labels, helping small models leverage large-scale data.The authors identify hard samples’ ground-truth labels as the main cause of the defects addressed by distillation.
  • 4 Analysis and Discussions: Teacher predictions reveal class relationships through correlated outputs, including high correlations among similar or related classes.Pearson-correlation heatmaps show related classes as a block-diagonal structure and capture uncorrelated classes as well.
  • 4 Analysis and Discussions: TinyViT-21M predictions have less obvious class-correlation structure without distillation, while distillation makes their correlations closer to the teacher’s.The comparison uses TinyViT-21M predictions with and without distillation.

5 Experiments

Experiments show that fast pretraining distillation improves small-transformer performance, while sparse logits reduce storage and preserve accuracy. TinyViT achieves strong ImageNet-1k results and transfers effectively to downstream tasks.

  • Ablation Study: Pretraining without distillation yields limited gains on IN-21k, while distillation increases DeiT-Ti, DeiT-S, and Swin-T accuracy by 2.2%, 2.1%, and 2.2%, respectively.Compared with training from scratch on IN-1k, non-distilled IN-21k pretraining improves accuracy by only 0.8%, 0.6%, and 0.7%.
  • Ablation Study: Accuracy increases with the number of saved logits K until saturation, while storage cost grows linearly on both IN-1k and IN-21k.The selected settings are K=10 on IN-1k and K=100 on IN-21k under storage constraints.
  • Ablation Study: Better teacher models yield better students, but Florence requires 11GB GPU memory and causes 2.4 times longer training without fast distillation.The fast framework loads stored logits from disk, avoiding additional GPU memory and matching the baseline training time.
  • Results on ImageNet: 86.5% top-1 accuracy is reached by TinyViT-21M at higher resolution, establishing state-of-the-art performance on ImageNet-1k for small models.With distillation, its gains over Swin-T and DeiT-S rise to 3.6% and 4.9%, respectively.
  • Transfer Learning Results: Pretraining with distillation produces the best representability, with more than 10% gains over CIFAR-100, Flowers, and Cars when using larger pretraining datasets.Linear-probe and few-shot results support improved representability, except for ChestX, where medical images differ from natural-image training data.
  • Transfer Learning Results: TinyViT achieves 1.5% higher AP than Swin-T under the same Cascade R-CNN training recipe on COCO object detection.The passage reports an additional improvement after applying pretraining distillation, but does not provide its completed value.

6 Conclusions

TinyViT combines efficient tiny architectures with fast pretraining distillation to improve ImageNet-1k performance and downstream transferability. The authors conclude that the framework is effective and plan to explore more data and stronger teachers.

  • Conclusions: TinyViT is a family of tiny, efficient vision transformers pretrained on large-scale datasets with fast distillation.The framework is evaluated on ImageNet-1k and downstream benchmarks.
  • Conclusions: Extensive experiments demonstrate TinyViT’s efficacy on ImageNet-1k and superior transferability across downstream benchmarks.
  • Conclusions: Future work will use more data and more powerful teachers to further improve the representability of small models.

TinyViT: Fast Pretraining Distillation for Small Vision Transformers —— Supplementary Material ——

The supplementary material details the model architectures and evaluates the speed of fast distillation, including an experiment using ground-truth labels.

  • Supplementary Material: The supplement elaborates the TinyViT model architectures described in Section 3.2.
  • Supplementary Material: An additional experiment compares the training cost of fast pretraining distillation with the conventional method.
  • Supplementary Material: A separate experiment investigates distillation with ground-truth labels and explains why soft labels are used alone.

A Model Architectures

TinyViT uses a hierarchical four-stage architecture designed for efficient dense prediction and scales model size through progressive contraction. Its design combines early convolutional inductive bias with adjustable dimensions, depths, expansion ratios, and window sizes.

  • Architecture: TinyViT is a hierarchical four-stage transformer designed to support dense prediction downstream tasks.The architecture includes attention biases and a 3 × 3 depthwise convolution to capture local information.
  • Scaling: Progressive contraction generates candidate models under parameter and throughput constraints, selecting the best validation model for further reduction.The process begins from a large model and iteratively adjusts contraction factors until the target model is achieved.
  • Architecture: Replacing the first-stage transformer block with MBConv blocks improves accuracy in small vision transformers.The authors conjecture that early convolution provides locality-related inductive bias and helps small models converge.
  • Scaling: Decreasing embedded dimensions substantially reduces parameters, while narrower models increase depth, especially in Stage 3, to satisfy parameter constraints.Embedded dimensions for TinyViT-21M, TinyViT-11M, and TinyViT-5M are specified as progressively smaller configurations.
  • Design choices: An MLP expansion ratio of 4 outperforms 3, and larger windows improve accuracy at higher computational cost, especially a 14 × 14 Stage 3 window.Window sizes do not affect model size, but the Stage 3 choice provides improved accuracy with little extra computation.

B How fast the distillation is?

The proposed fast pretraining distillation reduces the cost of distilling small models by storing sparse teacher logits and reusing them during training. With Florence as teacher, it is faster than conventional distillation while enabling larger training batches.

  • Efficiency: 29.8% faster than conventional distillation, the proposed method uses stored Florence logits to pretrain TinyViT-21M.Storing the top-100 logits takes 92.4 GPU days, while pretraining takes 140.0 GPU days for 90 ImageNet-21k epochs; conventional distillation takes 330.9 GPU days.

C Distillation with ground-truth

Distillation without ground-truth labels performs better than distillation with them in the evaluated TinyViT variant. The authors attribute the drop to mutually non-exclusive ImageNet-21k labels that can suppress related classes.

  • Comparison: Distillation with ground-truth labels causes a slight performance drop for the evaluated TinyViT-21M variant.The comparison uses a model pretrained for 90 epochs on ImageNet-21k and finetuned for 30 epochs.
  • Interpretation: ImageNet-21k contains correlated label pairs such as “chair” and “furniture,” making one-hot labels less precise for distillation.Pretrained foundation-model soft labels retain category-relation information that helps distill the small model.
Loading 2207.10666v1…