Putting Registers to Work: Task Registers for Token Pruning in Vision Transformers
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 , our jointly adapted model, TAP-J, reaches mIoU at encoder throughput on ADE20K and box AP at 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