🤖 machine learning

Adaptive Computation Depth via Learned Token Routing in Transformers

L'article présente l'attention sélective de jetons (TSA), un mécanisme de commutation léger, différentiable de bout en bout, qui permet aux transformateurs de sauter dynamiquement les calculs de couches redondants pour chaque jeton en fonction de la difficulté contextuelle, réalisant ainsi des gains d'efficacité significatifs avec une perte de qualité minimale et sans nécessiter de régularisation explicite de la profondeur.

Auteurs originaux : Ahmed Abdelmuniem Abdalla Mohammed

Publié 2026-05-08

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Ahmed Abdelmuniem Abdalla Mohammed

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 dirigez une usine où chaque produit unique qui descend le convoyeur reçoit exactement le même traitement. Qu'il s'agisse d'un objet simple et parfait ou d'un objet complexe et défectueux, chaque article passe exactement le même laps de temps à chaque station. C'est ainsi que fonctionnent actuellement les modèles d'IA standard (Transformers) : ils traitent chaque mot d'une phrase à travers exactement le même nombre de « couches de réflexion », indépendamment de la facilité ou de la difficulté à comprendre ce mot.

Ce papier présente un nouveau système appelé Token-Selective Attention (TSA). Considérez le TSA comme l'installation d'un gardien intelligent et automatique à chaque station de l'usine.

Le Problème : L'Usine « Taille Unique »

Dans un modèle d'IA standard, si vous lui demandez d'écrire « To be, or not to be », le modèle traite le mot « To » et le mot « question » avec la même quantité d'énergie cérébrale.

Les mots faciles (comme « the » ou « is ») sont comme des objets simples. Ils ne nécessitent pas beaucoup de traitement.
Les mots difficiles (comme des concepts complexes ou des noms rares) sont comme des objets défectueux. Ils ont besoin de temps et d'attention supplémentaires pour être corrigés.

Actuellement, l'usine gaspille de l'énergie en traitant les objets faciles avec la même intensité que les objets difficiles.

La Solution : Le Gardien Intelligent (TSA)

Les auteurs ont ajouté un minuscule et léger « gardien » (un petit réseau de neurones) entre chaque couche de l'IA. Ce gardien examine chaque mot (token) individuellement et se demande : « Ce mot a-t-il besoin de passer à la station de traitement suivante, ou peut-il l'éviter ? »

La Décision : Le gardien ne dit pas simplement « Oui » ou « Non ». Il attribue un score de probabilité (comme un variateur de lumière). Si un mot est facile, le gardien peut dire : « Tu peux sauter l'étape suivante », ou « Tu n'as besoin que de faire la moitié du travail ». Si un mot est difficile, il dit : « Va jusqu'au bout ».
La Magie : La meilleure partie est que le gardien s'enseigne lui-même. Il n'a pas besoin qu'un humain lui dise quels mots sont difficiles. Il apprend uniquement en essayant de minimiser les erreurs. Si sauter une étape pour un mot spécifique provoque une erreur, le gardien apprend à maintenir ce mot actif la prochaine fois. Si le fait de sauter l'étape fonctionne bien, il apprend à économiser de l'énergie.

Fonctionnement en Pratique

Le papier a testé cela sur deux aspects principaux :

Des Puzzles Logiques Simples : Lorsque l'IA devait copier une liste de nombres, le gardien a réalisé : « C'est facile », et a sauté environ 65 % du travail. Lorsqu'il devait trier des nombres (ce qui est plus difficile), il n'a sauté qu'environ 27 % du travail. Il a naturellement compris que les tâches plus difficiles nécessitent plus d'étapes.
L'Écriture d'Histoires : Lorsqu'on lui a demandé d'écrire comme Shakespeare ou de résumer des articles Wikipédia, l'IA a économisé 14 % à 23 % de sa puissance de calcul.
- Elle a appris que les signes de ponctuation, les espaces et les mots courants sont faciles et peuvent être traités rapidement.
- Elle a appris que les mots de contenu uniques nécessitent le traitement complet de l'« usine ».

Les Résultats : Plus Rapide et Plus Intelligent

Aucune Perte de Qualité : L'IA écrit aussi bien que le modèle standard, mais elle utilise considérablement moins d'énergie (puissance de calcul).
Mieux que la « Sortie Anticipée » : D'autres méthodes tentent d'arrêter toute la phrase tôt si l'IA se sent « confiante ». Le TSA est plus intelligent ; il empêche les mots individuels de passer par des étapes inutiles tout en laissant les mots difficiles continuer. Le papier a constaté que le TSA produisait de meilleurs résultats que ces anciennes méthodes lorsqu'elles étaient configurées pour économiser la même quantité d'énergie.
Vitesse Réelle : Lorsque les chercheurs ont réellement exécuté le code sur un ordinateur, ils ont observé une accélération réelle (environ 2,3 % plus rapide) car l'ordinateur n'avait littéralement pas à faire les calculs mathématiques pour les mots sautés.

