Joint MLP and Token Pruning for Personalizing Vision Transformers

Abstract

ViTs have achieved excellent performance in image recognition tasks, but their large parameter counts and high computational complexity limit their deployment on resource-constrained devices. Most existing ViT pruning methods adopt class-agnostic pruning strategies, which fail to distinguish the diverse structural requirements of different target classes. As a result, they are prone to removing critical features, leading to class-wise accuracy imbalance in practical deployment. To address this issue, this paper proposes a class-aware joint pruning framework for ViTs, which collaboratively compresses the model from two orthogonal dimensions: MLP neurons and visual tokens. Specifically, (1) based on first-order Taylor expansion, we quantify the contribution of each MLP neuron to the target classes and adaptively prune redundant neurons to achieve structured compression, followed by lightweight fine-tuning on the target class subset; (2) we propose a Class-Guided Token Selection (CGTS) method, which constructs class prototype vectors using a few support samples of the target classes and then dynamically selects patch tokens that are semantically highly relevant to the target classes during inference in a zero-shot manner, requiring no additional training or fine-tuning. The two modules complement each other, achieving dual compression from the parameter dimension and the inference data dimension. Experiments on CIFAR-100 and TinyImageNet datasets using DeiT-Tiny/Small models demonstrate that, compared with state-of-the-art pruning methods, our method reduces GMACs on target class subsets by up to 48%, improves inference speed by nearly 50%, and requires only 0.8 KB of additional storage overhead per subset, ultimately achieving a superior trade-off among accuracy, computational efficiency, and storage overhead.

1. Introduction

The rapid advancement of artificial intelligence (AI) has driven profound transformations in computer vision. Vision Transformer (ViT) [1], benefiting from its global self-attention mechanism and strong representation learning capability, has achieved superior performance over traditional convolutional neural networks (CNNs) in image classification [2,3,4], object detection [5,6,7,8], and semantic segmentation [9,10,11]. However, ViT models have enormous parameters and thus suffer from high computational complexity, which hinders their efficient lightweight deployment on resource-constrained devices [12].

Model pruning [13,14], as a widely used deep learning lightweight technique, removes redundant parameters and structures in neural networks to compress complex models and accelerate inference with minimal accuracy degradation. Most early research on model pruning focused on CNNs, which remove network redundancy by pruning local structures such as convolutional kernels and feature channels. With the widespread adoption of ViTs in computer vision, ViT pruning has gradually emerged as a new research hotspot. Unlike the pruning paradigm of CNNs, which only exploits local feature redundancy, Vision Transformers possess inherent global modeling capability, giving rise to dedicated pruning strategies and optimization schemes distinct from those designed for CNNs. Existing ViT pruning methods, including structured pruning [15,16,17] and token pruning [18,19,20], can effectively reduce model complexity and computational overhead.

To ensure their generality, models are typically trained with enormous data, allowing them to classify diverse classes. However, in practical deployment scenarios, models are usually used in a specific environment and only perceive a limited number of classes. Then, many components in a fully trained model that are not relevant to these predicted classes are redundant. Hence, in recent years, researchers have proposed class-aware pruning methods for CNNs. By accurately identifying the relevant components of different classes, these methods selectively retain key features of target classes and eliminate redundant structures irrelevant to these classes. They can compress models to improve performance while maintaining high accuracy. Representative methods such as CAPTOR [21], CRISP [22], and OCAP [23] have achieved promising results for CNNs.

However, these methods proposed for CNNs [24,25] are not directly applicable to ViTs, as ViTs introduce distinct pruning dimensions, such as tokens and multi-head attentions. In addition, the relevance of these new pruning dimensions to class-aware pruning remains largely unexplored. Most existing ViT pruning methods adopt a class-agnostic paradigm and follow a globally unified pruning criterion. They fail to distinguish the important features of target classes, likely removing the features related to target classes and hurting the final accuracy. As demonstrated in CNNs, class-aware pruning can reduce model complexity and memory footprint while retaining or improving prediction accuracy. Thus, it is worth exploring class-aware pruning in ViTs, as ViTs are gradually replacing CNNs to be the backbones of many CV tasks.

To the best of our knowledge, this paper is among the first to simultaneously investigate class-aware joint compression of ViTs from the two orthogonal dimensions (MLP neurons and visual tokens) and propose a corresponding joint pruning framework. Structured MLP pruning has been extensively studied for reducing model parameters, yet existing approaches lack adaptation to target classes. Token pruning has been considered a promising technique to reduce model complexity for transformer-based models, such as ViTs and large language models (LLMs). However, we find that many token pruning methods that use gates to control token propagation actually cannot improve a model’s performance due to the extra overhead caused by the control gate. In this paper, we propose a lightweight and zero-shot pruning method tailored for class-aware scenario.

The main contributions of this paper are as follows:

• We observe that MLP modules in ViTs are highly class-sensitive. Therefore, we propose a class-aware structured MLP pruning method based on first-order Taylor expansion, which computes the importance scores of neurons with respect to target classes and adaptively prunes redundant neurons, thereby achieving effective compression of the MLP modules, which requires lightweight fine-tuning on the target class subset to recover minor accuracy loss. Since the neurons in MLPs are pruned structuredly, dense patterns in MLPs are reserved and no special accelerator is required. This pruning stage compresses the feature space and provides a low-noise backbone network for subsequent token-level optimization.
• Building on MLP pruning, we propose zero-shot training-free Class-Guided Token Selection (CGTS), which filters class-relevant tokens via prototype vectors and adaptively determines the pruning layer and retention ratio without manual tuning. The two modules form a unified joint framework via progressive fusion, delivering orthogonal complementary gains across the parameter and inference dimensions.
• We conduct comprehensive experiments to evaluate our proposed method compared to some state-of-the-art methods. Experimental results demonstrate that, while preserving accuracy, our method reduces Giga Multiply-Accumulate Operations (GMACs) by up to 48%, improves inference speed by nearly 50%, and incurs only 0.8 KB of additional storage overhead per subset. Compared with existing pruning methods, our method achieves a better balance between computational cost, storage overhead, and inference latency.

