TinyViT

TinyViT is a family of tiny hierarchical vision transformers (5M / 11M / 21M parameters) together with a fast pretraining distillation framework.

• Paper: Kan Wu, Jinnian Zhang, Houwen Peng, Mengchen Liu, Bin Xiao, Jianlong Fu and Lu Yuan (Microsoft Research / Microsoft Cloud+AI), TinyViT: Fast Pretraining Distillation for Small Vision Transformers, ECCV 2022, arXiv:2207.10666.
• Two separable contributions: (1) a distillation scheme and (2) a small hybrid conv/transformer architecture.
• Distillation: makes pretraining-stage distillation on ImageNet-21k practically free by precomputing and storing the teacher’s sparse logits instead of running the teacher in the training loop.
• Architecture: obtained by progressively contracting a 21M seed model.
• Headline result: TinyViT-21M reaches 84.8% top-1 on ImageNet-1k, comparable to Swin-B pretrained on IN-21k (85.2% with 88M parameters) at $4.2\times$ fewer parameters, and 86.5% at $512^2$ resolution.
• Used by: TinyViT-5M is the backbone that MobileSAM distils SAM’s ViT-H image encoder into, which is where most people meet it.

Fig. 1: pretraining alone gains only ~0.7% for small models; pretraining with distillation gains 1.6–2.2%.

Motivation: small models saturate on large data

The prevailing direction in 2022 was scaling ViTs up — V-MoE at 14.7B parameters on 305M images, Swin at 3B parameters on 70M images. TinyViT asks the opposite question: can a small model absorb the knowledge in a massive pretraining corpus at all?

Empirically it cannot, not directly. Pretraining DeiT-Ti / DeiT-S / Swin-T on IN-21k and finetuning on IN-1k buys only $+0.8/+0.6/+0.7$ points over training from scratch on IN-1k. Small models are saturated by large data rather than helped by it.

The paper’s diagnosis (Sec. 4) is hard samples. IN-21k contains 14M images over 21,841 categories, roughly 10% of which are considered hard by prior work — wrong labels, or images containing several equally prominent objects with only one label. Large models can memorise these; small models cannot, and burn capacity trying: TinyViT-21M reaches only 53.2% training accuracy on IN-21k versus 57.1% for Swin-L (197M). Two interventions confirm the diagnosis:

#	Model	Pretraining dataset	IN-1k	IN-Real	IN-V2
0	Swin-T	from scratch on IN-1k	81.2	86.7	69.7
1	Swin-T	original IN-21k	81.9 (+0.7)	87.0	70.6
2	Swin-T	cleaned IN-21k	82.2 (+1.0)	87.3	71.1
3	Swin-T	original IN-21k w/ distillation	83.4 (+2.2)	88.0	72.6
4	TinyViT-21M	from scratch on IN-1k	83.1	88.1	73.1
5	TinyViT-21M	original IN-21k	83.8 (+0.7)	88.4	73.8
6	TinyViT-21M	cleaned IN-21k	84.2 (+1.1)	88.5	73.8
7	TinyViT-21M	original IN-21k w/ distillation	84.8 (+1.7)	88.9	75.1

“Cleaned” means: use the pretrained Florence model to find images whose ground-truth label falls outside its top-5 predictions and delete them — 2M images, ~14% of IN-21k. Cleaning helps, but distillation helps more, without touching the data, because it never uses the ground-truth labels that are the source of the problem in the first place. The teacher’s soft label also injects a class-relationship prior: Pearson correlation heatmaps over the 21,841 IN-21k classes (Fig. 3 of the paper) show a clear block-diagonal structure in the teacher’s predictions that a from-scratch small student fails to reproduce and a distilled student recovers.

Fast pretraining distillation

Fig. 2: top — the teacher branch that saves sparse logits and an encoded augmentation; middle — disk; bottom — the student branch, which replays both. The two branches are independent and asynchronous.

Naive pretraining distillation is expensive for the wrong reason: most of the compute and nearly all of the GPU memory go into forwarding the training data through the teacher each iteration, not into training the student, and the teacher’s memory footprint caps the student’s batch size. TinyViT removes the teacher from the training loop entirely.

