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

arXiv:2608.10989v1 Announce Type: cross 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 rho=0.5, our jointly adapted model, TAP-J, reaches 47.0 mIoU at 1.30imes encoder throughput on ADE20K and 53.7 box AP at 1.32imes encoder throughput on COCO while remaining competitive on ImageNet-1K.

Related

• Human-like Object Grouping in Self-supervised Vision Transformers
• bViT: Investigating Single-Block Recurrence in Vision Transformers for Image Recognition
• Weight Concentration Regularization for Improving Pruning Robustness Under High Sparsity

Source: arXiv cs.AI | 2026-08-12