All Tokens Matter: Token Labeling for Training Better Vision Transformers

In plain words

A vision transformer cuts an image into small square patches and turns each patch into a token. Normally, only one extra token, the one that stands for the whole image, is asked to say what the picture shows. The other tokens get no direct feedback, even though each one holds detail about its own small area.

Token labeling asks every patch to say what it sees too. A separate, strong image classifier first produces a map of likely labels for each small area of each training photo. During training, each patch is checked against the label for its own area, alongside the usual whole-image check. The paper also mixes two training photos by swapping rectangles of patches, so that the labels stay matched to the patches. The authors report that the same training recipe gives a small model higher accuracy on ImageNet than a standard baseline, and that a larger model beats the other transformers they compare with at a similar size.

TL;DR

Token labeling is an extra training loss for vision transformers. Every patch token is trained against a dense label map produced in advance by an ImageNet-only NFNet-F6 classifier, so the model gets location-specific supervision at almost no training cost §1 §3.2 §4.1. A CutMix-style augmentation on patch embeddings (MixToken) is used with it §3.3. The authors report 84.4% ImageNet top-1 for LV-ViT-S at 26M parameters and 86.4% for LV-ViT-L at 512 px tab. 4. Those numbers are author-reported. The released code differs from the paper in several training and inference details, listed under Paper vs code.

Problem

Vision transformers use a class token to aggregate the whole image, so the classification loss touches only that token. The patch tokens, which carry local information, get no direct supervision §1. Earlier data-augmentation methods such as CutMix mix pixels across images, which can leave a single patch containing content from two images, and then a per-patch label is unclear §3.3. Distillation, the other common route, needs a teacher model running during training §3.2. The paper targets this gap: it wants dense supervision for patch tokens that is cheap to train with, without a teacher running online.

Method

The baseline is a standard ViT: the image is split into patches, each patch is linearly projected to a token, a class token is added, and the classification loss is taken on the class token's output after the last transformer block §3.1 eq. 1.

Token labeling §3.2. Each training image also gets a dense target: one 1000-dimensional label per patch token, arranged on the token grid. The training objective is the class loss plus a weighted average of the per-token cross-entropy losses eq. 2 eq. 3. The labels come from a pretrained annotator, the NFNet-F6 classifier, trained on ImageNet only (86.3% top-1) §4.1. Its dense output for each training image is computed once, before training, so the training step itself adds little compute §3.2 §4.1.

Storage and cropping §4.1. The annotator's output map is 18 x 18 x 1000 at its 576 px input §4.1. The authors store only the five largest values per cell, in half precision, which they report as about 10 GB instead of 2 TB for full maps §4.1. During training the map is cropped with the same random crop as the image, pooled to the token grid (14 x 14 for 224 px) with RoIAlign, and turned into probabilities by a softmax over classes. A 1 x 1 pooling over the whole crop gives the class target mixup.py. The paper describes this as cropping and interpolating the score map §3.2.

