ToaSt: Token Channel Selection and Structured Pruning for Efficient ViT
Authors: Hyunchan Moon, Cheonjun Park, Steven L. Waslander
Organizations: LG Electronics, Seoul, Republic of Korea · Hankuk University of Foreign Studies, Yongin, Republic of Korea · University of Toronto, Toronto, Canada
Vision Transformers (ViTs) have achieved remarkable success across various vision tasks, yet their deployment is often hindered by prohibitive computational costs. While structured weight pruning and token compression have emerged as promising solutions, they suffer from prolonged retraining and inter-layer dependencies that complicate optimization, respectively. We propose ToaSt, a decoupled framework applying specialized strategies to distinct ViT components. We apply coupled head-wise structured pruning to Multi-Head Self-Attention modules, leveraging attention operation characteristics to enhance robustness. For Feed-Forward Networks (over 60% of FLOPs), we introduce Token Channel Selection (TCS), a training-free method that filters redundant noise channels at inference time. Extensive evaluations across nine diverse models, including DeiT, ViT-MAE, and Swin Transformer, demonstrate that ToaSt achieves superior trade-offs between accuracy and efficiency, consistently outperforming existing baselines. On ViT-MAE-Huge, ToaSt achieves 88.52% accuracy (+1.64%p) with 39.4% FLOPs reduction. ToaSt also transfers effectively to diverse downstream tasks (COCO detection, ADE20K segmentation, CIFAR-100 classification), achieving 52.2 versus 51.9 mAP on COCO. Code: github.com/SHANNonLab-HUFS/ToaSt
Figures & tables
Figure 1: ToaSt compression methodology. (a) Standard ViT block architecture. (b) Token compression propagates compression effects across layers due to inter-layer dependencies. (c) ToaSt independently compresses each layer through coupled weight pruning (MHSA) and token channel selection (FFN), preventing cross-layer propagation while reducing dk and D dimensions.
Figure 2: Overview of the ToaSt framework for layer-independent compression. (a) Structured Coupled MHSA Weight Pruning: Pruning indices are synchronized across coupled groups (Q-K and V-Proj) to reduce the internal head dimension dk while preserving the attention mechanism’s functional integrity. (b) Token Channel Selection for FFN: Redundant channels in the intermediate FFN layer are identified and eliminated based on feature importance analysis (Eq. 6 ), maintaining the global embedding dimension D at the block interface.
Figure 3: Impact of coupled index synchronization on accuracy. Under identical pruning ratios and the geometric median metric, Non-Align prunes Q/K/V/O with independent indices per matrix, while Ours (Align) enforces shared indices within Q-K and V-O pairs, significantly mitigating accuracy drop at high pruning ratios.
Figure 4: Layer-wise redundancy analysis of Swin-Base FFN. (Left) Sparsity increases in deeper stages, indicating many “dead neurons.” (Center) Linear Reconstruction R2 remains near 1.0, proving that feature channels are highly dependent. (Right) Effective Rank Ratio: the fraction of singular values needed to capture 90% of the total variance ( minkk/C s.t. ∑i=1kσi2≥0.9⋅∑j=1Cσj2 ). A lower ratio indicates higher redundancy. The ratio collapses in later stages, confirming that the 4D expansion contains massive redundancy.
Model
Method
Top-1 (%)
Top-5 (%)
GFLOPs
FLOPs ↓ (%)
Throughput (img/s)
Speedup
DeiT-Tiny
Baseline
72.20
91.10
1.3
–
2090.9
1.00 ×
ToMe ( Bolya et al., 2022 )
71.25
90.74
0.7
46.2
2484.6
1.19 ×
DiffRate ( Chen et al., 2023 )
71.78
90.87
0.9
30.8
2422.5
1.16 ×
ToaSt (Ours)
74.25
92.65
0.76
41.5
4249.7
2.03 ×
DeiT-Small
Baseline
79.82
94.95
4.6
–
2313.2
1.00 ×
ToMe ( Bolya et al., 2022 )
79.35
94.65
2.7
41.3
2737.1
1.18 ×
Table 1: Comparison with state-of-the-art methods on ImageNet-1K. Throughput measured on H100 GPU with batch size 128. Best results in bold . FFN TCS is applied training-free at inference time. FLOPs ↓ denotes percentage reduction from baseline. Speedup = compressed model throughput / baseline throughput.
Model
Method
GFLOPs
Top-1 (%)
Δ Acc
Throughput(imgs/s)
DeiT Tiny
Baseline
1.3
72.20
–
2090.95
DiffRate
0.8
71.67
- 0.53
2255.7
ToaSt (Ours)
0.8
74.30
+2.1
4249.68
DeiT Small
Baseline
4.6
79.82
–
2313.2
ToMe
2.7
79.35
- 0.47
2737.05
DiffRate
2.7
79.38
- 0.44
3259.85
Table 2: Comparison at similar FLOPs budgets. ToaSt achieves +1–4%p higher accuracy than token compression methods.
Figure 5: Layer-wise FFN TCS Sensitivity Analysis. Sensitivity analysis of FC1, FC2, and combined pruning across DeiT-Small layers at various ratios (10%-90%). (a) FC1 shows high sensitivity in early layers but robustness in later layers (L9-11), with L11 improving accuracy up to 80% pruning. (b) FC2 exhibits lower sensitivity, enabling aggressive pruning (50-90%) in later layers. (c) Combined pruning validates asymmetric layer-adaptive ratios exploiting distinct redundancy patterns between FC1 and FC2.
Backbone
Method
GFLOPs
Box
Mask
mAP
mAP
Swin Small
Baseline
194
51.9
45.0
ToaSt
144 (25.8% ↓ )
52.2
45.0
Swin Base
Baseline
343
51.9
45.0
ToaSt
251 (26.8% ↓ )
52.2
45.1
ToaSt
246 (28.3% ↓ )
51.8
44.9
Table 3: Object detection results on COCO val2017 with Cascade Mask R-CNN. Our compressed backbones maintain or improve detection performance. Detailed configuration is in Appendix Table 26 .
Setting
GFLOPs
Quality
Δ
ADE20K (Baseline)
87.0
48.13 mIoU
–
ADE20K (4-layer TCS)
82.8
48.13 mIoU
0.00
ADE20K (6-layer TCS)
80.9
47.98 mIoU
− 0.15
CIFAR-100 (Baseline)
15.4
89.37%
–
CIFAR-100 (ToaSt)
10.1 (34.4% ↓ )
89.44%
+0.07
Table 4: Cross-task generalization on ADE20K and CIFAR-100 with Swin-Base. TCS applied training-free on ADE20K; CIFAR-100 uses full ToaSt with fine-tuning.
Model
Method
Top-1 (%)
GFLOPs
Throughput
Speedup
DeiT Small
Baseline
79.82
4.6
2313.20
1.00 ×
MHSA Only
76.12
3.1
3684.59
1.59 ×
ToaSt
83.40
2.5
4783.32
2.07 ×
ViT-MAE Large
Baseline
85.96
61.6
349.03
1.00 ×
MHSA Only
81.81
42.8
492.07
1.41 ×
ToaSt
88.94
38.5
527.01
1.51 ×
Table 5: Ablation study on representative architectures. ToaSt consistently outperforms MHSA-only pruning across supervised (DeiT), self-supervised (ViT-MAE), and hierarchical (Swin) transformers. Full results in Appendix F .
Appendix figures & tables22 assets
Supplementary material from the paper’s appendix.
Appendix
Operation
Complexity
FLOPs (%)
Primary Factor
Q,K,V Projections
O(3ND2)
∼ 23.0%
Hidden dim D
Attention Score (QK⊤)
O(N2D)
∼ 3.9%
Sequence length N
Attention Output (AV)
O(N2D)
∼ 3.9%
Sequence length N
Output Projection
O(ND2)
∼ 7.7%
Hidden dim D
FC1
O(4ND2)
∼ 30.7%
Hidden dim D
FC2
O(4ND2)
∼ 30.7%
Hidden dim D
Appendix
Table 6: Complexity breakdown of ViT component operations (DeiT-Small: N=197 , D=384 ).
Figure 6: Scaling trends of standard ViT models. (Left) Total parameter count. (Right) Model width ( D ) vs. depth ( L ). Width increases dramatically while depth saturates in larger models.
Model
L1 -norm
L2 -norm
Geometric Median
DeiT-Small (90%)
76.12
76.17
76.24
Swin-Small (90%)
80.54
81.14
81.20
Appendix
Table 7: Comparison of importance metrics across architectures.
Block
Kept SNR
Pruned SNR
Ratio
17
0.198
0.065
3.04 ×
18
0.229
0.055
4.20 ×
19
0.282
0.051
5.50 ×
20
0.409
0.077
5.29 ×
21
0.562
0.153
3.67 ×
Appendix
Table 8: Channel SNR analysis on Swin-Base. Kept vs. pruned channel SNR in FC2-pruned blocks.
Operation
Original FLOPs
Pruned FLOPs
Reduction
Q,K Projection
2ND⋅dkH
2ND⋅0.1dkH
90%
QK⊤ (Attn)
N2dkH
N2⋅0.1dkH
90%
V Projection
ND⋅dvH
ND⋅0.1dvH
90%
AV (Weighted Sum)
N2dvH
N2⋅0.1dvH
90%
Output Projection
NdvH⋅D
N⋅0.1dvH⋅D
90%
Appendix
Table 9: Computational reduction in MHSA with 90% per-head pruning. The reduction applies uniformly across all linear and attention operations.
Model
Epochs
Wall-clock Time
DeiT-Small
290
∼ 1 day
ViT-MAE-Large
139
∼ 31 hours
ViT-MAE-Huge
15
∼ 15 hours
Appendix
Table 10: Fine-tuning cost for MHSA pruning. Measured on 4 × H100 GPUs. TCS requires no training.
Table 12: Layer-wise FFN pruning ratios for ViT-MAE-Huge (32 layers). MHSA: 90% channel reduction (except layer 1). Only FC2 is pruned in the final 8 layers.
Stage 0
Stage 1
Stage 2
Stage 3
Layer
1
2
1
2
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
1
2
FC1
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0
0.3
FC2
0
0
0
0
0
0
0
0
0
0
0
0
0
0.9
0.9
0.9
0.9
0.9
0.9
0.9
0.9
0.9
0.9
0.9
Appendix
Table 13: Layer-wise FFN pruning ratios for Swin-Base. Stage configuration: [2, 2, 18, 2] blocks. MHSA: 90% channel reduction (except first layer of each stage).
Model
Method
Top-1 (%)
GFLOPs
Throughput
Speedup
DeiT-Tiny
Baseline
72.20
1.3
2090.95
1.00 ×
MHSA Only
69.81
0.9
2947.26
1.41 ×
ToaSt
74.25
0.76
4249.68
2.03 ×
DeiT-Small
Baseline
79.82
4.6
2313.20
1.00 ×
MHSA Only
76.12
3.1
3684.59
1.59 ×
ToaSt
83.40
2.5
4783.32
2.07 ×
Appendix
Table 14: Complete ablation study across all evaluated architectures.
Importance Source
Top-1 (%)
Top-5 (%)
Coupled (Ours)
71.43
90.75
Q only → broadcast to K
71.28
90.53
K only → broadcast to Q
71.32
90.71
Proj only → broadcast to V
71.10
90.38
Appendix
Table 15: Ablation on importance source for SCWP. DeiT-Small with 90% head-dimension pruning, 50 epochs fine-tuning.
Model
With Attn (%)
Without Attn (%)
Δ
DeiT
Tiny
74.25
66.22
+8.03
Small
83.40
75.83
+7.57
Base
84.82
82.60
+2.22
ViT-MAE
Large
89.21
89.51
− 0.30
Huge
88.52
90.52
− 2.00
Swin
Tiny
81.76
81.85
− 0.09
Appendix
Table 16: Effect of attention-guided weighting ( λcls in Eq. ( 6 )) on Top-1 accuracy across all architectures.
λcls
λpatch
DeiT-T
DeiT-S
DeiT-B
MAE-B
MAE-H
0.0
1.0
65.54
75.39
83.26
85.40
90.57
1.0
1.0
74.58
83.37
84.75
84.94
89.50
2.0
1.0
74.25
83.40
84.82
84.13
88.52
3.0
1.0
73.51
83.13
85.03
83.01
87.33
5.0
1.0
72.31
82.84
84.84
80.98
85.12
1.0
2.0
73.85
82.40
84.40
85.45
89.90
Appendix
Table 17: Sensitivity of λcls and λpatch in Eq. ( 6 ). Top-1 accuracy across CLS-token architectures. Best per model in bold .
Model
Sampling
Top-1
GFLOPs
Throughput
Model
Sampling
Top-1
GFLOPs
Throughput
DeiT-Tiny
No Sampling
75.21
0.76
4100.7
ViT-MAE-Huge
No Sampling
88.64
101.4
201.3
20% (default)
74.25
0.76
4249.7
2% (default)
88.52
101.4
206.2
1%
72.65
0.76
4228.5
1%
88.06
101.4
206.0
1 Sample
72.65
0.76
4251.5
1 Sample
87.71
101.4
206.3
DeiT-Small
No Sampling
84.67
2.5
4564.2
Swin-Tiny
No Sampling
82.12
2.5
3073.6
2% (default)
83.40
2.5
4783.3
20% (default)
81.76
2.5
3096.2
Appendix
Table 18: Sampling strategy ablation across nine ViT architectures. Top-1 accuracy (%) and throughput (img/s) on ImageNet-1K (H100 GPU, batch size 128). Default configurations used in the main paper are in bold .
Pattern
DeiT-S
MAE-L
MAE-H
Swin-B
fc2_heavy (Ours)
82.84
88.24
88.79
85.85
stepwise
81.01
86.16
90.06
84.29
uniform
56.54
77.56
77.04
74.42
inverse
18.06
74.57
68.24
63.06
Appendix
Table 19: Pruning ratio pattern ablation. Top-1 accuracy (%) on ImageNet-1K across four representative architectures under matched FLOPs budgets. Best per model in bold , second-best underlined .
Method
Training
Top-1 (%)
Δ vs. TCS
SVD (truncated)
Free
69–73
− 10 to − 15
LoRA (rank-matched, 20 ep.)
20 epochs
76.13
− 7.27
TCS (Ours)
Free
83.40
—
Appendix
Table 20: Comparison with low-rank methods on DeiT-Small at matched FLOPs. TCS substantially outperforms both training-free SVD and learned LoRA, validating that FFN redundancy is best exploited through token-adaptive channel selection rather than fixed or learned weight decomposition.
Model
Method
Top-1 (%)
Top-5 (%)
GFLOPs
Throughput (img/s)
DeiT-Tiny
ToMe
71.25
90.74
0.7G
2484.6
ToaSt
74.25
92.65
0.76G
4249.7
ToaSt + ToMe (r=7)
72.30
–
0.64G
3636.0
DeiT-Small
ToMe
79.35
94.65
2.7G
2737.1
ToaSt
83.40
96.97
2.5G
4783.3
ToaSt + ToMe (r=13)
79.98
95.64
1.83G
6138.1
Appendix
Table 21: Combination of ToaSt with ToMe token compression. ToaSt and ToMe target orthogonal dimensions (channel D vs. sequence N ), enabling direct combination with additive throughput gains.
Model
Method
Top-1 (%)
Top-5 (%)
GFLOPs
FLOPs ↓ (%)
Latency (ms)
Throughput (img/s)
Speedup
DeiT-Tiny
Baseline
72.20
91.10
1.3
–
61.22
2090.95
–
ToMe
71.25
90.74
0.7
46.2
51.54
2484.59
1.19 ×
DiffRate
71.78
90.87
0.9
30.8
52.84
2422.54
1.16 ×
DiffRate
71.67
90.78
0.8
38.5
56.75
2255.70
1.08 ×
ToaSt
69.93
89.67
0.9
30.8
36.89
3469.48
1.66 ×
ToaSt
74.30
92.26
0.8
38.5
29.32
4364.95
2.09 ×
Appendix
Table 22: Detailed results for DeiT models on ImageNet-1K (H100 GPU, batch size 128).
Model
Method
Top-1 (%)
Top-5 (%)
GFLOPs
FLOPs ↓ (%)
Latency (ms)
Throughput (img/s)
Speedup
ViT-MAE-Base
Baseline
83.75
96.54
17.6
–
112.27
1140.16
–
ToMe (r=16)
81.09
95.61
8.8
50.0
62.00
2064.49
1.81 ×
ToMe (r=13)
81.87
96.02
10.4
40.9
71.78
1783.33
1.56 ×
DiffRate
82.90
96.14
11.5
34.7
82.43
1552.75
1.36 ×
ToaSt
83.44
95.49
11.5
34.7
78.80
1624.28
1.42 ×
ToaSt
84.13
96.39
11.0
37.5
75.62
1692.61
1.48 ×
Appendix
Table 23: Detailed results for ViT-MAE models on ImageNet-1K (H100 GPU, batch size 128).
Model
Method
Top-1 (%)
Top-5 (%)
GFLOPs
FLOPs ↓ (%)
Latency (ms)
Throughput (img/s)
Speedup
Swin-Tiny
Baseline
81.20
95.50
4.5
–
49.02
2610.94
–
ToaSt
81.76
95.70
3.1
31.3
45.58
2705.76
1.04 ×
Swin-Small
Baseline
83.20
96.20
8.7
–
83.42
1534.41
–
STViT-R
82.60
96.07
5.8
33.3
77.74
1646.58
1.07 ×
ToaSt
84.65
96.80
5.4
38.2
67.03
1909.46
1.24 ×
Swin-Base
Baseline
83.50
96.50
15.4
–
116.35
1100.10
–
Appendix
Table 24: Detailed results for Swin Transformer models on ImageNet-1K (H100 GPU, batch size 128).
Model
Method
Total (ms)
Attn (ms)
Selection (ms)
FFN (ms)
Overhead
DeiT Tiny
ToMe
125.81
50.90
38.93
35.98
30.9%
ToaSt
53.01
29.42
8.79
14.80
16.6%
DeiT Small
ToMe
37.90
13.80
10.56
13.54
27.9%
ToaSt
27.45
9.92
2.57
14.97
9.4%
DeiT Base
ToMe
82.80
34.12
5.31
43.37
6.4%
ToaSt
75.59
23.18
3.51
48.91
4.6%
Appendix
Table 25: Block-wise latency profiling on H100 GPU (batch size 128, fp32). ToMe’s bipartite matching overhead dominates in smaller models, while ToaSt maintains low overhead across all scales.
Backbone
MHSA Pruning
TCS Configuration
Swin-Small
60%
FC1: 30% last 1L, FC2: 90% last 3L
Swin-Base (Config 1)
60%
FC1: 30% last 1L, FC2: 90% last 4L
Swin-Base (Config 2)
60%
FC1: 30% last 1L, FC2: 90% last 6L
Appendix
Table 26: Detailed compression configurations for COCO object detection experiments.
Vision Transformers achieve strong performance but incur high computational costs from quadratic token interactions. Training-free token reduction methods improve efficiency, but their robustness remains underexplored. We propose Prune-and Merge ViT, a training-free framework for off-the-shelf models. It combines class attention-based pruning with similarity-based merging to reduce redundancy while preserving semantic information. To evaluate robustness we introduce two metrics. Relative Accuracy Drop (RAD) measures relative accuracy degradation under transformations, and Prediction Flip Rate (PFR) quantifies prediction consistency. On pretrained ViT and DeiT, Prune-and-Merge reduces computation competitively while maintaining robustness and better accuracy than existing token reduction methods. Our analysis suggests that hybrid token reduction does not introduce hallucination-like behavior and preserves prediction stability for common image transformations.
Vatsal, Soumil Malhotra, Priyanshu Chahal +1
Department of Mathematics, Indian Institute of Technology Roorkee
Vision Transformers (ViTs) are strong backbones for semantic segmentation, but their computational cost limits deployment. Recent token compression methods for efficient transformer-based segmentation reduce this cost by decreasing the number of tokens. However, existing evaluations primarily focus on low-to-moderate compression, leaving their behavior under aggressive compression and corrupted inputs unclear. Meanwhile, structural pruning provides an orthogonal route to efficiency by removing redundant components in the ViT architecture, but is rarely compared to token compression under a unified protocol. To bridge this gap, we benchmark representative token compression and structural pruning methods for ViT-based semantic segmentation under matched FLOPs on ADE20K and Cityscapes, together with their common-corruption variants ADE20K-C and Cityscapes-C. Our results reveal a consistent trend on both clean and corrupted inputs: token compression is highly effective at mild reductions but degrades sharply when compression becomes severe, consistent with substantial information loss from overly aggressive token reduction. In contrast, structural pruning exhibits a smoother degradation curve and is more stable at high compression. Motivated by these findings, we study a prune-then-merge pipeline that applies moderate token compression on top of a moderately pruned backbone. At comparable FLOPs, this combined strategy consistently achieves a better accuracy-robustness trade-off at high compression, offering a practical recipe for deployment-oriented ViT segmentation. Code is available at https://github.com/phatnguyencs/vit-seg-compression.
Tien-Phat Nguyen, Ngai-Man Cheung
Temasek Laboratories, Singapore University of Technology and Design, Singapore
Vision Transformers (ViTs) achieve strong performance but suffer from high computational costs due to quadratic self-attention complexity. Although token reduction techniques such as pruning and merging mitigate this, they typically overlook how representations evolve across network depth. We propose RAPID, a depth-aware token reduction framework that adapts reduction strategies to the layer-wise characteristics of token representations. The primary methodological contribution is a bifurcated strategy: in shallow-to-middle layers, RAPID employs a redundancy-similarity aware pruning metric to eliminate over-represented local patterns. As features transition to global semantic concepts in deeper layers, the framework shifts to an importance-similarity aware merging mechanism. This stage leverages classification (CLS) token attention weights to protect semantically critical tokens while fusing less important but similar neighbors. Empirical validation on ImageNet-1K using ViT and DeiT architectures demonstrates that RAPID establishes a superior accuracy-compression Pareto frontier compared to plug-and-play baselines such as ToMe and ToFu. RAPID is particularly robust in aggressive compression regimes, achieving up to 4.29% higher accuracy than ToMe at extreme reduction rates. Our framework provides a training-free template for optimizing vision models by aligning reduction strategies with hierarchical feature evolution.
Kyumin Choi, Ikbeom Jang
Hankuk University of Foreign Studies Seoul, South Korea