Chapter 27 · Section 27.4 · 90 min

Training ViTs

This section is about training a Vision Transformer well when you have only about a million labeled images. A Vision Transformer, or ViT, cuts an image into square patches, turns each patch into a vector, and runs a standard Transformer over that sequence of vectors (§27.1 and §27.2). A convolutional network, or CNN, slides small filters over the image, so it is built to know two facts about images. Nearby pixels belong together, and an object is the same object wherever it sits. A ViT is built to know neither, so it must learn both from examples.

§27.3 measured what that costs. Pretrained on 300 million labeled images, a ViT beats a ResNet, the standard deep CNN with residual connections, pretrained on the same data. Trained on ImageNet-1k alone with the original recipe, it loses. ImageNet-1k has 1.28 million labeled photos in 1,000 classes. Top-1 accuracy is the fraction of test images whose highest-scoring class is the correct one.

The fix is a training recipe, not a new network. ViT-B/16 is the base-size ViT, with 16 by 16 pixel patches, 12 blocks of width 768, and 86 million parameters. Trained on ImageNet-1k with the original recipe, it reaches 77.9% top-1, and that is after fine-tuning and testing at 384 by 384 pixels. DeiT-B is exactly the same network, trained on the same images with a different recipe. DeiT stands for data-efficient image Transformers (Touvron et al. 2021a). It reaches 81.8% at 224 pixels, and 83.1% after a short second stage of training at 384 (Table 5 of the paper). The first stage takes 53 hours on one machine with eight GPUs, and the second takes 20 hours more. Only the training changed.

This section takes that recipe apart one ingredient at a time, each with the published experiment that measures its worth. Some ingredients change each training image in ways that keep its label. Some mix two images and their labels. Some regularize the network, by skipping whole layers at random or by starting each layer almost switched off. A teacher network labels every image a second time. A few rules about image resolution and learning rates make fine-tuning work.

Two topics close the section. Large ViTs grow a few tokens whose vectors are about ten times longer than the rest, and extra tokens called registers remove them. And an optimizer that avoids sharp minima, points where the loss rises steeply in some direction, stands in for much of the augmentation. Every vision backbone you fine-tune at work was trained with some version of this recipe. Interviewers ask what each ingredient does, how to code it, and how it goes wrong.

What a recipe can supply instead of data

A model with little built-in knowledge needs more examples. A training recipe can make each example count for more, in three ways.

Augmentation changes a training image in ways that do not change its label. A cat flipped left to right, cropped, or with shifted colors is still a cat. A CNN gets part of this knowledge from its design, because shifting the input of a convolution shifts its output the same way (§27.3). A ViT can learn it only from examples, and augmentation manufactures those examples.

Regularization limits how closely the model can fit its training set. ViT-B has 86 million parameters and ImageNet-1k has 1.28 million images, so with enough passes over the data it can memorize them. Two regularizers from Chapter 0 are weight decay, which pulls every weight toward zero, and label smoothing, which softens the targets. For deep Transformers the one that matters most is stochastic depth, which skips whole layers at random.

A teacher is an already trained model whose predictions serve as extra labels. A ViT trained to agree with a CNN picks up some of what the CNN's design knows about images, without copying the design.

Each of the three makes a training example worth more, and each has a price. You pay in slower fitting, in compute, or in a teacher you must train first. So the trade pays only when data is scarce. Steiner et al. (2021) found that for ViTs trained on ImageNet-1k, the right augmentation and regularization were worth about as much as ten times more data. They also trained on ImageNet-21k, about 14 million images in about 21,800 classes. There, with a short schedule, the same additions made all but the largest models worse. So the first question to ask about any recipe is how much data it was tuned for.

The DeiT recipe and what each part is worth

The DeiT recipe is the baseline that later ViT recipes modify, so we state it exactly. An epoch is one pass over the training set, and the batch size $B$ is the number of images in one step. Most of the recipe comes from Chapter 0.

The optimizer is AdamW. It is Adam with the weight decay applied directly to the weights, rather than added to the gradient.
The learning rate rises linearly during a short warmup, then decays along a cosine curve (§0.6).
Label smoothing with $\epsilon = 0.1$ replaces the one-hot target $\mathbf{e}_y$ , the vector with a 1 at the true class $y$ and 0 elsewhere, by $(1-\epsilon)\mathbf{e}_y + \frac{\epsilon}{K}\mathbf{1}$ , for $K$ classes. The model is never asked for full certainty (§0.5).
Dropout zeroes random activations during training. Dropout and weight decay are both regularizers, covered in §0.7.

The remaining rows of the table are new, and the rest of the section explains each one.

📌 The DeiT-B recipe against the original ViT-B recipe on ImageNet-1k (DeiT, Table 9)

	ViT-B, as DeiT describes it	DeiT-B
epochs, batch size $B$	300, 4,096	300, 1,024
optimizer, peak learning rate $\eta$	AdamW, $3\times10^{-3}$	AdamW, $5\times10^{-4}\cdot B/512$
schedule, warmup	cosine, 3.4 epochs	cosine, 5 epochs
weight decay $\lambda$	0.3	0.05
label smoothing $\epsilon$	none	0.1
dropout	0.1	none
stochastic depth $p_{\text{dp}}$	none	0.1
repeated augmentation	none	yes, 3 copies
gradient clipping	yes	none
RandAugment	none	magnitude 9, noise 0.5
Mixup and CutMix $\alpha_{\text{mix}}$	none	0.8 and 1.0
random erasing	none	probability 0.25

Two rows need a comment now. The learning rate scales linearly with the batch size from a base of 512 images, so $B = 1{,}024$ gives $\eta = 10^{-3}$ . And DeiT's weight decay is six times smaller than ViT's, because the authors found that ViT's value hurt convergence in their setting.

Weight decay also skips some parameters, as in most ViT code. It skips every one-dimensional parameter (the biases, and the gains and shifts of LayerNorm), the class token, and the position table. Pulling those toward zero does not make the model simpler in any useful way.

An ablation removes one ingredient at a time and retrains, with every other setting fixed. The DeiT authors ran one on DeiT-B, and it is the most useful table in their paper.

📌 What each ingredient is worth (DeiT, Table 8, DeiT-B on ImageNet-1k)

change from the default recipe	top-1 at 224	after fine-tuning at 384
none, the default	81.8 ± 0.2	83.1 ± 0.1
SGD instead of AdamW for pretraining	74.5	77.3
SGD for the fine-tuning only	81.8	83.1
no RandAugment (and no AutoAugment)	79.6	80.4
AutoAugment instead of RandAugment	81.2	81.9
no Mixup	78.7	79.8
no CutMix	80.0	80.6
neither Mixup nor CutMix	75.8	76.7
no random erasing	4.3, did not train well	0.1
no stochastic depth	3.4, did not train well	0.1
no repeated augmentation	76.5	77.4
dropout added	81.3	83.1
exponential moving average of weights added	81.9	83.1

Read the table row by row, because it says more than "everything helps".

The optimizer matters most. SGD instead of AdamW costs 7.3 points in pretraining, yet SGD for the fine-tuning stage costs nothing. The last part of this section looks at why the optimizer matters so much for a ViT trained from scratch, through the shape of the loss surface.
Mixup and CutMix, the two ways of mixing images, overlap. Removing Mixup costs 3.1 points and removing CutMix 1.8, but removing both costs 6.0, more than the sum. When one is gone, the other partly covers for it.
Repeated augmentation, which shows each image several times per step with different augmentations, is worth 5.3 points. Few people have heard of it, yet it is worth about as much as any single ingredient.
Dropout, the standard regularizer of the early deep learning era, makes things slightly worse.
The two rows marked "did not train well" are a warning, not a measurement. With every other setting fixed at values tuned for the full recipe, those runs collapsed. They show that the recipe is a tuned whole, not that random erasing is worth 77 points.

So when an interviewer asks which ingredient matters most, the honest answer is the optimizer, then the two mixes together, then repeated augmentation. And no single row can be read in isolation from the rest.

Augmentations that keep the label

RandAugment

An augmentation policy says which image operations to apply, in what order, and how strongly. Before 2019 the strongest policies were found by search. AutoAugment (Cubuk et al. 2019) trained a controller network to choose them, out of about $10^{32}$ possible policies. The search was so expensive that it had to run on a small stand-in task, such as a subset of the data. RandAugment (Cubuk et al. 2020) replaced the whole search with two numbers.

In plain words, pick two random distortions for each image, apply both at one shared strength, and tune only how many and how strong. The paper also found that the best strength is not a constant. It grew with model size. It also grew with dataset size, which is surprising, since augmentation helped most on small datasets. The authors suggest that strong distortions leave too little signal when images are few. So tune the strength on the task and model you care about, not on a small stand-in. That lesson returns at the end of the section.

timm, the PyTorch image-models library that DeiT and most ViT code build on, implements a milder version under the same name. DeiT's setting is written rand-m9-mstd0.5-inc1. Here m9 is the magnitude 9, mstd0.5 its standard deviation, and inc1 selects versions of the operations in which a larger magnitude always means a stronger distortion. It differs from the paper in three ways.

timm draws $N = 2$ operations from a list of 15. The list has no identity, and it adds invert and a second kind of solarize.
Each drawn operation is applied only with probability 0.5.
Each application draws its own magnitude from a normal distribution with mean 9 and standard deviation 0.5, clipped to $[0, 10]$ .

The difference matters when you try to reproduce a number. The cell samples both versions.

Python · cell 1

The paper's version leaves an image untouched only when it draws identity twice, which happens $1/14^2 \approx 0.005$ of the time. It applies two real operations to $(13/14)^2 \approx 86\%$ of images. timm's version leaves a quarter of the images untouched and applies exactly one operation to half of them. It has the same name and the same $N$ and $M$ , and it is a much milder augmentation. If a paper says "RandAugment, $M = 9$ ", ask which code it ran.

Random erasing and repeated augmentation

Random erasing (Zhong et al. 2020) blanks out a random rectangle of the image. A close cousin, Cutout (DeVries & Taylor 2017), blanks a square of fixed size. In DeiT each image is erased with probability 0.25. The rectangle covers between 2% and a third of the image. Its aspect ratio is drawn log-uniformly between 0.3 and $1/0.3$ . timm's pixel mode, DeiT's default, fills it with independent standard normal noise in the normalized color space. The model learns that a class survives the loss of any one region, a kind of robustness to occlusion that §27.3 measured.

Repeated augmentation (Hoffer et al. 2020; Berman et al. 2019) is the surprising 5 points. Each image appears $n_{\text{rep}} = 3$ times in the same step, and each copy gets its own random augmentation. DeiT's sampler sends the three copies to three different GPUs. A batch of 1,024 therefore holds about 341 distinct images.

The sampler then cuts the epoch back to the dataset's size. So one DeiT "epoch" sees only about a third of the distinct images, and its 300 epochs are about 100 passes over the data with three views each. Why does it help? Averaging the loss over three views of one image gives a less noisy estimate of that image's augmented loss. That pushes the model harder toward predictions that agree across views. The price is fewer distinct images per batch.

