TAP-ViTs: Task-Adaptive Pruning for On-Device Deployment of Vision Transformers

TL;DR

TAP-ViTs tackles the challenge of deploying Vision Transformers on edge devices under privacy and hardware heterogeneity constraints. It introduces a GMM-based metric dataset construction to infer device-specific task requirements from private data without sharing raw samples, and pairs it with a dual-granularity pruning strategy that jointly evaluates composite neuron importance and adaptive layer importance. The framework demonstrates superior accuracy under comparable compression across multiple ViT backbones and datasets, and even surpasses the full model at low pruning ratios in some cases, while maintaining practical edge deployment costs. This privacy-preserving, task-adaptive pruning approach offers a scalable pathway for on-device ViT deployment in heterogeneous and privacy-sensitive environments, with potential extensions to generative data synthesis and multi-modal models.

Abstract

Vision Transformers (ViTs) have demonstrated strong performance across a wide range of vision tasks, yet their substantial computational and memory demands hinder efficient deployment on resource-constrained mobile and edge devices. Pruning has emerged as a promising direction for reducing ViT complexity. However, existing approaches either (i) produce a single pruned model shared across all devices, ignoring device heterogeneity, or (ii) rely on fine-tuning with device-local data, which is often infeasible due to limited on-device resources and strict privacy constraints. As a result, current methods fall short of enabling task-customized ViT pruning in privacy-preserving mobile computing settings. This paper introduces TAP-ViTs, a novel task-adaptive pruning framework that generates device-specific pruned ViT models without requiring access to any raw local data. Specifically, to infer device-level task characteristics under privacy constraints, we propose a Gaussian Mixture Model (GMM)-based metric dataset construction mechanism. Each device fits a lightweight GMM to approximate its private data distribution and uploads only the GMM parameters. Using these parameters, the cloud selects distribution-consistent samples from public data to construct a task-representative metric dataset for each device. Based on this proxy dataset, we further develop a dual-granularity importance evaluation-based pruning strategy that jointly measures composite neuron importance and adaptive layer importance, enabling fine-grained, task-aware pruning tailored to each device's computational budget. Extensive experiments across multiple ViT backbones and datasets demonstrate that TAP-ViTs consistently outperforms state-of-the-art pruning methods under comparable compression ratios.

Figures (6)

• Figure 1: Heatmaps of cross-task neuron-importance divergence.
• Figure 2: The Framework of TAP-ViTs.
• Figure 3: Performance vs. Pruning Ratio
• Figure 4: Kendall's $\tau$ correlation of neuron-importance rankings across tasks.
• Figure 5: Comparison of layer-importance rankings across tasks.