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

TL;DR

HEART-ViT introduces a second-order, input-adaptive token and head pruning framework for Vision Transformers, significantly reducing computational costs while maintaining or improving accuracy, and demonstrating practical efficiency on edge devices.

Contribution

It is the first unified, Hessian-guided framework for dynamic token and head pruning in ViTs, enabling explicit control over loss budgets and improved efficiency.

Findings

Achieves up to 49.4% FLOPs reduction on ImageNet datasets.

Reduces latency by 36% and increases throughput by 46%.

Maintains or surpasses baseline accuracy after pruning.

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.