DeiT — Training a Vision Transformer on ImageNet Alone

In the last post we saw that a vision transformer is hungry: ViT only beat convnets after it had eaten JFT-300M. Trained on ImageNet-1k alone, ViT-B/16 lands at 77.9 % top-1, well below a good ResNet, and its authors concluded that transformers "do not generalize well when trained on insufficient amounts of data" [TOUV21]. DeiT asks a simple question: is that a property of the architecture, or of the way it was trained?

The answer is the whole post, and it can be stated in three steps:

1. Keep the network. DeiT-B is ViT-B. Same patches, same 12 layers, same 86M parameters.
2. Change the training. A recipe borrowed from the convnet world (strong augmentation, regularisation, AdamW) takes the same model from 77.9 % to 83.1 %.
3. Add one token. A distillation token that learns to copy a convnet teacher takes it to 84.5 %, and to 85.2 % with a 1000-epoch schedule. The 300-epoch run fits on one 8-GPU machine in two to three days [TOUV21].

Brush up before you start ;)

• Patches, the class token and learned positions: An Image Is Worth 16×16 Words — ViT, Patch by Patch. I will not repeat them.
• Dot-product attention and softmax: Query, Key, Value — Dot-Product Attention from Scratch; the block itself: Attention Is All You Need — The Transformer, Block by Block.
• The teacher family, RegNetY: Simple, Powerful, and Fast— RegNet Architecture from Facebook AI.

1. The Problem Is Not the Network

• ViT's own ImageNet-1k number, 77.9 % for ViT-B/16 at 384 px, was already known to be soft. A footnote in DeiT points out that the timm library's training procedure had lifted the same model to 79.35 % on 8 V100s before DeiT existed [TOUV21]. So part of the gap was a recipe problem from day one.
• The second question the paper asks is how to distil a transformer at all. The only tool on the shelf was Hinton's soft distillation [HINT15], designed for convnets and recurrent nets, and nobody had checked whether a transformer student could listen to a convnet teacher.

So the paper has two halves, and I will take them in order: first the recipe, then the token.

2. What Changed in the Network: One Token

Nothing else. Let me say it plainly because the name suggests otherwise: DeiT is not a new architecture. DeiT-B has the ViT-B shapes of Table. 3 in An Image Is Worth 16×16 Words — ViT, Patch by Patch: D = 768, h = 12 heads of d = 64 (the d_k of Query, Key, Value — Dot-Product Attention from Scratch), 12 layers and 16×16 patches, so a 224 px image becomes N = 14 × 14 = 196 tokens [TOUV21]. The differences from ViT-B are a linear classifier instead of an MLP head during pre-training, the training recipe, and the one extra token of §5.

Two smaller models shrink D and h and keep d = 64:

Table. 1 The three DeiT sizes (Table 1 of [TOUV21]). Head dimension d = 64 throughout.

3. The Recipe

If the network is the same, every point of accuracy has to come from how it is trained. The paper puts ViT-B's hyper-parameters next to DeiT-B's in its Table 9, and the list reads like a convnet training manual:

Table. 2 ViT-B against DeiT-B, from Table 9 of [TOUV21]. Training is at 224 px; fine-tuning at 384 px interpolates the positional embeddings bicubically.

A few things to notice:

• Everything below the rule is data augmentation: RandAugment applies a random pair of photometric and geometric ops at a set strength; Mixup blends two images and their labels; CutMix pastes a crop of one image into another and mixes the labels by area; random erasing blanks a rectangle. Above the rule, stochastic depth drops whole residual branches at random during training, and EMA keeps an exponential moving average of the weights for evaluation. A transformer comes with none of a convnet's built-in priors, so the paper leans on augmentation to make up for it.
• Repeated augmentation puts three differently augmented copies of the same image in one batch. With 3 repeats, 300 "epochs" is really 100 passes over distinct images, each three times longer [TOUV21].
• Two NLP habits are dropped: dropout goes, weight decay falls from 0.3 to 0.05. Stochastic depth comes in instead.