The remainder of this paper is organized as follows. Section 2 covers related work. Section 3 details the CGTS and MLP pruning methods and their synergistic integration. Section 4 presents the experimental results, and Section 5 concludes this paper. Our code is available at https://github.com/lee95783/class_aware (accessed on 2 August 2026).

2. Related Work

Pruning ViT models is primarily achieved through token pruning and structured pruning. Token pruning reduces the quadratic complexity of self-attention by dynamically or statically discarding uninformative tokens. DynamicViT [18] inserts lightweight prediction modules between Transformer blocks to learn token retention probabilities end-to-end and progressively prunes redundant tokens, while training-free methods such as Zero-TPrune [26] leverage a weighted PageRank algorithm on attention graphs to assess token importance in a zero-shot setting. On the other hand, MLP modules bear the primary burden of feature transformation in ViTs and are thus the central target of structured pruning; approaches like AMP [27] and DGMR [28] achieve substantial parameter reduction by evaluating neuron importance. Moreover, SViTE [29] attempts to cover both token and structural dimensions within a unified framework by jointly imposing sparsity constraints on weights, attention heads, and tokens during training.

While these methods have distinct optimization directions, they all follow a task-agnostic pruning paradigm: the token retention decisions of DynamicViT [18] and Zero-TPrune [26] are independent of specific target classes, and the joint sparsity constraints of SViTE [29] are not customized for particular class subsets, making it impossible to achieve fine-grained targeted sparsification. Furthermore, except for a few training-free approaches, the vast majority still require expensive full-scale fine-tuning for each pruning configuration.

In contrast to task-agnostic strategies, class-aware pruning seeks to incorporate target-class priors into pruning decisions, enabling compressed models to focus more sharply on task-relevant discriminative features. Early efforts predominantly targeted CNN architectures: Hemmat et al. [30] employed a statistically driven “class-specific excitation rate” to prune neurons ineffective for the target classes; Qin et al. [21] clustered filters with similar response patterns to remove filter groups mismatched with the target classes; and Hu et al. [31] further formulated channel pruning as a submodular set function maximization problem, optimizing inter-class and intra-class scatter via the trace ratio criterion. Subsequent work extended class-aware pruning to practical deployment efficiency [22], 3D point cloud tasks [32], and on-device dynamic scenarios [23]. Although these studies have advanced the development of class-aware pruning, they all rest on the homogeneous structural assumptions of CNNs, and their filter-oriented importance metrics cannot be readily transferred to the heterogeneous components of ViTs, such as MLPs and tokens.

Recently, class-aware ideas have been preliminarily introduced into ViT lightweighting. X-Pruner [33] learns class-level explainability masks for each prunable unit, retaining structured components critical to the target classes through end-to-end training. Vulcan [34] reveals from a knowledge-disentanglement perspective that FFNs encode class-specific knowledge and proposes a “train-then-prune” paradigm to generate class-customized ViT models. However, these methods focus exclusively on structured component pruning without addressing token-level redundancy, and they still rely on full training or fine-tuning procedures, incurring substantial training overhead.

However, these methods still exhibit clear limitations: X-Pruner [33] adopts a soft gating mechanism that fails to achieve genuine structural compression and inference speedup; Vulcan [34] only performs structured pruning on feed-forward network (FFN) modules and completely neglects token-level redundancy. More critically, both methods rely on full end-to-end training, leading to high deployment costs for new class subsets.

To sum up, three key limitations remain unresolved in existing ViT pruning methods: (1) most methods follow a task-agnostic paradigm and cannot achieve targeted compression for specific class subsets; (2) existing class-aware ViT pruning only covers a single structural dimension, without joint optimization of MLP parameters and visual tokens; and (3) the vast majority of schemes depend on full-scale fine-tuning and lack lightweight zero-shot deployment solutions.

We propose a class-aware dual-dimension synergistic compression framework for Vision Transformers. Our work simultaneously covers both token and structural dimensions within a unified class-aware framework, addressing the methodological gap in existing studies that follow a task-agnostic paradigm or only optimize a single dimension. The entire framework avoids full-scale retraining, requiring only minimal fine-tuning on the target class subset for deployment adaptation, with training overhead substantially lower than that of existing full-scale fine-tuning schemes. The two compression dimensions operate in a progressively synergistic manner, effectively reducing edge-side inference latency while preserving accuracy on the target classes.

3. Proposed Method

This section presents the details of our proposed method. The overview of the method is shown in Figure 1. Overall, we adopt a backbone-first then fine-grained optimization strategy. After completing MLP backbone pruning and determining the fundamental network structure, we further conduct fine-grained pruning at the patch token level. The pruning pipeline starts with two offline tasks, as described in Section 3.1: class-aware pruning of MLP modules based on first-order Taylor expansion, followed by lightweight fine-tuning of the pruned model, thereby achieving the compression and reconstruction of the model backbone. Class prototypes are then constructed offline, and pruning layers are adaptively selected on the compressed model, followed by dynamic token screening in the online inference stage, as elaborated in Section 3.2. The progressive fusion and collaborative compression strategy of the two pruning schemes are explained in Section 3.3.