MixToken §3.3 eq. 5 eq. 6 eq. 7. Instead of mixing raw pixels, the method mixes patch embeddings. Two images are embedded, and a binary CutMix-style mask M decides for each token whether it comes from the image or from its partner. Token labels are mixed with the same mask eq. 6. The class label is mixed by the fraction of tokens taken from each image eq. 7. The paper does not state the mixing-weight distribution; the code uses Beta(1, 1) [code: tlt/models/lvvit.py#L137] lvvit.py.

LV-ViT architecture §4.2 app. B.1 tab. 1. The LV-ViT models change two things from ViT. The linear patch projection is replaced by a small convolutional stem: a 7 x 7 stride-2 convolution, two 3 x 3 stride-1 convolutions, then an 8 x 8 stride-8 projection to tokens app. B.1 layers.py. Each residual branch is divided by a constant (2 for S and M) app. B.1 layers.py. The four sizes are:

Training steps, in the order the code runs them:

1. Annotate each training image once, offline, keeping the top-5 raw outputs per cell §4.1 generate_label.py.
2. Per sample, crop and flip the stored maps with the image, and build the class and token targets (pseudocode below).
3. Embed the image, apply MixToken to the embeddings, and run the transformer (pseudocode below).
4. Compute the class loss on the mixed class target and the token loss on the unmixed token targets, then add them (pseudocode below).

Pseudocode for step 1 (offline labels):

# Offline dense labels from the annotator  [§4.1]
# [code: generate_label.py#L264-L272]
# Input:  one training image at 576 px
# Output: top-5 logits and class ids for each of the 18 x 18 cells

def annotate(image):
    # Run the annotator on the whole image; its output is raw logits  [§4.1]
    logits = annotator(image)                    # (18, 18, 1000)
    # Keep the 5 largest logits per cell, stored in half precision  [§4.1]
    top_logits, top_classes = topk(logits, k=5)  # (18, 18, 5) each
    return top_logits, top_classes

Pseudocode for step 2 (targets for one crop):

# Token targets for one training crop  [§3.2] [§4.1]
# [code: tlt/data/mixup.py#L21-L45] [code: tlt/data/mixup.py#L9-L18]
# Input:  stored top-5 maps and a crop box in relative coordinates
# Output: class target (1000,) and token targets (14, 14, 1000)

def crop_targets(top_logits, top_classes, box):
    # Scatter the top-5 logits into a 1000-channel map; others stay 0
    # [code: tlt/data/mixup.py#L9-L18]
    dense = scatter(top_logits, top_classes)     # (18, 18, 1000)
    # Pool over the box to the 14 x 14 token grid with RoIAlign
    # [code: tlt/data/mixup.py#L27-L33]
    token_logits = roi_align(dense, box, output_size=(14, 14))
    # Pool over the whole box for the class target  [code: L35-L37]
    class_logits = roi_align(dense, box, output_size=(1, 1))
    # Softmax over classes turns each logit map into a target  [code: L45]
    return softmax(class_logits), softmax(token_logits)

Pseudocode for step 3 (MixToken and the forward pass):

# MixToken on patch embeddings, with partner-based token outputs
# [§3.3] [eq. 5] [code: tlt/models/lvvit.py#L189-L221]
# Input:  patch embeddings (B, 14, 14, d) for a batch of images
# Output: class logits (B, 1000), token logits (B, 14, 14, 1000), box

def mixtoken_forward(patches):
    # Draw the mixing weight from Beta(1, 1), which is uniform
    # [code: tlt/models/lvvit.py#L137] [code: tlt/models/lvvit.py#L194]
    lam = sample_beta(1.0, 1.0)
    # One CutMix-style box for the whole batch  [code: lvvit.py#L43-L59]
    box = random_box(grid=(14, 14), lam=lam)
    # Copy the box region from the batch-flipped partner  [eq. 5]
    # [code: tlt/models/lvvit.py#L198]
    mixed = patches.copy()
    mixed[:, box] = patches.flip(0)[:, box]
    class_logits, token_logits = transformer(mixed)
    # Inside the box, take the partner's token outputs  [code: lvvit.py#L217]
    token_logits[:, box] = token_logits.flip(0)[:, box]
    return class_logits, token_logits, box

Pseudocode for step 4 (loss):

# Training loss with a class term and a token term  [§3.2] [eq. 3]
# [code: tlt/loss/cross_entropy.py#L56-L82]
# Input:  class_logits (B, 1000), token_logits (B, 196, 1000), box
# Output: scalar loss

def token_label_loss(class_logits, token_logits, class_t, token_t, box):
    # Fraction of grid cells outside the box  [code: cross_entropy.py#L75]
    lam = 1 - box_area(box) / 196
    # Mix the class target with the batch partner using lam; only when lam < 1
    # [code: cross_entropy.py#L76-L77]
    class_t = lam * class_t + (1 - lam) * class_t.flip(0)
    # Class term: soft-target cross-entropy, averaged over the batch  [eq. 1]
    class_loss = soft_ce(class_logits, class_t).mean()
    # Token targets are used unmixed  [code: cross_entropy.py#L73]
    # Token term: cross-entropy averaged over all tokens  [eq. 2]
    token_loss = soft_ce(token_logits, token_t).mean()
    # Weighted sum. Paper: 0.5 on the token term. Code default: 1.0
    # [§3.2] [code: main.py#L272-L275]
    return class_loss + token_weight * token_loss

Inference, as the released code does it. The paper does not describe this rule:

# Prediction in eval mode  [code: tlt/models/lvvit.py#L209-L211]
# Not described in the paper.
# Input:  patch embeddings (B, 14, 14, d) from one unmixed image
# Output: class scores (B, 1000)

def predict(patches):
    # Run the transformer with no MixToken  [code: lvvit.py#L189]
    class_logits, token_logits = transformer(patches)
    # Add half of each class's largest token logit to its class logit
    # [code: tlt/models/lvvit.py#L211]
    return class_logits + 0.5 * token_logits.max(over_tokens)

Hyperparameters

The ImageNet-1K recipe below is the one the paper states, with the code's value beside it when they differ. "Not stated" means the paper does not give the value.

Results

Author-reported. Setup: ImageNet-1K training only, no external data, 300 epochs, top-1 on the validation set; test resolution as shown §4.1 tab. 4 tab. 4 caption.

Ablations, author-reported:

• Ablation path from DeiT-S to LV-ViT-S: 79.9 to 84.4 in total. Token labeling with MixToken adds from 82.4 to 83.3 in that path tab. 8.
• Token labeling vs ImageNet labels: MixToken with token labeling 83.3; MixToken with ReLabel 83.0; CutMix with ReLabel 82.8; MixToken with hard labels 82.5; CutMix with hard labels 82.4 tab. 2.
• Augmentations on LV-ViT-S: full set 83.3; without MixToken 81.3; adding MixUp 83.1; removing CutOut 83.0; removing RandAugment 82.8 tab. 3.
• Token participation: using 20% of tokens for the token loss costs 0.5 points on LV-ViT-S and 0.4 on LV-ViT-M §4.2 fig. 4.
• Annotator: EfficientNet-B3 (81.6%) as annotator still helps LV-ViT-S; NFNet-F6 gives 83.3, and EfficientNet-L2 (pretrained on JFT-300M, which the paper excludes from main results) gives 83.5 §4.2 fig. 4.
• Segmentation with FCN, ADE20K: LV-ViT-S 46.1 to 47.2 (single scale) with token labeling tab. 5.
• Other ViTs and MLPs (re-implemented): Mixer-S/16 75.6 to 76.1, Mixer-B/16 78.3 to 79.5, Mixer-L/16 77.7 to 80.1, ResNeSt-50 80.9 to 81.5 tab. 11.
• Against CaiT-S36 at 384: same 85.4 top-1 with 56M vs 69M parameters and 42B vs 48B FLOPs, 300 vs 400 epochs tab. 12.

Paper vs code

Checked against the official repository at code_commit. Both sides are cited.

• Token-loss weight. The paper sets beta = 0.5 §3.2. The code's defaults for the class and token weights are both 1.0 main.py, and the documented ImageNet commands do not override them README.md.
• Score maps. The paper says a "score map" is produced by the annotator §4.1. The code stores the annotator's raw top-5 outputs, not probabilities generate_label.py. It converts them to probabilities after RoIAlign with a softmax over all 1000 classes, filling unlisted classes with zero mixup.py mixup.py.
• Token targets under MixToken. The paper mixes token labels with the mask M eq. 6. The code does not mix token targets: they are used unmixed cross_entropy.py. Inside the mask, the code takes the token outputs of the flipped partner instead lvvit.py. This is a different construction from eq. 6, and the card cannot confirm from reading that the two are equivalent.
• Class-target mixing. The paper mixes the class label by the mask's average eq. 7. The code computes the mixing weight from the model's box, area-corrected, and applies it only when the weight is below 1 cross_entropy.py.
• Inference rule. The paper does not describe it. The code's eval path returns the class logit plus 0.5 times the maximum token logit per class lvvit.py. The validation loop uses that output main.py.
• Label smoothing. The paper states no smoothing. The code applies smoothing of 0.1 to all targets by default main.py main.py.
• EMA. README training commands enable --model-ema README.md. The paper does not mention EMA, and does not say which weights its numbers use.
• Drop path. Paper's default table gives 0.1 tab. 7. README gives 0.1 for S, 0.2 for M and 0.3 for L README.md.
• LR for L. The formula gives 1.6e-3 at batch 1024 §4.1. The README's L command uses 1e-3 README.md.
• RandAugment for L. Table 7 gives 9/0.5 tab. 7. The README's L command uses N=3 README.md.
• Fine-tuning epochs. The paper gives 30 epochs app. A. The README fine-tuning command gives no epoch count README.md.
• Segmentation. The paper gives a minimum LR of 5e-6 §4.4. The S config uses poly decay with minimum LR 0 upernet_lvvit_s_512x512_160k_ade20k.py. The L config sets the backbone image size to 448 upernet_lvvit_l_512x512_160k_ade20k.py, while the paper says 512 §4.4.
• Checkpoint accuracy. README lists 84.0 for LV-ViT-M at 224 README.md; Table 4 gives 84.1 tab. 4.
• Patch-embedding width. The ablation table's caption describes 64 filters tab. 9. The code uses 128 filters for L layers.py. The paper does not state the width for L.
• Residual scaling for L. The code uses a factor of 3 lvvit.py. The paper ablates 2 and 3 but does not assign them to sizes tab. 10.

Pitfalls

• The annotator is the expensive part, and it is not timed. One forward pass of NFNet-F6 at 576 px over the whole training set is needed before training. The paper reports only storage (about 10 GB for top-5 maps) §4.1.
• Stored values are logits, not probabilities. Anyone who stores probabilities, or softmaxes before RoIAlign, gets different targets. The softmax over 1000 classes depends on how many classes are kept, because the rest are zero mixup.py mixup.py.
• Crop coordinates in the code. RoIAlign boxes are the crop coordinates scaled by the map size, minus 0.5 mixup.py. The paper does not state this offset.
• Flips must move the label maps too. Horizontal flips flip the label maps along with the image label_transforms_factory.py.
• One box per batch. The MixToken box is drawn once for the whole batch and shared by all samples lvvit.py. Sampling one box per sample is a different method.
• Image-level mixing double-mixes the token targets. Enabling --mixup or --cutmix adds image-level mixing and mixes the token targets again with the same weight mixup.py main.py. The documented commands leave both at 0.
• Ground-truth blend is off by default. The code can blend hard labels into the class target with a weight of 0.9 or 0.5 cross_entropy.py. The paper does not describe it, and the flag is not in the documented commands.
• Rounding of tables. The paper's tricks table gives per-step gains that do not chain: 79.9 to 81.2 is labeled +1.2, not +1.3 tab. 8. Use the endpoints.
• Small effects, no seeds reported. Several ablation gaps are 0.1 to 0.5 points tab. 2 tab. 3. The paper does not say how many runs each number averages.
• Cost of the augmentation. The paper says the token-label and MixToken steps add negligible compute §3.2 §3.3. It does not report wall-clock cost.

Lineage

• Builds on: ViT, arXiv:2010.11929 §3.1; DeiT, arXiv:2012.12877, for the training recipe and the distillation contrast §2 §4.1; CutMix and MixUp §3.3 §4.1; ReLabel, arXiv:2101.05022, for the dense-label idea §3.2 §4.1; NFNet, arXiv:2102.06171, as the annotator §4.1.
• Compared against: DeiT, T2T-ViT, CaiT (arXiv:2103.17239), Swin, CrossViT and other ViTs and CNNs tab. 4 tab. 6.
• Superseded by: not stated.