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

2026-08-11

AI Summary

Existing token pruning methods for vision Transformers struggle to generalize effectively across diverse tasks such as image classification, semantic segmentation, and object detection. This work addresses this limitation by first revealing through probing experiments the varying sensitivity of different tasks to pruning strategies. Building on this insight, we propose Task-Adaptive Pruning (TAP), which introduces task-specific registers to dynamically guide token pruning and feature restoration at each layer. TAP jointly optimizes the pruning criterion, depth budget allocation, and restoration scale in a task-aware manner. Under a token retention ratio of ρ = 0.5 and without compromising ImageNet-1K classification accuracy, TAP achieves 47.0 mIoU on ADE20K (1.30× encoder throughput gain) and 53.7 box AP on COCO (1.32× throughput gain).

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$, our jointly adapted model, TAP-J, reaches $47.0$ mIoU at $1.30\times$ encoder throughput on ADE20K and $53.7$ box AP at $1.32\times$ encoder throughput on COCO while remaining competitive on ImageNet-1K.

Problem

• token pruning
• Vision Transformers
• task adaptation
• multi-task transfer
• spatial demands

Innovation

• Task-Adaptive Pruning
• Vision Transformers
• Token Pruning
• Task Registers
• Multi-task Efficiency