TinyViT: Fast Pretraining Distillation for Small Vision Transformers

In plain words

Large vision transformers do well on image benchmarks, but they are too big for phones and small devices. A small model trained on the same data usually does not catch up, because it runs out of capacity to use the extra data. This paper asks whether a small model can still benefit from a huge pretraining dataset if a big model teaches it.

The usual way to teach is to run the big "teacher" model on every training batch, which is slow and needs a lot of memory. TinyViT avoids this. The teacher is run once over the dataset, and for each augmented image it keeps only its most likely answers, plus a small random seed that lets the exact same augmentation be rebuilt later. A small student is then trained against these saved answers, with no teacher in memory.

The authors also build a family of small models by starting from a large design and shrinking it step by step, keeping the versions that meet a size and speed budget. The resulting models reach strong ImageNet accuracy for their size, and the authors report that they also transfer well to other tasks.

TL;DR

• Pretraining with distillation lets small vision transformers use large datasets better than plain pretraining does §1 tab. 2.
• Teacher outputs are stored in advance as the top-K probabilities per augmented image, together with a 4-byte seed that replays the augmentation, so the teacher is not run during student training §3.1 eq. 1 eq. 2.
• The TinyViT family (5M, 11M, 21M parameters) is derived from a large seed by progressive model contraction §3.2.
• Author-reported: TinyViT-21M reaches 84.8% ImageNet-1k top-1 at 224px after distillation on ImageNet-21k, and 86.5% at 512px tab. 4.

Problem

Pretraining on very large datasets helps large models far more than small ones: plain pretraining on ImageNet-21k gave the small DeiT-Ti, DeiT-S and Swin-T only 0.6 to 0.8 points on ImageNet-1k tab. 2. Distilling from a large teacher during pretraining adds 2.1 to 2.2 points tab. 2.

The cost is in the teacher. Running a large teacher on every batch uses most of the compute and GPU memory, which forces small batch sizes and slows student training §1 §3.1. The authors also note that a small model struggles to fit hard, noisy samples in ImageNet-21k, and that ground-truth labels in that dataset overlap (for example "chair" and "furniture") §4 app. C.

The setting assumes access to a strong pretrained teacher, a dataset of this size, and disk space for the stored labels (16 GB for ImageNet-1k with 300 epochs, 481 GB for ImageNet-21k with 90 epochs, as reported by the authors §5.2).

Method

TinyViT has two parts that are independent of each other: a fast distillation procedure §3.1 and a family of architectures §3.2.

Part 1: fast pretraining distillation §3.1.

1. Before student training, pass each augmented training image through the teacher once. The augmentation is a strong pipeline (RandAugment, CutMix and others) §3.1.
2. Keep the top-K values of the teacher's output and their class indices, and discard the rest. The authors store these sparse labels instead of all C logits §3.1 eq. 2.
3. Store the augmentation as one seed. A pseudo-random generator, for which the paper names PCG as a common choice, turns the seed back into the same augmentation parameters during training §3.1.
4. During student training, load the stored indices, values and seed, rebuild the augmented image, rebuild the sparse teacher distribution with the formula in eq. 2, and minimise the cross-entropy between the teacher distribution and the student output eq. 1.

The student loss uses no ground-truth labels. The authors explain this as a choice that avoids the overlap between ImageNet-21k labels §3.1 app. C.

Part 2: model family by progressive model contraction §3.2.

5. Use a hierarchical base with four stages. The stem is two 3×3 stride-2 convolutions. Stage 1 uses MBConv blocks, and the last three stages use transformer blocks with window attention §3.2 app. A.
6. Define contraction factors: stage embedding widths, stage depths, window sizes, the MBConv expansion ratio, the MLP expansion ratio and the head dimension §3.2.
7. From the current model, generate smaller candidates by changing the factors. Keep the candidates that meet the parameter and throughput limits, train and validate them, and continue from the most accurate one until the target size is reached §3.2 app. A.

The authors describe step 7 as a constrained local search over the contraction factors §3.2. The paper does not give the search budget, the number of candidates per step, or the validation data beyond the description in app. A.