Why augmentation is a prior, stated exactly

Why should training on flipped images make a model better at unflipped ones? For one clean case there is an exact answer. Average a model's predictions over an image and its flip. The averaged model cannot tell an image from its flip, and its loss on the augmented data is never higher. So the augmented objective always prefers a model that ignores the flip. Here is that claim with real numbers first.

✏️ Averaging over a flip lowers the augmented loss

A two-class model, cat or dog, outputs a score $z$ , and its probability of cat is $\sigma(z) = 1/(1 + e^{-z})$ . On a photo of a cat it scores $z = 2.2$ , so $\sigma(z) = 0.90$ . On the same photo flipped left to right it scores $z = 0$ , so $\sigma(z) = 0.50$ . Training with random flips shows each version half the time. The augmented loss is the average cross-entropy, $\frac12\big(-\ln 0.90 - \ln 0.50\big) = \frac12(0.105 + 0.693) = 0.399$ .

Now average the scores over the two versions, $z = (2.2 + 0)/2 = 1.1$ , and give that score to both. Then $\sigma(1.1) = 0.750$ on each, and the augmented loss is $-\ln\sigma(1.1) = 0.287$ . The averaged model ignores the flip, and its augmented loss is lower.

The loss went down because cross-entropy is convex in the score, so the loss of an average is at most the average of the losses. Nothing about the numbers was special. Here is the precise version. Write $\mathcal{A}$ for a distribution over transformations $t$ that map images to images and are assumed to keep the label. Training on augmented data minimizes the augmented risk

R_{\text{aug}}(\theta) = \mathbb{E}_{(\mathbf{x}, y)}\,\mathbb{E}_{t\sim\mathcal{A}}\;\ell\big(\mathbf{z}_\theta(t(\mathbf{x})),\, y\big),

where $\mathbf{z}_\theta(\mathbf{x}) \in \R^K$ are the model's logits (the scores before the softmax) and $\ell$ is the loss, cross-entropy here. In plain words, the augmented risk is the average loss over images and over random transformations of them.

The theorem needs one more term. A group of transformations is a set that contains the identity, contains the inverse of each member, and contains the composition of any two members. The identity together with the horizontal flip is a group of two. So are the circular shifts of §27.3, which move an image by whole pixels and wrap what falls off one edge around to the other.

∂ Proof that the orbit average is invariant and has lower augmented risk

Invariance first. As $g$ runs over $G$ , so does $gh$ , because multiplying by $h$ on the right is a bijection of a group. So $\bar{\mathbf{z}}(h\mathbf{x}) = \frac{1}{|G|}\sum_g \mathbf{z}(gh\mathbf{x}) = \frac{1}{|G|}\sum_{g'}\mathbf{z}(g'\mathbf{x}) = \bar{\mathbf{z}}(\mathbf{x})$ .

Then fix one pair $(\mathbf{x}, y)$ . Invariance makes every term on the left equal, and Jensen's inequality for the convex $\ell$ bounds the loss of an average by the average of the losses,

\frac{1}{|G|}\sum_{g}\ell\big(\bar{\mathbf{z}}(g\mathbf{x}), y\big) = \ell\big(\bar{\mathbf{z}}(\mathbf{x}), y\big) \le \frac{1}{|G|}\sum_{g}\ell\big(\mathbf{z}(g\mathbf{x}), y\big).

Taking the expectation over $(\mathbf{x}, y)$ gives $R_{\text{aug}}(\bar{\mathbf{z}}) \le R_{\text{aug}}(\mathbf{z})$ .

In plain words, if a model treats a flipped image differently from the original, averaging its predictions over the flips makes it invariant and lowers the augmented loss at the same time. So whenever the model class can represent the averaged predictions, the augmented objective always has an invariant minimizer. Training on augmented data says "be invariant to these changes" through the loss instead of through the architecture. This is the core of the group-theoretic analysis of Chen, Dobriban & Lee (2020). Random crops and RandAugment are not groups, since a crop cannot be undone. For them the theorem is an idealization that points the same way.

That is the exact sense in which augmentation stands in for an inductive bias, a preference built into a model before it sees data. A CNN responds to shifts in a fixed way by construction. A ViT trained on shifted, cropped, and flipped images is pushed toward the same behavior by its objective, and it spends data and compute to get there. An interviewer who asks "why does augmentation help a ViT more than a CNN" wants this argument.

Augmentations that mix labels, Mixup and CutMix

The ingredients so far keep each label intact. Two of the most valuable rows of the ablation do something stranger. They train on images that belong to two classes at once.

📐 Mixup (Zhang et al. 2018)

Given two training examples $(\mathbf{x}_i, \mathbf{y}_i)$ and $(\mathbf{x}_j, \mathbf{y}_j)$ , with $\mathbf{y}$ a target distribution over the $K$ classes (one-hot or smoothed), draw $\lambda_{\text{mix}} \sim \operatorname{Beta}(\alpha_{\text{mix}}, \alpha_{\text{mix}})$ and train on

\tilde{\mathbf{x}} = \lambda_{\text{mix}}\mathbf{x}_i + (1 - \lambda_{\text{mix}})\mathbf{x}_j, \qquad \tilde{\mathbf{y}} = \lambda_{\text{mix}}\mathbf{y}_i + (1 - \lambda_{\text{mix}})\mathbf{y}_j .

In DeiT's code (timm's Mixup in batch mode) one $\lambda_{\text{mix}}$ is drawn per batch and image $i$ of a batch of $B$ is paired with image $B - 1 - i$ , the batch reversed.

In plain words, overlay two images, the first at opacity $\lambda_{\text{mix}}$ and the second at $1 - \lambda_{\text{mix}}$ . Then ask the model to put probability $\lambda_{\text{mix}}$ on the first class and $1 - \lambda_{\text{mix}}$ on the second (Zhang et al. 2018). With $\lambda_{\text{mix}} = 0.7$ , a cat photo mixed with a dog photo has the target 0.7 cat and 0.3 dog. We write subscripts because $\lambda$ alone is the book's weight decay.

The Beta parameter $\alpha_{\text{mix}}$ shapes the mixing. $\operatorname{Beta}(1, 1)$ is uniform on $[0, 1]$ . For $\alpha_{\text{mix}} < 1$ the density piles up near 0 and 1, so many mixed images are nearly one of the two originals. With DeiT's $\alpha_{\text{mix}} = 0.8$ , 26% of draws land within 0.1 of an end. With $\alpha_{\text{mix}} = 0.2$ , 67% do. The Mixup paper found that $\alpha_{\text{mix}}$ in $[0.1, 0.4]$ helped ImageNet ResNets, while much larger values made them underfit.

One identity makes every Mixup implementation short, and it explains what the mixed label does.

∴ Cross-entropy is linear in its target

For any predicted distribution $\mathbf{p}$ and targets $\mathbf{y}_i, \mathbf{y}_j$ ,

H\big(\lambda_{\text{mix}}\mathbf{y}_i + (1-\lambda_{\text{mix}})\mathbf{y}_j,\ \mathbf{p}\big) = \lambda_{\text{mix}}\,H(\mathbf{y}_i, \mathbf{p}) + (1 - \lambda_{\text{mix}})\,H(\mathbf{y}_j, \mathbf{p}),

and the left side, as a function of $\mathbf{p}$ , is minimized exactly at $\mathbf{p} = \lambda_{\text{mix}}\mathbf{y}_i + (1-\lambda_{\text{mix}})\mathbf{y}_j$ .

Proof. $H(\mathbf{y}, \mathbf{p}) = -\sum_k y_k\log p_k$ is a sum of terms each linear in $y_k$ , which gives the identity. The minimizer is Gibbs' inequality from §0.5, $H(\mathbf{q}, \mathbf{p}) \ge H(\mathbf{q}, \mathbf{q})$ with equality only at $\mathbf{p} = \mathbf{q}$ .

In plain words, a mixed label costs the same as two weighted losses, one per class, which is how many codebases implement it. And the best possible prediction on a mixed image is the mix itself. So between two training images the model is trained to report intermediate, honest probabilities instead of a confident guess. Zhang et al. argue that this favors functions that behave linearly between training examples. Thulasidasan et al. (2019) showed that it improves calibration, the match between confidence and accuracy, which we measure below.

CutMix, and why its label needs an area correction

Mixup's images are see-through overlays that never occur in nature. CutMix (Yun et al. 2019) pastes a rectangle instead, so every pixel is a real pixel of one of the two images.

📐 CutMix (Yun et al. 2019)

For images of height $H$ and width $W$ , draw $\lambda_{\text{mix}} \sim \operatorname{Beta}(\alpha_{\text{mix}}, \alpha_{\text{mix}})$ , a box center $r_y \sim U(0, H)$ , $r_x \sim U(0, W)$ , and give the box sides $H\sqrt{1 - \lambda_{\text{mix}}}$ and $W\sqrt{1 - \lambda_{\text{mix}}}$ , clipped to the image. With $\mathbf{M}_{\text{cut}} \in \{0, 1\}^{H\times W}$ equal to 0 inside the clipped box and 1 outside (1 keeps image $i$ ),

\tilde{\mathbf{x}} = \mathbf{M}_{\text{cut}}\odot\mathbf{x}_i + (\mathbf{1} - \mathbf{M}_{\text{cut}})\odot\mathbf{x}_j, \qquad \tilde{\mathbf{y}} = \lambda'\mathbf{y}_i + (1 - \lambda')\mathbf{y}_j, \qquad \lambda' = 1 - \frac{\text{area of the clipped box}}{HW},

applied to every channel. Current timm works in whole pixels. The sides are $h = \lfloor H\sqrt{1-\lambda_{\text{mix}}}\rfloor$ and $w = \lfloor W\sqrt{1-\lambda_{\text{mix}}}\rfloor$ , and the center is an integer. The box covers rows $[r_y - \lfloor h/2\rfloor,\ r_y - \lfloor h/2\rfloor + h)$ and the matching columns, both clipped to the image. The weight is always recomputed from the clipped area. The older timm that DeiT's code pins ends the box at $r_y + \lfloor h/2\rfloor$ instead, one row or column short when $h$ or $w$ is odd.

In plain words, cut a box out of one image, fill it with the same region of another, and label the result by how much of each image you can see. The drawn $\lambda_{\text{mix}}$ only sets the box size before clipping. When the box hangs off the edge of the image, less of image $j$ is pasted than $1 - \lambda_{\text{mix}}$ says. So the label must move toward image $i$ .

How large is that error on average? Clipping happens whenever the center lies within half a box side of an edge. For big boxes that is almost always. The average effect has a closed form.

∂ Derivation of the expected clipped area

Work on one axis first. A box of side $s = \rho H \le H$ centered at $c \sim U(0, H)$ covers $[c - s/2,\ c + s/2]$ . Clipping to $[0, H]$ removes $(s/2 - c)_+$ at the bottom edge and $(c + s/2 - H)_+$ at the top, where $(u)_+ = \max(u, 0)$ . The two cannot both happen, since that would need $c < s/2$ and $c > H - s/2$ , which requires $s > H$ . So the clipped length is $s - (s/2 - c)_+ - (c + s/2 - H)_+$ .