For an input image $x$ with strong augmentation $\mathcal{A}$ (RandAugment, CutMix, …), the teacher branch stores the pair $(\mathcal{A},\hat{\mathbf{y}})$ with $\hat{\mathbf{y}}=T(\mathcal{A}(x))$. Both halves must be stored because augmentation is stochastic — the same image through the same pipeline gives a different crop every time, so the augmentation actually used and the logits it produced have to be kept together, per image per epoch. Training then optimises

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

where $S(\cdot )$ is the student and $CE$ is cross entropy. The framework is label-free: no ground-truth labels enter training at all, so unlabelled web-scale data is usable directly. It is also, per-iteration, exactly as cheap as training without distillation — the teacher $T(\cdot )$ has vanished from the equation.

Two encodings make the storage tractable.

Sparse soft labels. Storing dense logits over $C=21,841$ classes for every augmented image in every epoch is hopeless. Instead only the top-$K$ values $\{{\hat{y}}_{\mathcal{I}(k)}{\}}_{k=1}^K$ and their indices $\{\mathcal{I}(k){\}}_{k=1}^K$ are saved, and the dense vector is reconstructed at training time with label smoothing over the remainder:

${\hat{y}}_c=\{\begin{array}{ll}{\hat{y}}_{\mathcal{I}(k)}& \text{if }c=\mathcal{I}(k),\\ \frac{1-{\sum }_{k=1}^K{\hat{y}}_{\mathcal{I}(k)}}{C-K}& \text{otherwise.}\end{array}$

With $K\ll C$ this cuts storage by orders of magnitude at no measurable accuracy cost — and sparsification arguably denoises the teacher, keeping the class relationships while dropping the long tail.

Data augmentation encoding. The augmentation parameters $\mathbf{d}$ (crop coordinates, rotation degrees, …) would cost hundreds of bytes per image. TinyViT instead encodes them as a single seed $d_0=\mathcal{E}(\mathbf{d})$ — 4 bytes — and recovers $\mathbf{d}={\mathcal{E}}^{-1}(d_0)$ at training time by re-seeding a PCG pseudo-random generator. In implementation the “encoder” is simply a generator for $d_0$ that reuses the decoder: it draws $d_0$, runs $\mathbf{d}={\mathcal{E}}^{-1}(d_0)$ for the teacher, and saves $d_0$.

This is what distinguishes the method from FKD, the closest prior work: FKD stores the explicit transformation parameters, TinyViT stores a seed; and TinyViT’s scheme supports arbitrary augmentations including Mixup and CutMix, whereas FKD does not. FKD also targets finetune-stage distillation of CNNs, not pretraining-stage distillation of transformers.

Cost. With Florence (682M parameters, 97.9 GFLOPs) as the teacher, storing the top-100 logits for all 90 epochs of IN-21k takes 92.4 GPU days, and pretraining TinyViT-21M (4.4 GFLOPs) on the stored logits takes 140.0 GPU days — 232.4 total, against 330.9 GPU days for conventional in-loop distillation, a 29.8% saving. Two structural advantages compound this: the per-epoch teacher logits are independent, so they can be generated in parallel rather than epoch-by-epoch; and the saved logits are reusable for arbitrary students, so the 92.4 days is paid once for an entire model family.

Architecture

Supplementary Table 1: the elastic base architecture. Every $\gamma$ is a contraction factor searched over.

TinyViT is a hierarchical hybrid, four stages with progressively reduced resolution ($56^2\to 56^2\to 28^2\to 14^2\to 7^2$ for $224^2$ input), in the style of Swin and LeViT — hierarchical rather than plain-ViT so that dense-prediction downstream tasks get multi-scale features.

• Patch embedding: two stacked convolutions, kernel $3\times 3$, stride 2, padding 1.
• Stage 1: lightweight MBConv blocks (MobileNetV3-style inverted residuals). Convolutions at early layers learn low-level representations efficiently thanks to their inductive bias; the ablation notes that replacing the stage-1 transformer block with MBConv straightforwardly improves accuracy, and conjectures the locality prior helps small models converge.
• Stages 2–4: transformer blocks with window attention to bound the cost, plus attention biases (LeViT-style) and a $3\times 3$ depthwise convolution between attention and MLP to capture local information.
• Downsampling between stages: MBConv blocks with stride 2.
• Residual connections on every block in stage 1 and on all attention and MLP blocks; GELU activations; BatchNorm after convolutions, LayerNorm after linear layers.
• Classifier: AvgPool + LayerNorm + Linear.

Note that MobileSAM uses this architecture with one modification: the stride of the last downsampling convolution is set to 1, so that the output feature resolution matches what SAM’s frozen mask decoder expects.

Scaling down by progressive model contraction

Rather than searching an architecture from scratch, the paper starts from a large seed model and shrinks it. The contraction factors are:

Factor	Meaning
${\gamma }_{D_{1-4}}$	embedded dimension of the four stages (fewer heads in MHSA when reduced)
${\gamma }_{N_{1-4}}$	number of blocks per stage (depth)
${\gamma }_{W_{2-4}}$	window size in the last three stages
${\gamma }_R$	channel expansion ratio of the MBConv block
${\gamma }_M$	MLP expansion ratio for all transformer blocks
${\gamma }_E$	dimension of each attention head (smaller ⇒ more heads ⇒ lower cost)

The procedure is a constrained local search in this factor space: from the current model, generate candidates by perturbing the contraction factors; keep those satisfying the parameter-count and throughput constraints; evaluate them on a 99%/1% train/val split of the IN-1k training set; the best-validating model becomes the seed for the next reduction step, until the target size is reached.

All three released models share $\{{\gamma }_{N_1},{\gamma }_{N_2},{\gamma }_{N_3},{\gamma }_{N_4}\}=\{2,2,6,2\}$, $\{{\gamma }_{W_2},{\gamma }_{W_3},{\gamma }_{W_4}\}=\{7,14,7\}$ and $\{{\gamma }_R,{\gamma }_M,{\gamma }_E\}=\{4,4,32\}$, and differ only in embedded dimensions:

Model	$\{{\gamma }_{D_1},{\gamma }_{D_2},{\gamma }_{D_3},{\gamma }_{D_4}\}$
TinyViT-21M	{96, 192, 384, 576}
TinyViT-11M	{64, 128, 256, 448}
TinyViT-5M	{64, 128, 160, 320}

Four observations from the search, offered as manual-design guidance:

1. Replacing the stage-1 transformer block with MBConv improves accuracy for small ViTs.
2. Reducing the embedded dimensions ${\gamma }_{D_{1-4}}$ cuts parameters fastest, so it is the first knob to turn; as the model narrows, its depth (especially stage 3, ${\gamma }_{N_3}$) is increased to keep hitting the parameter constraint — narrow-and-deep beats wide-and-shallow at fixed size.
3. MLP expansion ratio ${\gamma }_M=4$ beats 3.
4. Window sizes ${\gamma }_{W_{2-4}}$ do not affect model size at all, only compute; larger windows improve accuracy, and stage 3 in particular gets a good return from a $14\times 14$ window at little extra cost — hence the asymmetric $\{7,14,7\}$.

Hyperparameters

Transcribed from §5.1 and the supplementary material.

IN-21k pretraining

Epochs	90
Optimizer	AdamW
Weight decay	0.01
Initial LR / schedule	0.002, cosine
Warm-up	5 epochs
Batch size	4,096
Gradient clipping	max norm 5
Stochastic depth	0 (5M / 11M), 0.1 (21M)
Augmentation	random resize & crop, horizontal flip, color jittering, random erasing, RandAugment, Mixup, CutMix

IN-1k finetuning from the pretrained model

Epochs	30
Batch size	1,024
LR / schedule	$5\times 10^{-4}$, cosine, 5-epoch warm-up
Weight decay	$10^{-8}$
Layer-wise LR decay	0.8, output layer → input layer
BatchNorm	running statistics frozen
Mixup / CutMix	disabled

IN-1k finetuning at higher resolution

Window sizes	$\{7,7,14,7\}$ at $224^2$; $\{12,12,24,12\}$ at $384^2$; $\{16,16,32,16\}$ at $512^2$
Attention biases	bilinearly interpolated to the new window size
Epochs	30
Accumulated batch size	1,024
LR / schedule	$4\times 10^{-5}$, cosine, 5-epoch warm-up
Weight decay	$10^{-8}$
BatchNorm	frozen; Mixup/CutMix disabled

IN-1k training from scratch

Epochs	300
Optimizer	AdamW, weight decay 0.05
LR / schedule	0.001, cosine, 20 warm-up epochs
Batch size	1,024
Gradient clipping	max norm 5.0
Stochastic depth	0.0 / 0.1 / 0.2 for 5M / 11M / 21M

Knowledge distillation

Stored logits	top-100 for IN-21k, for all 90 epochs
Teachers evaluated	Swin-L, BEiT-L, CLIP-ViT-L/14, Florence (the last two finetuned on IN-21k for 30 epochs first)
Distillation temperature	1.0
Mixup / CutMix	disabled for pretraining distillation on TinyViT
Student hyperparameters	identical to the non-distilled runs above
Framework	PyTorch + timm

Ablations

Pretraining distillation on existing small ViTs. The framework is architecture-agnostic; applied to off-the-shelf DeiT and Swin it roughly triples the payoff of IN-21k pretraining:

Model	params (M)	IN-1k from scratch	IN-21k pretrain, no distill.	IN-21k pretrain w/ distill.
DeiT-Ti	5	72.2	73.0 (+0.8)	74.4 (+2.2)
DeiT-S	22	79.9	80.5 (+0.6)	82.0 (+2.1)
Swin-T	28	81.2	81.9 (+0.7)	83.4 (+2.2)

Pretraining data scale.

Fig. 4: distillation gains hold across the whole data-scale sweep for both TinyViT-21M (left) and TinyViT-5M (right), but both curves flatten — capacity, not data, is the eventual bottleneck.

Number of saved logits $K$.

Fig. 5: accuracy saturates in $K$ while storage grows linearly. Left: distilling on IN-1k. Right: distilling on IN-21k then finetuning on IN-1k.

Accuracy rises with $K$ until saturation while storage grows linearly, so the operating point is chosen just past the knee: $K=10$ (1.0% of logits) on IN-1k for 300 epochs at 16 GB, and $K=100$ (0.46% of logits) on IN-21k for 90 epochs at 481 GB.

Teacher choice. Better teachers give better students, and because the teacher forward pass is offline the training-time cost of a better teacher is zero:

#	IN-21k pretrained teacher	IN-1k	IN-Real	IN-V2	Training time (GPU h)	Memory (GB)
0	w/o distillation	83.8	88.4	73.8	3,360	0
1	BEiT-L (326M, 84.1)	84.1	88.4	73.8	6,415 (1.9×)	3.9
2	Swin-L (229M, 84.4)	84.2	88.6	73.9	5,804 (1.7×)	6.8
3	CLIP-ViT-L/14 (321M, 85.2)	84.8	88.9	75.1	7,087 (2.1×)	2.7
4	Florence (682M, 86.2)	84.8	89.0	74.8	7,942 (2.4×)	10.7

The training times and memory in this table are measured without the fast framework, on V100s — they are what the fast framework eliminates. With it, distillation costs the same GPU hours and the same GPU memory as row #0.

Ground truth hurts. Mixing ground-truth labels into the distillation loss consistently lowers accuracy (CLIP teacher: 84.3 → 84.5 without GT; Florence: 84.2 → 84.9 without GT), because IN-21k’s labels are not mutually exclusive — “chair”/“furniture”, “horse”/“animal” — so a one-hot target actively suppresses correct parent or child classes.

Results

ImageNet-1k, MACs and throughput measured on a V100. ✻ marks IN-21k pretraining with the proposed fast distillation, ↑ higher-resolution finetuning.

Model	Top-1	params (M)	MACs (G)	Throughput (img/s)	Input	Arch.
MobileViT-S	78.4	6	1.8	2,661	256	Hybrid
TinyViT-5M	79.1	5.4	1.3	3,060	224	Hybrid
TinyViT-5M✻	80.7	5.4	1.3	3,060	224	Hybrid
LeViT-256	81.6	19	1.1	7,386	224	Hybrid
TinyViT-11M	81.5	11	2.0	2,468	224	Hybrid
TinyViT-11M✻	83.2	11	2.0	2,468	224	Hybrid
DeiT-S	79.9	22	4.6	2,276	224	Trans
Swin-T	81.2	28	4.5	1,393	224	Trans
EffNet-B5	83.6	30	9.9	330	456	CNN
TinyViT-21M	83.1	21	4.3	1,571	224	Hybrid
TinyViT-21M✻	84.8	21	4.3	1,571	224	Hybrid
TinyViT-21M✻ ↑384	86.2	21	13.8	394	384	Hybrid
TinyViT-21M✻ ↑512	86.5	21	27.0	167	512	Hybrid

Trained from scratch on IN-1k, TinyViT-21M is already $+1.9$ / $+3.2$ over Swin-T / DeiT-S, which validates the contracted architecture on its own; with pretraining distillation the gaps widen to $+3.6$ / $+4.9$. TinyViT also beats the NAS-searched AutoFormer and GLiT at comparable size.

Transfer

Linear probe (TinyViT-21M, head replaced by a linear layer, everything else frozen):

Training data	CIFAR-10	CIFAR-100	Flowers	Cars	Pets
IN-1k	91.7	75.2	80.9	56.3	86.5
IN-1k ✻	91.7	74.5	82.4	61.7	85.5
IN-21k	96.3	84.7	99.7	67.7	92.6
IN-21k ✻	96.9	86.6	99.7	75.1	93.8

Pretraining on the larger dataset is worth >10 points on CIFAR-100, Flowers and Cars, and distillation adds on top of that — pretraining distillation on large data gives the best representation. Few-shot results (5-way, 5/20/50-shot on CropDisease, EuroSAT, ISIC 2018, ChestX, 600 episodes) follow the same pattern, with the exception of ChestX, whose grayscale medical images are unlike anything in either the teacher’s or the student’s training data.

Object detection — COCO 2017 val, Cascade Mask R-CNN, 12 epochs, Swin-T’s training recipe, parameter counts are backbone only:

Backbone	params	IN-1k	AP	AP${}_{50}$	AP${}_{75}$	AP${}_S$	AP${}_M$	AP${}_L$
Swin-T	28M	81.2	48.1	67.1	52.1	31.1	51.2	63.5
TinyViT-21M	21M	83.1	49.6 (+1.5)	68.5	54.2	32.3	53.2	64.8
TinyViT-21M ✻	21M	84.8	50.2 (+2.1)	69.4	54.4	32.9	53.9	65.2

The split is informative: $+1.5$ AP comes from the architecture, a further $+0.6$ from pretraining distillation — i.e. the distilled representation transfers, it is not merely a classification-head artefact.

Limitations

• Capacity, not data, is the ceiling. Every data-scale curve saturates as epochs increase; distillation raises the plateau but does not remove it. The conclusion explicitly names “more data with more powerful teachers” as future work.
• Storage is the new cost. 481 GB for the IN-21k top-100 logits over 90 epochs. The logits are also tied to a specific number of epochs and a specific augmentation pipeline — change either and the 92.4 GPU days of teacher forward passes must be paid again. The saving is real only when the stored logits are amortised over several students.
• The teacher pass itself is not free: 92.4 GPU days for Florence, which merely moves cost out of the training loop rather than eliminating it.
• Label-free means teacher-bound. Since the ground truth is deliberately discarded, the student inherits every bias and error of the teacher, and the method’s ceiling is set by teacher quality (Table 3 shows exactly this monotonic dependence).
• The scaling-down method is a constrained local search over hand-chosen contraction factors, not a full architecture search; the conclusion calls a better scaling-down method an open direction.
• Window attention bounds the receptive field per block; the accuracy/compute trade-off of window size is handled by hand ($\{7,14,7\}$), not searched jointly.
• All headline numbers are ImageNet classification plus one COCO detection setting; no segmentation or dense-prediction ablation of the hierarchical design is reported.

See also

• MobileSAM — distils SAM’s ViT-H encoder into TinyViT-5M; the main reason TinyViT shows up in segmentation work
• SAM — the foundation model whose encoder cost motivates tiny backbones
• LW-DETR — the same efficiency question asked of detection transformers rather than backbones
• GPUs for Local ML — what “92.4 GPU days of teacher logits” and “481 GB of storage” mean in practice
• Implementing Research Papers