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

Technische Samenvatting: Registers aan het werk zetten: Task Registers voor Token Pruning in Vision Transformers

Probleemstelling

Vision Transformers (ViT's) vertonen een kwadratische computationele schaling met het aantal tokens, wat token pruning motiveert om de inferentiekosten te verlagen. Hoewel pruning-criteria (bijv. aandachtsscores, ruimtelijke dekking) vaak worden ontwikkend voor beeldclassificatie, worden pretrained ViT-backbones steeds vaker aangepast aan taken met dichte voorspellingen (dense prediction) zoals semantische segmentatie en objectdetectie.

De auteurs identificeren een kritieke kloof: pruning-policies die ontworpen zijn voor de ene taak, vertalen zich niet noodzakelijkerwijs effectief naar andere taken. Via gecontroleerde "probe"-experimenten toont het artikel aan dat:

1. Divergentie in de rangschikking van criteria: Segmentatie- en detectietaken rangschikken pruning-criteria anders. Zo veroorzaakt aandacht-gebaseerde pruning bijvoorbeeld aanzienlijke nauwkeurigheidsverliezen in vroege lagen voor classificatie, terwijl ruimtelijke dekking stabieler is. Omgekeerd presteert het criterium dat de detectiefout minimaliseert vaak slecht voor segmentatie.
2. Dieptegevoeligheid: De effectiviteit van een pruning-criterium varieert aanzienlijk over de netwerkdiepte en verschilt tussen taken.
3. Herstelconflicten: Taken met dichte voorspellingen vereisen feature-reconstructie op verwijderde posities. Het artikel stelt vast dat segmentatie en detectie tegenovergestelde herstel-eindpunten verkiezen: segmentatie profiteert van het herstellen van de volledige feature offset (het verschil tussen de verwijderde en de behouden tokens), terwijl detectie beter presteert door alleen de behouden "stand-in" token te gebruiken.

Bestaande methoden passen doorgaans een enkele, statische pruning-policy toe of gebruiken taak-agnostische selectors, waardoor ze er niet in slagen rekening te houden met deze taakspecifieke en diepteafhankelijke variaties.

Methodologie: Task-Adaptive Pruning (TAP)

Het artikel stelt Task-Adaptive Pruning (TAP) voor, een framework dat Task Registers gebruikt om token pruning, budgetallocatie en feature-herstel dynamisch te controleren.

Kernmechanisme:

• Task Registers: In plaats van registers enkel te gebruiken als taak-agnostische opslag voor feature-artefacten (zoals in eerder werk van Darcet et al., 2024), introduceert TAP één leerbaar initieel register per taak. Tijdens een forward pass is alleen het register dat bij de huidige taak hoort actief.
• Evoluerende Staat: Het actieve register evolueert samen met de beeldtokens door de backbone. De staat ervan op elk niveau conditioneert drie belangrijke beslissingen:
  1. Token Selectie: Het register berekent een score voor elke kandidaat-patch door de bestaande query- en key-projecties van de block te hergebruiken. Tokens worden gerangschikt en een subset wordt behouden voor self-attention.
  2. Exacte Budgetallocatie: Een gedeelde lineaire readout, geconditioneerd op de registerstaat, bepaalt het exacte integer aantal tokens dat op elk pruning-niveau moet worden verwijderd. Dit zorgt ervoor dat het globale verwijderingsbudget exact wordt gehaald zonder dat een rate loss nodig is.
  3. Herstel Schaling: Voor taken met dichte voorspellingen controleert het register een schaleringsfactor ($\alpha$) die bepaalt hoeveel van de opgeslagen feature offset (het verschil tussen de verwijderde token en de stand-in) wordt hersteld bij de readout.

Implementatievarianten:

• TAP-J (Joint Fine-tuning): Elke taak heeft een aparte backbone, een taakregister en readouts. Dit maakt volledige specialisatie mogelijk.
• TAP-F (Frozen Base): Een enkele bevroren backbone wordt gedeeld over taken. Taakspecifieke registers, low-rank (LoRA) updates en heads worden gezamenlijk getraind. Dit vermindert de opslag van parameters terwijl de aanpasbaarheid behouden blijft.

Feature Herstel:
Verwijderde tokens worden deterministisch gekoppeld aan de behouden token met de hoogste cosine-gelijkenis in de key-ruimte (de "stand-in"). Het feature-verschil (offset) wordt opgeslagen. Bij dichte readout-punten wordt de reconstructie berekend als $\hat{x}_i = x_{\text{stand-in}} + \alpha \cdot \delta_i$, waarbij $\alpha$ per taak en diepte wordt geleerd. Classificatietaken slaan deze herstelstap over.

Belangrijkste Bijdragen

1. Gecontroleerd Framework: De auteurs ontwikkelen een framework dat no-pruning checkpoints bevriest en parameter-vrije reductie-criteria laag voor laag toepast om de impact van pruning-keuzes te isoleren over classificatie, segmentatie en detectie. Dit onthult de discrepantie in de rangschikking van criteria en herstelvoorkeuren.
2. Task Registers als Controllers: Het artikel herdefinieert de rol van registers van passieve opslag naar actieve, taakspecifieke controllers van sparse computation. Het register coördineert token selectie, de exacte distributie van het budget over de diepte en de herstel-schaling.
3. Verenigde Pruning en Herstel: TAP behandelt token-verwijdering en feature-herstel als een verenigde operatie. Het koppelt gediscardeerde tokens aan stand-ins en schaalt de herstel van offsets op basis van de taakregister, waardoor de gereconstrueerde features buiten de encoder blijven om de sparsity te behouden.

Experimentele Resultaten

De auteurs evalueren TAP op ImageNet-1K (classificatie), ADE20K (segmentatie) en COCO (detectie/instance segmentatie) met een ViT-B backbone met $\rho = 0.5$ (50% tokens behouden).

• Prestaties:
  • TAP-J bereikt de kleinste nauwkeurigheidsdalingen onder de volledig fine-tuned methoden. Op ADE20K bereikt het 47.0 mIoU (slechts 0.2 daling van no-pruning) met 1.30× encoder throughput. Op COCO bereikt het 53.7 box AP (0.3 daling) met 1.32× throughput. Op ImageNet-1K bereikt het 83.2 top-1 (0.4 daling).
  • TAP-F (gedeelde backbone) volgt TAP-J licht, maar presteert aanzienlijk beter dan statische regels en andere baselines, met 45.8 mIoU op ADE20K en 52.5 box AP op COCO met minimale parameter overhead (~1.2M taakspecifieke parameters).
• Efficiëntie: TAP-J biedt encoder versnellingen van 1.26× tot 1.32× terwijl het concurrerende nauwkeurigheid behoudt. De controller overhead (register query, budget readout, bookkeeping) beslaat slechts ~0.7% van de encoder runtime.
• Geheugen: TAP vermindert de piek inferentiegeheugens met 9–18% vergeleken met dichte baselines, waarbij segmentatie de grootste winst boekt door de verminderde activatiegroottes in tussenliggende blokken.
• Ablatieonderzoeken: Het verwijderen van de taakregister (door een gedeelde te gebruiken) of het vastzetten van de allocatie/herstel naar statische waarden verslechtert de prestaties aanzienlijk, wat de noodzaak van taak- en diepte-adaptieve controle bevestigt.

Betekenis en Claims

Het artikel claimt dat token reductie niet kan worden behandeld als een enkele overdraagbare regel. Naarmate ViT-representaties evolueren van lokale structuren naar hoogwaardige semantiek, verandert de waarde van tokens met zowel de diepte als de specifieke taak die de backbone leest.

De betekenis van TAP ligt in het vermogen om sparse computation aan te passen aan de specifieke behoeften van verschillende visuele taken met behulp van een gedeelde backbone. Door de evoluerende taakregisters aan het werk te zetten, lost TAP:

• Conflicten in de rangschikking van pruning-criteria tussen taken op.
• Optimaliseert de distributie van het verwijderingsbudget over de netwerkdiepte.
• Conditioneert feature-herstel op de specifieke vereisten van dichte versus sparse readouts.

De auteurs positioneren TAP als een stap naar verenigde visuele systemen die meerdere taken kunnen dienen via een gedeelde backbone met minimale taakspecifieke opslag, waarbij de balans tussen voorspellende prestaties en encoder throughput effectiever wordt gehouden dan met statische pruning-policies. Het artikel blijft bescheiden en merkt op dat TAP-J de beste classificatie-alleen methode (DiffRate) met 0.1 punt achterblijft, waarschijnlijk omdat classificatie vertrouwt op een class token en de dichte herstelstap overslaat, waardoor een van de mechanismen van TAP voor die specifieke taak onderbenut blijft. Toekomstig werk suggereert het verkennen van beeld-geconditioneerde budgetten en meerdere stand-ins.