The expected bottom loss is $\frac1H\int_0^{s/2}(s/2 - c)\,dc = \frac{1}{H}\cdot\frac12\big(\frac s2\big)^2 = \frac{s^2}{8H}$ , and the top loss is the same by symmetry. So the expected clipped length is $s - \frac{s^2}{4H} = \rho H\,(1 - \rho/4)$ .

The row and column centers are independent, so the expected area is the product of the expected lengths, $\rho H(1 - \rho/4)\cdot\rho W(1 - \rho/4) = \rho^2(1-\rho/4)^2HW$ . Then $\mathbb{E}[\lambda'] = 1 - \rho^2(1 - \rho/4)^2$ , and $\rho^2 = 1 - \lambda_{\text{mix}}$ gives the result. Since $(1 - \rho/4)^2 \le 1$ , the expected weight is never below the nominal one.

At $\lambda_{\text{mix}} = 0.5$ the expected weight is 0.661, so without the correction the average label is off by 0.16. At $\lambda_{\text{mix}} = 0$ the box is as large as the image, yet after clipping it leaves on average $1 - (3/4)^2 = 0.4375$ of image $i$ in view. The first cell shows both mixes on two drawn images, a disk and stripes. The second checks the formula against timm's integer boxes on $224\times224$ images.

Python · cell 2

Python · cell 3

The simulation and the formula agree to about 0.002, and the gap comes from timm's integer centers and sides. At $\lambda_{\text{mix}} = 0.2$ the box is clipped in 99% of draws, and the true weight averages 0.52, not 0.2. DeiT draws CutMix's $\lambda_{\text{mix}}$ uniformly. Over those draws, a CutMix without the correction would mislabel the average image by 0.18 of its probability mass. That is a bug an interviewer can ask you to find in ten lines of code.

Which of the two mixes is better depends on the model. For a ResNet-50 on ImageNet, Yun et al. report 78.6% with CutMix, 77.4% with Mixup, and 76.3% with neither. CutMix's pasted regions are real image patches, so the network never sees see-through overlays. To put the right probability on each class it has to find where each object is, and Yun et al. found that this improves localization. For DeiT-B the ablation points the other way. Removing Mixup costs 3.1 points and removing CutMix 1.8. So within the full recipe, a ViT trained on ImageNet-1k leans more on Mixup, which blends every pixel. DeiT uses both. For each batch it flips a fair coin (mixup_switch_prob 0.5) and applies one or the other.

Interview

Mixup and CutMix on a batch, with the area correction

Implement mixup_cutmix(x, y, lam, box=None, mode="mixup") the way current timm does it in batch mode. The batch x has shape (B, C, H, W), and y has shape (B, K), one target distribution per image (one-hot or already smoothed). Image i is paired with image B - 1 - i, the batch reversed.

With mode="mixup", return lam * x + (1 - lam) * x_rev and the same mix of the targets.
With mode="cutmix", box = (cy, cx) is the integer box center (row, column). The sides are h = int(H * sqrt(1 - lam)) and w = int(W * sqrt(1 - lam)). The box covers rows from cy - h // 2 up to but not including cy - h // 2 + h, and the matching columns, both clipped to the image. Inside the box, every channel of image i is replaced by image B - 1 - i. The targets use the corrected weight lam_c = 1 - area / (H * W) of the clipped box.

Return (x_mixed, y_mixed, lam_used), where lam_used is lam for Mixup and lam_c for CutMix. Do not modify the inputs.

import numpy as np

def mixup_cutmix(x, y, lam, box=None, mode="mixup"):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

The widget mixes the disk and the stripes of the cell above. In Mixup mode the image and the label blend by the same weight. In CutMix mode you drag the box, and the label bar follows the pixels you can see. In both, the histogram shows where draws of $\lambda_{\text{mix}}$ land for the chosen $\alpha_{\text{mix}}$ .

What do mixed labels do to a model's confidence? A model is calibrated when its confidence matches its accuracy. Among the predictions it makes with confidence 0.8, about 80% are right. The expected calibration error (ECE) of §0.11 measures the gap. It sorts predictions into bins by confidence. In each bin it takes the absolute difference between accuracy and mean confidence, and it averages those differences weighted by bin size.

The cell trains the same small network twice on a 3-class problem, with and without Mixup, using the Adam update of §0.6. The classes overlap, so the best possible accuracy is well below 1. A model that fits its 200 training points is then bound to be overconfident.

Python · cell 4

✏️ Mixup calibrates a small classifier

The network is a two-layer MLP with 64 hidden units, trained for 800 full-batch steps on 200 points and tested on 3,000 fresh ones.

	train accuracy	test accuracy	mean confidence	ECE
no Mixup	0.89	0.70	0.87	0.17
Mixup, $\lambda_{\text{mix}}\sim\operatorname{Beta}(1, 1)$	0.77	0.72	0.75	0.03

Without Mixup the network fits 89% of its training points but is right on only 70% of test points, while reporting an average confidence of 0.87. With Mixup it fits fewer training points, generalizes slightly better, and its confidence nearly matches its accuracy. Mixing labels kept the model from memorizing the noise where the classes overlap.

The same mechanism has a cost when it is overdone. Mixup on top of label smoothing softens the targets twice. On clean data, a model trained to output intermediate probabilities can end up underconfident. So measure calibration rather than assume it. In practice this matters wherever a model's probabilities are used as probabilities, for example to decide when to defer to a human.

Dropping whole residual branches, stochastic depth

Stochastic depth skips whole layers of the network at random during training. Two ideas from earlier chapters make it precise. The first is the residual structure of a Transformer block (§5.1). Each sublayer reads the running vector $\mathbf{x}_\ell$ and adds its output back, $\mathbf{x}_{\ell+1} = \mathbf{x}_\ell + f_\ell(\mathbf{x}_\ell)$ . The sublayer $f_\ell$ is called a residual branch, and skipping it leaves $\mathbf{x}_\ell$ unchanged.

The second is inverted dropout (§0.7). Each activation is zeroed with probability $p$ , and the survivors are divided by $1 - p$ , so the expected activation is unchanged. Stochastic depth (Huang et al. 2016) applies the same idea to a whole residual branch at once.

📐 Stochastic depth, or drop path (Huang et al. 2016)

In a residual network $\mathbf{x}_{\ell+1} = \mathbf{x}_\ell + f_\ell(\mathbf{x}_\ell)$ , training with stochastic depth replaces each update by

\mathbf{x}_{\ell+1} = \mathbf{x}_\ell + \frac{\xi_\ell}{1 - p_\ell}\,f_\ell(\mathbf{x}_\ell), \qquad \xi_\ell \sim \operatorname{Bernoulli}(1 - p_\ell),

with one independent $\xi_\ell$ per training example, the same for every token and channel of that example. Evaluation uses the plain update. In a ViT the attention branch and the MLP branch of each block are dropped independently, each with the block's rate. The rates grow with depth, up to a maximum $p_{\text{dp}}$ . timm, and so DeiT, uses $p_\ell = \frac{\ell - 1}{L - 1}\,p_{\text{dp}}$ for block $\ell = 1, \dots, L$ , which is linspace(0, p_dp, L), so the first block is never dropped. Huang et al.'s linear rule is $p_\ell = \frac{\ell}{L}\,p_{\text{dp}}$ .

In plain words, for each image, flip a biased coin for every residual branch. Skip the branches that lose, and scale up the ones that run so that each block's output is right on average. Every image trains a network of a different random depth. Huang et al. originally made one decision per mini-batch and rescaled at test time instead. The per-example form, rescaled during training, is what modern code calls drop path.

∴ Stochastic depth preserves each block's expected update

If $\xi_\ell \sim \operatorname{Bernoulli}(1 - p_\ell)$ is independent of $\mathbf{x}_\ell$ , then $\mathbb{E}[\mathbf{x}_{\ell+1}\mid\mathbf{x}_\ell] = \mathbf{x}_\ell + f_\ell(\mathbf{x}_\ell)$ , the evaluation-mode update.

Proof. Given $\mathbf{x}_\ell$ , the only random quantity is $\xi_\ell$ , and $\mathbb{E}[\xi_\ell] = 1 - p_\ell$ , so $\mathbb{E}[\mathbf{x}_{\ell+1}\mid\mathbf{x}_\ell] = \mathbf{x}_\ell + \frac{1 - p_\ell}{1 - p_\ell}f_\ell(\mathbf{x}_\ell)$ .

The theorem is about one block, and it is easy to overread. Over two blocks, $\mathbb{E}[\mathbf{x}_{\ell+2}\mid\mathbf{x}_\ell] = \mathbb{E}[\mathbf{x}_{\ell+1}\mid\mathbf{x}_\ell] + \mathbb{E}[f_{\ell+1}(\mathbf{x}_{\ell+1})\mid\mathbf{x}_\ell]$ . The expectation of a nonlinear function is not the function of the expectation. So the evaluation network matches the average training network exactly only for one block, or when the branches are affine. Beyond that it is an approximation, the same one dropout makes. The cell checks both statements and prints the two depth schedules for ViT-B.

Python · cell 5

Why one decision per example, and not per element or per batch? This is a common interview follow-up.

Dropping single elements of $f_\ell(\mathbf{x}_\ell)$ is ordinary dropout on the branch output. It never skips a branch, and the ablation shows that dropout hurts DeiT.
Dropping per batch, as Huang et al. did, makes every example in a step see the same network, which raises the variance of the gradient.
Per example, every batch already averages over many networks of different depths, the implicit ensemble that motivated Huang et al.

The price is that the dropped computation is still done and then multiplied by zero. So drop path saves no compute in a batched implementation.

Core

Drop path and its schedule

Implement drop_path(x, p, rng, training=True) for an array x whose first axis indexes examples (with any number of further axes). In training mode with $0 \le p \lt 1$ , draw one keep decision per example with probability $1 - p$ using rng (a NumPy Generator). Zero the dropped examples entirely, and divide the kept ones by $1 - p$ . With $p = 1$ return zeros. In evaluation mode, or with $p = 0$ , return x unchanged. Do not modify x.

Then implement drop_path_rates(L, p_max), the list of $L$ per-block rates of timm's rule, rising linearly from 0 for the first block to p_max for the last.

import numpy as np

def drop_path(x, p, rng, training=True):
    return x                        # placeholder: never drops anything

def drop_path_rates(L, p_max):
    return [p_max] * L              # placeholder: a uniform rate

Tests, verified solutions, and AI help, open once you add your API key.

LayerScale

Deep Transformers are hard to train when every residual branch starts at full strength. The sum of many unit-sized updates makes the residual stream grow, and the early updates are noisy. §5.7 met two fixes that shrink each branch at the start. One scales each branch's output weights by $1/\sqrt{2L}$ at initialization. The other multiplies each branch by a learned scalar that starts at zero. LayerScale (Touvron et al. 2021b) uses a learned gain for every channel instead. The same paper introduced CaiT, for class-attention in image Transformers, a family of deep ViTs.

In plain words, every branch starts almost switched off, channel by channel. The network starts close to the identity map, and each branch's contribution grows at its own pace. In CaiT's experiments at the width of DeiT-S, a 36-block network without LayerScale failed before the end of training. It reached 81.9% only after its stochastic depth was tuned up to 0.25, while LayerScale reached 82.9%. So LayerScale is what you reach for when a deep ViT diverges. We meet it again in ConvNeXt (§27.5).

🔬 LayerScale adds no expressive power, yet it changes training

The gain can be folded into the branch's last linear layer. For row-vector tokens, $(\mathbf{u}\mathbf{W}_O + \mathbf{b}_O)\operatorname{diag}(\boldsymbol\gamma) = \mathbf{u}\,\mathbf{W}_O\operatorname{diag}(\boldsymbol\gamma) + \mathbf{b}_O\odot\boldsymbol\gamma$ . So every function a LayerScale network computes, a plain network can compute too. What changes is the optimization. §4.3 made the same point about attention's $1/\sqrt{d_k}$ , which can be folded into the query weights without changing the function but not without changing training. Adam takes steps of roughly the same size whatever the scale of a weight. So the way a function is parametrized changes how fast each part of it moves.

Learning from a convnet teacher

The third substitute for data is a teacher. Knowledge distillation trains a student model to match the outputs of a trained teacher model. §20.4 derives the classic soft form, recapped here. Dividing the logits by a temperature $\tau > 1$ flattens both distributions, so the teacher's ranking of the wrong classes shows. A teacher that gives a husky photo 0.7 husky, 0.25 wolf, and 0.05 cat says that huskies look like wolves, which a one-hot label never says. The student's softened distribution is pulled toward the teacher's by a KL divergence term, scaled by $\tau^2$ so that its gradient keeps the same size at any temperature,

\mathcal{L}_{\text{soft}} = (1 - \lambda_{\text{KD}})\,H\big(\mathbf{y}, \operatorname{softmax}(\mathbf{z}_s)\big) + \lambda_{\text{KD}}\,\tau^2\,D_{\mathrm{KL}}\big(\operatorname{softmax}(\mathbf{z}_t/\tau)\,\big\Vert\,\operatorname{softmax}(\mathbf{z}_s/\tau)\big).

Here $\mathbf{z}_s$ are the student's logits, $\mathbf{z}_t$ the teacher's, $\mathbf{y}$ the target (the label, smoothed or mixed), and $\lambda_{\text{KD}} \in [0, 1]$ the weight on the teacher. In plain words, the student is trained partly on the true label and partly to copy the teacher's softened probabilities. DeiT's soft runs used $\tau = 3$ and $\lambda_{\text{KD}} = 0.1$ . DeiT's contribution was to find that, in its setting, a cruder target works better.

📐 Hard-label distillation (Touvron et al. 2021a, Eq. 3)

Let $y_t = \arg\max_c \mathbf{z}_t(c)$ be the teacher's predicted class on the same augmented image the student sees. The loss is

\mathcal{L}_{\text{hard}} = \tfrac12\,H\big(\mathbf{y}^{\text{LS}}, \operatorname{softmax}(\mathbf{z}_{\text{cls}})\big) + \tfrac12\,H\big(\mathbf{e}_{y_t}, \operatorname{softmax}(\mathbf{z}_{\text{dist}})\big),

where $\mathbf{y}^{\text{LS}} = (1 - \epsilon)\mathbf{e}_y + \frac{\epsilon}{K}\mathbf{1}$ is the smoothed true label with $\epsilon = 0.1$ (mixed by Mixup or CutMix when they are on), and the teacher's label is used without smoothing. With a single output head, $\mathbf{z}_{\text{cls}} = \mathbf{z}_{\text{dist}}$ . With the distillation token below, they are two heads.

In plain words, treat the teacher's top guess as a second label, as if a second annotator had labeled the image, and weight the two labels equally. The teacher looks at the augmented image, so its label can differ from the true one. A random crop that cut the object out, or a CutMix whose pasted region dominates, gets a label that matches what is actually visible. The DeiT authors also note that hard labels need no tuning, with no $\tau$ or $\lambda_{\text{KD}}$ to choose. DeiT's teacher is RegNetY-16GF, a CNN from the RegNet family with 84 million parameters and 82.9% top-1. It was trained on ImageNet-1k with the same augmentation as DeiT.

Soft distillation gave DeiT no gain at all (81.8% either way). That does not mean soft distillation fails for images. Beyer et al. (2022) distilled a large ResNet teacher into a ResNet-50 with the KL term alone ( $\lambda_{\text{KD}} = 1$ , no true labels). It worked very well under two conditions.

The teacher must see exactly the image the student sees, the same crop and the same Mixup. They call this consistent teaching.
The student must train far longer than supervised training usually does. They call this patient teaching.

Their best recipe, function matching, adds aggressive Mixup with $\lambda_{\text{mix}}$ uniform on $[0, 1]$ . The student then matches the teacher's function on many images that lie between training images.

📌 A good teacher is patient and consistent (Beyer et al. 2022, Tables 1 and 3)

ResNet-50 student, ImageNet top-1, batch 4,096. The teacher is a ResNet-152 twice as wide, pretrained on ImageNet-21k and fine-tuned to 83.0% on ImageNet (a BiT model, from Big Transfer).

epochs	300	1,200	4,800	9,600
teacher outputs precomputed on a fixed view	77.8	78.1	77.6	77.0
consistent views, function matching	80.3	81.5	82.2	82.3

The best ResNet-50 they trained on the true labels in the same setup reached 78.2%. With a stronger teacher (the teacher at 224 and at 384 pixels, logits averaged), the student reaches 82.8% after 9,600 epochs, a ResNet-50 record at the time.

The fixed teacher stops improving after about a thousand epochs and then overfits, because its target for an image is the same whatever view the student sees. The consistent teacher gives a fresh target for every new view. Its student is still improving at 9,600 epochs, about three million steps.

DeiT's soft run met the first condition, since its code feeds the teacher the same mixed batch as the student. It missed the second. It lasted 300 epochs, where even function matching is 2.7 points below its teacher. And it put only $\lambda_{\text{KD}} = 0.1$ of the loss on the teacher, against one half for hard distillation. So DeiT's hard-label win is a finding about short schedules and its own weighting, not a law. An interviewer who asks why hard labels won wants these two reasons, the short schedule and the small weight.

🔬 Two more details behind DeiT's soft-distillation result

In the released DeiT code the KL term is averaged over the $K$ classes as well as over the images. That shrinks it by a further factor of $K = 1{,}000$ relative to the formula above, and a comment in the code says the authors also tried the per-image normalization. Beyer et al. tested a ResNet student, so whether patient soft distillation also beats hard labels for a ViT student at equal compute is an open question.

A second, smaller suspect comes from §0.5, which noted that networks trained with label smoothing make worse teachers. Smoothing erases what their soft outputs say about similar classes (Müller, Kornblith & Hinton 2019). Whether DeiT's RegNetY-16GF teacher carried that handicap, the paper does not say or test.

The distillation token

DeiT then gives the teacher's label its own token. Recall the input of a ViT from §27.1. It is a sequence of $T$ vectors of width $d$ . First comes the class token $\mathbf{x}_{\text{cls}} \in \R^d$ , a learned vector whose final state summarizes the image. Then come the $N_p$ patch embeddings, one per patch, and a learned position row $\mathbf{P}_{\text{pos}}$ is added to every token. For ViT-B/16 at 224 pixels, $N_p = 14\cdot14 = 196$ and $T = 197$ .

📐 The distillation token (Touvron et al. 2021a)

Add a second learned prefix token $\mathbf{x}_{\text{dist}} \in \R^d$ right after the class token. The sequence becomes the class token, the distillation token, and the $N_p$ patch embeddings, so $T = N_p + 2$ , with a position row for each token.

After the final LayerNorm, the class token's output feeds a linear head with logits $\mathbf{z}_{\text{cls}}$ , trained on the true label.
The distillation token's output feeds a second linear head with logits $\mathbf{z}_{\text{dist}}$ , trained on the teacher's label $y_t$ , through $\mathcal{L}_{\text{hard}}$ .
At test time the paper adds the two heads' softmax outputs, $\hat{\mathbf{p}} = \frac12\big(\operatorname{softmax}(\mathbf{z}_{\text{cls}}) + \operatorname{softmax}(\mathbf{z}_{\text{dist}})\big)$ , and predicts its argmax.

In plain words, one token learns to answer like the human annotators and one learns to answer like the CNN, and at test time the two vote. The extra cost is one token in 198, about half a percent of the compute. The two tokens really do learn different things. Their learned input embeddings have cosine similarity 0.06, and their outputs after the last layer reach 0.93, close but not equal. When the authors instead trained two class tokens on the same true label, the tokens converged to cosine 0.999 and gave no gain.

📌 Distillation results on ImageNet-1k (DeiT, Tables 2, 3, and 5)

method, 300 epochs	DeiT-Ti at 224	DeiT-B at 224	DeiT-B at 384
no distillation	72.2	81.8	83.1
soft distillation	72.2	81.8	83.2
hard distillation, class token only	74.3	83.0	84.0
distillation token, class head alone	73.9	83.0	84.2
distillation token, distillation head alone	74.6	83.1	84.4
distillation token, both heads fused	74.5	83.4	84.5

The teacher, RegNetY-16GF, has 82.9%, so the fused DeiT-B student beats its own teacher. With 1,000 epochs the distilled DeiT-B reaches 84.2% at 224 and 85.2% at 384, while the undistilled one saturates after about 400 epochs. A convnet teacher is better than a Transformer teacher of equal accuracy. A DeiT-B teacher (81.8%) gives a student of 81.9%, and a RegNetY-8GF teacher (81.7%) gives 82.7%.

The last line is the key evidence for the section's theme. Two teachers with the same accuracy produce students almost a point apart, and the better teacher is the one with a different built-in bias. DeiT's Table 4 adds a matching observation. The distilled ViT disagrees with the CNN teacher on 10.2% of images, against 13.3% for a ViT trained without distillation. So the student has moved toward the CNN's way of deciding. The authors read this as the CNN's inductive bias passing to the student through its labels, following Abnar, Dehghani & Zuidema (2020).

One detail matters when you reproduce the numbers. The released code averages the two heads' logits at inference, (x + x_dist) / 2, while the paper's fusion adds probabilities. The two rules usually choose the same class, but not always. The exercise asks for the paper's rule, and its tests include a case where the two disagree.

Interview

DeiT's hard distillation loss and late fusion

Implement deit_loss(z_cls, z_dist, y, z_teacher, eps=0.1), the batch mean of $\mathcal{L}_{\text{hard}}$ . All logits have shape (B, K), and y holds integer labels of shape (B,). Train the class head on the smoothed true label $(1 - \epsilon)\mathbf{e}_y + \frac{\epsilon}{K}\mathbf{1}$ (the convention of §0.5 and of timm). Train the distillation head on the teacher's argmax, with no smoothing. Use a numerically stable log-softmax.

Then implement deit_predict(z_cls, z_dist), which returns the fused probabilities $\frac12(\operatorname{softmax}(\mathbf{z}_{\text{cls}}) + \operatorname{softmax}(\mathbf{z}_{\text{dist}}))$ of shape (B, K).

import numpy as np

def deit_loss(z_cls, z_dist, y, z_teacher, eps=0.1):
    raise NotImplementedError

def deit_predict(z_cls, z_dist):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

Train small, test big, the resolution discrepancy

Here is a puzzle every vision practitioner meets. Touvron et al. (2019) took a ResNet-50 trained at 224 pixels and, without changing a single weight, evaluated it on larger test images. Its top-1 rose from 77.0% at 224 to 78.4% at 288, then fell again at larger sizes. Why would a model do better at a resolution it never trained on? Because the standard pipelines show objects at different sizes in training and at test, and the mismatch can be computed exactly. The paper calls its fix FixRes.

The two pipelines differ in one step. In training, RandomResizedCrop picks a random region of the image covering a fraction $a$ of its area and stretches it to $K_{\text{train}}\times K_{\text{train}}$ pixels. The fraction is uniform, $a \sim U[a_-, a_+]$ , with $[a_-, a_+] = [0.08, 1]$ by default. The aspect ratio is drawn log-uniformly from $[3/4, 4/3]$ . At test time the image is resized so that its shorter side is $K^{\text{img}}_{\text{test}}$ pixels, and a $K_{\text{test}}\times K_{\text{test}}$ center crop is taken. The classic protocol resizes to 256 and crops 224. Write $r_{\text{train}}$ and $r_{\text{test}}$ for the apparent size in pixels of the same object as the network sees it in the two cases. In this part, $r$ is an apparent size, not a grid row.

Take one crop first. It keeps half of the image's area, $a = 0.5$ , so its side is $\sqrt{0.5} = 0.71$ of the image's side. Stretching it to 224 pixels makes everything in it $1/0.71 = 1.41$ times larger than if the whole image had been resized to 224. The test pipeline resizes the whole image to 256 instead, which makes objects $256/224 = 1.14$ times larger than that same baseline. So in this crop an object looks $1.41/1.14 = 1.24$ times larger in training than at test.

∴ The train-test apparent size ratio (Touvron et al. 2019, Eqs. 4 and 5)

For a square image and a crop of aspect ratio 1, an object has

\frac{r_{\text{test}}}{r_{\text{train}}} = \sqrt a\;\frac{K^{\text{img}}_{\text{test}}}{K_{\text{train}}}, \qquad\text{so}\qquad \mathbb{E}\Big[\frac{r_{\text{test}}}{r_{\text{train}}}\Big] = \mathbb{E}[\sqrt a]\,\frac{K^{\text{img}}_{\text{test}}}{K_{\text{train}}}, \qquad \mathbb{E}[\sqrt a] = \frac23\cdot\frac{a_+^{3/2} - a_-^{3/2}}{a_+ - a_-} .

For $[a_-, a_+] = [0.08, 1]$ , $\mathbb{E}[\sqrt a] = 0.708$ (FixRes calls this number $F$ ), and with $K^{\text{img}}_{\text{test}} = 256$ and $K_{\text{train}} = 224$ the expected ratio is $0.708\cdot256/224 = 0.809$ .

Proof. Let the image have side $D$ pixels and the object an apparent side $r$ in it. In training the region has side $\sqrt a\,D$ and is stretched to $K_{\text{train}}$ , a scale factor of $K_{\text{train}}/(\sqrt a\,D)$ , so $r_{\text{train}} = rK_{\text{train}}/(\sqrt a\,D)$ . At test time the image is resized to side $K^{\text{img}}_{\text{test}}$ , a factor $K^{\text{img}}_{\text{test}}/D$ , and the center crop removes pixels without rescaling, so $r_{\text{test}} = rK^{\text{img}}_{\text{test}}/D$ . The ratio is $\sqrt a\,K^{\text{img}}_{\text{test}}/K_{\text{train}}$ , and only $a$ is random. For $a$ uniform on $[a_-, a_+]$ ,

\mathbb{E}[\sqrt a] = \frac{1}{a_+ - a_-}\int_{a_-}^{a_+} a^{1/2}\,da = \frac{1}{a_+ - a_-}\cdot\frac23\Big(a_+^{3/2} - a_-^{3/2}\Big).

With $a_- = 0.08$ , $a_-^{3/2} = 0.02263$ , so $\mathbb{E}[\sqrt a] = \frac23\cdot\frac{0.97737}{0.92} = 0.7082$ .

In plain words, random crops zoom in. On average, training shows objects $1/0.809 = 1.24$ times larger than the test pipeline does, about the same as the single crop with half the area. A network is not scale invariant, so it is tested on object sizes it rarely trained on. Testing at a larger size undoes the zoom, which is where the 77.0% to 78.4% comes from.

Python · cell 6

✏️ The training zoom at real scale

Train at $K_{\text{train}} = 224$ with RandomResizedCrop on $[0.08, 1]$ , and test with resize 256 and center crop 224.

By the formula, $\mathbb{E}[\sqrt a] = 0.708$ . The expected ratio is 0.809, so training zooms objects by 1.24.
With torchvision's sampler, which rejects crops that do not fit, $\mathbb{E}[\sqrt a]$ falls to 0.668. The real zoom is 1.31.
To make the expected ratio exactly 1, resize test images to $K_{\text{train}}/\mathbb{E}[\sqrt a]$ and keep the usual crop ratio of 0.875. That gives a test crop of 277 pixels by the formula and 293 by the simulation.
FixRes measured the best test size for ResNet-50 at 288, between the two.

These targets summarize a whole distribution of zooms. Equalizing the mean of the inverse ratio, $\mathbb{E}[1/\sqrt a]$ , instead would give a different size. Treat them as the right neighborhood, not an exact answer.

The widget draws the whole spread of zooms behind these means, one bar per range of size ratios. Raise the test resize until the mean meets the line of equal size, at about 316 pixels for square crops and 335 for torchvision's.

Testing at a larger size fixes the apparent size, but it shifts something else. For a CNN, the statistics of the activations before the final pooling change with input size. FixRes repairs that by fine-tuning the classifier and the BatchNorm statistics at the test size. For a ViT the shift is bigger. A 288-pixel image at patch size $P = 16$ has 324 patches instead of 196. The position table must be resized to the new grid (§27.1), and every attention layer mixes more tokens than it ever saw in training. So ViTs do not just test at a higher resolution. They fine-tune there.

Core

The FixRes factor and the test size that removes the zoom

Implement fixres_factor(a_min, a_max), the expectation $\mathbb{E}[\sqrt a]$ for $a$ uniform on [a_min, a_max]. Here $a$ is the area fraction, not the side fraction. When a_min == a_max, return $\sqrt{a}$ itself.

Then implement fixres_test_size(K_train, a_min, a_max, crop_ratio=0.875). It returns the test crop size $K_{\text{test}}$ that makes $\mathbb{E}[r_{\text{test}}/r_{\text{train}}] = 1$ under the theorem, when the test pipeline resizes to $K^{\text{img}}_{\text{test}}$ and center-crops $K_{\text{test}} =$ crop_ratio $\cdot K^{\text{img}}_{\text{test}}$ . Return a float, without rounding to a multiple of the patch size.

import numpy as np

def fixres_factor(a_min, a_max):
    return np.sqrt((a_min + a_max) / 2)      # placeholder: the square root of the mean area

def fixres_test_size(K_train, a_min, a_max, crop_ratio=0.875):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

Fine-tuning at the higher resolution

Standard ViT practice combines both ideas. Pretrain at a low resolution, which is cheap because the cost of attention and of the MLPs grows with the number of tokens (§27.2). Then fine-tune briefly at a higher one. DeiT-B fine-tunes at 384 for about 25 epochs by the paper's count, and for 30 in the fine-tuning command of the released code. Before that it resizes the position table by bicubic interpolation. The authors chose bicubic because bilinear interpolation shrinks the norms of the interpolated vectors, which costs accuracy unless the model is fine-tuned.

DeiT's Table 10 shows what the fine-tuning resolution buys a model pretrained at 224. Fine-tuned at 160 it reaches 79.9%, at 224 81.8%, at 320 82.7%, and at 384 83.1%. DeiT III pushes the idea further. It pretrains a ViT-H at 126 pixels, 81 tokens at patch size 14, and fine-tunes at 224. With an 800-epoch schedule that beats pretraining at 224 directly (85.1% against 84.6%), while using about 70% fewer tokens in pretraining. At 400 epochs the two are within 0.1 points. The authors note that the lower resolution also regularizes the largest models. In practice, if a model will be deployed at 384, pretrain it lower and fine-tune it at 384, rather than paying for 384 from the start.

Fine-tuning a pretrained ViT, layer-wise learning-rate decay

Most ViTs you will train start from pretrained weights. There are three common sources.

Supervised pretraining on ImageNet-21k.
CLIP, an image encoder trained to match images with their captions (Chapter 30).
A self-supervised method that learns from unlabeled images. MAE learns by reconstructing masked patches (masked image modeling), and DINOv2 learns by matching a teacher network's outputs across views of the same image (self-distillation).

The lower blocks of such a model compute general features, edges and textures and parts, that a new task can reuse almost unchanged. The head is new and must learn fast. A single learning rate is then either too large for the bottom or too small for the top.

The fix is to shrink the learning rate geometrically with depth below the head. It was introduced for text models as discriminative fine-tuning (Howard & Ruder 2018) and used in the text encoder ELECTRA (Clark et al. 2020). It has been standard for ViTs since BEiT, a masked-image pretraining method (Bao et al. 2022).

In plain words, the head gets the full learning rate, and each step down the network multiplies it by $\zeta$ . With $\zeta = 0.75$ , the last block gets 0.75 of the head's rate, the block below it 0.56, and so on down. A smaller $\zeta$ keeps more of the pretrained features.

✏️ Layer-wise rates for ViT-B

ViT-B has $L = 12$ blocks, so the patch embedding is 13 layers below the head and gets $\zeta^{13}$ of the head's rate.

$\zeta$	where it is used	patch embedding's share of the head's rate
0.65	BEiT, base models	0.0037
0.75	BEiT, large models, and the MAE paper's default (MAE's released ViT-B recipe uses 0.65)	0.024
0.8	Segment Anything, fine-tuning an MAE-pretrained ViT-H	0.055