The factors shared by all three TinyViT sizes are listed in the paper §3.2 app. A. Each size uses its own embedding widths:

# Shared factors and per-size widths  [§3.2] [app. A]
# Depths (2, 2, 6, 2); window sizes (7, 7, 14, 7); MBConv expansion 4;
# MLP expansion 4; head dimension 32.
# TinyViT-21M:  embed dims (96, 192, 384, 576)
# TinyViT-11M:  embed dims (64, 128, 256, 448)
# TinyViT-5M:   embed dims (64, 128, 160, 320)

The distillation step. The pseudocode below follows the description in §3.1 and the code path in main.py. It is written to be read, not run.

# Store teacher labels once, before pretraining  [§3.1] [eq. 2]
# [code: TinyViT/save_logits.py#L103-L160]
# Input:  images from ImageNet-21k; teacher T; sparsity K (K = 100)
# Output: per augmented view: a seed and K (index, probability) pairs

def store_teacher_labels(dataset, teacher, K):
    for image in dataset:
        # Draw one seed; it fixes the augmentation  [§3.1]
        seed = draw_seed()                           # int32, 4 bytes
        view = augment(image, seed)                  # strong augmentation
        # Teacher output after softmax, temperature 1.0  [§5.1]
        probs = softmax(teacher(view))               # (C,)
        # Keep only the K largest probabilities  [§3.1] [eq. 2]
        values, indices = top_k(probs, K)            # (K,), (K,)
        save(seed, indices, values)

# One student step from stored labels  [§3.1] [eq. 1] [eq. 2]
# [code: TinyViT/main.py#L284-L330]
# Shapes: images (B, 3, H, W); values, indices (B, K); C classes

def distill_step(student, images, seeds, values, indices):
    # Rebuild the same augmented view from the stored seed  [§3.1]
    views = replay_augment(images, seeds)            # (B, 3, H, W)
    logits = student(views)                          # (B, C)
    # Spread the leftover mass evenly over the other C - K classes  [eq. 2]
    rest = (1 - sum(values, axis=-1, keepdims=True)) / (C - K)  # (B, 1)
    teacher = broadcast(rest, (B, C))                # (B, C)
    teacher = scatter(teacher, indices, values)      # (B, C)
    # Soft cross-entropy against the teacher; no ground-truth labels  [eq. 1]
    loss = mean(-sum(teacher * log_softmax(logits), axis=-1))
    return loss

Key equations as stated in the paper:

\mathcal{L} = CE\left(\hat{\mathbf{y}}, S(\mathcal{A}(x))\right) eq. 1

where \hat{\mathbf{y}} = T(\mathcal{A}(x)) is the teacher output for the augmented image, S is the student, and CE is cross-entropy §3.1.

\hat{y}_c = \begin{cases} \hat{y}_{\mathcal{I}(k)} & \text{if } c = \mathcal{I}(k), \\ \dfrac{1 - \sum_{k=1}^{K} \hat{y}_{\mathcal{I}(k)}}{C - K} & \text{otherwise} \end{cases} eq. 2

where \mathcal{I}(k) are the stored indices of the top-K entries §3.1.

Architecture details §3.2 app. A.

• Stem: two 3×3 convolutions, stride 2, padding 1, each followed by batch norm, with GELU between them tiny_vit.py.
• Stage 1 uses MBConv blocks with expansion ratio 4 §3.2 tiny_vit.py tiny_vit.py. Each downsampling layer is a PatchMerging block in the code: 1×1 convolution, depthwise 3×3 convolution with stride 2, then 1×1 convolution, all at the output width. The paper's table labels these layers MBConv app. A tiny_vit.py.
• Stages 2 to 4 use transformer blocks with window attention. Each block has a learned attention bias per head and relative offset tiny_vit.py. A 3×3 depthwise convolution sits between attention and the MLP, with a residual connection after each sub-block §3.2 tiny_vit.py tiny_vit.py.
• Normalisation: batch norm for convolutions and LayerNorm for linear layers. Activation: GELU §3.2.
• Head: average over tokens, then LayerNorm, then a linear classifier app. A tiny_vit.py.