3.1 What Each Ingredient Is Worth

Table 8 of the paper removes the ingredients one at a time from DeiT-B (81.8 % at 224 px, 83.1 % after fine-tuning at 384 px) [TOUV21]:

• Switching AdamW for SGD: −7.3.
• Removing Mixup: −3.1. Removing CutMix: −1.8. Removing both: −6.0.
• Removing repeated augmentation: −5.3.
• Removing RandAugment (and not replacing it): −2.2. AutoAugment in its place: −0.6.
• Removing random erasing, or removing stochastic depth: the run does not train at all (4.3 % and 3.4 %) with the fixed hyper-parameters.
• Adding dropout back: −0.5. Adding EMA of the weights: +0.1.

Flip the resolution switch in the figure and the same bars shrink a little after fine-tuning at 384 px: SGD and "no Mixup, no CutMix" swap places at the top, and dropout's −0.5 disappears entirely. The second switch adds the two runs that collapsed.

resolution (0 = 224 px, 1 = fine-tuned at 384 px)

0

show the two runs that collapsed (0 = no, 1 = yes)

0

NOTE: dropout hurts and EMA barely helps. The point of the ablation is not that the recipe is clever; it is that NLP defaults do not carry over to pixels and augmentation does the heavy lifting.

Resolution matters too, and in the direction you would hope. Fine-tuning the same DeiT-B at 160 / 224 / 320 / 384 px gives 79.9 / 81.8 / 82.7 / 83.1 % at 609 / 291 / 134 / 86 images per second on a V100 (Table 10 of the paper) [TOUV21]. Train cheap at 224, fine-tune once at 384.

4. Distillation: Soft and Hard

With the recipe in place, DeiT-B is at 81.8 % at 224 px. The second half of the paper asks a teacher for help.

Hinton's picture of distillation is a biological one: "Many insects have a larval form that is optimized for extracting energy and nutrients from the environment and a completely different adult form that is optimized for the very different requirements of traveling and reproduction" [HINT15]. The large, slow model you train is the larva; the model you deploy is the adult; and the point is that they need not be the same shape.

The reason a teacher's answers carry more than the labels do is Hinton's other example: "An image of a BMW … may only have a very small chance of being mistaken for a garbage truck, but that mistake is still many times more probable than mistaking it for a carrot" [HINT15]. The small probabilities on the wrong classes say how the teacher sees the world. Let's devise the equations step-by-step.

Throughout: $Z_s$ and $Z_t$ are the student's and teacher's logits over K = 1000 classes, $\psi$ is the softmax, and $y$ is the true label (one-hot, label-smoothed with ε = 0.1 in every DeiT run).

4.1 The Baseline Loss

Nothing new here: pay for the log-probability the student gives the right class.

4.2 Soft Distillation

Hinton's method [HINT15], DeiT's Eq. 2: add a second term that compares the student's and teacher's softened distributions.

Three changes from Eq. 1, one at a time:

• τ, the temperature. The sharpness knob of Query, Key, Value — Dot-Product Attention from Scratch §4, turned up: at τ > 1 the small wrong-class probabilities, the garbage-truck ones, carry signal instead of being rounded to zero. DeiT uses the usual τ = 3 [TOUV21].
• τ². Gradients from the softened targets shrink as 1/τ², so Hinton multiplies the term by τ² to keep it on the same scale as the label term [HINT15]. It is not optional.
• λ, the mix. DeiT uses λ = 0.1 [TOUV21]: 90 % "match the label", 10 % "match the teacher's whole opinion, including its doubts".

NOTE: the paper writes the KL with the student first. The released code (losses.py) computes kl_div(log_softmax(student/T), log_softmax(teacher/T)), which is KL(teacher ‖ student), i.e. the student's cross-entropy against the teacher's soft targets minus the teacher's constant entropy. Eq. 2 is written the way the code runs, because the direction changes the gradient.