Figure 1. Overview of the MLP pruning and CGTS framework.

3.1. MLP Pruning

In Vision Transformer architectures, multi-head attention and MLP feed-forward networks exhibit distinct differences in pruning adaptability. Attention heads occupy a low proportion of total parameters and present class-agnostic properties, resulting in limited pruning gains. Experimental results further demonstrate that pruning attention heads can barely reduce inference overhead and is prone to inducing significant accuracy degradation when numerous classes are retained. In contrast, MLP modules have high parameter and computational ratios with prominent neuron redundancy. Meanwhile, they possess strong class-aware modeling ability and deliver remarkable pruning potential and benefits, which makes them our primary pruning target.

To quantify the contribution of each MLP neuron to the target class, we adopt the first-order Taylor expansion to approximate the impact of neuron removal on the loss function. Compared with weight-magnitude-based pruning criteria, the first-order Taylor approximation more accurately reflects each neuron’s contribution to model outputs. Meanwhile, unlike Hessian-based pruning that depends on second-order curvature, it requires only a small set of samples for low-overhead importance ranking, with better numerical stability and lower deployment costs on target class subsets. For the j-th neuron in the k-th MLP layer, its importance score is defined as the average absolute value of the product of activation and gradient computed on the target class subset ${\mathcal{T}}_c$:

$$S_k\left(j\right)={\displaystyle \frac{1}{|{\mathcal{T}}_c|}}\sum _{(\mathbf{x},y)\in {\mathcal{T}}_c}\left|f_k^j\left(\mathbf{x}\right)·{\displaystyle \frac{\partial \mathcal{L}(\mathbf{x},y)}{\partial f_k^j\left(\mathbf{x}\right)}}\right|$$

(1)

where $f_k^j\left(\mathit{x}\right)$ denotes the activation value of input $\mathit{x}$ at the corresponding neuron, and $\mathcal{L}$ represents the cross-entropy loss function. This score jointly captures the activation intensity of neurons and their sensitivity to the loss function. A higher score indicates a greater contribution of the neuron to target class recognition. The entire calculation only requires limited forward and backward propagation on the target class subset, without accessing the full training dataset.

After obtaining the importance scores of all neurons in each layer, we determine the number of reserved neurons according to a predefined pruning ratio $\alpha \in (0,1)$. Specifically, for the k-th layer, we sort the $N_k$ neurons in descending order based on $S_k\left(j\right)$ and retain the top $\lceil (1-\alpha )·N_k\rceil$ high-contribution neurons while discarding the remaining redundant ones. Accordingly, we reconstruct compact weight matrices with reduced dimensions. The original input weight matrix ${\mathit{W}}_1^k\in {\mathbb{R}}^{N_k\times D}$ is replaced by ${\mathit{W}}_1^{k\prime }\in {\mathbb{R}}^{\lceil (1-\alpha )N_k\rceil \times D}$, and the output weight matrix ${\mathit{W}}_2^k\in {\mathbb{R}}^{D\times N_k}$ is updated to ${\mathit{W}}_2^{k\prime }\in {\mathbb{R}}^{D\times \lceil (1-\alpha )N_k\rceil }$. We migrate the weights corresponding to reserved neurons from the original matrices to the newly constructed ones. This operation effectively reduces model parameters and FLOPs, leading to lower inference latency under standard inference frameworks.

Finally, we perform minimal fine-tuning on the pruned compact model using the target class subset to compensate for the slight accuracy degradation caused by structured pruning. We summarize the class-aware MLP pruning pipeline in Algorithm 1.

Algorithm 1 Class-aware MLP structured pruning.
Require: Pre-trained ViT $\mathcal{M}$ with MLP layers; Weights ${\mathbf{W}}_1^k,{\mathbf{W}}_2^k$ and neuron number $N_k$ for
layer k; Target class dataset ${\mathcal{T}}_c$; Pruning ratio $\alpha \in (0,1)$; Fine-tune epochs E Ensure: Compact pruned & fine-tuned model ${\mathcal{M}}_{\mathrm{pruned}}$
1:	for each MLP layer k of $\mathcal{M}$ do
2:	for each neuron j in layer k do
3:	$S_k\left(j\right)\leftarrow {\displaystyle \frac{1}{|{\mathcal{T}}_c|}}{\sum }_{(\mathbf{x},y)\in {\mathcal{T}}_c}\left|f_k^j\left(\mathbf{x}\right)·{\displaystyle \frac{\partial \mathcal{L}(\mathbf{x},y)}{\partial f_k^j\left(\mathbf{x}\right)}}\right|$
4:	end for
5:	Sort ${\left\{S_k\left(j\right)\right\}}_{j=1}^{N_k}$ in descending order
6:	$M\leftarrow \lceil (1-\alpha )·N_k\rceil$	▹ Number of reserved neurons
7:	$\mathrm{idx}\leftarrow$ Top-M indices of sorted scores
8:	${\mathbf{W}}_1^{k\prime }\leftarrow {\mathbf{W}}_1^k[\mathrm{idx},:]$	▹ Slice input weight rows
9:	${\mathbf{W}}_2^{k\prime }\leftarrow {\mathbf{W}}_2^k[:,\mathrm{idx}]$	▹ Slice output weight columns
10:	Update ${\mathbf{W}}_1^k\leftarrow {\mathbf{W}}_1^{k\prime },{\mathbf{W}}_2^k\leftarrow {\mathbf{W}}_2^{k\prime }$
11:	end for
12:	for epoch $=1$ to E do
13:	Fine-tune ${\mathcal{M}}_{\mathrm{pruned}}$ on ${\mathcal{T}}_c$ with cross-entropy loss
14:	end for
15:	return ${\mathcal{M}}_{\mathrm{pruned}}$

