💻 computer science

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

Cet article propose la méthode Task-Adaptive Pruning (TAP), qui utilise des registres spécifiques à la tâche pour classer dynamiquement les jetons et allouer les budgets d'élagage à travers différentes tâches de Vision Transformer, atteignant ainsi des améliorations significatives du débit tout en maintenant des performances compétitives en classification d'images, en segmentation sémantique et en détection d'objets.

Auteurs originaux : Hongsen Cao, Mona Jaber, Shanxin Yuan, Ahmed Sayed

Publié 2026-08-12

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Hongsen Cao, Mona Jaber, Shanxin Yuan, Ahmed Sayed

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous avez un cerveau de robot super intelligent qui regarde des images et essaie de comprendre ce qu'elles contiennent. Ce cerveau est construit comme une équipe géante de petits travailleurs, chacun tenant un morceau du puzzle (un « token »). Pour prendre une décision, chaque travailleur doit discuter avec tous les autres. Le problème ? Si l'image est grande, le nombre de discussions explose, et le cerveau est submergé, lent et gourmand en électricité. Les scientifiques ont essayé de résoudre cela en disant au cerveau d'ignorer les morceaux ennuyeux de l'image et de ne garder que les plus importants, un peu comme un professeur qui dirait à une classe de n'écouter que les élèves qui ont quelque chose d'intelligent à dire. C'est ce qu'on appelle l'« élagage de tokens » (token pruning). Mais voici le plus délicat : ce qui est considéré comme « intelligent » ou « important » change selon ce que le robot essaie de faire. Un morceau d'image crucial pour repérer un chat peut être inutile pour déterminer la couleur du ciel. Pendant longtemps, les scientifiques ont essayé d'utiliser la même « liste d'ignorance » pour chaque tâche, mais il s'avère qu'une solution unique convient rarement à tout.