La Conclusion

Ce papier présente un moyen de rendre les modèles d'IA « paresseux » de la manière la plus intelligente possible. Au lieu de forcer chaque mot à faire la même quantité de travail, le TSA permet à l'IA de décider, à la volée, quels mots sont faciles et lesquels ont besoin d'aide supplémentaire. C'est comme avoir une usine où les produits faciles filent sur le convoyeur, tandis que les difficiles reçoivent le traitement VIP complet, économisant du temps et de l'énergie sans sacrifier la qualité.

Résumé Technique : Profondeur de Calcul Adaptative via Routage de Tokens Apprenti dans les Transformers

Énoncé du Problème

Les architectures Transformer standard appliquent un nombre fixe de couches à chaque token d'une séquence, indépendamment de la difficulté contextuelle ou de la densité d'information de ce token. Cette conception privilégie la simplicité architecturale à l'adaptabilité par token, entraînant une inefficacité computationnelle significative. En pratique, les tokens courants dans des contextes prévisibles nécessitent beaucoup moins de traitement que les tokens rares dans des constructions nouvelles, pourtant ils reçoivent les mêmes ressources de calcul lors de chaque passage avant. Cette inefficacité est particulièrement coûteuse à l'échelle de l'inférence, où le coût dominant est le passage complet avant à travers toutes les couches pour tous les tokens. Les solutions existantes, telles que le Temps de Calcul Adaptatif (ACT) pour les RNN, les Transformers Universels, les Mélanges de Profondeurs (MoD) et les mécanismes de Sortie Anticipée (Early Exit), tentent de résoudre ce problème mais reposent souvent sur une sélection rigide, des budgets accumulés ou des seuils de confiance a posteriori qui peuvent ne pas être entièrement différentiables ou optimaux.

Méthodologie : Attention Sélective de Tokens (TSA)

L'article propose l'Attention Sélective de Tokens (TSA), un mécanisme qui introduit une porte apprise, par token, sur les mises à jour résiduelles entre les blocs transformateurs consécutifs.

Architecture et Routage

La TSA insère un routeur léger $r_l$ après chaque bloc $f_l$ (pour $l = 0, \dots, L-2$ ) dans un transformateur en mode décodeur uniquement avec normalisation pré-activation. Le bloc initial $f_0$ s'exécute toujours inconditionnellement.

Conception du Routeur : Chaque routeur est un Perceptron Multicouche (MLP) à deux couches avec une sortie sigmoïde, produisant une probabilité d'arrêt continue $p_l \in (0, 1)$ pour chaque token. La dimension cachée est $d/4$ (arrondie à 16), ajoutant environ 1,7 % de surcharge paramétrique au modèle de base.
Mise à Jour Gérée : Les mises à jour résiduelles pour le bloc suivant $f_{l+1}$ $f_{l + 1}$ sont mises à l'échelle par $(1 - p_l)$ $(1 - p_{l})$ .
- Si $p_l = 0$ , la mise à jour est standard.
- Si $p_l = 1$ , l'état reste inchangé, sautant efficacement le bloc.
- L'interpolation est lisse, garantissant que le mécanisme est entièrement différentiable sans nécessiter d'estimateurs en ligne droite, d'échantillonnage Gumbel ou d'apprentissage par renforcement.

Entraînement et Régularisation

Pendant l'entraînement, toutes les couches s'exécutent pleinement ; la porte met à l'échelle les mises à jour résiduelles mais ne saute pas le calcul. La perte totale est définie comme $L = L_{tâche} + L_{profondeur}$ .

Régularisation de Profondeur : Un terme $L_{profondeur} = \lambda \cdot \frac{1}{L-1} \sum (1 - p_l)$ est ajouté pour encourager doucement un arrêt précoce.
Comportement Émergent : Il est notable que l'article démontre que le mécanisme de routage apprend à sauter des opérations même lorsque $\lambda = 0$ (aucune pression explicite sur la profondeur), entraîné uniquement par le gradient de la perte de tâche.

Inférence

À l'inférence, le mécanisme bascule en mode « épars ». Les tokens avec $p_l > 0.5$ sont exclus du Réseau de Neurones Feed-Forward (FFN) via des opérations de collecte/dissémination, tandis que l'attention reste dense pour préserver la sémantique exacte des clés-valeurs. Cela génère des économies réelles de FLOPs et des accélérations en temps réel.

Contributions Clés