Segment Anything, a promptable segmentation model (§27.6), paired $\zeta = 0.8$ with drop path 0.4 and weight decay 0.1 (Kirillov et al. 2023). Larger models with better features usually want a smaller learning rate overall, but a $\zeta$ closer to 1, since they have more layers to spread the decay over.

Core

Layer-wise learning rates from parameter names

Implement layer_id(name, L) for parameter names of an $L$ -block ViT in timm's style.

Names starting with patch_embed, and the names cls_token, dist_token, reg_token, and pos_embed, are layer 0.
A name blocks.k.… is layer k + 1. Blocks are numbered from 0, and k can have several digits.
Everything else (norm.weight, head.bias, and so on) is layer L + 1.

Then implement layer_decay_lrs(L, base_lr, zeta), the list of the $L + 2$ learning rates $\eta\,\zeta^{L+1-i}$ for $i = 0, \dots, L+1$ , and param_lr(name, L, base_lr, zeta), the learning rate of one parameter.

def layer_id(name, L):
    raise NotImplementedError

def layer_decay_lrs(L, base_lr, zeta):
    return [base_lr * zeta ** i for i in range(L + 2)]    # placeholder: exponent the wrong way round

def param_lr(name, L, base_lr, zeta):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

Averaging the weights

One more ingredient appears in many recipes. An exponential moving average (EMA) of the weights keeps a second copy of the parameters $\boldsymbol\theta$ . After every step it is updated by $\boldsymbol\theta_{\text{EMA}} \leftarrow \beta\,\boldsymbol\theta_{\text{EMA}} + (1 - \beta)\,\boldsymbol\theta$ , with a decay $\beta$ just below 1. Evaluating that copy averages away the noise of the last stretch of training. The weight on the parameters from $k$ steps ago is $(1-\beta)\beta^k$ , so the average age of what it remembers is $\sum_k k(1 - \beta)\beta^k = \beta/(1 - \beta)$ steps.

