Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

Résumé Technique : Reg4Pru – Régularisation par Routage Aléatoire de Tokens pour l'Élagage de Tokens

Énoncé du Problème
Les Vision Transformers (ViT) sont devenus fondamentaux pour les modèles de vision modernes en raison de leurs capacités de mise à l'échelle et de généralisation. Cependant, leur dépendance aux mécanismes d'auto-attention entraîne une complexité computationnelle qui croît de manière quadratique ($O(N^2)$) avec le nombre de tokens $N$. Ce goulot d'étranglement est particulièrement aigu dans la segmentation d'imagerie médicale, où des entrées à haute résolution sont nécessaires pour préserver les détails diagnostiques critiques.

Bien que les méthodes d'élagage de tokens (suppression des tokens redondants) et les stratégies de réactivation (réintroduction des tokens élagués lors de l'inférence pour des prédictions denses pixel par pixel) offrent une voie vers l'efficacité, les auteurs observent une dégradation significative des performances dans les couches profondes. Plus précisément, lorsque les tokens élagués sont réactivés lors de l'inférence, ils présentent des représentations radicalement différentes de leur état lors de l'entraînement. Ce décalage de distribution entre l'entraînement et le test entraîne une instabilité et de mauvaises performances de prédiction dense, provoquant de manière contre-intuitive une baisse de performance à mesure que la profondeur du modèle augmente.

Méthodologie
Le papier propose Reg4Pru, une technique de régularisation lors de l'entraînement conçue pour atténuer l'instabilité associée à l'élagage et à la réactivation de tokens. Le mécanisme central repose sur le routage aléatoire de tokens pendant la phase d'entraînement pour simuler les effets de l'élagage sans supprimer définitivement les tokens.

1. Routage Aléatoire de Tokens : Lors de la passe directe (forward pass) de l'entraînement, le modèle sélectionne une plage de routage aléatoire $[l, n]$ entre les blocs transformeurs intermédiaires. Un sous-ensemble aléatoire de tokens ($\tau_r$) est choisi pour sauter le calcul des blocs de $l$ à $n$, tandis que les autres tokens ($\tau_k$) sont traités normalement.
2. Flux de Gradient et Représentation : Les tokens routés contournent des blocs spécifiques, coupant efficacement leurs gradients par rapport à ces couches. Cependant, en randomisant à la fois le sous-ensemble de tokens routés et la plage de routage $[l, n]$ à chaque étape, différents tokens contribuent aux gradients de différents blocs au fil des itérations. Cela force le modèle à apprendre des représentations invariantes à la profondeur, garantissant que les tokens peuvent être réactivés à n'importe quelle profondeur lors de l'inférence sans souffrir de décalage de représentation.
3. Stratégie d'Inférence : À l'inférence, un réseau de politique distinct prédit quels tokens élaguer en se basant sur un calendrier de ratio de maintien hiérarchique. Les tokens élagués sont temporairement retirés pour économiser le calcul et sont réactivés au dernier bloc pour générer des masques de segmentation denses.
4. Fonctions de Perte : L'objectif d'entraînement combine une perte standard basée sur le ratio (pour imposer des budgets de tokens) et, optionnellement, une perte de politique informée ($L_{pol}$) qui utilise les masques de segmentation de vérité terrain pour guider les décisions d'élagage. Le mécanisme de routage agit comme un régularisateur aux côtés de ces pertes.

Contributions Clés

• Analyse de l'Instabilité liée à la Profondeur : Les auteurs identifient et caractérisent une instabilité spécifique dans l'élagage de tokens pour la prédiction dense, où les performances déclinent dans les couches plus profondes en raison de la réactivation de tokens qui n'ont pas été exposés à la pleine profondeur du réseau pendant l'entraînement.
• Cadre Reg4Pru : Ils introduisent une stratégie de routage de tokens lors de l'entraînement avec des bornes aléatoires. Contra�à les méthodes de routage précédentes qui utilisent des bornes fixes ou des politiques de routage apprises, Reg4Pru randomise la plage de routage pour approximer la distribution des caractéristiques à travers toutes les profondeurs d'élagage, agissant ainsi comme un régularisateur.
• Validation Empirique : La méthode est validée sur le jeu de données de segmentation des vaisseaux sanguins rétiniens FIVES, démontrant que le routage aléatoire stabilise significativement les performances par rapport au routage à bornes fixes ou sans routage.

Résultats
Les expériences ont été menées sur le jeu de données FIVES (images de 2048$\times$2048, redimensionnées en quadrants pour les tests de débit) en utilisant une architecture EoMT initialisée par DINOv2-small.

• Gains de Performance : Comparé à un modèle élagué entraîné sans routage, Reg4Pru a amélioré la précision moyenne (AP) de 46 % en valeur absolue (passant de 25,60 % à 71,67 % pour la configuration $L_{ratio}$).
• Efficacité : La configuration Reg4Pru a atteint une accélération relative de 29 % du temps d'exécution (wall-clock time) par rapport à la ligne de base non élaguée tout en maintenant une haute qualité de segmentation.
• Stabilité : L'analyse par profondeur (Figure 5) a montré que si les modèles élagués sans routage souffraient d'une chute brutale de l'AP dans les blocs plus profonds, Reg4Pru maintenait des performances stables. Le routage à bornes fixes a atténué une partie de l'instabilité mais n'a pas réussi à empêcher la chute des performances après la fin de la plage fixe, confirmant la nécessité de la randomisation.
• Comparaison : Reg4Pru a surpassé les méthodes de fusion de tokens (ToMe) en termes d'AP lors de la comparaison de configurations de débit similaires, particulièrement lorsque ToMe était appliqué à la fois sur les couches MHSA et MLP.

Signification et Revendications
Le papier affirme que Reg4Pru est un régularisateur précieux pour les stratégies de réduction de tokens dans les tâches de prédiction dense. Les auteurs soutiennent que l'instabilité observée dans l'élagage est le résultat d'un décalage de distribution entraînement-test concernant les représentations de tokens à différentes profondeurs. En randomisant la plage de routage pendant l'entraînement, Reg4Pru force le modèle à apprendre des représentations qui sont robustes au fait d'être sautées à diverses profondeurs, comblant ainsi l'écart entre l'entraînement et l'inférence.

Les auteurs restent modestes quant à la généralisabilité de leurs résultats, notant que les expériences ont été menées sur un seul jeu de données médical à haute résolution. Ils suggèrent que, bien que la méthode soit efficace pour la segmentation, d'autres évaluations sont nécessaires pour tester son applicabilité à d'autres tâches de prédiction dense (par exemple, l'estimation de la profondeur) et son potentiel en tant que mécanisme de régularisation général lorsqu'elle est combinée avec des techniques telles que le dropout ou la profondeur stochastique. Ils reconnaissent également que, bien qu'ils aient exploré des politiques informées, les politiques basées sur le ratio produisent actuellement des performances supérieures dans leur configuration spécifique.