Data-Efficient Training: DeiT & Augmentation for ViTs

"My first trainer fed me three hundred million pictures and called me a prodigy. My second trainer had one million, a bag of tricks, and a wise old CNN sitting next to me whispering answers. I learned just as much. The second trainer had the better story."

A Vision Transformer Who Found a Mentor

Section 22.2 left us with a working ViT and an unanswered question: it needs "enough data", but how much, and what do you do if you do not have it? The honest answer from the 2020 paper was sobering. On ImageNet alone, trained from scratch, ViT-Base scored several points below a comparable ResNet of the kind you met in Chapter 20. It only pulled ahead after pretraining on JFT-300M, Google's internal 300-million-image dataset, which almost no one outside a handful of labs can touch. For a year that made the ViT look like a beautiful idea gated behind data most teams will never have. This section is the story of how that gate came down.

1. Why ViTs Are Data-Hungry Beginner

The data appetite is a direct consequence of the missing inductive bias from Section 22.1. A convolution assumes locality and translation equivariance, and those assumptions are correct for natural images, so a CNN starts partway to the answer and needs comparatively little data to finish. A ViT assumes nothing about spatial structure; it must learn, purely from labeled examples, that nearby patches are related and that an object looks the same wherever it appears. Learning a true fact about the world from data, rather than having it hard-coded, costs examples. With a small dataset the ViT has enough freedom to fit the training set in ways that do not generalize, so it overfits where the CNN, hemmed in by its biases, does not.

Figure 22.3.1 sketches the qualitative picture the ViT paper reported: at small data scales the CNN wins because its biases substitute for data, but the curves cross, and at large scales the ViT's flexibility lets it keep improving past the point where the CNN's biases become a ceiling. The exact crossover depends on the architectures and recipe, and pinning it down quantitatively is the subject of Section 22.5; here it motivates the problem DeiT solves.

2. The DeiT Recipe: Augmentation and Regularization Intermediate

DeiT (Data-efficient image Transformer, Touvron et al., 2021) made one central claim: most of the ViT's data hunger can be fed by augmentation and regularization instead of by more images. The recipe is essentially the modern training bundle of Chapter 21, applied unusually aggressively. It combines RandAugment (automated, stacked photometric and geometric ops), MixUp and CutMix (the label-mixing augmentations of Section 21.2 that blend both pixels and targets), random erasing, strong stochastic depth (randomly dropping whole transformer blocks during training), label smoothing, weight decay, and a long cosine schedule with warmup. Each technique fights overfitting from a different angle; together they let ViT-Base reach competitive ImageNet accuracy trained only on ImageNet's $1.28$ million images.

The point worth internalizing is that the augmentation is doing the job the missing inductive bias used to do. MixUp and CutMix manufacture a near-endless supply of varied training images, forcing the model to rely on robust, generalizable features rather than memorizing the finite training set, which is exactly the freedom-without-data problem subsection 1 described. The code below assembles the DeiT-style transform stack with the torchvision v2 transforms from Chapter 21.

# Assemble the DeiT training augmentation: per-image RandAugment and random
# erasing in the Compose, plus batch-level MixUp/CutMix that mix images and
# labels together. This manufactured variety stands in for more real images.
import torch
from torchvision.transforms import v2

# DeiT-style training augmentation: aggressive by CNN standards.
train_tf = v2.Compose([
    v2.RandomResizedCrop(224, antialias=True),
    v2.RandomHorizontalFlip(),
    v2.RandAugment(num_ops=2, magnitude=9),          # automated stacked ops
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=[0.485, 0.456, 0.406],         # ImageNet statistics
                 std=[0.229, 0.224, 0.225]),
    v2.RandomErasing(p=0.25),                         # cut out a random region
])

# MixUp / CutMix operate on a whole batch (mixing images AND labels):
mixup_cutmix = v2.RandomChoice([
    v2.MixUp(num_classes=1000),
    v2.CutMix(num_classes=1000),
])
# in the loop:  images, labels = mixup_cutmix(images, labels)

3. Distillation Through a Token Advanced

DeiT's signature contribution is a new way to do knowledge distillation. The classical idea (Hinton et al.) is to train a small "student" to match the softened output probabilities of a strong "teacher", transferring the teacher's knowledge. The word "softened" matters: a one-hot label says only "this is a cat", but a teacher's full probability vector says "mostly cat, a little lynx, almost never car", and those relative confidences encode which classes look alike. Learning to reproduce that richer signal teaches the student more per image than the bare label does. DeiT adapts this to the transformer in an architecturally native way: it adds a second special token, the distillation token, alongside the class token from Section 22.2. The distillation token flows through the encoder just like the class token, but its output is trained to predict the teacher's label rather than the ground-truth label. At inference the predictions from both tokens are combined.

The teacher in DeiT is, pointedly, a strong CNN (a RegNet). The ViT student therefore learns from a network that already has the locality and translation biases the ViT lacks, in effect importing those biases through the teacher's predictions rather than rediscovering them from scratch. This explains why CNN teachers worked better than ViT teachers in their experiments: the student is borrowing exactly the inductive bias it is missing. The training objective combines a standard cross-entropy on the class token with a distillation loss on the distillation token,

where $z_{\text{cls}}$ and $z_{\text{dist}}$ are the two tokens' output logits, $\psi$ is the softmax, $y$ is the true label, and $y_{\text{teacher}}$ is the CNN teacher's predicted (hard) label. Figure 22.3.2 shows the two-token architecture, and the code sketches the training step. The illustration below gives the mentor-and-student picture behind it.

# DeiT's two-token training: the class token is supervised on ground truth and
# the distillation token on a CNN teacher's hard label, so the student imports
# the teacher's locality bias. At inference the two predictions are averaged.
import torch.nn.functional as F

def deit_distillation_loss(cls_logits, dist_logits, targets, teacher_logits):
    """Hard-label distillation: average of true-label CE and teacher-label CE."""
    teacher_labels = teacher_logits.argmax(dim=1)          # CNN teacher's prediction
    loss_cls = F.cross_entropy(cls_logits, targets)        # supervise on ground truth
    loss_dist = F.cross_entropy(dist_logits, teacher_labels)  # supervise on teacher
    return 0.5 * loss_cls + 0.5 * loss_dist

# at inference, fuse the two predictions:
def deit_predict(cls_logits, dist_logits):
    return (F.softmax(cls_logits, -1) + F.softmax(dist_logits, -1)) / 2

# Skip the hand-built recipe: load a distilled DeiT-Base with both tokens
# already trained and fused, then fine-tune with timm's script that bundles
# the full augmentation and regularization stack shown above.
import timm
# distilled DeiT-Base: both tokens already trained, fused at inference
model = timm.create_model("deit_base_distilled_patch16_224", pretrained=True).eval()
# fine-tune on your data with timm's training script, which already includes
# RandAugment, MixUp, CutMix, random erasing, stochastic depth, and EMA.