Routage de Tokens Différentiable : Une porte douce continue simple sur les mises à jour résiduelles permettant un routage par token et par couche sans modifications architecturales du Transformer de base.
Émergence Pilotée par le Gradient : Preuve qu'un routage proportionnel à la difficulté émerge naturellement du seul gradient de la perte de tâche. Même sans régularisation de profondeur ( $\lambda=0$ ), le routeur apprend à sauter environ 20 % des opérations token-couche.
Compromis Efficacité-Qualité : Validation sur la modélisation de langage au niveau des caractères montrant des économies de 14 à 23 % des opérations token-couche (TLOps) avec une perte de qualité inférieure à 0,5 %.
Supériorité sur la Sortie Anticipée : À efficacité équivalente, la TSA atteint une perte de validation inférieure de 0,71 % par rapport aux bases de référence de Sortie Anticipée, suggérant que des décisions de routage apprises de bout en bout sont supérieures au seuillage de confiance a posteriori.
Accélération Réelle : Démonstration d'une accélération réelle en temps (jusqu'à 2,3 %) sur du matériel standard (Apple M1 Pro) via une inférence éparses.

Résultats Expérimentaux

Tâches Synthétiques

Sur les tâches de copie et de tri (séquences de longueur 10) :

Tâche de Copie : Le routeur a appris que la tâche était une application identité, sautant 54,9 % des TLOps ( $\alpha = 0.341$ ) sans perte de précision.
Tâche de Tri : Le routeur a alloué plus de calcul, sautant 22,5 % des TLOps ( $\alpha = 0.730$ ), reflétant la difficulté accrue.

Modélisation de Langage au Niveau des Caractères

Tiny-Shakespeare : La TSA a économisé 22,8 % des TLOps ( $\alpha = 0.726$ ) avec une augmentation de la perte de validation de seulement 0,006 nats (+0,4 %).
enwik8 : Sur un corpus plus diversifié (Wikipedia), la TSA a économisé 13,9 % des TLOps ( $\alpha = 0.833$ ) sans dégradation de qualité (perte de validation marginalement meilleure).
Robustesse : Le mécanisme est robuste au paramètre d'hyper-régularisation $\lambda$ sur deux ordres de grandeur ( $[0, 0.1]$ ). À $\lambda=0.05$ , il a réalisé des économies de TLOps de 50,4 % avec un coût de qualité négligeable.

Comparaison avec la Sortie Anticipée

À une fraction active équivalente ( $\alpha \approx 0.726$ ), la TSA a atteint une perte de validation inférieure de 0,71 % par rapport à une base de référence de Sortie Anticipée. Cela indique que l'apprentissage de décisions de routage de bout en bout produit de meilleures décisions que l'application d'un seuil de confiance au moment de l'inférence.

Débit en Temps Réel

Soft-TSA (Entraînement) : Ajout d'une surcharge d'environ 1 % indépendamment de $\alpha$ .
Sparse-TSA (Inférence) : A atteint une accélération de 2,3 % sur Tiny-Shakespeare ( $\alpha=0.726$ ) et une performance à l'équilibre sur enwik8 ( $\alpha=0.833$ ). Les accélérations nécessitaient une taille de lot de $\ge 64$ pour surmonter la surcharge de synchronisation CPU-GPU.

Importance et Revendications

L'article affirme que la TSA offre une solution minimale et architecturalement simple à l'inefficacité des Transformers à profondeur fixe. Son importance principale réside dans la démonstration que la profondeur de calcul adaptative peut être apprise directement à partir de l'objectif de tâche sans supervision explicite concernant la difficulté ou des signaux complexes d'apprentissage par renforcement.

Les auteurs soulignent que le routeur agit comme une « porte de bruit apprise », atténuant les mises à jour lorsque les contributions résiduelles sont redondantes. Cela permet au modèle d'allouer dynamiquement le calcul en fonction de la difficulté du token, réalisant des gains d'efficacité significatifs (jusqu'à 55 % sur les tâches synthétiques, 14–23 % sur la modélisation de langage) avec des compromis de qualité minimes. Le travail suggère que le routage différentiable de bout en bout est une stratégie plus efficace pour la profondeur adaptative que la sélection discrète ou le seuillage a posteriori, ouvrant la voie à une inférence plus efficace sur du matériel standard.

L'article reste modeste concernant l'échelle, notant que toutes les expériences ont été menées sur des modèles d'environ 5 à 6 millions de paramètres, et que le comportement de mise à l'échelle vers des modèles plus grands (10M–100M+) reste une question ouverte. De plus, les accélérations actuelles dépendent de la taille du lot, et la parcimonie de l'attention n'a pas encore été mise en œuvre.