4.3 Hard-Label Distillation

DeiT's Eq. 3 throws the KL term away and keeps only the teacher's decision:

• $y_t$ is the teacher's top class for this augmented crop. There is no τ and no λ; the teacher's answer, in the paper's words, "plays the same role as the true label" [TOUV21].
• Why would the teacher's label differ from the dataset's? The paper's example: an image labelled "cat" that is "a large landscape and a small cat in a corner. If the cat is no longer on the crop of the data augmentation it implicitly changes the label of the image" [TOUV21]. With Mixup, CutMix and random erasing in the recipe, this happens all the time. The teacher re-labels every crop, and that is the mechanism, not noise.

Numbers again. Take two classes, true label A, temperature 1, and compare what each term charges:

Table. 3 The two distillation terms on a two-class toy problem, computed by the figure below at τ = 1.

• Row 1: the student copies the teacher exactly, and the soft term is zero even though the student is still 10 % unsure. The hard term keeps pushing.
• Row 2: the teacher is barely sure (0.55) and the student is confident; soft distillation punishes the student for being more certain than the teacher. The hard term does not care how sure the teacher was.
• Row 3: both terms agree the student should move.

Now drag the student's $p_s$ in the figure and watch the bottom of the soft bowl follow it; raise τ and the bowl flattens while the hard step stays exactly where it was.

student p_s (class A)

0.7

temperature τ

3.0

The soft curve is a bowl with its bottom at $p_s=p_t$; raising τ flattens the bowl. The hard curve is a step at $p_t=0.5$ and flat on either side: the teacher's confidence never enters. And here is the surprise of the paper (Table 3 of the paper, DeiT-B, 224 px, 300 epochs) [TOUV21]:

• No distillation: 81.8 %.
• Soft distillation: 81.8 %. The textbook method did nothing.
• Hard distillation: 83.0 %. Parameter-free, and +1.2.

5. The Distillation Token

Both losses above are applied to the class token's output. The paper's own contribution is to give the teacher a token of its own.

• A distillation token is a learned D-vector appended to the input next to the class token, so the network processes N + 2 tokens. It goes through the same self-attention layers as everything else, exactly as the class token does.
• After the last layer L, each special token is projected by its own linear head, and each head has its own target: class token → true label, distillation token → teacher's hard label.

• $z_L^{cls},z_L^{dist}\in {ℝ}^D$ are the two special tokens after the last layer; $W_{cls},W_{dist}\in {ℝ}^{K\times D}$ are the two heads. In the code the ½ / ½ is alpha = 0.5. Soft distillation with the token uses Eq. 2's KL on $Z_s^{dist}$ instead.
• Compared with Eq. 3 the only change is where the two cross-entropies land: on two different tokens instead of the same one. Two students in the same lecture hall, hearing the same lecture, graded against two different answer keys.

Is it just a second class token? The paper checked. Two class tokens with the same target converge to the same vector (cosine similarity 0.999) and add nothing. The learned class and distillation input tokens sit at cosine 0.06, and their outputs reach 0.93 at the last layer: similar, because they are predicting the same thing, but not the same [TOUV21].

5.1 Inference: Late Fusion

At test time there are three options: use the class head, use the distillation head, or add the two softmax outputs and take the argmax.

The paper calls this late fusion, and it is the best of the three. The whole ladder, DeiT-B, 300 epochs, at 224 px and after fine-tuning at 384 px (Table 3 of the paper) [TOUV21]:

Table. 4 Distillation variants for DeiT-B, ImageNet-1k top-1 (Table 3 of [TOUV21]). The alembic ⚗ in the paper marks models trained with the token.

6. Which Teacher?

Here is the second surprise. You would expect the best teacher for a transformer to be a bigger transformer. Table 2 of the paper says otherwise (teacher → DeiT-B⚗ student, 224 px / ↑384 px) [TOUV21]:

