Later Is Better: Token Reduction for ViTs Under Distribution Shift
Organizations: KAIST
Abstract
Training-free token reduction accelerates vision transformers by removing redundant tokens across layers, recovering most of the original accuracy at a fraction of the compute. These methods, however, are designed and evaluated primarily on clean data, and under real-world distribution shift their accuracy gap to the uncompressed model widens with the removal rate. We show that this gap is governed by the reduction schedule, the depth profile of removal, usually left fixed as an implementation detail. Concretely, we introduce a one-parameter late-concentrated power-law schedule that consistently improves out-of-distribution accuracy over flat at no extra inference cost. On ImageNet-C with DeiT-S, the late schedule closes 83% of that gap at a 26% compute reduction (+1.17pp), and 99% of it at a lighter 7% reduction (+0.26pp). The gain cannot be attributed to retaining more tokens or using extra compute: held to flat's compute, the late schedule removes more tokens in total and leaves fewer tokens at the end, yet still wins. Single-layer probes point to a mechanism: earlier reductions perturb features that pass through more remaining layers, front-loading reduction error in depth. The effect is broad, holding across five token-reduction methods (ToMe, EViT, ATS, ATC, PiToMe), nine backbones, all ImageNet-C corruption types, eight further shift suites, and two further modalities, video and vision-language QA. It is also specific to shift, still positive on clean and rising monotonically to ~4x that at the highest severity 5. The schedule keeps its gain under six test-time adaptation methods, and needs no per-input or per-domain tuning.
Figures & tables
| late-concentrated ( vs. flat) | ||||||||
|---|---|---|---|---|---|---|---|---|
| GFLOPs (cut) | flat | |||||||
| 3.97G ( ) | 39.14 | |||||||
| 3.69G ( ) | 38.83 | |||||||
| 3.34G ( ) | 38.44 | |||||||
| 3.15G ( ) | 37.99 | |||||||
| 2.90G ( ) | 37.41 | |||||||
| Sev. | flat | late | ||
|---|---|---|---|---|
| 0 (clean) | 79.40 | 79.68 | 1.0 | |
| 1 | 70.52 | 71.04 | 1.9 | |
| 2 | 64.33 | 64.99 | 2.4 | |
| 3 | 59.06 | 59.89 | 3.0 | |
| 4 | 50.06 | 51.12 | 3.8 | |
| 5 | 37.99 | 39.16 | 4.2 |
| ( ) | ( ) | ( ) | |||||||
| Method | flat | late | flat | late | flat | late | |||
| DeiT-S | |||||||||
| ToMe | 38.83 | 39.37 | 38.44 | 39.32 | 37.99 | 39.16 | |||
| EViT | 37.50 | 38.97 | 36.69 | 38.91 | 35.98 | 38.50 | |||
| ATS | 34.95 | 35.43 | 34.50 | 35.21 | 33.92 | 34.88 | |||
| ATC | 38.28 | 38.73 | 37.96 | 38.74 | 37.55 | 38.65 | |||
| ViT-S | ViT-B | |||||
| Method | flat | late | flat | late | ||
| Source (inference-only) | 38.05 | 39.46 | 53.39 | 54.96 | ||
| Tent ( Wang et al., 2021 ) | 48.47 | 52.04 | 54.49 | 55.72 | ||
| EATA ( Niu et al., 2022 ) | 52.65 | 53.66 | 59.09 | 60.11 | ||
| SAR ( Niu et al., 2023 ) | 51.52 | 51.96 | 59.71 | 61.07 | ||
| FOA ( Niu et al., 2024 ) | 42.42 | 42.81 | 59.80 | 60.41 | ||
Appendix figures & tables30 assets
Supplementary material from the paper’s appendix.
Appendix
| Sev | flat | ||||
|---|---|---|---|---|---|
| 0 (clean) | 79.40 | 79.67 | 79.68 | ||
| 1 | 70.52 | 70.95 | 71.04 | ||
| 2 | 64.33 | 64.82 | 64.99 | ||
| 3 | 59.06 | 59.71 | 59.89 | ||
| 4 | 50.06 | 50.88 | 51.12 | ||
| 5 | 37.99 | 38.89 | 39.16 |
| Backbone | flat | late ( ) | late ( ) |
|---|---|---|---|
| DeiT-S | 36.67 | 37.88 ( ) | 37.77 ( ) |
| DeiT-B | 41.44 | 42.96 ( ) | 42.74 ( ) |
| ViT-S | 36.07 | 37.48 ( ) | 36.92 ( ) |
| ViT-B | 50.93 | 52.71 ( ) | 52.15 ( ) |
| ViT-B | ViT-S | |||||
|---|---|---|---|---|---|---|
| Corruption | flat | late | flat | late | ||
| gaussian_noise | 53.64 | 56.30 | 2.66 | 31.51 | 33.38 | 1.87 |
| shot_noise | 53.85 | 56.16 | 2.31 | 30.93 | 32.54 | 1.61 |
| impulse_noise | 54.11 | 56.92 | 2.82 | 31.08 | 32.93 | 1.85 |
| defocus_blur | 44.74 | 46.98 | 2.23 | 30.26 | 31.78 | 1.52 |
| glass_blur | 35.04 | 35.57 | 0.53 | 22.61 | 23.00 | 0.39 |
| ImageNet | DN target domains | DN in-domain | ||||
|---|---|---|---|---|---|---|
| Backbone | A | V2 | sketch | clipart | painting | real |
| ViT-S | ||||||
| ViT-B | ||||||
| DeiT-S | ||||||
| DeiT-B | ||||||
| CLIP ViT-B | ||||||
| DeiT-S | ViT-S | |||||
|---|---|---|---|---|---|---|
| Corruption | flat | late | flat | late | ||
| gaussian_noise | 23.41 | 24.88 | 1.47 | 14.85 | 16.04 | 1.19 |
| shot_noise | 23.78 | 24.95 | 1.17 | 16.23 | 17.83 | 1.60 |
| impulse_noise | 26.02 | 27.47 | 1.45 | 15.60 | 16.76 | 1.16 |
| defocus_blur | 16.49 | 16.70 | 0.21 | 6.98 | 7.03 | 0.05 |
| glass_blur | 16.92 | 17.53 | 0.61 | 17.29 | 17.40 | 0.11 |
| Severity 3 | Severity 5 | |||||
|---|---|---|---|---|---|---|
| Corruption | flat | late | flat | late | ||
| Shot noise | 67.40 | 68.28 | 50.04 | 51.68 | ||
| Fog | 70.92 | 71.36 | 66.80 | 67.96 | ||
| Contrast | 76.08 | 76.08 | 73.20 | 73.08 | ||
| Brightness | 75.84 | 75.52 | 71.72 | 71.80 | ||
| Saturate | 75.64 | 75.76 | 65.60 | 65.52 | ||
| ImageNet-R | ImageNet-Sketch | |||||
|---|---|---|---|---|---|---|
| Backbone | flat | late | flat | late | ||
| DeiT-S | 42.14 | 42.45 | 29.40 | 29.68 | ||
| DeiT-B | 44.56 | 44.63 | 32.21 | 32.42 | ||
| ViT-S | 45.90 | 46.57 | 32.89 | 33.33 | ||
| ViT-B | 58.77 | 58.85 | 44.92 | 45.09 | ||
| DeiT-S | ViT-S | |||||
|---|---|---|---|---|---|---|
| Corruption | flat | late | flat | late | ||
| blue noise | 50.36 | 50.62 | 53.08 | 55.10 | ||
| brownish noise | 66.40 | 67.36 | 60.42 | 62.34 | ||
| caustic refr. | 53.60 | 54.46 | 51.32 | 53.12 | ||
| checkerboard | 62.66 | 63.58 | 64.06 | 65.20 | ||
| sine waves | 31.42 | 31.70 | 24.62 | 26.42 | ||
| DeiT-S | ViT-S | |||||
|---|---|---|---|---|---|---|
| Corruption | flat | late | flat | late | ||
| bit error | 15.50 | 16.04 | 16.70 | 16.98 | ||
| color quant | 59.82 | 59.80 | 41.56 | 41.44 | ||
| far focus | 50.96 | 51.62 | 53.30 | 54.14 | ||
| flash | 39.28 | 39.84 | 38.98 | 39.70 | ||
| fog 3d | 38.72 | 38.52 | 31.34 | 31.96 | ||
| Corruption | flat | late | |
|---|---|---|---|
| Speckle noise | 41.67 | 43.98 | |
| Gaussian blur | 23.38 | 24.09 | |
| Spatter | 47.15 | 49.72 | |
| Saturate | 62.10 | 63.17 | |
| Mean | 43.58 | 45.24 |
| Method | flat | late | flat | late | flat | late | |||
| DeiT-S | |||||||||
| ToMe | 38.44 | 39.13 | 37.99 | 38.89 | 37.41 | 38.58 | |||
| EViT | 36.69 | 38.10 | 35.98 | 37.87 | 35.16 | 37.62 | |||
| ATS | 34.50 | 34.99 | 33.92 | 34.61 | 33.35 | 34.38 | |||
| ATC | 37.96 | 38.53 | 37.55 | 38.35 | 37.04 | 38.26 | |||
| Method | full | flat | late | flat | late | ||
|---|---|---|---|---|---|---|---|
| ViT-S | |||||||
| Tent | 50.19 | 49.27 | 50.84 | 48.47 | 52.04 | ||
| EATA | 55.20 | 54.32 | 54.84 | 52.65 | 53.66 | ||
| SAR | 53.39 | 52.77 | 53.22 | 51.52 | 51.96 | ||
| FOA | 45.27 | 43.97 | 44.68 | 42.42 | 42.81 | ||
| gauss | shot | imp | defoc | glass | motn | zoom | snow | frost | fog | brt | cont | elas | pix | jpeg | mean | ||
| ViT-S | |||||||||||||||||
| Tent | flat | 42.4 | 44.3 | 44.5 | 46.5 | 45.2 | 51.3 | 18.2 | 10.9 | 55.8 | 60.8 | 73.4 | 48.9 | 57.5 | 64.2 | 63.3 | 48.5 |
| late | 45.5 | 47.1 | 47.2 | 47.7 | 46.1 | 52.2 | 45.1 | 21.7 | 56.8 | 61.3 | 73.5 | 51.0 | 57.4 | 64.5 | 63.5 | 52.0 | |
| EATA | flat | 42.1 | 44.1 | 44.2 | 46.3 | 45.6 | 50.5 | 46.8 | 55.5 | 54.9 | 59.8 | 73.6 | 41.6 | 57.8 | 63.9 | 63.1 | 52.7 |
| late | 45.0 | 46.5 | 46.6 | 46.7 | 46.1 | 51.1 | 47.4 | 55.8 | 55.8 | 60.6 | 73.8 | 44.1 | 57.7 | 64.1 | 63.5 | 53.7 | |
| SAR | flat | 40.9 | 42.5 | 43.1 | 45.1 | 43.9 | 50.0 | 45.0 | 53.5 | 53.8 | 58.1 | 71.9 | 45.1 | 56.0 | 62.5 | 61.5 | 51.5 |
| GFLOPs / step | ms / step | |||||
| Method | full | late | cut | full | late | cut |
| ViT-S | ||||||
| Tent | 1216 | 890 | 156.9 | 124.5 | ||
| EATA | 1216 | 890 | 171.5 | 139.2 | ||
| SAR | 2431 | 1781 | 318.5 | 256.1 | ||
| FOA | 15555 | 11533 | 1927.4 | 1590.8 | ||
| Reducer | Family | noise | blur | weather | digital |
|---|---|---|---|---|---|
| ToMe | merging | ||||
| EViT | pruning | ||||
| ATS | sampling | ||||
| ATC | clustering | ||||
| PiToMe | merging |
| Condition | early | late | Spearman |
|---|---|---|---|
| clean | 0.60 | 0.25 | |
| gaussian noise (noise) | 0.92 | 0.20 | |
| defocus blur (blur) | 0.74 | 0.20 | |
| fog (weather) | 0.83 | 0.17 | |
| jpeg compression (digital) | 0.77 | 0.23 |
| own-path (flat) | oracle clean-decision | late schedule ( ) | |
|---|---|---|---|
| Schedule | Family | GFLOPs | mean | vs. flat |
|---|---|---|---|---|
| flat (const ) | – | – | ||
| power law ( ) | power law (convex) | |||
| exponential ( ) | exponential (convex) | |||
| two-level step | step (non-smooth) |
| gauss | shot | imp | defoc | glass | motn | zoom | snow | frost | fog | brt | cont | elas | pix | jpeg | mean | |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| uncompressed | 33.89 | 33.79 | 34.92 | 27.73 | 18.92 | 32.23 | 28.41 | 47.53 | 53.16 | 54.22 | 70.25 | 38.55 | 34.75 | 29.87 | 52.72 | 39.40 |
| flat | 31.62 | 31.70 | 32.40 | 26.53 | 18.42 | 31.11 | 27.31 | 46.37 | 51.10 | 51.43 | 69.57 | 36.59 | 34.58 | 28.88 | 52.19 | 37.99 |
| power law ( ) | 33.79 | 33.58 | 34.85 | 27.57 | 18.72 | 32.24 | 28.75 | 47.62 | 52.45 | 52.97 | 70.18 | 37.31 | 34.79 | 29.74 | 52.82 | 39.16 |
| exponential | 33.43 | 33.28 | 34.53 | 27.42 | 18.72 | 32.11 | 28.36 | 47.49 | 52.42 | 52.85 | 70.05 | 37.43 | 35.07 | 29.63 | 52.77 | 39.04 |
| two-level step | 33.58 | 33.45 | 34.59 | 27.61 | 18.73 | 32.20 | 28.67 | 47.71 | 52.53 | 52.94 | 70.09 | 37.87 | 35.13 | 29.72 | 52.79 | 39.17 |
| layer | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | total | last 3 |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| DeiT-S, searched | 0 | 2 | 6 | 12 | 11 | 12 | 10 | 12 | 18 | 5 | 31 | 72 | 191 | 57% |
| DeiT-B, searched | 0 | 1 | 2 | 2 | 7 | 8 | 11 | 11 | 17 | 49 | 29 | 55 | 192 | 69% |
| power law, (both) | 0 | 1 | 3 | 5 | 7 | 11 | 14 | 19 | 24 | 30 | 36 | 42 | 192 | 56% |
| layer | clean | gaussian |
|---|---|---|
| 1 | 16.4 | 25.8 |
| 2 | 12.7 | 18.4 |
| 3 | 8.3 | 10.9 |
| 4 | 6.3 | 5.9 |
| 5 | 4.1 | 5.7 |
| 6 | 3.4 | 4.7 |
| corruption | flat | late | full | excl-1 | excl-2 |
|---|---|---|---|---|---|
| gaussian noise | 31.62 | 33.79 | +2.17 | +2.16 | +1.89 |
| shot noise | 31.70 | 33.58 | +1.88 | +1.92 | +1.70 |
| impulse noise | 32.40 | 34.85 | +2.45 | +2.33 | +2.21 |
| defocus blur | 26.53 | 27.57 | +1.04 | +1.11 | +1.13 |
| glass blur | 18.42 | 18.72 | +0.30 | +0.27 | +0.29 |
| motion blur | 31.11 | 32.24 | +1.13 | +1.16 | +0.98 |
| MAE ViT-B | MAE ViT-L | |||||
| flat | late | flat | late | |||
| Budget | ||||||
| GFLOPs | 12.67 ( ) | 12.62 ( ) | – | 44.98 ( ) | 44.90 ( ) | – |
| tokens removed | 96 | 169 | – | 96 | 180 | – |
| final tokens | 100 | 27 | – | 100 | 16 | – |
| Top-1 accuracy (%) by corruption | ||||||
| backbone | pair ( , ) | net | recovery | new correctness |
|---|---|---|---|---|
| DeiT-S | , | |||
| DeiT-S | , | |||
| DeiT-S | , | |||
| DeiT-S | , | |||
| DeiT-S | , | |||
| DeiT-S | , |
| flat (GFLOPs) | |||||||
|---|---|---|---|---|---|---|---|
| 2 | 2 (3.99G) | 3 ( ) | 3 ( ) | 4 ( ) | 4 ( ) | 5 ( ) | 5 ( ) |
| 4 | 4 (3.72G) | 5 ( ) | 6 ( ) | 7 ( ) | 8 ( ) | 9 ( ) | 10 ( ) |
| 6 | 6 (3.46G) | 8 ( ) | 9 ( ) | 10 ( ) | 13 ( ) | 16 ( ) | 20 ( ) |
| 8 | 8 (3.20G) | 10 ( ) | 12 ( ) | 14 ( ) | 16 ( ) | 18 ( ) | 21 ( ) |
| 10 | 10 (2.94G) | 13 ( ) | 15 ( ) | 17 ( ) | 21 ( ) | 26 ( ) | 33 ( ) |
| 12 | 12 (2.67G) | 15 ( ) | 18 ( ) | 22 ( ) | 27 ( ) | 34 ( ) | 44 ( ) |
| schedule | gaussian top-1 | ||
|---|---|---|---|
| 6 | const | 32.35 | – |
| 6 | late | 33.84 | 1.49 |
| 8 | const | 31.62 | – |
| 8 | late | 33.12 | 1.50 |
| 8 | late | 33.49 | 1.86 |
| 8 | late | 33.79 | 2.17 |
| backbone | schedule | GFLOPs | ms / batch | img / s |
|---|---|---|---|---|
| (DeiT-S, ViT-S) | uncompressed | 4.25 | 67.1 | 954 |
| flat (const ) | 3.20 | 54.6 | 1172 | |
| late ( ) | 3.15 | 54.6 | 1172 | |
| (DeiT-B, ViT-B) | uncompressed | 16.86 | 212.0 | 302 |
| flat (const ) | 12.67 | 167.0 | 383 | |
| late ( ) | 12.46 | 167.5 | 382 |