3.2. CGTS

CGTS consists of two components: layer selection and token selection. Token pruning is performed only once: once some patch tokens are removed at a certain layer, all subsequent layers benefit. The selection of the optimal pruning layer requires a trade-off between semantic richness and speedup gain, and varies with the subset size K. Therefore, we first select the optimal layer offline in Section 3.2.1 and then dynamically select tokens at this layer using the class prototype in Section 3.2.2. The complete workflow of CGTS is summarized in Algorithm 2.

Algorithm 2 CGTS.

Require: Target class set $\mathcal{C}$, candidate layers ${\mathcal{L}}_{\mathrm{cand}}$, minimum layer $L_{\min }$, support sample

number $N_c$, retention ratio $\rho$, input token sequence $\mathbf{t}$ Ensure: Pruned token sequence ${\mathbf{t}}_{\mathrm{pruned}}$   1: Offline: Adaptive Layer Selection   2: for $L\in {\mathcal{L}}_{\mathrm{cand}}$ do   3:       Compute class-wise prototype:   4:       ${\displaystyle {\mathbf{e}}_c^{\left(L\right)}=\frac{1}{N_c}{\displaystyle \sum _{i=1}^{N_c}}{\mathbf{h}}_{\mathrm{CLS}}^{(c,i)}}$   5:       Calculate layer discriminability $D\left(L\right)$ via Equation (2)   6: end for   7: $\tau \leftarrow \left({\min }_{L\in {\mathcal{L}}_{\mathrm{cand}}}D\left(L\right)+{\max }_{L\in {\mathcal{L}}_{\mathrm{cand}}}D\left(L\right)\right)/2$   8: $L^{*}\leftarrow \min \left\{L\mid L\ge L_{\min },D\left(L\right)>\tau \right\}$   9: ${\mathbf{e}}_{\mathrm{shared}}\leftarrow {\displaystyle \frac{1}{\left|\mathcal{C}\right|}}{\sum }_{c\in \mathcal{C}}{\mathbf{e}}_c^{\left(L^{*}\right)}$ 10: Inference: Class-Guided Token Selection 11: Extract patch tokens $\left\{{\mathbf{x}}_j\right\}$ at input of layer $L^{*}$, $N\leftarrow \left|\right\{{\mathbf{x}}_j\left\}\right|$ 12: for $j=1$ to N do 13:       $s_j\left(c\right)\leftarrow {\mathbf{x}}_j^{\top }{\mathbf{e}}_{\mathrm{shared}}$ 14: end for 15: Sort $\left\{s_j\left(c\right)\right\}$ in descending order 16: $\mathcal{I}\leftarrow$ indices of top $\lceil \rho N\rceil$ scores 17: Sort $\mathcal{I}$ in ascending order 18: ${\mathbf{t}}_{\mathrm{pruned}}\leftarrow \left[{\mathbf{h}}_{\mathrm{CLS}};{\mathbf{x}}_j\mid j\in \mathcal{I}\right]$ 19: return ${\mathbf{t}}_{\mathrm{pruned}}$

3.2.1. Layer Selection

It is crucial to balance accuracy and inference efficiency when selecting the pruning layer L for both single-class and multi-class tasks. If L is set too shallow, features are not fully aggregated, class prototypes lack semantic discrimination, and similarity scores are nearly random, easily causing erroneous removal of critical tokens. If L is set too deep, most computational overhead has already been incurred, resulting in limited compression gains.

To achieve adaptive optimal pruning layer selection, we propose a prototype discriminability-based layer selection strategy, which only relies on support samples and does not require test set labels. The core goal is to preferentially select layers with strong capability to distinguish effective semantic tokens from irrelevant background tokens, so as to accurately remove redundant tokens and preserve critical semantic information. To quantitatively measure the semantic discrimination significance of different layers, we use the score difference between high-quality and low-quality tokens to represent the basic discrimination amplitude and adopt standard deviation normalization to eliminate feature scale differences across layers, enabling fair comparison of discriminability and constructing a standardized prototype discriminability score. Specifically, we extract patch token features of support samples at the input of candidate layer L and calculate their similarity scores with ${\mathbf{e}}_c$. For each image, all scores are sorted in descending order, where ${\mu }_{\mathrm{top}}$ denotes the mean of the top 50% scores, ${\mu }_{\mathrm{bot}}$ denotes the mean of the bottom 50% scores, and $\sigma$ denotes the standard deviation of all scores. The prototype discriminability score of layer L for class c is defined as

$$D_c\left(L\right)={\displaystyle \frac{{\mu }_{\mathrm{top}}-{\mu }_{\mathrm{bot}}}{\sigma }}$$

(2)

A larger $D_c\left(L\right)$ indicates that the layer can better distinguish class-relevant tokens from irrelevant ones. For multi-class tasks, the comprehensive layer discriminability score is calculated as $D\left(L\right)={\displaystyle \frac{1}{\left|\mathcal{C}\right|}}{\sum }_{c\in \mathcal{C}}{D}_c\left(L\right)$, representing the average discrimination capability for all target classes. After obtaining $D\left(L\right)$ for all candidate layers, the optimal pruning layer $L^{*}$ is determined as follows:

$$L^{*}=\min \{L\mid L\ge L_{\min }\mathrm{and}D\left(L\right)>\tau \}$$

(3)

where $L_{\min }$ is the minimum allowable pruning layer, and $\tau$ is set to the midpoint of the value range of all candidate $D\left(L\right)$. This rule skips unreliable shallow layers and selects the earliest layer with discriminability above the average level, maximizing computational savings under controllable accuracy loss.

3.2.2. Token Selection

ViT divides an input image into a fixed number of patch tokens, together with a class (CLS) token for classification. In practical inference, only a small number of patch tokens carry discriminative semantic information to support the final prediction, while most background or low-information tokens contribute little to the classification decision. This motivates us to dynamically prune redundant tokens to reduce computational cost while maintaining model accuracy. The core of effective token pruning lies in exploring a training-free, lightweight and efficient discriminative criterion to accurately identify semantically critical tokens corresponding to the target classes.

Our token selection scheme is built upon two core assumptions. First, within a single category, the key visual patch tokens representing target objects share highly similar semantics. Second, for different classes within the same deployment subset, their foreground discriminative features follow common patterns and approximate distributions, providing a theoretical basis for fusion into a unified representation.

In ViT, the CLS token continuously aggregates global semantic information through multi-layer self-attention. After iterative feature reinforcement layer by layer within Transformer, the CLS features of images from the same class form stable cluster distributions in the high-dimensional feature space. Accordingly, we take the semantic center of these feature clusters as the class-specific representative vector, which is used to accurately measure the semantic similarity between any image patch token and the target class. Under the constraint of limited support samples, the mean of feature samples serves as the optimal unbiased center estimator that minimizes the intra-class mean square error. Therefore, for each target class c, we take the mean CLS feature of a small number of support samples at pruning layer L as the prototype vector ${\mathbf{e}}_c\in {\mathbb{R}}^D$ for this class:

$${\mathbf{e}}_c^{\left(L\right)}={\displaystyle \frac{1}{N_c}}{\displaystyle \sum _{i=1}^{N_c}}{\mathbf{h}}_{\mathrm{CLS}}^{(c,i)}$$

(4)

where ${\mathbf{h}}_{\mathrm{CLS}}^{(c,i)}\in {\mathbb{R}}^D$ denotes the CLS feature of the i-th sample, and $N_c$ is the number of support samples for class c. The construction of class prototypes only requires forward propagation without any parameter update.

During the inference stage, given an input image and the target class c, ViT first divides the image into N patch embeddings with D dimensions, which are concatenated with the CLS token and fed into the Transformer layers. When the feature flow reaches the preset pruning layer L, we extract the feature representation of all image patch tokens at the input of this layer, denoted as ${\mathbf{x}}_j\in {\mathbb{R}}^D,(j=1,\dots ,N)$. To measure the correlation between each token and the target class c, we compute the inner product between ${\mathbf{x}}_j$ and the prototype vector ${\mathbf{e}}_c$ as the similarity score:

$$s_j\left(c\right)={\mathbf{x}}_j^{\top }{\mathbf{e}}_c$$

(5)

This score reflects the directional consistency between token features and the target class prototype in the feature space; a higher score indicates that the semantics of the token are more consistent with class c. All N tokens are sorted in descending order according to $s_j\left(c\right)$, and only the top-k tokens are retained, where $k=\lceil \rho ·N\rceil$ and $\rho \in (0,1]$ is the preset retention ratio. The retained k tokens are concatenated with the CLS token and further fed into layer L and subsequent layers for computation, while the remaining $N-k$ tokens are directly discarded. The CLS token is not involved in similarity calculation and is always retained unconditionally. Since the pruning operation only involves feature extraction, inner product computation, and sorting without any learnable parameters, CGTS can be directly applied to any pre-trained ViT model without additional fine-tuning.

For multi-class tasks ($K>1$), we aggregate the CLS feature means of all K classes of support samples at the target layer to construct a global shared prototype vector. During inference, token screening is completed in a single pass relying solely on this single prototype, eliminating repeated forward passes and multi-branch result fusion, thus achieving lightweight and efficient pruning inference. Although different classes within the subset have visual differences, the shared prototype preserves the foreground discriminative semantics common to all target classes and can still effectively capture the overall discriminative semantic region of the target subset. In constrained inference scenarios where the input belongs to one of the K classes, the dot-product operation automatically highlights tokens relevant to the current class. Therefore, the global averaging strategy does not impair token screening effectiveness.

3.3. Integration of MLP Pruning and CGTS

MLP pruning and CGTS adopt a progressive fusion paradigm: we first conduct class-aware structured MLP pruning on target subsets to obtain a low-noise compact backbone and then construct class prototypes based on purified features for CGTS token filtering. The two pruning modules operate on disjoint dimensions, and either standalone strategy bears obvious drawbacks. Sole MLP pruning fails to cut the quadratic computational cost introduced by patch tokens, while standalone CGTS suffers degraded prototype discriminability caused by redundant neurons in the unrefined network. This serial architecture remedies the flaws of both single pruning branches simultaneously, achieving a better balance between classification accuracy and compression efficiency and reducing overall inference latency collaboratively.

4. Evaluation

4.1. Experimental Setup

Data and models:

This paper targets practical edge deployment scenarios that require recognition of a fixed small set of categories. With 100 and 200 classes, respectively, CIFAR-100 and TinyImageNet are well aligned with the category scale of typical real-world edge applications. They also enable us to build multiple nested class subsets to systematically evaluate the performance of our class-aware pruning under varying numbers of target categories.