DeiT's code uses $\beta = 0.99996$ , an average over about 25,000 steps. DeiT-B takes about 1,250 steps per epoch, so that is roughly the last 20 epochs. In the ablation, EMA added only 0.1 points, which vanished after fine-tuning. EMA matters most as the teacher in self-supervised methods, where the averaged network provides the targets (DINO).

Registers for the artifacts of large ViTs

The recipe so far makes ViTs accurate. The next problem is of a different kind. It shows up only in the features, and only in large models trained for a long time. Many uses of a ViT read its patch features directly, for example to segment an image or find objects without labels, so the features matter as much as the accuracy.

Darcet et al. (2024) noticed the problem while using DINOv2 to find objects without labels. DINOv2 is a ViT trained on unlabeled images by self-distillation, matching a teacher network's outputs across views. In the original DINO, the class token's attention over the patches outlines objects cleanly. In DINOv2 it was full of bright spots in the background.

Each spot is a single patch token whose output vector has about ten times the norm of the others. In DINOv2's ViT-g, 2.37% of the output tokens have a norm above 150, while most are below 100. The same artifacts appear in supervised DeiT III models and in OpenCLIP, an open reproduction of the image-text model CLIP, already at the B and L sizes.

The authors' reading is this. The network needs somewhere to store and process information that belongs to no particular patch, and the class token alone is not enough. So it takes over patches whose own content it can afford to lose. The fix gives it dedicated space.