• DeiT-B as teacher (81.8 %): student 81.9 / 83.1.
• RegNetY-4GF as teacher (80.0 %): student 82.7 / 83.6.
• RegNetY-16GF as teacher (82.9 %): student 83.1 / 84.2.

A RegNetY-4GF is 1.8 points worse than DeiT-B as a classifier and still produces a better student. The paper's explanation is that distillation "can transfer inductive biases in a soft way" [TOUV21]: the convnet's built-in translation habits reach the transformer through the labels it hands over. Table 4 of the paper backs it up: the distilled model disagrees with the convnet on 10.2 % of images and with a from-scratch DeiT on 10.7 %, and the distillation head alone is the closest to the convnet of all (10.0 %) [TOUV21]. The student has learned to think like its teacher.

The default teacher is therefore a RegNetY-16GF [RADO20], 84M parameters, 82.9 % top-1, trained with DeiT's own recipe. And the student overtakes it: DeiT-B⚗ reaches 83.4 % at 224 px while running at 291 images per second against the teacher's 335, with a smaller memory footprint [TOUV21].

7. Results

The headline table, ImageNet-1k only, no outside data (Table 5 of the paper) [TOUV21]:

Table. 5 ImageNet-1k top-1 with no external data (Table 5 of [TOUV21]). ⚗ = trained with the distillation token, which adds one token and one head, hence 87M; ↑384 = fine-tuned at 384 px. ViT-B pre-trained on JFT-300M scores 84.15 at 384 px for comparison.

Reading it as one story, the same ViT-B throughout:

• 77.9 → 83.1 at 384 px: +5.2 points from the recipe alone.
• 83.1 → 84.5: +1.4 from the distillation token.
• 84.5 → 85.2: +0.7 from training 1000 epochs instead of 300. Without distillation, accuracy saturates after about 400 epochs; with the token it keeps climbing (Fig. 3 of the paper; 83.4 → 84.2 at 224 px, Table 5) [TOUV21].
• The last line beats ViT-B pre-trained on 300 million JFT images (84.15) using one million.

NOTE: 85.2 % is not the three-day model. Three hundred epochs of DeiT-B take 53 hours on one 8-GPU node (37 hours on two), plus an optional 20 hours of fine-tuning at 384 px; that buys 81.8 %, or 83.4 % with the token. A RegNetY-16GF is 20 % slower to train for its 82.9 % [TOUV21].

Transfer holds up as well: DeiT-B⚗↑384 reaches 99.2 % on CIFAR-10, 91.4 % on CIFAR-100, 98.9 % on Flowers, 93.9 % on Cars and 80.1 / 83.0 % on iNaturalist 2018 / 2019 (Table 7 of the paper) [TOUV21].

8. What Came Next

DeiT settled the data question for ImageNet-sized problems, and its recipe became the default: the next paper in this series, Swin Transformer — Windows, Shifts and a Pyramid for Dense Prediction, trains with DeiT's hyper-parameters minus repeated augmentation and EMA, and beats DeiT-S by 1.5 points at the same FLOPs. But it does so by changing the architecture, which DeiT deliberately left alone. ViT and DeiT produce one feature map at 1/16 resolution with attention that costs quadratically in the number of tokens, and neither is what a detector or a segmenter wants to see. Swin's answer is a pyramid built from windows that shift. Stay tuned.

9. References

1. [TOUV21] Touvron, Cord, Douze, Massa, Sablayrolles, Jégou. Training data-efficient image transformers & distillation through attention. arXiv:2012.12877.
2. [DOSO20] Dosovitskiy et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv:2010.11929.
3. [HINT15] Hinton, Vinyals, Dean. Distilling the Knowledge in a Neural Network. arXiv:1503.02531.
4. [RADO20] Radosavovic, Kosaraju, Girshick, He, Dollár. Designing Network Design Spaces. arXiv:2003.13678.