We use DeiT-Tiny and DeiT-small [2] pre-trained on CIFAR-100 and TinyImageNet as our backbones. The all-classes accuracy of DeiT-Tiny and DeiT-small is 83.79% and 86.18%. We evaluate all methods with four class subsets, $K\in \{5,10,20,50\}$, randomly drawn from the full datasets. All subset results use restricted evaluation: argmax over the K target logits, which aligns with our core deployment premise that all input samples belong to one of the predefined K categories.

It is worth noting that this closed-set restricted evaluation protocol differs from the full-category standard evaluation adopted by general ViT benchmarks. This constrained setup highly matches our targeted edge deployment scenario with fixed target classes and no out-of-distribution inputs in real operation, so we adopt this evaluation rule to precisely fit the inference requirements of the dedicated device studied in this paper.

Implementation details:

During the fine-tuning stage after MLP structured pruning, we adopt the AdamW optimizer with an initial learning rate of $1\times {10}^{-4}$ and weight decay of $1\times {10}^{-5}$. The batch size is set to 32, and we apply a cosine annealing scheduler without warm-up for a total of five training epochs. For CGTS prototype construction, we randomly draw 10 training images per target class as support samples, and test data is never involved. The index lists of support samples for all subsets are pre-generated offline and loaded directly in inference, which eliminates redundant real-time sampling operations.

Baselines:

We select two types of baselines for comprehensive evaluation: mainstream class-agnostic ViT pruning approaches and the representative class-aware ViT work X-Pruner. We compare against five methods: (1) DynamicViT [18], a learned token predictor inserted at layer 6; (2) SViTE [29], a static class-aware token mask computed at layer 6; (3) X-Pruner [33], class-conditional soft head gating trained with augmented Lagrangian multipliers (ALMs); (4) X-Pruner (hard), where trained X-Pruner gates are thresholded at 0.5 and heads are structurally removed via weight compaction; and (5) Zero-TPrune [26], a zero-shot token pruning method that scores tokens via Weighted PageRank (WPR) on the self-attention graph, requiring no training or calibration data. All fine-tuned baselines use five epochs on K-class data. Our method additionally applies class-aware Taylor MLP pruning before fine-tuning.

Metrics: We report top-1 accuracy, GMACs per image, backbone size (MB), per-deployment storage overhead, and GPU latency at batch size 64 (NVIDIA Tesla P100, Santa Clara, CA, USA). We run all experiments multiple times with different random seeds for model initialization and random subset sampling and report averaged metrics to guarantee statistical reliability, following standard evaluation protocols for Vision Transformer pruning. Specifically, we used five random class subsets per K and three independent runs per subset.

4.2. Main Results

Table 1 reports the experimental results of DeiT-Tiny evaluated on the CIFAR-100 dataset, covering a comprehensive comparison of classification accuracy and inference efficiency across all methods.

Table 1. Comprehensive comparison on CIFAR-100 with DeiT-Tiny. ^† Token keep (physical sequence reduction). ^‡ Soft head gating, no FLOPs reduction. * Hard structural head removal after gate thresholding; accuracy varies by K (see text). ^§ Zero-shot WPR scoring; no training or calibration data needed, but adds 8.9% latency overhead and increases peak memory vs CGTS (see Section 4.3). — denotes not applicable.

		Accuracy (%)				Efficiency
Method	Train	$\mathit{K}\mathbf{=}\mathbf{5}$	$\mathit{K}\mathbf{=}\mathbf{10}$	$\mathit{K}\mathbf{=}\mathbf{20}$	$\mathit{K}\mathbf{=}\mathbf{50}$	GMACs	MB	Per-Subset	ms/B64	img/s
Unpruned	—	96.20	95.50	94.35	89.00	1.225	22.17	—	101.6	630
Ours MLP30% + CGTS0.6 †	5ep	97.60	96.10	94.85	88.28	0.844	17.91	0.8 KB	74.8	856
Ours MLP50% + CGTS0.6 †	5ep	96.80	95.40	94.25	87.56	0.636	15.08	0.8 KB	67.8	944
CGTS only † (zero-shot)	none	97.00	95.60	93.95	88.80	1.022	22.17	0.8 KB	83.6	766
Zero-TPrune †,§ (keep = 0.7)	none	96.60	95.80	93.95	88.88	1.030	22.17	0 B	91.1	703
DynamicViT † (keep = 0.7)	5ep × N	97.60	95.20	94.05	89.62	1.022	22.17	22.2 MB	90.3	709
SViTE † h0 (keep = 0.7)	5ep × N	96.20	94.70	92.85	86.86	1.022	22.17	22.2 MB	85.1	751
X-Pruner ‡ ALM (keep = 0.7)	5ep × N	98.60	97.40	95.60	90.10	1.225	22.17	14.1 KB	111.5	574
X-Pruner * hard (keep = 0.7)	5ep × N	97.80	96.60	93.90	87.06	1.225	22.17	14.1 KB	111.5	574
X-Pruner * hard (keep = 0.5)	5ep × N	88.20	94.90	83.15	56.82	≈1.2	19–22	14.1 KB	—	—

First, compared with the unpruned baseline, our MLP30% + CGTS0.6 method achieves better accuracy on small and medium-scale subsets with $K=5,10,20$ and obtains prominent efficiency gains. Specifically, it reduces GMACs by 31%, cuts model size by 19%, and increases inference throughput by 36% relative to the unpruned baseline.

Further compared with existing mainstream pruning methods, our approach achieves competitive classification accuracy, which is comparable to or even better than most peer schemes. It also delivers consistently superior efficiency: compared with existing pruning methods, our method can further reduce GMACs by an additional 15%, compress model size by 19%, and boost inference throughput by 14%.