In plain words, give the model four blank notepads that are not part of the image, and it stops scribbling on the background. The same idea, learned memory tokens added to a Transformer's input, had been proposed for text as the Memory Transformer (Burtsev et al. 2020). Registers bring it to vision.

Language models have an exact counterpart. They park attention they do not need on the first tokens of a sequence, the attention sinks of §4.4. §12.6 showed that one dedicated learnable sink token at the start of every sequence is enough to absorb it. In both cases the model wants a place to put attention or computation that is not content. Interviewers like this connection, because it shows you see one mechanism across two fields.

Registers are cheap. Four tokens added to a few hundred raise the cost of the linear layers by about the ratio of token counts, and the attention cost by about the square of that ratio. The cell uses the multiply-add count of §27.2.

Python · cell 7

One register is enough to remove the visible artifacts. More registers kept improving ImageNet accuracy, while dense tasks had an optimum, with one register giving most of the benefit. DINOv3, the 2025 successor of DINOv2, uses four registers in its 7-billion-parameter teacher (Siméoni et al. 2025). Registers have to be added before training.

For the many models released without them there is a fix that needs no training. Jiang et al. (2025) traced the artifacts to a handful of MLP neurons, which they call register neurons, that write the high norms into the outlier tokens. They append one extra token at test time and move those neurons' high activations into it. That gives cleaner attention maps, and downstream results comparable to trained registers.

Warm-up

Find the high-norm tokens and recover the patch grid

A ViT's output X has shape (T, d). Its first n_prefix rows are prefix tokens (the class token, then any registers or distillation token), followed by the $N_p = H_p W_p$ patch rows in raster order, row by row.

Implement high_norm_tokens(X, n_prefix, thresh), returning (idx, frac). Here idx holds the sorted patch indices, counted from 0 among the patches only, whose Euclidean norm exceeds thresh, as an integer array, and frac is the fraction of patches they make up. Then implement strip_prefix(X, n_prefix, H_p, W_p), returning the patch features as an (H_p, W_p, d) grid, with patch $i$ at row $\lfloor i/W_p\rfloor$ and column $i \bmod W_p$ .

import numpy as np

def high_norm_tokens(X, n_prefix, thresh):
    raise NotImplementedError

def strip_prefix(X, n_prefix, H_p, W_p):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

Sharpness, and an optimizer that stands in for augmentation

Go back to the largest single row of DeiT's ablation. Replacing AdamW by SGD costs 7.3 points, although SGD with momentum is the standard optimizer for ResNets and trains them well. For a ViT, where the optimizer ends up seems to matter far more than for a CNN. Chen, Hsieh & Gong (2022) measured where it ends up.

Their main measure is the sharpness of §0.6. It is the largest eigenvalue $\lambda_{\max}$ of the Hessian $\mathbf{H} = \nabla^2_{\boldsymbol\theta}\mathcal{L}$ , the matrix of second derivatives of the training loss $\mathcal{L}$ , at the final weights $\boldsymbol\theta$ . In plain words, sharpness is the curvature of the loss along its steepest direction. In one dimension, a loss $\frac12\lambda\theta^2$ has sharpness $\lambda$ , and moving $\theta$ by 0.1 from the minimum raises the loss by $0.005\lambda$ . Sharpness is found by power iteration on Hessian-vector products (§0.10), each of which costs about two gradient evaluations (§0.3).

Their baseline is the original ViT recipe, which augments only with Inception-style crops, random crops of varying size and shape plus horizontal flips. It leaves ViT-B/16 at 74.6% at 224 pixels on ImageNet-1k, below a ResNet-152 trained with the same crops. They compared it with MLP-Mixer, a model that mixes patches with MLPs instead of attention and so has even fewer built-in priors than a ViT.

📌 ViTs converge to sharp minima (Chen, Hsieh & Gong 2022, Tables 1 and 4)

Trained from scratch on ImageNet-1k with Inception-style crops only (the original recipes, 300 epochs of AdamW for the ViT and the Mixer, 90 of SGD for the ResNet), at 224 pixels.

model	training loss	loss with Gaussian noise added to the weights	Hessian $\lambda_{\max}$	top-1
ResNet-152 (60M parameters)	0.86	2.39	179.8	78.5
ViT-B/16 (87M)	0.65	6.66	738.8	74.6
MLP-Mixer-B/16 (59M)	0.45	7.78	1,644.4	66.4

The fewer priors the architecture has, the lower its training loss, the sharper its minimum, and the worse its test accuracy. Split into blocks of the Hessian, the ViT's sharpest part is the patch embedding ( $\lambda_{\max} = 300.4$ for that block alone), and early blocks are sharper than late ones.

That flat minima generalize better than sharp ones is an empirical regularity with a long history, not a theorem. Dinh et al. (2017) showed why it cannot be a theorem as stated. In a ReLU network you can scale one layer up and the next down. The function does not change, yet the minimum becomes as sharp as you like. So a comparison of sharpness across architectures, like the table's, is suggestive rather than conclusive. The stronger evidence is the experiment below, which lowers the sharpness of the same ViT and watches its accuracy.

With that caveat, the table suggests a remedy. If the optimizer lands in sharp minima, change the objective so that sharp minima score badly.

📐 Sharpness-aware minimization, SAM (Foret et al. 2021)

For a minibatch loss $\mathcal{L}(\boldsymbol\theta)$ with $\boldsymbol\theta\in\R^n$ and a radius $\rho_{\text{SAM}} > 0$ , SAM minimizes the worst loss in a ball around the weights,

\mathcal{L}^{\text{SAM}}(\boldsymbol\theta) = \max_{\lVert\boldsymbol\epsilon\rVert_2 \le \rho_{\text{SAM}}} \mathcal{L}(\boldsymbol\theta + \boldsymbol\epsilon),

with weight decay added by the base optimizer as usual. One SAM step on a minibatch is

\mathbf{g} = \nabla\mathcal{L}(\boldsymbol\theta), \qquad \hat{\boldsymbol\epsilon} = \rho_{\text{SAM}}\,\frac{\mathbf{g}}{\lVert\mathbf{g}\rVert_2}, \qquad \mathbf{g}_{\text{SAM}} = \nabla\mathcal{L}(\boldsymbol\theta + \hat{\boldsymbol\epsilon}),

with both gradients on the same minibatch. Then the base optimizer (SGD or AdamW) takes an ordinary step from $\boldsymbol\theta$ , not from $\boldsymbol\theta + \hat{\boldsymbol\epsilon}$ , with $\mathbf{g}_{\text{SAM}}$ in place of $\mathbf{g}$ . The norm is taken over all parameters together.

In data-parallel training each device computes its own $\hat{\boldsymbol\epsilon}$ from its own shard of the batch, and the devices average their $\mathbf{g}_{\text{SAM}}$ . Foret et al. call this variant $m$ -sharpness, with $m$ the per-device batch, and found that smaller $m$ tends to generalize better. The perturbation $\boldsymbol\epsilon$ is a vector in weight space, unrelated to the label-smoothing $\epsilon$ above.

In plain words, take a step uphill of length $\rho_{\text{SAM}}$ along the gradient, measure the gradient there, go back, and descend with that gradient instead (Foret et al. 2021). Each step costs two forward and backward passes, so SAM doubles the compute per step. What the two gradients compute can be stated exactly.

∴ What SAM's second gradient measures

Let $\mathcal{L}$ be three times continuously differentiable, with gradient $\mathbf{g} = \nabla\mathcal{L}(\boldsymbol\theta) \ne \mathbf{0}$ and Hessian $\mathbf{H}$ at $\boldsymbol\theta$ .