Hyperparameters

Values marked "code" were not in the paper text and are read from the official repository at code_commit. Code learning rates are written per 512 images and scaled linearly by the total batch size main.py.

Pretraining on ImageNet-21k with distillation

Fine-tuning on ImageNet-1k from the pretrained model

Higher-resolution fine-tuning

Training from scratch on ImageNet-1k

Architecture and teacher-storage hyperparameters

Results

Author-reported. All TinyViT numbers are from the paper; the setup is ImageNet-21k pretraining (about 14M images, 21,841 classes §4) followed by ImageNet-1k fine-tuning or training, with 224px input unless stated §5.1 tab. 4.

Headline ImageNet-1k classification tab. 4

Selected rows. The full table has more baselines (ResNet-18, PVT-Tiny, ResT-Small, CrossViT, GLiT, Mobile-Former, ViTAS, T2T-ViT, Swin-T and others). Throughput was measured by the authors on a V100 GPU tab. 4. "Distilled" means pretraining on ImageNet-21k with the proposed distillation using CLIP-ViT-L/14 as the teacher §5.3.

The text compares the distilled 21M model with Swin-B pretrained on ImageNet-21k (85.2% with 88M parameters) §1.

Pretraining distillation on existing small models tab. 2

Data cleaning and distillation on ImageNet-21k tab. 1

Top-1 on ImageNet-1k, ImageNet-Real and ImageNet-V2 after fine-tuning on ImageNet-1k tab. 1.

Choice of teacher tab. 3

Teacher numbers in brackets are parameters and linear-probe accuracy on ImageNet-1k tab. 3. Training time and memory were measured without the fast distillation method tab. 3.

Linear probe and few-shot transfer, TinyViT-21M tab. 5

Linear probe accuracy (%), followed by 5-, 20- and 50-shot 5-way accuracy (%) on four datasets tab. 5.

COCO object detection tab. 6

Cascade Mask R-CNN, 12 epochs, backbone only, COCO 2017 validation tab. 6.

Scaling and teacher-logit sparsity fig. 4 fig. 5. The authors report that distillation keeps helping as the number of ImageNet-21k images seen grows, while the models tend to saturate with more epochs fig. 4 §5.2. Accuracy rises with K until it saturates, and storage grows linearly with K fig. 5 §5.2.

Analysis §4 fig. 3. TinyViT-21M reaches 53.2% training accuracy on ImageNet-21k against 57.1% for Swin-L-197M §4. The predicted class correlations of the distilled student look more like those of the teacher than those of the student trained without distillation fig. 3.

Paper vs code

Compared against the official repository at code_commit . The code is the Microsoft Cream repository, under the TinyViT folder that the paper's abstract links.

Differences:

1. Stored values are probabilities, not logits. The paper calls the stored values "logits" §3.1. The code takes a softmax over the outputs and stores the top-K probabilities, which are then used directly as targets save_logits.py main.py. The sum-to-one form of eq. 2 only works with probabilities, so the code matches the formula.
2. Attention-bias interpolation is bicubic, not bilinear. The paper says the attention biases are bilinear-interpolated for higher resolution §5.1. The code interpolates the bias tables with bicubic mode when loading a checkpoint whose table size does not match utils.py. The paper's description of the interpolation mode disagrees with the code.
3. Teacher set. The paper uses Swin-L, BEiT-L, CLIP-ViT-L/14 and Florence as teachers §5.2. The repository only has a configuration for the CLIP teacher clip_vit_large_patch14_22k.yaml. The README calls the public teacher "CLIP-ViT-Large/16" README.md, while the config and the paper use ViT-L/14.
4. Default K. The paper uses K = 10 for ImageNet-1k and K = 100 for ImageNet-21k §5.2. The default in the code is 100 config.py. The 1k run must set K explicitly; no 1k distillation config is in the repository.
5. Learning rates are scaled. The paper states the effective learning rates §5.1. The code takes a base rate per 512 images and scales it linearly with the total batch main.py. The two agree at the batch sizes in the paper (4096 for pretraining, 1024 for fine-tuning), so the paper's numbers are the code's numbers after scaling. An implementer who uses a different batch size must scale the base rate.
6. Mixup, cutmix and label smoothing are not in the paper for ImageNet-1k. The paper says nothing about these for from-scratch training or fine-tuning except that fine-tuning disables mixup and cutmix §5.1. The 1k config does not override them, so the code defaults apply: mixup 0.8 and cutmix 1.0 config.py. The fine-tuning configs set both to zero, and label smoothing 0.1 then applies config.py main.py.