Notably, our method achieves slightly lower accuracy than competitors on the large subset $K=50$. We attribute this performance gap to the single-prototype limitation of CGTS when dealing with heterogeneous large-K subsets.

Meanwhile, X-Pruner with soft gating yields the highest accuracy across all K settings but fails to yield efficiency gains. The extra gating overhead slows down its inference speed by about 9% compared with the unpruned baseline.

By contrast, the hard-gated X-Pruner under the aggressive keep = 0.5 setting suffers catastrophic accuracy degradation, dropping to 83.2% at $K=20$ and 56.8% at $K=50$. This phenomenon verifies that soft gating mechanisms cannot be extended to structural model pruning and removal. In addition, under the more aggressive pruning setting of MLP50% + CGTS0.6, the model still maintains stable classification accuracy with only marginal performance decline, while the efficiency gain is further enhanced. Figure 2 intuitively illustrates the accuracy trends of all compared pruning methods across different target class sizes on the CIFAR-100 dataset.

Figure 2. Accuracy of various token/head pruning methods under different numbers of target classes K on CIFAR-100 with DeiT-Tiny ($K=5,10,20,50$).

In addition, we further conduct generalization validation on the TinyImageNet-200 dataset with high class complexity using the DeiT-Small model, as reported in Table 2. Compared with CIFAR-100, TinyImageNet possesses stronger semantic heterogeneity and higher task difficulty, leading to distinct accuracy trends among various pruning methods. Figure 3 plots the corresponding accuracy curves on the TinyImageNet dataset, showing the performance change rule under a 200-class complex task.

Table 2. Comprehensive comparison on TinyImageNet with DeiT-Small. ^† Token keep (physical sequence reduction). ^‡ Soft head gating, no FLOPs reduction.

Method	$\mathit{K}\mathbf{=}\mathbf{5}$	$\mathit{K}\mathbf{=}\mathbf{10}$	$\mathit{K}\mathbf{=}\mathbf{20}$	$\mathit{K}\mathbf{=}\mathbf{50}$
Unpruned (no fine-tune)	99.60	98.20	96.10	94.12
Fine-tuned only (pr = 0.0)	99.60	98.40	95.60	92.92
Ours MLP30% + CGTS0.6 †	98.80	97.80	93.70	90.32
Ours MLP50% + CGTS0.6 †	98.40	96.80	91.80	89.20
DynamicViT † (keep = 0.7)	99.20	97.20	95.20	92.00
SViTE † h0 (keep = 0.7)	98.80	98.00	94.40	92.44
X-Pruner ‡ ALM (keep = 0.7)	99.20	98.60	95.30	93.00

Figure 3. Accuracy of various token/head pruning methods under different numbers of target classes K on TinyImageNet with DeiT-Small ($K=5,10,20,50$).

Overall, the X-Pruner soft gating scheme achieves the best accuracy across all K settings without physical structural pruning. DynamicViT and SViTE also maintain competitive performance. Although our method delivers slightly lower accuracy than the above competitors, the performance gap is controllable.

At $K=5$, MLP30% + CGTS0.6 achieves comparable accuracy to SViTE, merely 0.40% lower than DynamicViT. In the cases of $K=10,20,50$, the accuracy difference between our method and mainstream structural pruning methods remains stable within 0.6–1.5%. The core reason lies in the higher complexity of TinyImageNet. The single-prototype mechanism of CGTS suffers more serious representation degradation on heterogeneous large-K subsets. Even so, our method still maintains acceptable classification accuracy, and its core advantage of balancing accuracy and efficiency is well retained.

Overall, compared with the mainstream pruning methods covered in this paper, our approach achieves simultaneous reductions in GMACs, model size, deployment storage cost, and inference latency while maintaining acceptable classification accuracy.

4.3. Efficiency Analysis

Under the batch size setting of $B=64$, structural MLP compression (physically reducing the dimensions of fc1/fc2) is essential to obtain real latency gains. Soft head gating (X-Pruner) and soft token masking cannot yield actual acceleration under the standard PyTorch 2.5.1 execution paradigm. The token predictor of DynamicViT introduces additional overhead at $B=64$, resulting in an 8% slower inference speed compared with CGTS under the same GMACs. Zero-TPrune also exhibits considerable extra overhead. Although Zero-TPrune requires no calibration data with negligible per-subset storage consumption, its WPR scoring mechanism instantiates a complete $[B,H,N,N]$ attention matrix and performs three power iteration steps during each forward propagation. At $B=64$, its inference latency increases by 8.9% compared with CGTS, and the peak GPU memory rises by 6% from 185 MB of the unpruned baseline to 197 MB. This is because the buffer of $B\times H\times N\times N$ dominates the activation memory consumption. By shortening the sequence length at layers 7–11, CGTS reduces the peak memory by 10%. Our method is the only one that achieves simultaneous reductions across all four efficiency dimensions.

We further analyze the inference scenario of $B=1$ for edge deployment, with the corresponding results summarized in Table 3. All token pruning methods (CGTS, DynamicViT, and Zero-TPrune) fail to bring efficiency benefits at $B=1$, due to low GPU utilization and dominant kernel launch overhead. Zero-TPrune performs the worst at $B=1$, with inference latency increased by 4% compared with the baseline, verifying that its per-sample WPR iteration cost becomes a notable burden for edge deployment. In contrast, structural MLP compression remains effective and achieves a 16% speedup. Even under low arithmetic intensity conditions, smaller weight matrices effectively reduce memory bandwidth consumption. This further demonstrates the deployment advantage of our proposed method.