The first-order inner problem $\max_{\lVert\boldsymbol\epsilon\rVert_2\le\rho_{\text{SAM}}}\ \mathcal{L}(\boldsymbol\theta) + \boldsymbol\epsilon^\top\mathbf{g}$ has the value $\mathcal{L}(\boldsymbol\theta) + \rho_{\text{SAM}}\lVert\mathbf{g}\rVert_2$ , attained only at $\hat{\boldsymbol\epsilon} = \rho_{\text{SAM}}\,\mathbf{g}/\lVert\mathbf{g}\rVert_2$ .
$\mathbf{g}_{\text{SAM}} = \mathbf{g} + \rho_{\text{SAM}}\,\mathbf{H}\mathbf{g}/\lVert\mathbf{g}\rVert_2 + O(\rho_{\text{SAM}}^2)$ , which is the gradient of $\mathcal{L}(\boldsymbol\theta) + \rho_{\text{SAM}}\lVert\nabla\mathcal{L}(\boldsymbol\theta)\rVert_2$ up to $O(\rho_{\text{SAM}}^2)$ .
If $\mathcal{L}(\boldsymbol\theta) = \mathcal{L}^\star + \frac12(\boldsymbol\theta - \boldsymbol\theta^\star)^\top\mathbf{H}(\boldsymbol\theta - \boldsymbol\theta^\star)$ with $\mathbf{H}$ symmetric positive semidefinite, then at the minimum $\max_{\lVert\boldsymbol\epsilon\rVert_2\le\rho_{\text{SAM}}}\mathcal{L}(\boldsymbol\theta^\star + \boldsymbol\epsilon) - \mathcal{L}^\star = \frac12\rho_{\text{SAM}}^2\lambda_{\max}(\mathbf{H})$ .

The proof is in the derivation below.

∂ Proof of the three statements about SAM

(1) By Cauchy-Schwarz, $\boldsymbol\epsilon^\top\mathbf{g} \le \lVert\boldsymbol\epsilon\rVert_2\lVert\mathbf{g}\rVert_2 \le \rho_{\text{SAM}}\lVert\mathbf{g}\rVert_2$ , with equality in both only when $\boldsymbol\epsilon$ points along $\mathbf{g}$ with length $\rho_{\text{SAM}}$ .

(2) Three continuous derivatives make Taylor's remainder for the gradient of second order, so $\nabla\mathcal{L}(\boldsymbol\theta + \boldsymbol\epsilon) = \mathbf{g} + \mathbf{H}\boldsymbol\epsilon + O(\lVert\boldsymbol\epsilon\rVert^2)$ . Setting $\boldsymbol\epsilon = \hat{\boldsymbol\epsilon}$ gives the first expression. For the second, the chain rule gives $\nabla_{\boldsymbol\theta}\lVert\mathbf{g}(\boldsymbol\theta)\rVert_2 = \mathbf{J}_{\mathbf{g}}^\top\mathbf{g}/\lVert\mathbf{g}\rVert_2$ . The Jacobian of the gradient is the Hessian, which is symmetric, so this is $\mathbf{H}\mathbf{g}/\lVert\mathbf{g}\rVert_2$ .

(3) $\mathcal{L}(\boldsymbol\theta^\star + \boldsymbol\epsilon) - \mathcal{L}^\star = \frac12\boldsymbol\epsilon^\top\mathbf{H}\boldsymbol\epsilon \le \frac12\lambda_{\max}\lVert\boldsymbol\epsilon\rVert_2^2 \le \frac12\rho_{\text{SAM}}^2\lambda_{\max}$ by the Rayleigh quotient bound of §0.10, with equality when $\boldsymbol\epsilon$ is $\rho_{\text{SAM}}$ times a top unit eigenvector.

In plain words, away from a minimum, SAM is to first order gradient descent on the loss plus $\rho_{\text{SAM}}$ times the gradient's norm, a penalty for sitting on a steep slope. Its second gradient is a Hessian-vector product in disguise, $\mathbf{g}_{\text{SAM}} - \mathbf{g} \approx \mathbf{H}\hat{\boldsymbol\epsilon}$ . That is a one-sided version of the finite difference that §0.10 uses to measure sharpness, now used inside training. Near a minimum, where the gradient vanishes, the worst loss within distance $\rho_{\text{SAM}}$ is half of $\rho_{\text{SAM}}^2$ times the sharpness. So SAM pays one extra gradient per step to prefer minima whose steepest curvature is small.

Take the one-dimensional loss $\frac12\lambda\theta^2$ again, with $\rho_{\text{SAM}} = 0.2$ . At a sharp minimum with $\lambda = 700$ , the worst loss within the ball is $\frac12\cdot0.04\cdot700 = 14$ . At a flat one with $\lambda = 20$ it is 0.4. Plain gradient descent cannot tell the two minima apart, since both have zero loss and zero gradient. SAM's objective scores the sharp one 35 times worse.

📌 SAM on ViTs, with no new data and no new augmentation (Chen, Hsieh & Gong 2022, Tables 1, 2, 3, 5, 6, and 13)

ViT-B/16 with the same crops and hyperparameters, AdamW plus SAM at $\rho_{\text{SAM}} = 0.2$ , rises from 74.6% to 79.9%. Its $\lambda_{\max}$ falls from 738.8 to 20.9 while its training loss rises from 0.65 to 0.82. Over $\rho_{\text{SAM}} = 0.05, 0.1, 0.2, 0.25$ it reaches 77.5, 78.8, 79.9, and 79.3%.
With SGD instead of AdamW, ViT-B/16 reaches 71.5%, and SGD plus SAM 79.1%, so most of the gap between the two optimizers disappears.
Training the plain ViT for 600 epochs instead of 300 changes nothing, so the gain is not the doubled compute.
The gain shrinks as priors are added. ResNet-152 gains 0.8 points, the hybrid R50-B/16 (a ViT-B/16 reading ResNet-50 features) 2.7, ViT-B/16 5.3, and Mixer-B/16 11.0.
Strong augmentation (Mixup and RandAugment) gives the ViT 79.6%, and augmentation plus SAM 81.5%. On a tenth of ImageNet-1k, plain training gives 32.8%, augmentation 38.5%, and SAM 46.1%.

The last line is where SAM meets the theorem on augmenting over a group. Augmentation asks for a model that does not change along the transformations, a smoothness in input space, taught through augmented examples. SAM asks for a loss that does not change much in its worst direction in weight space, paid for with compute. They are different things.

Chen et al. found that the strongly augmented ViT-B/16 has a $\lambda_{\max}$ of 1,659.3, larger than the plain model's. Yet its loss under random weight noise is low, 1.23 against 6.66. Augmentation flattens the landscape on average but not in its sharpest direction, which is why the two add up.

The ordering of the gains, smallest for the ResNet and largest for the Mixer, is this section's theme once more. An architecture without priors leaves more of the work to training, and the optimizer is part of training. Two facts keep SAM out of large-scale pretraining, and none of the large recipes in this chapter use it. It doubles the compute of every step, and Chen et al. note that its gain shrinks as the training set grows.

Python · cell 8

✏️ SAM on the toy classifier

$\rho_{\text{SAM}}$	training loss	test accuracy	sharpness
0 (plain Adam)	0.285	0.699	28.2
0.05	0.343	0.716	8.2
0.2	0.395	0.709	1.8

At $\rho_{\text{SAM}} = 0.2$ the sharpness is fifteen times lower and the training loss is higher, the pattern of Chen et al.'s table in miniature. The test accuracy moves by about as much as the noise of a 3,000-point test set, whose standard error is about 0.8 points. A toy MLP on two-dimensional data cannot show the five-point gain. That needs a model without priors on a hard problem.

The initialization here matches the Mixup experiment's. So train_sam(0.0) is that experiment's network without Mixup, written with flat parameters so that one vector holds them all.

Interview

One SAM step, and SAM across devices

You are given grad_fn, a function that returns the minibatch gradient (a vector of shape (n,)) at any parameter vector. Write rho for $\rho_{\text{SAM}}$ , and implement three functions. Do not modify theta in any of them.

sam_grad(grad_fn, theta, rho) returns $\mathbf{g}_{\text{SAM}} = \nabla\mathcal{L}(\boldsymbol\theta + \rho_{\text{SAM}}\,\mathbf{g}/\lVert\mathbf{g}\rVert_2)$ , with $\mathbf{g} = \nabla\mathcal{L}(\boldsymbol\theta)$ . If $\mathbf{g} = \mathbf{0}$ , use $\hat{\boldsymbol\epsilon} = \mathbf{0}$ .
sam_step(grad_fn, theta, rho, lr) takes one SAM step with plain SGD as the base optimizer and returns the new parameters.
sam_grad_sharded(grad_fns, theta, rho) is the data-parallel version. Here grad_fns[k] is device $k$ 's gradient function on its own shard. Each device perturbs with its own gradient, and the result is the mean of the devices' SAM gradients.

import numpy as np

def sam_grad(grad_fn, theta, rho):
    return grad_fn(theta)                  # placeholder: the plain gradient

def sam_step(grad_fn, theta, rho, lr):
    raise NotImplementedError

Tests, verified solutions, and AI help, open once you add your API key.

In PyTorch, SAM is a thin wrapper around any optimizer. The first backward pass fills the gradients. The wrapper moves the weights to $\boldsymbol\theta + \hat{\boldsymbol\epsilon}$ , and a second forward and backward pass on the same batch refills the gradients. Then the wrapper moves the weights back before the base optimizer steps. The block repeats the experiment above with AdamW and measures the sharpness with double-backward Hessian-vector products.

🔥 In PyTorchfor a machine with PyTorch

SAM again ends at a higher training loss and a sharpness more than ten times lower. One detail matters in CNNs. With BatchNorm, the second forward pass sees the same batch again and should not update the running statistics a second time. A ViT, with its LayerNorms, does not have that problem.

Choosing the ingredients

No recipe is right for every model and dataset. The trade-offs follow from the three substitutes we started with. Augmentation and regularization help when data is scarce relative to the model's capacity, and they hurt when it is not. Steiner et al. found that on ImageNet-1k almost any augmentation and regularization helped. On ten times more data with a short schedule, they hurt all but the largest models. Augmentation helped in many more settings than regularization did, so their rule of thumb is to prefer augmentation to regularization. The table collects each ingredient's job, its cost, and where it stops paying.

ingredient	what it stands in for	cost	where it stops helping
RandAugment, or 3-Augment	invariance to photometric and geometric changes	slower fitting, CPU time in the data loader	for ViTs, DeiT III found the simpler 3-Augment (one of grayscale, solarization, blur, plus color jitter and flips) better, 83.1% against 82.7% for ViT-B
Mixup and CutMix	a smoother, better calibrated function between examples	slower fitting, mixed labels	short schedules and large datasets; DeiT III drops Mixup on ImageNet-21k
random erasing	robustness to occlusion	a little	DeiT III does without it
repeated augmentation	more views of each image per step	fewer distinct images per batch	ImageNet-21k, where DeiT III turns it off
stochastic depth	an ensemble of shallower networks	nothing at evaluation	small models (DeiT III uses 0 for Ti and S)
LayerScale	a near-identity start for deep stacks	$2Ld$ parameters	shallow models, where it changes little
distillation	the teacher's inductive bias	training and running a teacher	when the student is already as strong as the teacher
fine-tuning at higher resolution	removing the training zoom	a short second stage	none, almost every ViT does it
layer-wise learning-rate decay	keeping pretrained features	one hyperparameter	training from scratch, where it does not apply
registers	clean dense features	under 2% compute	small or short-trained models, which do not grow artifacts
sharpness-aware minimization	the flat minima that a convnet's priors give it, sought by the optimizer	a second forward and backward pass every step	strong augmentation and large datasets, where its gain shrinks; architectures with priors (0.8 points for ResNet-152)

