HEART-VIT: Hessian-Guided Efficient Dynamic Attention and Token Pruning in Vision Transformer

TL;DR

HEART-ViT addresses the high inference cost of Vision Transformers by proposing a unified, second-order, input-adaptive pruning framework that jointly removes tokens and attention heads. By computing curvature-weighted sensitivity scores $S_z(x) = z(x)^T \mathcal{H}_z(x) z(x)$ via efficient Hessian-vector products, it prunes components under an explicit loss budget, with normalization and dynamic hard/soft gating to stabilize pruning. The approach reveals that token pruning dominates FLOPs reduction due to the quadratic cost of self-attention, while head pruning refines representational capacity; together they yield strong Pareto-optimal gains on ImageNet benchmarks and translate to real-world speedups on edge devices. Empirical results show substantial FLOPs and latency reductions with minimal or positive accuracy recovery after fine-tuning, and ablations highlight the key role of per-layer normalization and asymmetric pruning schedules in achieving best efficiency-accuracy trade-offs.

Abstract

Vision Transformers (ViTs) deliver state-of-the-art accuracy but their quadratic attention cost and redundant computations severely hinder deployment on latency and resource-constrained platforms. Existing pruning approaches treat either tokens or heads in isolation, relying on heuristics or first-order signals, which often sacrifice accuracy or fail to generalize across inputs. We introduce HEART-ViT, a Hessian-guided efficient dynamic attention and token pruning framework for vision transformers, which to the best of our knowledge is the first unified, second-order, input-adaptive framework for ViT optimization. HEART-ViT estimates curvature-weighted sensitivities of both tokens and attention heads using efficient Hessian-vector products, enabling principled pruning decisions under explicit loss budgets.This dual-view sensitivity reveals an important structural insight: token pruning dominates computational savings, while head pruning provides fine-grained redundancy removal, and their combination achieves a superior trade-off. On ImageNet-100 and ImageNet-1K with ViT-B/16 and DeiT-B/16, HEART-ViT achieves up to 49.4 percent FLOPs reduction, 36 percent lower latency, and 46 percent higher throughput, while consistently matching or even surpassing baseline accuracy after fine-tuning, for example 4.7 percent recovery at 40 percent token pruning. Beyond theoretical benchmarks, we deploy HEART-ViT on different edge devices such as AGX Orin, demonstrating that our reductions in FLOPs and latency translate directly into real-world gains in inference speed and energy efficiency. HEART-ViT bridges the gap between theory and practice, delivering the first unified, curvature-driven pruning framework that is both accuracy-preserving and edge-efficient.

Figures (7)

• Figure 1: Our method pushes the Pareto frontier of FLOPs vs. accuracy on ImageNet-1K. We compare ViT-B/16 and DeiT-B/16 under symmetric (Sym) and asymmetric (Asym) pruning against strong baselines and state-of-the-art transformer variants. Our pruned models consistently achieve higher accuracy at significantly reduced FLOPs, surpassing both dense ViT/DeiT baselines and competitive efficient transformers. Detailed results are on Appendix.Table \ref{['tab:SOTA comparison']}. Notes: Sym = Symmetric pruning; 20/20 = 20% Tokens + 20% Heads.
• Figure 2: Accuracy decomposition under symmetric (Sym) and asymmetric (Asym) pruning. (a)–(b) show results on ImageNet-1K for ViT-B/16 and DeiT-B/16, while (c)–(d) present the corresponding results on ImageNet-100. Bars indicate pruned retention (bottom) and accuracy recovered by fine-tuning (FT, top); the dashed line marks the dense baseline; $\Delta$ annotations indicate the accuracy change relative to the baseline.detailed results has been shown in (Appendix: Table \ref{['tab:imagenet100']} & \ref{['tab:fullimagenet']}.)
• Figure 3: Accuracy vs FLOPs reduction curves for Symmetric and Asymmetric pruning on ImageNet-1K (a–b) and ImageNet-100 (c–d). Baseline accuracy is shown as dashed lines.
• Figure 4: Layerwise analysis of ViT-B/16 under 50% symmetric pruning (tokens + heads) on ImageNet-1K. CKA reveals mid-layer representational shifts that are partially recovered after fine-tuning. CLS cosine shows pruning drives the model toward an alternative semantic trajectory, while residual ratios highlight temporary suppression of transformations that fine-tuning restores.
• Figure 5: Layerwise analysis of DeiT-B/16 under 50% symmetric pruning on ImageNet-1K. Compared to ViT, DeiT shows stronger representational shifts in CKA and CLS similarity, and larger fluctuations in residual ratios, indicating that its distilled training makes it more sensitive to pruning perturbations.