Evo-ViT: Slow-Fast Token Evolution for Dynamic Vision Transformer

In plain words

A vision transformer cuts an image into many small patches and spends work on every patch in every layer. Evo-ViT lets the model decide, layer by layer, which patches carry the important information. The model's own summary token already points at the parts of the image it attends to, and Evo-ViT uses that signal to sort the patches.

The important patches get the full, expensive treatment. The others are not thrown away. They are folded into one shared stand-in token, which gets a cheap update, and that update is handed back to every unimportant patch. The image keeps its layout, and the model still has access to the background. The result is a faster network that is trained this way from the start, rather than one that is slimmed down after a finished training run.

TL;DR

Evo-ViT is an efficiency method for ViT-style image classifiers that works during training, not only at inference. At chosen layers it splits patch tokens into informative ones (high CLS attention, updated fully) and placeholders (summarised into one token, updated cheaply). The paper reports large throughput gains with small top-1 losses on ImageNet-1k for DeiT and LeViT. Note the gap between the paper and its DeiT code, listed under Paper vs code: the representative-token update is implemented only for LeViT.

Problem

Transformers spend quadratic cost on the number of tokens, so cutting tokens is the main lever. Token pruning removes tokens, but the paper argues that this (a) breaks the 2D spatial structure that deep-narrow models with spatial pooling rely on, and (b) relies on a pre-trained model, so it cannot be trained from scratch § Introduction. The paper also argues that tokens in shallow and middle layers are hard to prune, because their features are still changing fast and are weakly correlated with each other [§ Methodology: Structure preserving token selection] fig. 4. The goal is a method that works from the start of training and keeps all tokens in play § Introduction.

Method

Evo-ViT works in two parts, applied at selection layers.

Structure preserving token selection. The CLS token's attention to each patch serves as the importance score eq. 3 fig. 2. A global version of this score is carried across layers with a residual blend eq. 4. At each selection layer the top-k patches by this score are informative and the rest are placeholders. Unlike pruning, no token is dropped [§ Methodology: Structure preserving token selection].

Slow-fast token updating. Placeholders are merged into one representative token by a weighted sum using their global scores eq. 5. Informative tokens and the representative token go through MSA and then FFN, with residual connections eq. 6. The placeholders then receive the representative token's two residuals, copied to every placeholder position eq. 7. Informative tokens get the slow, full update; placeholders get a fast, coarse one [§ Methodology: Slow-fast token updating].

Training strategies. The paper uses a layer-to-stage schedule: selection and slow-fast updating happen at every layer for the first 200 epochs, then selection happens only at the start of each stage for the last 100 epochs. For DeiT, a stage is four layers § Methodology: Training Strategies. An auxiliary loss adds a classification loss on the CLS token alongside the loss on the average-pooled features eq. 8; at inference only the average-pooled features are used § Methodology: Training Strategies.

The paper's settings: selection from the fifth layer for DeiT and the third non-convolutional layer for LeViT, a keep ratio of 0.5 at each selected layer, and α = 0.5 for the global attention § Experiments: Setup.

Background used by the method, with the paper's notation:

ViT encoder layer § Preliminaries:

\begin{aligned} x_0 &= \lceil x_{cls} \mid x_{patch} \rceil + x_{pos}, \\ y_k &= x_{k-1} + MSA(LN(x_{k-1})), \\ x_k &= y_k + FFN(LN(y_k)) \end{aligned} eq. 1

Self-attention and CLS attention § Preliminaries:

SA(Q,K,V)=Softmax\left(\frac{Q K^T}{\sqrt{C}}\right) V eq. 2

A_{cls} = Softmax\left(\frac{q_{cls} K^T}{\sqrt{C}}\right) eq. 3

Global class attention over layers, with α the trade-off [§ Methodology: Structure preserving token selection]:

A_{cls,g}^{k} = \alpha\, A_{cls,g}^{k-1} + (1-\alpha)\, A_{cls}^{k} eq. 4

Representative token aggregation § Methodology: Slow-fast token updating:

x_{rep} = \phi_{agg}(x_{ph}) eq. 5

Slow path, with residuals recorded as x^(1) and x^(2) [§ Methodology: Slow-fast token updating]:

\begin{aligned} x_{inf}^{(1)}, \; x_{rep}^{(1)} &= MSA(x_{inf}, x_{rep}), \\ x_{inf} &\leftarrow x_{inf} + x_{inf}^{(1)}, \quad x_{rep} \leftarrow x_{rep} + x_{rep}^{(1)}, \\ x_{inf}^{(2)}, \; x_{rep}^{(2)} &= FFN(x_{inf}, x_{rep}), \\ x_{inf} &\leftarrow x_{inf} + x_{inf}^{(2)} \end{aligned} eq. 6

Fast path for placeholders § Methodology: Slow-fast token updating:

x_{ph} \leftarrow x_{ph} + \phi_{exp}(x_{rep}^{(1)}) + \phi_{exp}(x_{rep}^{(2)}) eq. 7

Assisted CLS loss § Methodology: Training Strategies:

\begin{aligned} \hat{y}_{cls}, \hat{y} &= m(x_{cls}, x_{patch}), \\ \mathcal{L} &= \phi(\hat{y}_{cls}, y) + \phi(Avg(\hat{y}), y) \end{aligned} eq. 8

The pseudocode below is written from the equations above. The CLS token's placement in the slow path is not stated in the paper; step 2 follows the code evo_deit.py.

# Slow-fast token evolution at one selection layer  [eq. 3-7]
# Input:  x (B, 1+N, C) with CLS at index 0; g (B, N) global scores
# Output: x (B, 1+k+(N-k), C), informative tokens first; g (B, N)

def slow_fast_layer(x, g, msa, ffn, keep_ratio, alpha):
    # 1. Rank patch tokens by global class attention  [§ Structure preserving]
    num_keep = int(g.shape[1] * keep_ratio)  # k, informative count
    order = argsort(g, descending=True)      # (B, N)
    cls_token = x[:, :1]                     # (B, 1, C), never pruned
    patch = gather(x[:, 1:], order)          # (B, N, C)
    g = gather(g, order)                     # (B, N)
    # 2. Split into informative and placeholder tokens  [§ Structure preserving]
    x_inf = concat(cls_token, patch[:, :num_keep], axis=1)  # (B, 1+k, C)
    x_ph = patch[:, num_keep:]               # (B, N-k, C)
    # 3. Merge placeholders into one representative token  [eq. 5]
    weights = g[:, num_keep:] / sum(g[:, num_keep:])  # sums to 1 [code]
    x_rep = sum(weights * x_ph, axis=1, keepdims=True)  # (B, 1, C)
    # 4. Slow path: MSA over informative and representative tokens  [eq. 6]
    tokens = concat(x_inf, x_rep, axis=1)    # (B, 2+k, C)
    msa_out = msa(tokens)                    # residual branch, (B, 2+k, C)
    tokens = tokens + msa_out
    # 5. Slow path: FFN on the same tokens  [eq. 6]
    ffn_out = ffn(tokens)                    # residual branch, (B, 2+k, C)
    tokens = tokens + ffn_out
    # 6. Fast path: copy the representative residuals to placeholders  [eq. 7]
    rep_msa = msa_out[:, -1:]                # x_rep^(1), (B, 1, C)
    rep_ffn = ffn_out[:, -1:]                # x_rep^(2), (B, 1, C)
    x_ph = x_ph + rep_msa + rep_ffn          # broadcast over N-k tokens
    # 7. Blend this layer's CLS attention into informative scores  [eq. 3-4]
    cls_attn = cls_attention(msa, tokens)    # (B, k), informative only
    g_inf = alpha * g[:, :num_keep] + (1 - alpha) * cls_attn
    g = concat(g_inf, g[:, num_keep:], axis=1)  # placeholders unchanged
    # 8. Output: CLS and informative tokens, then placeholders  [§ Slow-fast]
    x_out = concat(tokens[:, :1 + num_keep], x_ph, axis=1)
    return x_out, g

Hyperparameters

Paper values come from the setup and training sections. Anything marked code is the released code's default or README value; the paper does not state it.

Results

Author-reported, ImageNet-1k classification, 300 epochs. Top-1 accuracy and throughput in images per second. Throughput for the DeiT comparison is one NVIDIA V100 at batch 256 § Main Results. Other tables do not state the hardware.

Token pruning methods on DeiT tab. 1:

Comparison with state-of-the-art ViTs, 224 px unless marked * (384 px) tab. 2:

Module ablation on DeiT-T and LeViT-128S, each row adds one change tab. 3:

Token selection strategies on DeiT-T, all at the fifth layer with the structure preserving strategy and without layer-to-stage training. Sub-sampling methods run at the seventh layer tab. 4:

Other reported results:

• Keeping ratio: a consistent ratio of 0.5 in all layers gave the best result among the variants compared § Experiments: Consistent keeping ratio fig. 6.
• Selection stays consistent across layers in a well-trained model, and tokens picked early can be recovered later § Experiments: Visualization fig. 5.
• Last-layer CLS attention used for selection in every layer lost accuracy (69.7% vs 72.0%). The paper offers a possible reason: networks need some background context, and restricting every layer to object tokens may underfit background features [§ Experiments: Different Token Selection Strategy] tab. 4.

Paper vs code

The paper itself links https://github.com/YifanXu74/Evo-ViT (in its abstract section), and the repository's MIT copyright line names the first author, so no further search was needed. Pinned at code_commit. Models are in deit/ and levit/; training is in main_deit.py, main_levit.py and their engines. Discrepancies:

• The representative-token update is in LeViT but not in the DeiT model. Paper: placeholders are summarised and fast-updated by eq. 5 eq. 7. LeViT code merges them with normalised global weights and adds the representative residuals evo_levit.py. DeiT code keeps placeholders unchanged, and a code comment says that keeping only the placeholders already works well on DeiT without a representative token evo_deit.py. The DeiT results in tab. 1 and the DeiT "fast updating" row in tab. 3 therefore cannot be reproduced by the released DeiT model as written.
• Trade-off convention is flipped in code. Paper: α weights the previous scores eq. 4. Code: tradeoff weights the new scores, so α = 1 − tradeoff evo_deit.py. At α = 0.5 both are the same; copying a non-0.5 value from one into the other will invert the blend.
• DeiT: the first four layers do not use α = 0.5. The tradeoff list sets those layers to 1, which in the paper's convention is α = 0 [code: deit/evo_deit.py#L556]. Selection starts at layer index 4 [code: deit/evo_deit.py#L345]. The paper says α = 0.5 for all layers [§ Experiments: Setup].
• LeViT: α is learned after epoch 200. The code sets a learnable scalar passed through a sigmoid, starting at 0.5, and turns it on only after epoch 200 evo_levit.py evo_levit.py evo_levit.py engine_levit.py. The paper fixes α = 0.5. The code's convention for tradeoff is the flipped one above.
• Head aggregation differs. DeiT sums CLS attention over heads [code: deit/evo_deit.py#L369]; LeViT averages over heads [code: levit/evo_levit.py#L394]. The paper does not say. Inferred, not verified by running: the scale difference should not change the ranking, since it is the same for all tokens and the blend is linear.
• Loss is weighted. Paper: unweighted sum eq. 8. Code: 0.5 times the pooled-feature loss plus 0.5 times the CLS loss engine_deit.py. The criterion is label-smoothed cross-entropy, or soft-target cross-entropy when mixup is active main_deit.py. The paper says "cross-entropy".
• LeViT distillation is hard-label, two-head. Paper: "trained with distillation", no details § Experiments: Setup. Default code mode cls uses the teacher's argmax as the target on both heads, with α = 0.5 [code: levit/losses_levit.py#L68-L70] losses_levit.py losses_levit.py. The README names RegNetY-160 as the teacher for "all distillation settings" README.md. DeiT defaults to no distillation main_deit.py.
• Layer-to-stage switch matches. Both engines use layer-wise selection for epochs below 200 and stage-wise after engine_deit.py engine_levit.py. This is the one schedule detail where paper and code agree.
• Training recipe is code-only. The paper states only 300 epochs. The optimizer, learning rate, schedule, augmentation and batch size in the Hyperparameters table come from the code and the README, which lists different batch sizes per model. The paper's batch size of 256 applies to the throughput measurement, not to training. The README also says Evo-ViT-B can produce NaN losses early on and recommends a 9-epoch warm-up checkpoint and batch size 128 README.md.

Pitfalls

• The DeiT code does not implement the representative-token update. The paper's method is the spec. If you build the DeiT model from the paper, the LeViT merge and copy (see Paper vs code) is the closest code to follow. If you reproduce the released DeiT numbers, expect the released behaviour.
• Equation, figure and table numbers come from the LaTeX source, not a rendered PDF. The AAAI-22 class prints no section numbers (secnumdepth is 0), so sections are cited by title.
• CLS handling in the slow path is not stated. The code always keeps the CLS token as the first token and never prunes it evo_deit.py.
• Placeholders are not updated in the DeiT model. The DeiT code concatenates them back unchanged at the end of each stage [code: deit/evo_deit.py#L416]. Any layer-wise feature change they would get in the paper's design is absent there.
• Throughput comparisons are not all comparable. Only the DeiT comparison states V100 and batch 256 § Main Results. Tables 2 and 4 and the ablations give no hardware; the authors do not say whether they reran baselines.
• Accuracy costs differ by model. The paper reports larger drops for LeViT than for DeiT. Its explanation is that LeViT's deeper layers have few tokens and so less redundancy to remove § Main Results.
• Dense input. The paper says the method accelerates LeViT at 384 px with less accuracy loss and a larger speed-up than at 224 px § Main Results. Both 384 px DeiT rows in tab. 1 use a different baseline throughput (87 img/s), so compare within a resolution.
• Training cost is not reported. Only inference throughput is measured. The cost of the 300-epoch training run with these layers is not stated.

Lineage

• Builds on: DeiT (flat ViT training recipe and code base) [§ Experiments: Setup] README.md; LeViT (deep-narrow model, distillation and code base) § Experiments: Setup; CKA similarity, following Refiner, for layer analysis § Methodology: Structure preserving token selection.
• Compared against: DynamicViT, PatchSlimming (called PS-ViT in the table), SViTE, IA-RED² tab. 1.
• Superseded by: not known from content already on this wiki.