Cet article, intitulé « Putting Registers to Work », s'attaque précisément à ce problème. Les auteurs, une équipe de l'Université Queen Mary de Londres, ont découvert que les règles pour décider quels morceaux d'image jeter sont totalement différentes selon les tâches. Ils ont découvert qu'une stratégie qui fonctionne très bien pour repérer des objets (comme trouver un chien dans un parc) pourrait en fait nuire à la capacité du robot à comprendre une scène (comme déterminer si ce chien est sur de l'herbe ou sur du pavé). Pour résoudre cela, ils ont inventé un nouveau système ingénieux appelé Task-Adaptive Pruning (TAP) (Élagage Adaptatif à la Tâche).

Imaginez le cerveau du robot comme une cuisine très occupée. Habituellement, le chef (la politique d'élagage) utilise la même recette pour décider quels ingrédients jeter, qu'il prépare une soupe ou un gâteau. Les auteurs ont réalisé que c'était une erreur. Au lieu de cela, ils ont doté la cuisine d'un ensemble de « Registres de Tâches » spéciaux — imaginez-les comme des cartes de recettes intelligentes et magiques qui changent selon le plat en cours de préparation. Quand le robot doit identifier un chat, il sort la « Carte Chat », qui indique au cerveau exactement quelles parties de l'image garder et lesquelles ignorer. Quand il doit cartographier une pièce entière, il échange la carte pour la « Carte Pièce ». Ces cartes ne décident pas seulement de ce qu'il faut jeter ; elles déterminent aussi quand le jeter pendant que le cerveau traite l'image, couche par couche.

Les chercheurs ont testé cela en « figeant » le cerveau d'un robot standard et en essayant différentes « règles d'ignorance » sans le réentraîner. Ils ont découvert trois grandes surprises. Premièrement, la meilleure façon de choisir les morceaux importants pour trouver des objets est l'exact opposé de la meilleure façon pour cartographier des scènes. Deuxièmement, le robot est très sensible à la manière dont il gère les toutes premières couches de traitement, surtout pour des tâches simples comme dire si une image est un chat ou un chien. Troisièmement, lorsqu'un robot jette un morceau de l'image, il doit se souvenir de ce à quoi il ressemblait pour combler le vide plus tard. Mais voici le hic : pour trouver des objets, il est préférable de faire comme si le morceau n'avait jamais existé et d'utiliser un voisin proche pour combler le vide. Pour cartographier des scènes, il est préférable de se souvenir de la différence exacte et de la recoller plus tard.

Pour corriger tout cela, le TAP utilise une seule carte de registre « active » pour le travail en cours. Cette carte traverse le cerveau, observant l'image et prenant trois décisions à la volée :

Sélection : Quels tokens garder et lesquels supprimer.
Budget : Combien de tokens supprimer à chaque étape (en supprimer quelques-uns tôt, plus tard, ou vice versa).
Récupération : Quelle quantité de détail « oublié » ramener lorsque le robot a besoin de dessiner une carte détaillée.

Les résultats sont impressionnants. En utilisant ces cartes spécifiques aux tâches, le robot est devenu beaucoup plus rapide sans perdre son intelligence. Sur un test standard pour la détection d'objets (COCO), il est devenu 1,32 fois plus rapide tout en conservant presque la même précision (ne perdant que 0,3 point). Sur un test de cartographie de scènes (ADE20K), il est devenu 1,30 fois plus rapide avec une baisse minime de qualité. Même pour la classification d'images simple, il est resté compétitif.

Les auteurs suggèrent que cette approche est un pas vers un véritable cerveau de robot unifié — un cerveau capable de passer de la détection d'une voiture à la peinture d'une rue ou à la lecture d'un panneau, tout en utilisant le même moteur central mais avec différents « styles de pensée » activés à la volée. Ils n'ont pas seulement supposé que cela fonctionnerait ; ils l'ont prouvé en menant des milliers d'expériences, en figeant le cerveau pour tester chaque règle individuellement, et en montrant que leur système de « registre » surpasse systématiquement les anciennes méthodes « taille unique ». C'est un rappel que dans le monde de l'IA, parfois, la chose la plus intelligente que l'on puisse faire est de savoir exactement quand ignorer le bruit, et que la définition du « bruit » dépend entièrement de ce que l'on cherche à voir.

Résumé Technique : Mettre les registres au travail : Registres de tâches pour l'élagage des jetons dans les Vision Transformers

Énoncé du Problème

Les Vision Transformers (ViT) présentent une mise à l'échelle computationnelle quadratique avec le nombre de jetons (tokens), ce qui motive l'élagage des jetons pour réduire les coûts d'inférence. Bien que les critères d'élagage (par exemple, les scores d'attention, la couverture spatiale) soient souvent développés pour la classification d'images, les modèles de base ViT pré-entraînés sont de plus en plus adaptés à des tâches de prédiction dense comme la segmentation sémantique et la détection d'objets.

Les auteurs identent une lacune critique : les politiques d'élagage conçues pour une tâche ne se transfèrent pas nécessairement de manière efficace à d'autres. À travers des expériences de « sondage » contrôlées, l'article démontre que :

Divergence du classement des critères : Les tâches de segmentation et de détection classent les critères d'élagage différemment. Par exemple, l'élagage basé sur l'attention provoque des chutes de précision significatives dans les couches précoces pour la classification, alors que la couverture spatiale est plus stable. Inversement, le critère qui minimise l'erreur de détection performe souvent mal pour la segmentation.
Sensibilité à la profondeur : L'efficacité d'un critère d'élagage varie considérablement selon la profondeur du réseau et diffère entre les tâches.
Conflits de récupération : Les tâches de prédiction dense nécessitent une reconstruction des caractéristiques aux positions supprimées. L'article constate que la segmentation et la détection préfèrent des points de terminaison de récupération opposés : la segmentation bénéficie de la restauration de l'offset complet des caractéristiques (différence entre les jetons supprimés et conservés), tandis que la détection performe mieux en utilisant uniquement le jeton « substitut » conservé.

Les méthodes existantes appliquent généralement une politique d'élagage unique et statique ou utilisent des sélecteurs agnostiques à la tâche, échouant ainsi à prendre en compte ces variations spécifiques à la tâche et dépendantes de la profondeur.

Méthodologie : Élagage Adaptatif à la Tâche (TAP)

L'article propose le Task-Adaptive Pruning (TAP), un cadre qui utilise des Registres de Tâches pour contrôler dynamiquement l'élagage des jetons, l'allocation du budget et la récupération des caractéristiques.

Mécanisme Central :

Registres de Tâches : Au lieu d'utiliser les registres simplement comme un stockage agnostique de l'artéfact de caractéristiques (comme dans les travaux précédents de Darcet et al., 2024), TAP introduit un registre initial apprenable par tâche. Pendant une passe avant (forward pass), seul le registre associé à la tâche actuelle est actif.
État Évolutif : Le registre actif évolue aux côtés des jetons d'image à travers le backbone. Son état à n'importe quelle couche conditionne trois décisions clés :
1. Sélection de Jetons : Le registre calcule un score pour chaque patch candidat en réutilisant les projections de requête (query) et de clé (key) existantes du bloc. Les jetons sont classés, et un sous-ensemble est conservé pour l'auto-attention.
2. Allocation Exacte du Budget : Une lecture linéaire partagée, conditionnée par l'état du registre, détermine le nombre entier exact de jetons à supprimer à chaque couche d'élagage. Cela garantit que le budget global de suppression est exactement respecté sans nécessiter de perte de taux (rate loss).
3. Mise à l'Échelle de la Récupération : Pour les tâches denses, le registre contrôle un facteur d'échelle ( $\alpha$ ) qui détermine quelle part de l'offset de caractéristique stocké (différence entre le jeton supprimé et son substitut) est restaurée lors de la lecture (readout).

Variantes d'Implémentation :

TAP-J (Fine-tuning Conjoint) : Chaque tâche possède un backbone, un registre de tâche et des lectures (readouts) distincts. Cela permet une spécialisation complète.
TAP-F (Base Gelée) : Un seul backbone gelé est partagé entre les tâches. Les registres spécifiques à la tâche, les mises à jour de bas rang (LoRA) et les têtes sont entraînés conjointement. Cela réduit le stockage des paramètres tout en maintenant l'adaptabilité.

Récupération de Caractéristiques :
Les jetons supprimés sont appariés de manière déterministe au jeton conservé ayant la plus haute similitude cosinus dans l'espace des clés (le « substitut »). La différence de caractéristique (offset) est stockée. Aux points de lecture dense, la reconstruction est calculée comme $\hat{x}_i = x_{\text{substitut}} + \alpha \cdot \delta_i$ , où $\alpha$ est appris par tâche et par profondeur. Les tâches de classification sautent cette étape de récupération.

Contributions Clés

Cadre Contrôlé : Les auteurs développent un cadre qui gèle les points de contrôle sans élagage et applique des critères de réduction sans paramètres couche par couche pour isoler l'impact des choix d'élagage à travers la classification, la segmentation et la détection. Cela révèle le désalignement dans les classements des critères et les préférences de récupération.
Registres de Tâches comme Contrôleurs : L'article redéfinit le rôle des registres, passant d'un stockage passif à des contrôleurs actifs et spécifiques à la tâche du calcul parcimonieux. Le registre coordonne la sélection de jetons, la distribution du budget de suppression à travers la profondeur et la mise à l'échelle de la récupération.
Élagage et Récupération Unifiés : TAP traite la suppression de jetons et la récupération de caractéristiques comme une opération unifiée. Il apparie les jetons écartés à leurs substituts et met à l'échelle la restauration des offsets en fonction du registre de la tâche, gardant les caractéristiques reconstruites hors de l'encodeur pour maintenir la parcimonie.

Résultats Expérimentaux

Les auteurs évaluent TAP sur ImageNet-1K (classification), ADE20K (segmentation) et COCO (détection/segmentation d'instance) en utilisant un backbone ViT-B avec $\rho = 0.5$ (50 % de jetons conservés).

Performance :
- TAP-J obtient les baisses de précision les plus faibles parmi les méthodes de fine-tuning complet. Sur ADE20K, il atteint 47.0 mIoU (seulement 0.2 de baisse par rapport au sans-élagage) avec un débit de l'encodeur de 1.30×. Sur COCO, il atteint 53.7 box AP (0.3 de baisse) avec un débit de 1.32×. Sur ImageNet-1K, il atteint 83.2 top-1 (0.4 de baisse).
- TAP-F (backbone partagé) suit légèrement TAP-J mais surpasse significativement les règles statiques et les autres lignes de base, atteignant 45.8 mIoU sur ADE20K et 52.5 box AP sur COCO avec un surcoût de paramètres minime (~1.2M de paramètres spécifiques à la tâche).
Efficacité : TAP-J fournit des accélérations de l'encodeur de 1.26× à 1.32× tout en maintenant une précision compétitive. Le surcoût du contrôleur (requête de registre, lecture du budget, comptabilité) ne représente qu'environ 0.7 % du temps d'exécution de l'encodeur.
Mémoire : TAP réduit la mémoire d'inférence de pointe de 9 à 18 % par rapport aux lignes de base denses, la segmentation bénéficiant des gains les plus importants grâce à la réduction des tailles d'activation dans les blocs intermédiaires.
Études d'Ablation : Supprimer le registre de tâche (en utilisant un registre partagé) ou fixer l'allocation/la récupération à des valeurs statiques dégrade considérablement les performances, confirmant la nécessité d'un contrôle adaptatif à la tâche et à la profondeur.

Signification et Revendications

L'article affirme que la réduction de jetons ne peut pas être traitée comme une règle unique transférable. À mesure que les représentations ViT évoluent des structures locales vers des sémantiques de haut niveau, la valeur des jetons change à la fois avec la profondeur et avec la tâche spécifique qui lit le backbone.

La signification de TAP réside dans sa capacité à adapter le calcul parcimonieux aux besoins spécifiques de différentes tâches visuelles en utilisant un backbone partagé. En mettant les registres de tâches évolutifs au travail, TAP :

Résout les conflits dans les classements des critères d'élagage entre les tâches.
Optimise la distribution du budget de suppression à travers la profondeur du réseau.
Conditionne la récupération des caractéristiques aux exigences spécifiques des lectures denses versus parsemées.

Les auteurs positionnent TAP comme une étape vers des systèmes de vision unifiés capables de servir plusieurs tâches via un backbone partagé avec un stockage spécifique à la tâche minimal, équilibrant plus efficacement la performance prédictive avec le débit de l'encodeur que les politiques d'élagage statiques. L'article reste modeste, notant que TAP-J suit de peu la meilleure méthode de classification uniquement (DiffRate) de 0.1 point, probablement parce que la classification repose sur un jeton de classe et saute la récupération dense, laissant l'un des mécanismes de TAP sous-utilisé pour cette tâche spécifique. Des travaux futurs sont suggérés pour explorer les budgets conditionnés par l'image et les multiples substituts.