💻 computer science

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

Questo articolo propone il Task-Adaptive Pruning (TAP), un metodo che utilizza registri specifici per il compito per classificare dinamicamente i token e allocare i budget di pruning attraverso diversi compiti di Vision Transformer, ottenendo significativi miglioramenti del throughput pur mantenendo prestazioni competitive nella classificazione di immagini, nella segmentazione semantica e nel rilevamento di oggetti.

Autori originali: Hongsen Cao, Mona Jaber, Shanxin Yuan, Ahmed Sayed

Pubblicato 2026-08-12

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Hongsen Cao, Mona Jaber, Shanxin Yuan, Ahmed Sayed

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di avere un cervello robotico super intelligente che guarda le immagini e cerca di capire cosa c'è dentro. Questo cervello è costruito come una gigantesca squadra di piccoli operai, ognuno dei quali tiene in mano un pezzo del puzzle (un "token"). Per prendere una decisione, ogni operaio deve chiacchierare con tutti gli altri operai. Il problema? Se l'immagine è grande, il numero di chiacchiere esplode, e il cervello si sente sopraffatto, rallenta e ha fame di elettricità. Gli scienziati hanno cercato di risolvere questo problema dicendo al cervello di ignorare i pezzi noiosi dell'immagine e di tenere solo quelli importanti, un po' come un insegnante che dice a una classe di ascoltare solo gli studenti che hanno qualcosa di intelligente da dire. Questo è chiamato "token pruning" (potatura dei token). Ma ecco la parte complicata: ciò che conta come "intelligente" o "importante" cambia a seconda di ciò che il robot sta cercando di fare. Un pezzo di un'immagine che è cruciale per individuare un gatto potrebbe essere inutile per capire il colore del cielo. Per molto tempo, gli scienziati hanno cercato di usare la stessa "lista di esclusione" per ogni singolo lavoro, ma si è scoperto che un modello unico raramente si adatta a tutto.

Questo articolo, intitolato "Putting Registers to Work", affronta esattamente questo problema. Gli autori, un team della Queen Mary University of London, hanno scoperto che le regole per decidere quali pezzi dell'immagine buttare via sono totalmente diverse per compiti differenti. Hanno scoperto che una strategia che funziona benissimo per individuare oggetti (come trovare un cane in un parco) potrebbe in realtà danneggiare la capacità del robot di comprendere una scena (come capire se quel cane si trova sull'erba o sul marciapiede). Per risolvere questo problema, hanno inventato un nuovo e astuto sistema chiamato Task-Adaptive Pruning (TAP).

Pensa al cervello del robot come a una cucina frenetica. Di solito, lo chef (la politica di potatura) usa la stessa ricetta per decidere quali ingredienti scartare, indipendentemente dal fatto che stia preparando una zuppa o una torta. Gli autori si sono resi conto che questo era un errore. Invece, hanno fornito alla cucina un set di speciali "Task Registers" — pensa a loro come a delle smart, magiche schede ricetta che cambiano a seconda del piatto che si sta cucinando. Quando il robot ha bisogno di identificare un gatto, estrae la "Carta del Gatto", che dice al cervello esattamente quali parti dell'immagine tenere e quali ignorare. Quando deve mappare un'intera stanza, sostituisce la scheda con la "Carta della Stanza". Queste schede non decidono solo cosa buttare via; capiscono anche quando buttare via, mentre il cervello elabora l'immagine, strato dopo strato.

I ricercatori hanno testato questo approccio congelando un cervello robotico standard e provando diversi "regole di esclusione" senza riaddestrarlo. Hanno scoperto tre grandi sorprese. Primo, il modo migliore per scegliere i pezzi importanti per l'individuazione di oggetti è l'esatto opposto del modo migliore per la mappatura di scene. Secondo, il robot è molto sensibile a come gestisce i primissimi strati di elaborazione, specialmente per compiti semplici come dire se un'immagine è un gatto o un cane. Terzo, quando il robot butta via un pezzo dell'immagine, deve ricordare com'era per colmare il vuoto in seguito. Ma ecco il colpo di scena: per l'individuazione di oggetti, è meglio far finta che il pezzo non sia mai esistito e usare un vicino nelle vicinanze per colmare il vuoto. Per la mappatura di scene, è meglio ricordare effettivamente la differenza esatta e incollarla di nuovo più tardi.

Per risolvere tutto questo, TAP utilizza un singolo registro "attivo" per il lavoro corrente. Questa scheda attraversa il cervello, osservando l'immagine e prendendo tre decisioni al volo:

Selezione: Quali token tenere e quali scartare.
Budget: Quanti token scartare ad ogni passaggio (forse scartarne alcuni presto, di più più tardi, o viceversa).
Recupero: Quanto dettaglio "dimenticato" riportare indietro quando il robot ha bisogno di disegnare una mappa dettagliata.