Two later recipes show where the field went. The first is a short report on better plain ViT baselines (Beyer, Zhai & Kolesnikov 2022). It trained ViT-S/16 on ImageNet-1k to 76.5% in 90 epochs, about 6.5 hours on one TPUv3 machine with 8 cores, and to 80.0% in 300 epochs. The changes were small.

Mild RandAugment and Mixup.
The average of the patch tokens instead of a class token.
Fixed 2-D sine-cosine positions instead of learned ones.
A batch of 1,024 instead of 4,096.

The original recipe gave 66.8% in 90 epochs, so a handful of small changes was worth ten points. The second is DeiT III, which trained ViT-H to 85.2% on ImageNet-1k alone. It used binary cross-entropy instead of softmax cross-entropy, the LAMB optimizer, LayerScale, uniform stochastic depth growing with size, 3-Augment, and a low pretraining resolution. The deeper box explains the two unfamiliar pieces.

🔬 Binary cross-entropy and LAMB in DeiT III

Binary cross-entropy treats each of the $K$ classes as an independent yes-or-no question. Each class has its own logistic output $\sigma(z_k)$ , trained toward the target $\tilde y_k$ . With a Mixup or CutMix image, every class present in the mix can get a high probability, without the classes competing for one unit of probability. Wightman et al. (2021) adopted it for ResNet-50 because it fits Mixup and CutMix and was slightly better than cross-entropy. They set the target of every class in the mix to 1, not to its share. DeiT III found it a real gain for larger ViTs on ImageNet-1k. ViT-B at 300 epochs reached 80.9% with softmax cross-entropy and 82.2% with binary. On ImageNet-21k, DeiT III kept softmax cross-entropy.

LAMB (You et al. 2020) is Adam with a per-layer trust ratio. It computes Adam's update $\mathbf{u}$ for a layer's weights $\mathbf{W}$ and adds the decoupled weight decay. Then it rescales the result by $\lVert\mathbf{W}\rVert/\lVert\mathbf{u} + \lambda\mathbf{W}\rVert$ . So every layer moves by the same fraction $\eta$ of its own norm per step, which keeps very large batches with large learning rates stable. In practice the weight norm in the ratio is clipped to a fixed range, so the fraction is exact only inside that range.

On ImageNet-21k, DeiT III also replaced RandomResizedCrop with a simple random crop. With ten times more data, the regularization of aggressive crops was worth less than keeping the labeled object in view.

The strongest substitute for labeled data is not in the table at all. Pretraining without labels, by reconstructing masked patches or by matching views of the same image, gives ViT backbones that fine-tune better than the supervised recipes of this section. MAE's ViT-H reaches 87.8% at 448 pixels using ImageNet-1k images alone (He et al. 2022). That is the subject of Chapter 28. Several of its methods take the augmentations of this section and use them as the supervision itself (contrastive views).

In PyTorch

The block below is a complete DeiT-style training run on a toy problem, written without timm so that every ingredient is visible. The task is to say which quadrant of a noisy 32 by 32 image holds a bright square. The teacher is a hand-built stand-in for a CNN that averages brightness per quadrant, a model with exactly the right prior. The student is a small ViT with a class token and a distillation token, LayerScale, and per-example drop path on every branch. It trains with one Mixup or CutMix per batch on smoothed targets, the hard distillation loss, AdamW with weight decay on matrices only, DeiT's learning-rate rule, and warmup plus cosine decay. The learning rate is far above DeiT's because the run is 160 steps, not 375,000.

🔥 In PyTorchfor a machine with PyTorch

The loss barely moves for the first forty steps. The warmup takes the first ten, and the model sits near chance on the hard mixed and pasted images for a while after. Then the loss falls. It cannot approach zero, because a mixed target has entropy of its own, but the student ends above 90% accuracy on clean held-out images, against 25% for chance. At real scale the same ingredients come from timm, and DeiT's training script exposes each as a flag. The command below is shown, not run.

# DeiT-B with hard distillation from a RegNetY-16GF teacher, 8 GPUs x 128 images = batch 1,024.
# The script scales lr by batch/512, so 5e-4 becomes 1e-3.
torchrun --nproc_per_node=8 main.py \
  --model deit_base_distilled_patch16_224 --batch-size 128 --epochs 300 \
  --lr 5e-4 --weight-decay 0.05 --warmup-epochs 5 --smoothing 0.1 --drop-path 0.1 \
  --aa rand-m9-mstd0.5-inc1 --reprob 0.25 --repeated-aug \
  --mixup 0.8 --cutmix 1.0 --mixup-switch-prob 0.5 \
  --distillation-type hard --teacher-model regnety_160 --teacher-path /path/to/regnety_160.pth \
  --data-path /path/to/imagenet --output_dir /path/to/out
# --mixup 0.8 and --cutmix 1.0 are the Beta parameters alpha_mix, not probabilities.

⚠️ Where ViT training recipes go wrong

DeiT's Table 9 lists "Mixup prob. 0.8" and "Cutmix prob. 1.0", but in the code these are the Beta parameters $\alpha_{\text{mix}}$ . The choice between the two is made with probability 0.5 per batch. Reading them as probabilities gives a different and weaker recipe.
A CutMix written from Equation 2 of its paper alone, without the area correction that the paper's Algorithm A1 applies, mislabels the average image by about 0.18 of its probability mass.
Drop path schedules differ. timm's starts at 0 in the first block and Huang's at $p_{\text{dp}}/L$ . A mask per element is dropout, not stochastic depth.
Repeated augmentation changes what an "epoch" means and how many distinct images a batch holds. Changing the number of GPUs, or turning it off without adjusting the schedule, silently changes the run.
Layer-wise decay depends on the counting convention, such as whether the head is its own layer and where the final norm goes.
Registers and distillation tokens are prefix tokens. Strip them before reshaping patch features into a grid, or every dense feature map is shifted.
The inference rule of a distilled DeiT is not unique. The paper adds probabilities, while the code averages logits.
A SAM step must take both gradients on the same batch and normalize by the gradient norm over all parameters together. It must also move the weights back to $\boldsymbol\theta$ before the base optimizer steps, or it descends from the perturbed point.

🧭 What to remember

A ViT lacks a CNN's priors, so on ImageNet-1k it needs a recipe that stands in for data. DeiT keeps the ViT-B network and lifts it from 77.9% to 81.8% at 224 and 83.1% at 384, with AdamW, strong augmentation, repeated augmentation, stochastic depth, and label smoothing. DeiT's ablation prices each ingredient.
Augmentation is a prior written as data. For a group of label-preserving transformations, averaging any model over the group gives an invariant model with no higher augmented loss, so augmented images teach invariance through the objective instead of the architecture. The optimizer can carry a prior too. Sharpness-aware minimization cuts ViT-B/16's sharpness from 739 to 21 and lifts it from 74.6% to 79.9% with no new data or augmentation, for twice the compute per step.
Mixup and CutMix mix labels. Cross-entropy is linear in its target, so the best prediction on a mixed image is the mixed label, which improves calibration. CutMix must recompute its label weight from the clipped box, and on average $\mathbb{E}[\lambda'] = 1 - (1-\lambda_{\text{mix}})(1 - \sqrt{1-\lambda_{\text{mix}}}/4)^2$ .
Stochastic depth drops each residual branch per example, at a rate rising with depth, and rescales by $1/(1-p)$ , which preserves each block's expected update exactly. LayerScale starts every branch near zero, channel by channel.
Hard distillation from a convnet teacher, through a separate distillation token, beats soft distillation in DeiT's 300-epoch setting and transfers the teacher's inductive bias, so an 83.4% student beats its 82.9% teacher. Soft distillation is not worse in general. With identical views for teacher and student and a very long schedule, it keeps improving for thousands of epochs.
RandomResizedCrop zooms objects by about 1.24 on average ( $\mathbb{E}[\sqrt a] = 0.708$ for scale $[0.08, 1]$ ), so models do better at a higher test resolution, and ViTs fine-tune there with interpolated positions. Fine-tuning a pretrained ViT uses layer-wise learning-rate decay $\eta\zeta^{L+1-i}$ . Large, long-trained ViTs need a few register tokens to keep high-norm artifacts out of their patch features.
Next, §27.5 puts back the prior that dense, high-resolution vision needs most, features at several scales built from local operations. It does so in Swin's shifted windows, ConvNeXt, and hybrid backbones, most of them trained with this section's recipe.

Touvron et al. (2021a). Training data-efficient image transformers & distillation through attention. The DeiT recipe, the ablation of Table 8 that prices every ingredient, hard distillation, and the distillation token.
Touvron, Cord & Jégou (2022). DeiT III: Revenge of the ViT. The recipe revisited, with 3-Augment, binary cross-entropy, LayerScale, simple random crops for ImageNet-21k, and low-resolution pretraining.
Steiner et al. (2021). How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers. A systematic study of when augmentation and regularization substitute for data, and when they hurt.
Yun et al. (2019). CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features. The box sampling of Equation 2 and the evidence that pasting beats blending for localization.
Touvron et al. (2019). Fixing the train-test resolution discrepancy. Section 3 derives the apparent-size ratio of RandomResizedCrop, and the experiments show what testing at higher resolution buys.
Darcet et al. (2024). Vision Transformers Need Registers. How high-norm artifact tokens arise, what they contain, and how four register tokens remove them.
Beyer et al. (2022). Knowledge distillation: A good teacher is patient and consistent. Why soft distillation for images needs identical views for teacher and student and very long schedules, with the ImageNet runs up to 9,600 epochs.
Chen, Hsieh & Gong (2022). When Vision Transformers Outperform ResNets without Pretraining or Strong Data Augmentations. The sharp minima of plain ViTs, the Hessian evidence, and what SAM buys compared with strong augmentation.

Made for larger screens

Training ViTs

What a recipe can supply instead of data

The DeiT recipe and what each part is worth

Augmentations that keep the label

RandAugment

Random erasing and repeated augmentation

Why augmentation is a prior, stated exactly

Augmentations that mix labels, Mixup and CutMix

CutMix, and why its label needs an area correction

Mixup and CutMix on a batch, with the area correction

Dropping whole residual branches, stochastic depth

Drop path and its schedule

LayerScale

Learning from a convnet teacher

The distillation token

DeiT's hard distillation loss and late fusion

Train small, test big, the resolution discrepancy

The FixRes factor and the test size that removes the zoom

Fine-tuning at the higher resolution

Fine-tuning a pretrained ViT, layer-wise learning-rate decay

Layer-wise learning rates from parameter names

Averaging the weights

Registers for the artifacts of large ViTs

Find the high-norm tokens and recover the patch grid

Sharpness, and an optimizer that stands in for augmentation

One SAM step, and SAM across devices

Choosing the ingredients

In PyTorch

Read more

Watch next