💻 computer science

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

This paper proposes Task-Adaptive Pruning (TAP), a method that utilizes task-specific registers to dynamically rank tokens and allocate pruning budgets across different Vision Transformer tasks, achieving significant throughput improvements while maintaining competitive performance on image classification, semantic segmentation, and object detection.

Original authors: Hongsen Cao, Mona Jaber, Shanxin Yuan, Ahmed Sayed

Published 2026-08-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Hongsen Cao, Mona Jaber, Shanxin Yuan, Ahmed Sayed

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a super-smart robot brain that looks at pictures and tries to understand what's inside them. This brain is built like a giant team of tiny workers, each holding a piece of the puzzle (a "token"). To make a decision, every worker has to chat with every other worker. The problem? If the picture is big, the number of chats explodes, and the brain gets overwhelmed, slow, and hungry for electricity. Scientists have been trying to fix this by telling the brain to ignore the boring pieces of the picture and only keep the important ones, a bit like a teacher telling a class to only listen to the students who have something smart to say. This is called "token pruning." But here's the tricky part: what counts as "smart" or "important" changes depending on what the robot is trying to do. A piece of a picture that is crucial for spotting a cat might be useless for figuring out what color the sky is. For a long time, scientists tried to use the same "ignore list" for every single job, but it turns out that one size rarely fits all.

This paper, titled "Putting Registers to Work," tackles that exact problem. The authors, a team from Queen Mary University of London, discovered that the rules for deciding which picture pieces to throw away are totally different for different tasks. They found that a strategy that works great for spotting objects (like finding a dog in a park) might actually hurt the robot's ability to understand a scene (like figuring out if that dog is on grass or pavement). To solve this, they invented a clever new system called Task-Adaptive Pruning (TAP).

Think of the robot's brain as a busy kitchen. Usually, the chef (the pruning policy) uses the same recipe to decide which ingredients to toss out, no matter if they are making a soup or a cake. The authors realized this was a mistake. Instead, they gave the kitchen a set of special "Task Registers"—think of them as smart, magical recipe cards that change depending on the dish being cooked. When the robot needs to identify a cat, it pulls out the "Cat Card," which tells the brain exactly which parts of the image to keep and which to ignore. When it needs to map out a whole room, it swaps in the "Room Card." These cards don't just decide what to throw away; they also figure out when to throw it away as the brain processes the image, layer by layer.

The researchers tested this by freezing a standard robot brain and trying out different "ignore rules" without retraining it. They found three big surprises. First, the best way to pick important pieces for finding objects is the exact opposite of the best way for mapping scenes. Second, the robot is very sensitive to how it handles the very first layers of processing, especially for simple tasks like telling if an image is a cat or a dog. Third, when the robot throws a piece of the picture away, it needs to remember what it looked like to fill in the gap later. But here's the kicker: for finding objects, it's better to just pretend the piece was never there and use a nearby neighbor to fill the gap. For mapping scenes, it's better to actually remember the exact difference and paste it back in later.

To fix all this, TAP uses a single "active" register card for the current job. This card travels through the brain, watching the image and making three decisions on the fly:

Selection: Which tokens to keep and which to drop.
Budget: How many tokens to drop at each step (maybe drop a few early, more later, or vice versa).
Recovery: How much of the "forgotten" detail to bring back when the robot needs to draw a detailed map.

The results are impressive. By using these task-specific cards, the robot became much faster without losing its smarts. On a standard test for finding objects (COCO), it became 1.32 times faster while keeping its accuracy almost the same (dropping only 0.3 points). On a test for mapping scenes (ADE20K), it became 1.30 times faster with a tiny drop in quality. Even for simple image classification, it stayed competitive.

The authors suggest that this approach is a step toward a truly unified robot brain—one that can switch between spotting a car, painting a picture of a street, or reading a sign, all using the same core engine but with different "thinking styles" activated on the fly. They didn't just guess this would work; they proved it by running thousands of experiments, freezing the brain to test individual rules, and showing that their new "register" system consistently outperforms the old "one-size-fits-all" methods. It's a reminder that in the world of AI, sometimes the smartest thing you can do is know exactly when to ignore the noise, and that the definition of "noise" depends entirely on what you're trying to see.

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

Problem Statement

Vision Transformers (ViTs) exhibit quadratic computational scaling with the number of tokens, motivating token pruning to reduce inference costs. While pruning criteria (e.g., attention scores, spatial coverage) are often developed for image classification, pretrained ViT backbones are increasingly adapted to dense prediction tasks like semantic segmentation and object detection.

The authors identify a critical gap: pruning policies designed for one task do not necessarily transfer effectively to others. Through controlled "probe" experiments, the paper demonstrates that:

Criterion Ranking Divergence: Segmentation and detection tasks rank pruning criteria differently. For instance, attention-based pruning causes significant accuracy drops in early layers for classification, whereas spatial coverage is more stable. Conversely, the criterion that minimizes detection error often performs poorly for segmentation.
Depth Sensitivity: The effectiveness of a pruning criterion varies significantly across network depths and differs between tasks.
Recovery Conflicts: Dense prediction tasks require feature reconstruction at removed positions. The paper finds that segmentation and detection prefer opposite recovery endpoints: segmentation benefits from restoring the full feature offset (difference between removed and retained tokens), while detection performs better using only the retained "stand-in" token.

