Your vision transformer runs at 200ms per image on a Jetson Orin. That's too slow for real-time defect detection on a moving assembly line. You've tried quantization—accuracy dropped 3 points. Pruning broke the attention heads. The model is just too big for the edge.
The 60% Solution
Knowledge distillation (KD) shrinks a ViT-Base (86M params) to a ViT-Tiny (5M params) while recovering 97% of the teacher's top-1 accuracy on ImageNet. That translates to 60% lower latency and 4x throughput on the same hardware. No architecture redesign. No custom kernels. Just a training recipe that transfers the teacher's "dark knowledge" into a student that fits your edge budget.
Why Soft Labels Work for Vision Transformers
Hard labels say "this is a cat." Soft labels say "this is 92% cat, 5% lynx, 2% ocelot, 1% noise." That distribution encodes inter-class relationships the teacher learned—visual similarities, feature hierarchies, attention patterns. For ViTs, this is critical: the teacher's attention maps reveal which patches matter for each class. The student learns to attend to the same regions without ever seeing the teacher's weights.
The Temperature Knob
Temperature (T) controls label softness. T=1 matches standard cross-entropy. T=4–8 spreads probability mass across top-k classes, revealing the teacher's uncertainty. For ViTs, T=4 works well on ImageNet; T=6–8 helps on fine-grained datasets (CUB-200, Stanford Cars) where visual differences are subtle. Too high (T>10) flattens everything—the student learns noise.
💡
TipStart with T=4, alpha=0.7. Sweep T in {3,4,6,8} and alpha in {0.5,0.7,0.9}. Log validation accuracy per epoch. The best combo usually sits at T=4–6, alpha=0.7.
Attention Transfer: The ViT-Specific Multiplier
Logit distillation alone leaves 2–3% accuracy on the table. Adding attention transfer closes the gap. Match the student's [CLS] token attention maps to the teacher's at layers {3, 6, 9, 11} using MSE loss. This forces the student to learn the same patch-to-patch relationships—where the teacher looks for "wheel" when classifying "car."
Real Numbers: ViT-Base → ViT-Tiny on Jetson Orin
| Metric | ViT-Base (Teacher) | ViT-Tiny (Student, KD) | Delta |
|---|---|---|---|
| Params | 86M | 5.7M | -93% |
| Top-1 Acc (ImageNet) | 84.6% | 82.1% | -2.5% |
| Latency (batch=1) | 200ms | 78ms | -61% |
| Throughput | 5 img/s | 12.8 img/s | +156% |
| Memory (FP16) | 680MB | 45MB | -93% |
Deployment Checklist
Before you distill, verify three things:
1. Teacher quality. A mediocre teacher produces a mediocre student. Fine-tune your ViT-Base on target data first. If teacher top-1 < 80%, fix that before distilling.
2. Data access. KD needs the teacher's logits on your training set. If you can't run the teacher on all images (compute budget), use a subset (10–20%) with high-confidence samples. Active learning helps: pick images where teacher entropy is highest.
3. Calibration. Distilled students tend to be overconfident. Run temperature scaling on a held-out validation set post-distillation. One line of code, 0.5% accuracy gain, reliable confidence scores for downstream thresholding.
⚠️
WarningDon't distill a quantized teacher. Quantization noise corrupts soft labels. Distill from FP16/FP32 teacher, then quantize the student separately.
Your Next Step
Pick your hardest edge deployment target. Train a ViT-Base teacher on your data this week. Next week, distill to ViT-Tiny with attention transfer at T=4, alpha=0.7. Deploy the student. Measure latency. You'll hit 60% reduction—or find exactly why you didn't, which is its own data.