Table 3. Batch size B = 1 latency (edge deployment scenario).

Method	B = 1 (ms)	vs. Unpruned
Unpruned	7.18	—
Ours MLP30% + CGTS0.7	6.19	+16%
Ours MLP50% + CGTS0.5	6.82	+5%
CGTS only (keep = 0.7)	7.29	−2%
CGTS only (keep = 0.5)	6.97	+3%
Zero-TPrune (keep = 0.7)	7.45	−4%
Zero-TPrune (keep = 0.5)	7.20	0%
DynamicViT (keep = 0.7)	7.16	0%
SViTE (keep = 0.7)	5.39	+33% †
X-Pruner (keep = 0.7)	7.34	−2%

^† Static precomputed index gather; requires 22.2 MB checkpoint per deployment and delivers lower accuracy (−1.5% at $K=20$).

4.4. Memory Overhead Analysis

CGTS only needs to store a 768-byte prototype vector for the deployment of each class subset, without retaining a separate complete model. In contrast, both DynamicViT and SViTE require saving an independent full model checkpoint for each class subset. When the number of deployed subsets is $N=100$, the overall storage overhead of our method is only 22.2 MB (backbone + 100 × 0.8 KB), which is 1/100 of that of DynamicViT (2.22 GB). In addition, the deployment of a new K-class subset only requires one lightweight forward pass over support samples to generate prototypes, with no retuning required; prototypes only need recalculation when the full target category set is replaced. For multiple independent target subsets, our shared fixed backbone incurs negligible extra storage and computation overhead, showing remarkable advantages in deployment efficiency and rapid iteration.

4.5. Ablation Study

In this section, we conduct ablation experiments on the TinyImageNet dataset (Table 4). Taking the fine-tuned model without pruning and CGTS as the baseline, standalone MLP pruning with a ratio of 30% yields an accuracy drop below $0.5\%$ at $K=5,10$, achieving near-lossless compression. For standalone CGTS with a keep ratio of 0.6, accuracy degradation increases with K, rising from 0.80% at $K=5$ to 0.96% at $K=50$, which verifies the representation degradation law of the single-prototype mechanism. The losses are approximately additive when the two modules are combined. Taking $K=50$ as an example, the accuracy loss of MLP30% + CGTS0.6 is 2.60%, slightly higher than the summed individual loss of 1.88%, while no abrupt performance collapse is observed, demonstrating satisfactory compatibility. Under the aggressive configuration of MLP50% + CGTS0.6, the accuracy reaches 89.20% at $K=50$, with a drop of 3.72% compared with the baseline. More importantly, joint deployment can compress both MLP width and sequence length simultaneously, achieving multiplicative efficiency gains unavailable with any single module (FLOPs reduced by 31%, and throughput improved by nearly 50%). It is validated that the collaboration of the two modules can achieve prominent efficiency improvement at the cost of acceptable accuracy loss, thus enabling a flexible trade-off between model accuracy and compression ratio.

Table 4. Ablation study of our components on TinyImageNet (DeiT-Small, 5 fine-tuning epochs, CGTS at layer 6).

Method	$\mathit{K}\mathbf{=}\mathbf{5}$	$\mathit{K}\mathbf{=}\mathbf{10}$	$\mathit{K}\mathbf{=}\mathbf{20}$	$\mathit{K}\mathbf{=}\mathbf{50}$
Unpruned (Fine-tuned)	99.60	98.40	95.60	92.92
Only MLP Prune 30%	99.20	98.40	94.30	92.00
Only MLP Prune 50%	99.20	97.60	92.70	91.04
Only CGTS (keep = 0.6)	98.80	97.80	94.70	91.96
MLP30% + CGTS0.6	98.80	97.80	93.70	90.32
MLP50% + CGTS0.6	98.40	96.80	91.80	89.20

5. Conclusions

This paper proposes a class-aware joint pruning framework for Vision Transformers, which performs structured MLP pruning and dynamic token pruning simultaneously to tailor pre-trained ViT models into dedicated models for target class subsets. Extensive experiments demonstrate that the proposed method can reduce GMACs by up to 48% and improve inference throughput by nearly 50% while maintaining high accuracy on target classes, with only 0.8 KB storage overhead per single class. It outperforms several mainstream pruning methods in the accuracy–efficiency trade-off and enables flexible multi-subset deployment with a single backbone network.

Nevertheless, several directions remain for optimizing our framework in more complex real-world scenarios. Our single shared prototype is built under the assumption of consistent feature patterns within a class subset, which may lead to performance degradation when facing highly heterogeneous categories. Moreover, the current prototype construction pipeline leverages clean, identically distributed support samples; future work could further boost token pruning robustness to handle noisy labels and cross-domain distribution mismatches. Additionally, our method has substantial potential to deliver better compression gains when tailored for fine-grained classification tasks with subtle inter-class differences.

In the future, we aim to improve the single-prototype mechanism of CGTS by introducing multi-prototype learning or adaptive prototype representation strategies to mitigate feature degradation of the model on large-scale and highly heterogeneous class subsets and further reduce the accuracy gap in complex scenarios. Meanwhile, we can combine joint pruning with quantization and knowledge distillation to adapt to edge devices and extend the class-aware pruning to more ViT variants and high-level vision tasks as well as large-scale benchmarks including ImageNet-1K and COCO to verify the universality and scalability of the method, providing a more comprehensive technical solution for lightweight Transformer deployment in practical edge scenarios.