Existing methods typically apply a single, static pruning policy or use task-agnostic selectors, failing to account for these task-specific and depth-dependent variations.

Methodology: Task-Adaptive Pruning (TAP)

The paper proposes Task-Adaptive Pruning (TAP), a framework that utilizes Task Registers to dynamically control token pruning, budget allocation, and feature recovery.

Core Mechanism:

Task Registers: Instead of using registers merely as task-agnostic storage for feature artifacts (as in prior work like Darcet et al., 2024), TAP introduces one learnable initial register per task. During a forward pass, only the register associated with the current task is active.
Evolving State: The active register evolves alongside image tokens through the backbone. Its state at any layer conditions three key decisions:
1. Token Selection: The register computes a score for each candidate patch by reusing the block's existing query and key projections. Tokens are ranked, and a subset is retained for self-attention.
2. Exact Budget Allocation: A shared linear readout, conditioned on the register state, determines the exact integer number of tokens to remove at each pruning layer. This ensures the global removal budget is met exactly without requiring a rate loss.
3. Recovery Scaling: For dense tasks, the register controls a scaling factor ( $\alpha$ ) that determines how much of the stored feature offset (difference between removed token and its stand-in) is restored at the readout.

Implementation Variants:

TAP-J (Joint Fine-tuning): Each task has a separate backbone, task register, and readouts. This allows full specialization.
TAP-F (Frozen Base): A single frozen backbone is shared across tasks. Task-specific registers, low-rank (LoRA) updates, and heads are trained jointly. This reduces parameter storage while maintaining adaptability.

Feature Recovery:
Removed tokens are deterministically matched to the retained token with the highest key-space cosine similarity (the "stand-in"). The feature difference (offset) is stored. At dense readout points, the reconstruction is computed as $\hat{x}_i = x_{\text{stand-in}} + \alpha \cdot \delta_i$ , where $\alpha$ is learned per task and depth. Classification tasks skip this recovery step.

Key Contributions

Controlled Framework: The authors develop a framework that freezes no-pruning checkpoints and applies parameter-free reduction criteria layer-by-layer to isolate the impact of pruning choices across classification, segmentation, and detection. This reveals the misalignment in criterion rankings and recovery preferences.
Task Registers as Controllers: The paper redefines the role of registers from passive storage to active, task-specific controllers of sparse computation. The register coordinates token selection, exact budget distribution across depth, and recovery scaling.
Unified Pruning and Recovery: TAP treats token removal and feature recovery as a unified operation. It matches discarded tokens to stand-ins and scales the restoration of offsets based on the task register, keeping reconstructed features out of the encoder to maintain sparsity.

Experimental Results

The authors evaluate TAP on ImageNet-1K (classification), ADE20K (segmentation), and COCO (detection/instance segmentation) using a ViT-B backbone with $\rho = 0.5$ (50% tokens retained).

Performance:
- TAP-J achieves the smallest accuracy drops among fully fine-tuned methods. On ADE20K, it reaches 47.0 mIoU (only 0.2 drop from no-pruning) with 1.30× encoder throughput. On COCO, it achieves 53.7 box AP (0.3 drop) with 1.32× throughput. On ImageNet-1K, it reaches 83.2 top-1 (0.4 drop).
- TAP-F (shared backbone) trails TAP-J slightly but significantly outperforms static rules and other baselines, achieving 45.8 mIoU on ADE20K and 52.5 box AP on COCO with minimal parameter overhead (~1.2M task-specific parameters).
Efficiency: TAP-J provides encoder speedups of 1.26× to 1.32× while maintaining competitive accuracy. The controller overhead (register query, budget readout, bookkeeping) accounts for only ~0.7% of encoder runtime.
Memory: TAP reduces peak inference memory by 9–18% compared to dense baselines, with segmentation seeing the largest gains due to reduced activation sizes in intermediate blocks.
Ablation Studies: Removing the task register (using a shared one) or fixing the allocation/recovery to static values significantly degrades performance, confirming the necessity of task- and depth-adaptive control.

Significance and Claims

The paper claims that token reduction cannot be treated as a single transferable rule. As ViT representations evolve from local structures to high-level semantics, the value of tokens changes with both depth and the specific task reading the backbone.

The significance of TAP lies in its ability to adapt sparse computation to the specific needs of different visual tasks using a shared backbone. By putting evolving task registers to work, TAP:

Resolves conflicts in pruning criteria rankings between tasks.
Optimizes the distribution of the removal budget across network depth.
Conditions feature recovery on the specific requirements of dense vs. sparse readouts.

The authors position TAP as a step toward unified vision systems that can serve multiple tasks through a shared backbone with minimal task-specific storage, balancing predictive performance with encoder throughput more effectively than static pruning policies. The paper remains modest, noting that TAP-J trails the best classification-only method (DiffRate) by 0.1 point, likely because classification relies on a class token and skips dense recovery, leaving one of TAP's mechanisms underutilized for that specific task. Future work is suggested to explore image-conditioned budgets and multiple stand-ins.