I risultati sono impressionanti. Usando queste schede specifiche per il compito, il robot è diventato molto più veloce senza perdere la sua intelligenza. In un test standard per l'individuazione di oggetti (COCO), è diventato 1,32 volte più veloce mantenendo quasi la stessa precisione (scendendo solo di 0,3 punti). In un test per la mappatura di scene (ADE20K), è diventato 1,30 volte più veloce con una minima perdita di qualità. Anche per la classificazione semplice delle immagini, è rimasto competitivo.

Gli autori suggeriscono che questo approccio sia un passo verso un vero cervello robotico unificato — uno che può passare dal individuare un'auto, al dipingere un quadro di una strada, o al leggere un cartello, tutto usando lo stesso motore centrale ma con diversi "stili di pensiero" attivati al volo. Non hanno solo ipotizzato che questo funzionasse; lo hanno dimostrato eseguendo migliaia di esperimenti, congelando il cervello per testare le singole regole e mostrando che il loro nuovo sistema di "registri" supera costantemente i vecchi metodi "taglia unica". È un promemoria del fatto che, nel mondo dell'IA, a volte la cosa più intelligente che si possa fare è sapere esattamente quando ignorare il rumore, e che la definizione di "rumore" dipende interamente da ciò che si sta cercando di vedere.

Sintesi Tecnica: Mettere i Registri al Lavoro: Task Register per il Pruning dei Token nei Vision Transformer

Problema

I Vision Transformer (ViT) presentano una scalabilità computazionale quadratica rispetto al numero di token, il che motiva il pruning dei token per ridurre i costi di inferenza. Mentre i criteri di pruning (ad esempio, punteggi di attenzione, copertura spaziale) sono spesso sviluppati per la classificazione delle immagini, i backbone ViT pre-addestrati sono sempre più adattati a compiti di predizione densa come la segmentazione semantica e il rilevamento degli oggetti.

Gli autori identificano un divario critico: le policy di pruning progettate per un compito non si trasferiscono necessariamente in modo efficace ad altri. Attraverso esperimenti di controllo "probe", il paper dimosta che:

Divergenza nel Ranking dei Criteri: I compiti di segmentazione e rilevamento classificano i criteri di pruning in modo differente. Ad esempio, il pruning basato sull'attenzione causa cali significativi di accuratezza nei primi strati per la classificazione, mentre la copertura spaziale è più stabile. Al contrario, il criterio che minimizza l'errore di rilevamento spesso performa male per la segmentazione.
Sensibilità alla Profondità: L'efficacia di un criterio di pruning varia significativamente tra le profondità della rete e differisce tra i compiti.
Conflitti di Recupero: I compiti di predizione densa richiedono il recupero delle feature nelle posizioni rimosse. Il paper trova che la segmentazione e il rilevamento preferiscono endpoint di recupero opposti: la segmentazione beneficia dal ripristino del full feature offset (differenza tra i token rimossi e quelli mantenuti), mentre il rilevamento performa meglio utilizzando solo il token "stand-in" mantenuto.

I metodi esistenti tipicamente applicano una singola policy di pruning statica o utilizzano selettori agnostici rispetto al compito, fallendo nel tenere conto di queste variazioni specifiche per compito e dipendenti dalla profondità.

Metodologia: Task-Adaptive Pruning (TAP)

Il paper propone il Task-Adaptive Pruning (TAP), un framework che utilizza i Task Registers per controllare dinamicamente il pruning dei token, l'allocazione del budget e il recupero delle feature.

Meccanismo Core:

Task Registers: Invece di usare i registri solo come archiviazione task-agnostic per gli artefatti delle feature (come nel lavoro precedente di Darcet et al., 2024), TAP introduce un registro iniziale apprendibile per ogni compito. Durante un passaggio in avanti, solo il registro associato al compito corrente è attivo.
Stato Evolutivo: Il registro attivo evolve insieme ai token dell'immagine attraverso il backbone. Il suo stato in qualsiasi strato condiziona tre decisioni chiave:
1. Selezione dei Token: Il registro calcola un punteggio per ogni patch candidata riutilizzando le proiezioni query e key esistenti del blocco. I token vengono classificati e un sottoinsieme viene mantenere per l'auto-attenzione.
2. Allocazione Esatta del Budget: Un readout lineare condiviso, condizionato dallo stato del registro, determina il numero esatto di interi token da rimuovere in ogni strato di pruning. Ciò assicura che il budget globale di rimozione sia soddisfatto esattamente senza richiedere una perdita di tasso (rate loss).
3. Scaling del Recupero: Per i compiti densi, il registro controlla un fattore di scala ( $\alpha$ ) che determina quanto del feature offset memorizzato (differenza tra il token rimosso e il suo stand-in) venga ripristinato al readout.

Varianti di Implementazione:

TAP-J (Joint Fine-tuning): Ogni compito ha un backbone, un task register e dei readout separati. Ciò consente una specializzazione completa.
TAP-F (Frozen Base): Un singolo backbone congelato è condiviso tra i compiti. I registri specifici per compito, gli aggiornamenti a basso rango (LoRA) e le teste sono addestrati congiuntamente. Ciò riduce l'archiviazione dei parametri pur mantenendo l'adattabilità.

Recupero delle Feature:
I token rimossi sono accoppiati deterministicamente con il token mantenuto con la maggiore similarità del coseno nello spazio delle chiavi (lo "stand-in"). La differenza delle feature (offset) viene memorizzata. Nei punti di readout densi, la ricostruzione è calcolata come $\hat{x}_i = x_{\text{stand-in}} + \alpha \cdot \delta_i$ , dove $\alpha$ è appreso per compito e profondità. I compiti di classificazione saltano questo passaggio di recupero.

Contributi Chiave

Framework Controllato: Gli autori sviluppano un framework che congela i checkpoint senza pruning e applica criteri di riduzione privi di parametri strato per strato per isolare l'impatto delle scelte di pruning tra classificazione, segmentazione e rilevamento. Ciò rivela il disallineamento nei ranking dei criteri e nelle preferenze di recupero.
Task Registers come Controller: Il paper ridefinisce il ruolo dei registri da archiviazione passiva a controller attivi e specifici per compito della computazione sparsa. Il registro coordina la selezione dei token, la distribuzione del budget di rimozione attraverso la profondità e lo scaling del recupero.
Pruning e Recupero Unificati: TAP tratta la rimozione dei token e il recupero delle feature come un'operazione unificata. Accoppia i token scartati agli stand-in e scala il ripristino degli offset in base al task register, mantenendo le feature ricostruite fuori dall'encoder per mantenere la sparsità.

Risultati Sperimentali

Gli autori valutano TAP su ImageNet-1K (classificazione), ADE20K (segmentazione) e COCO (rilevamento/segmentazione delle istanze) utilizzando un backbone ViT-B con $\rho = 0.5$ (50% di token mantenuti).

Performance:
- TAP-J ottiene i cali di accuratezza più piccoli tra i metodi completamente fine-tuned. Su ADE20K, raggiunge 47.0 mIoU (solo 0.2 di calo rispetto a no-pruning) con un throughput dell'encoder di 1.30×. Su COCO, raggiunge 53.7 box AP (0.3 di calo) con un throughput di 1.32×. Su ImageNet-1K, raggiunge 83.2 top-1 (0.4 di calo).
- TAP-F (backbone condiviso) segue leggermente TAP-J ma supera significativamente le regole statiche e altri baseline, ottenendo 45.8 mIoU su ADE20K e 52.5 box AP su COCO con un overhead di parametri minimo (~1.2M parametri specifici per compito).
Efficienza: TAP-J fornisce accelerazioni dell'encoder da 1.26× a 1.32× mantenendo un'accuratezza competitiva. L'overhead del controller (query del registro, readout del budget, bookkeeping) incide solo per circa lo 0.7% del tempo di esecuzione dell'encoder.
Memoria: TAP riduce la memoria di picco durante l'inferenza del 9–18% rispetto ai baseline densi, con la segmentazione che vede i guadagni maggiori grazie alle dimensioni ridotte delle attivazioni nei blocchi intermedi.
Studi di Ablazione: Rimuovere il task register (usando uno condiviso) o fissare l'allocazione/recupero a valori statici degrada significativamente le performance, confermando la necessità di un controllo adattivo per compito e profondità.

Significato e Rivendicazioni

Il paper sostiene che la riduzione dei token non può essere trattata come una singola regola trasferibile. Man mano che le rappresentazioni ViT evolvono da strutture locali a semantiche di alto livello, il valore dei token cambia sia con la profondità che con il compito specifico che legge il backbone.

La significanza di TAP risiede nella sua capacità di adattare la computazione sparsa alle esigenze specifiche di diversi compiti visivi utilizzando un backbone condiviso. Mettere al lavoro i task register evolutivi, TAP:

Risolve i conflitti nei ranking dei criteri di pruning tra i compiti.
Ottimizza la distribuzione del budget di rimozione attraverso la profondità della rete.
Condiziona il recupero delle feature alle specifiche necessità di readout densi o sparsi.

Gli autori posizionano TAP come un passo verso sistemi visivi unificati che possono servire molteplici compiti attraverso un backbone condiviso con un minimo storage specifico per compito, bilanciando le prestazioni predittive con il throughput dell'encoder in modo più efficace rispetto alle policy di pruning statiche. Il paper rimane modesto, notando che TAP-J segue il miglior metodo di sola classificazione (DiffRate) di 0.1 punti, probabilmente perché la classificazione si affida a un token di classe e salta il recupero denso, lasciando uno dei meccanismi di TAP sottoutilizzato per quel compito specifico. Il lavoro futuro suggerito esplora budget condizionati dall'immagine e molteplici stand-in.