English

Putting Registers to Work: Task Registers for Token Pruning in Vision Transformers

Computer Vision and Pattern Recognition 2026-08-11 v1 Artificial Intelligence

Abstract

Token-pruning policies are usually designed for a single recognition pipeline, but pretrained Vision Transformers are reused across tasks with different spatial demands. We ask which parts of a pruning policy transfer across image classification, semantic segmentation, and object detection. For each pipeline, controlled probes freeze the no-pruning checkpoint and apply a series of parameter-free reduction criteria at one eligible layer at a time without retraining. The probes reveal three differences: segmentation and detection rank the criteria differently, classification is especially sensitive to attention-based pruning in the earliest layers, and the dense tasks prefer opposite recovery endpoints. These findings motivate Task-Adaptive Pruning (TAP). Existing register tokens serve as task-agnostic storage for feature artifacts. TAP instead introduces one task register per task and activates only the current one. Its evolving state ranks tokens, distributes an exact removal budget over depth, and sets the recovery scale for dense features. At a final keep rate of ρ=0.5\rho=0.5, our jointly adapted model, TAP-J, reaches 47.047.0 mIoU at 1.30×1.30\times encoder throughput on ADE20K and 53.753.7 box AP at 1.32×1.32\times encoder throughput on COCO while remaining competitive on ImageNet-1K.

Keywords

Cite

@article{arxiv.2608.10989,
  title  = {Putting Registers to Work: Task Registers for Token Pruning in Vision Transformers},
  author = {Hongsen Cao and Mona Jaber and Shanxin Yuan and Ahmed Sayed},
  journal= {arXiv preprint arXiv:2608.10989},
  year   = {2026}
}

Comments

24 pages, 9 figures. Includes supplementary material