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

arXiv: https://arxiv.org/abs/2608.10989 · HTML: https://mlsanity.com/paper/2608.10989/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Token pruning and merging for ViTs, one of 12 directions in Efficient LLM inference and serving.
- Builds on 25 of its 34 references that are in the map.
- First posted on arXiv 11 Aug 2026.

## Placed in

- Efficient LLM inference and serving / [Token pruning and merging for ViTs](https://mlsanity.com/inference/token-pruning-and-merging-for-vits/)

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space](https://mlsanity.com/paper/1706.02413/) (2017)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [Cascade R-CNN: Delving into High Quality Object Detection](https://mlsanity.com/paper/1712.00726/) (2017)
- [Token Merging: Your ViT But Faster](https://mlsanity.com/paper/2210.09461/) (2022)
- [Vision Transformers Need Registers](https://mlsanity.com/paper/2309.16588/) (2023)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [Exploring Plain Vision Transformer Backbones for Object Detection](https://mlsanity.com/paper/2203.16527/) (2022)
- [Not All Patches are What You Need: Expediting Vision Transformers via Token Reorganizations](https://mlsanity.com/paper/2202.07800/) (2022)
- [AdaViT: Adaptive Tokens for Efficient Vision Transformer](https://mlsanity.com/paper/2112.07658/) (2021)
- [Evo-ViT: Slow-Fast Token Evolution for Dynamic Vision Transformer](https://mlsanity.com/paper/2108.01390/) (2021)
- [SPViT: Enabling Faster Vision Transformers via Soft Token Pruning](https://mlsanity.com/paper/2112.13890/) (2021)
- [Joint Token Pruning and Squeezing Towards More Aggressive Compression of Vision Transformers](https://mlsanity.com/paper/2304.10716/) (2023)
- [DiffRate : Differentiable Compression Rate for Efficient Vision Transformers](https://mlsanity.com/paper/2305.17997/) (2023)
- [Zero-TPrune: Zero-Shot Token Pruning through Leveraging of the Attention Graph in Pre-Trained Transformers](https://mlsanity.com/paper/2305.17328/) (2023)
- [Dynamic Token Pruning in Plain Vision Transformers for Semantic Segmentation](https://mlsanity.com/paper/2308.01045/) (2023)
- [Which Tokens to Use? Investigating Token Reduction in Vision Transformers](https://mlsanity.com/paper/2308.04657/) (2023)
- [Revisiting Token Pruning for Object Detection and Instance Segmentation](https://mlsanity.com/paper/2306.07050/) (2023)
- [Beyond Attentive Tokens: Incorporating Token Importance and Diversity for Efficient Vision Transformers](https://mlsanity.com/paper/2211.11315/) (2022)
- [Dynamic Tuning Towards Parameter and Inference Efficiency for ViT Adaptation](https://mlsanity.com/paper/2403.11808/) (2024)
- [Token Cropr: Faster ViTs for Quite a Few Tasks](https://mlsanity.com/paper/2412.00965/) (2024)
- [VLTP: Vision-Language Guided Token Pruning for Task-Oriented Segmentation](https://mlsanity.com/paper/2409.08464/) (2024)

## Abstract (arXiv)

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× encoder throughput on ADE20K and 53.7 box AP at 1.32× encoder throughput on COCO while remaining competitive on ImageNet-1K.