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.