Consistent with the paper:

• Pretraining schedule, optimiser, weight decay, clipping, stochastic depth, and the disabled mixup and cutmix in distillation tiny_vit_21m_22k_distill.yaml.
• Soft-label loss without ground truth main.py.
• Embedding widths and depths of the three sizes .
• Seed stored as a 4-byte integer save_logits.py.
• The mean-over-tokens pooling then LayerNorm and a linear head tiny_vit.py.

Pitfalls

• The stored labels fix the augmentation. A student trained on saved labels sees only the augmentations that were saved with them. Any change to the augmentation pipeline (RandAugment policy, interpolation, colour jitter, random erasing, crop or the random generator) makes the saved seeds replay a different image. The paper names PCG as a common generator §3.1; this compile did not verify which generator the code uses for the replay.
• Seeds depend on rank and epoch when labels are generated. The code derives the seed for the labelling run from the base seed, the process rank and the start epoch save_logits.py. Training replays the stored seeds as saved, so the launch configuration matters only when labels are regenerated.
• Storage is large and one-off. The paper reports 481 GB for ImageNet-21k with K = 100 over 90 epochs, and 16 GB for ImageNet-1k with K = 10 over 300 epochs §5.2. The per-epoch labels are independent and can be computed in parallel §3.1.
• The 29.8% speed-up includes the one-off labelling. The appendix gives 92.4 GPU days to store the Florence labels and 140.0 GPU days to train with them, against 330.9 GPU days for conventional distillation app. B. The 29.8% figure matches (92.4 + 140.0) compared with 330.9, a derived calculation from the paper's numbers. The per-run saving on its own is larger.
• Teacher data overlap. The main teacher, CLIP-ViT-L/14, was fine-tuned on ImageNet-21k for 30 epochs §5.1, the same data the student is pretrained on. The paper does not discuss this overlap.
• Hard-sample cleaning is an analysis step. The paper removes about 2M images (about 14%) from ImageNet-21k using a Florence top-5 filter, to test the effect of noisy labels §4. The cleaned set is not the one used for the headline TinyViT numbers.
• Ground-truth labels slightly hurt. In the authors' test, adding ground-truth to the distillation loss lowered ImageNet-1k top-1 by 0.2 points with the CLIP teacher and 0.7 points with Florence app. C. The paper gives this as its reason for label-free distillation.
• Single runs. The paper reports one number per setting and no variance across seeds in the text.
• Caption and text disagree on the detector. The text says Cascade R-CNN §5.4; the caption of the detection table says Cascade Mask R-CNN tab. 6.
• Window padding is code-only. When a feature map is not a multiple of the window size, the code pads it tiny_vit.py. The paper does not mention padding. At the stated resolutions (224, 384 and 512) every stage divides evenly into windows, so padding does not trigger.
• Higher-resolution cost. The 512px model needs 27.0 GMACs and runs at 167 images per second on the authors' V100 tab. 4; the 21M model at 224px runs at 1,571. Resolution scaling is a large cost increase, not a free gain.
• Throughput depends on the device. The authors measured throughput on one V100 GPU tab. 4. Other hardware can reorder the models.

Lineage

• Builds on: the hierarchical window-attention design of Swin, and the attention-bias and stem choices of LeViT, both cited in §3.2. The distillation idea follows Hinton-style teacher-student training §2.
• Closest related work: FKD, which stores teacher logits for CNN fine-tuning. TinyViT differs by targeting pretraining for transformers, and by storing one seed per image instead of the explicit augmentation parameters §2.
• Compared against: DeiT (including its distillation token), Swin, LeViT, MobileViT, AutoFormer, GLiT and other small models tab. 4.