🤖 machine learning

Sticky Routing: Training MoE Models for Memory-Efficient Inference

L'article propose StickyMoE, une perte de cohérence de routage différentiable qui entraîne les modèles Mixture-of-Experts à maintenir les assignations d'experts à travers des segments de jetons sémantiquement cohérents, réduisant ainsi de manière significative l'échange de poids gourmand en mémoire sur les appareils de bord avec une dégradation minimale de la perplexité.

Auteurs originaux : Ali Kayyam

Publié 2026-07-13

📖 6 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Ali Kayyam

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 possédez une bibliothèque géante de 64 bibliothécaires experts différents, mais que votre petit appareil local (comme un téléphone ou une enceinte connectée) n'a de la place sur son bureau que pour en garder 2 à la fois. Le reste est stocké dans un sous-sol lent et poussiéreux.

Dans un modèle classique de « Mixture-of-Experts » (MoE), chaque fois que l'IA lit un nouveau mot, elle lance une pièce pour décider quel bibliothécaire solliciter. Le problème ? Elle lance une pièce pour chaque mot. Ainsi, elle demande au Bibliothécaire A pour « pomme », puis au Bibliothécaire B pour « tarte », puis au Bibliothécaire C pour « croûte ». À chaque changement, elle doit descendre au sous-sol, mettre le Bibliothécaire B dehors et remonter le Bibliothécaire C sur le bureau. Ce « va-et-vient » est si lent que l'IA passe plus de temps à chercher les bibliothécaires qu'à lire l'histoire.

L'article introduit StickyMoE, une nouvelle façon d'entraîner ces modèles pour qu'ils arrêtent de lancer la pièce de manière aussi erratique. Au lieu de cela, le modèle apprend à être « collant » (sticky). Si elle demande au Bibliothécaire A pour « pomme », elle est encouragée à continuer de demander au Bibliothécaire A pour « tarte » et « croûte » aussi, tant que le sujet est cohérent.

La grande découverte : Entraînez-le correctement, ne le réparez pas plus tard

Les auteurs ont testé plusieurs approches pour résoudre ce problème de « va-et-vient » et voici ce qu'ils ont découvert :

La correction « Post-Hoc » a échoué : Ils ont tenté de prendre un modèle déjà entraîné et de simplement « affiner » (fine-tuning) le mécanisme de lancer de pièce (le routeur) pour qu'il soit moins aléatoire. Ils ont appelé cela ReMoE. Le résultat ? Cela n'a pas fonctionné. Le taux de changement était presque inchangé (une variation absolue de moins de 0,5 %). Les auteurs soutiennent que c'est parce que les bibliothécaires (les experts) ont déjà été entraînés pour être des spécialistes de mots spécifiques et isolés. On ne peut pas simplement dire au lanceur de pièce d'être plus « collant » si les bibliothécaires eux-mêmes ne sont pas prêts à gérer des sujets longs et continus. C'est comme essayer d'apprendre à un sprinteur à devenir un marathonien en changeant seulement sa posture de départ ; les muscles (les poids des experts) sont bâtis pour des accélérations courtes, pas pour de longues courses.
La correction « Hard » était trop rigide : Ils ont également testé une méthode qui forçait le modèle à s'en tenir à un bibliothécaire pour un nombre défini de mots en utilisant une règle stricte (un « Hard-Window »). Bien que cela ait réduit les changements, cela a considérablement dégradé les réponses de l'IA (augmentant la « perplexité » jusqu'à 6,8 % sur le modèle moyen). C'était comme forcer un bibliothécaire à lire un roman entier même quand le sujet change pour un genre complètement différent ; la qualité de l'histoire en souffrait.
La solution « Sticky » a fonctionné : La méthode gagnante, StickyMoE, ajoute une légère « pénalité de cohérence » lors de la toute première étape de l'entraînement. Elle dit au modèle : « Hé, si tu changes d'expert entre deux mots qui se suivent immédiatement, c'est un peu pénalisant. » Cela encourage le modèle à apprendre la « collosité » tout en apprenant la langue.

Les résultats : Le côté « collant » est meilleur

Lorsqu'ils ont testé cela sur des modèles petits et moyens en utilisant le jeu de données WikiText-2 (une collection d'articles Wikipédia), les résultats étaient clairs :

Moins de changements : En ajustant un seul curseur (un nombre appelé λ), ils ont pu réduire le nombre de fois où le modèle devait changer d'expert jusqu'à 59 %. Par exemple, sur le modèle moyen, le changement est passé de 0,71 (changement presque à chaque fois) à 0,29.
Meilleure vitesse (simulée) : Comme le modèle changeait moins souvent, ils ont simulé un cache (l'étagère du bureau) capable de contenir 2 experts. Le « Taux de succès du cache » (la fréquence à laquelle le bon bibliothécaire était déjà sur le bureau) est passé de 0,54 à 0,88. Cela signifie que le modèle aurait besoin de courir au sous-sol 3,92 fois de moins qu'auparavant.
Meilleure qualité : Étonnamment, sur le modèle moyen, utiliser une quantité modérée de « collosité » (λ = 0,05) a en fait rendu l'IA plus intelligente, abaissant sa « perplexité » (une mesure de confusion) de 4,1 %. Le modèle a appris que rester sur un sujet l'aidait à mieux comprendre le contexte.
Pas d'effondrement : Une grande crainte était que le modèle ne devienne « paresseux » et n'utilise qu'un seul bibliothécaire pour tout. Mais l'« Entropie d'utilisation » est restée élevée (au-dessus de 1,92 bits sur un maximum de 2,0), ce qui signifie que les 4 experts étaient toujours utilisés équitablement. Le modèle n'est pas devenu paresseux ; il est juste devenu organisé.

Ce qui n'a pas fonctionné (et pourquoi)

L'article écarte explicitement l'idée que l'on puisse régler ce problème après la construction du modèle. L'approche ReMoE (affinement d'un modèle existant) a montré son inefficacité car les « représentations d'experts » (le cerveau des bibliothécaires) avaient déjà été façonnées par un processus d'entraînement qui ne se souciait pas de la continuité thématique. On ne peut pas rétrofitter la localité ; il faut l'intégrer dès le départ.

Ils ont également noté que la toute première couche du modèle (Couche 0) est la plus difficile à rendre « collante ». Même avec les meilleurs réglages, le taux de changement y est resté supérieur à 0,49. C'est parce que la première couche voit les mots bruts avant qu'ils ne soient mélangés au contexte, donc « pomme » et « tarte » peuvent paraître très différents à ce stade.

L'essentiel

L'article suggère que si vous voulez qu'une IA fonctionne rapidement sur un appareil doté d'une mémoire limitée, vous ne devriez pas seulement essayer d'être astucieux avec des règles de cache ou essayer de réparer le modèle plus tard. Au lieu de cela, vous devez entraîner le modèle à être « collant » dès le premier jour. En ajoutant une simple pénalité pour les changements d'experts trop fréquents entre des mots adjacents, vous obtenez un modèle qui reste naturellement sur le sujet, nécessite moins de trajets vers le sous-plancher lent, et comprend en fait mieux l'histoire.

Les auteurs admettent que cela est basé sur des simulations et un entraînement sur des jeux de données spécifiques (WikiText-2) avec des modèles petits et moyens (jusqu'à 22M de paramètres). Ils suggèrent qu'à mesure que les modèles deviennent plus grands, cette « collosité » pourrait encore mieux fonctionner, mais ils ne l'ont pas encore testée sur les modèles massifs de plusieurs milliards de paramètres. Ils notent également que leur méthode ne sait pas actuellement quand une phrase se termine, ce qui pourrait la rendre « trop collante » à travers une rupture de paragraphe, un problème à résoudre pour les travaux futurs.

Résumé Technique : Routage Persistant pour une Inférence de MoE Efficace en Mémoire

1. Énoncé du Problème : Le Goulot d'Étranglement de l'Échange d'Experts

Les modèles de Mélange d'Experts (MoE) sont conçus pour faire évoluer les modèles de langage de grande taille (LLM) en n'activant qu'un sous-ensemble épars d'experts par jeton (token). Théoriquement, cela permet des comptes de paramètres massifs avec une faible computation par jeton, ce qui les rend attractifs pour le déploiement sur périphériques (edge) où seuls quelques experts doivent résider dans une mémoire rapide (VRAM/SRAM) à un instant donné.

Cependant, l'article identifie un mode de défaillance critique en pratique : l'incohérence de routage temporel. Les routeurs MoE standards prennent des décisions indépendantes pour chaque jeton, basées uniquement sur la représentation cachée de ce jeton. Par conséquent, des jetons consécutifs activent fréquemment des experts entièrement différents. Sur un matériel à mémoire contrainte, chaque changement nécessite l'éviction des poids des experts mis en cache et le chargement de nouveaux experts depuis un stockage lent (RAM CPU/Flash). Comme la bande passante PCIe/NVMe est de plusieurs ordres de grandeur inférieure à celle de la mémoire GPU, ces « échanges d'experts » fréquents dominent la latence d'inférence, rendant souvent les modèles MoE plus lents que des modèles denses de taille comparable.

Les solutions existantes sont jugées insuffisantes :

Les approches au niveau système (ex: heuristiques de mise en cache, pré-lecteurs) sont réactives ; elles exploitent la localité existante mais ne peuvent pas corriger des routeurs intrinsèquement incohérents.
Les approches post-hoc (ex: fine-tuning uniquement du routeur) échouent car les représentations des experts ont été façonnées par un régime d'entraînement qui n'a jamais encouragé la localité. Le routeur ne peut pas facilement réaligner les décisions dans un espace de représentation qui n'a pas été conçu pour cela.
Les refontes architecturales (ex: Oracle-MoE) nécessitent souvent des changements structurels importants du modèle ou du pipeline d'entraînement.

2. Méthodologie : StickyMoE

Les auteurs proposent StickyMoE, une approche lors de l'entraînement qui traite la localité temporelle de routage comme un objectif explicite. L'innovation centrale est une perte de cohérence de routage différentiable ( $L_{cons}$ ) qui pénalise les changements brusques d'assignation d'experts entre jetons adjacents.

La Perte de Cohérence de Routage

La fonction de perte minimise la distance $\ell_2$ entre les vecteurs de probabilité de la porte softmax ( $g_t$ ) de jetons consécutifs :
$L_{cons} = \frac{1}{T-1} \sum_{t=2}^{T} \| g_t - g_{t-1} \|_2^2$
Où $g_t \in \mathbb{R}^N$ est la sortie du routeur pour le jeton $t$ .

Cette perte est intégrée à l'objectif d'entraînement total aux côtés des pertes standard de l'entropie croisée ( $L_{CE}$ ) et d'équilibrage de charge ( $L_{bal}$ ) :
$L = L_{CE} + \lambda L_{cons} + \mu L_{bal}$

Mécanismes Clés

Co-adaptation : Contra irement aux méthodes post-hoc, StickyMoE entraîne simultanément le routeur et les représentations des experts. Minimiser $L_{cons}$ met à jour les poids du routeur ( $W_r$ ) et effectue une rétropropagation à travers les états cachés ( $h_t$ ), encourageant le modèle à apprendre des représentations naturellement cohérentes à travers des segments sémantiquement cohérents.
Variante Soft-Hard : Pour traiter la « dérive à longue portée » (où de petits changements étape par étape s'accumulent pour devenir un changement distant), les auteurs introduisent une contrainte d'ancre. La séquence est divisée en fenêtres ; le premier jeton d'une fenêtre définit une ancre, et les jetons suivants sont pénalisés pour leur déviation par rapport à cette ancre, avec une pénalité dont le poids augmente linéairement avec la distance par rapport au début de la fenêtre.
Compatibilité : La méthode ne nécessite aucun changement architectural et ajoute un surcoût de calcul négligeable. Elle est compatible avec n'importe quel MoE top- $k$ standard et les techniques d'équilibrage de charge existantes.

3. Contributions Clés

Perte de Cohérence de Routage : Un objectif d'entraînement simple et différentiable qui pénalise directement le changement d'expert entre jetons adjacents, ne nécessitant qu'un seul hyperparamètre ( $\lambda$ ).
Complémentarité avec l'Équilibrage de Charge : Les preuves empiriques montrent que la perte de cohérence et la perte d'équilibrage de charge sont complémentaires. La perte de cohérence façonne quand les experts sont utilisés sans provoquer d'« effondrement d'expert » (l'entropie d'utilisation reste élevée, proche du maximum théorique).
Variante Soft-Hard : Une combinaison de pénalités douces par étape et de contraintes d'ancrage par segment qui garantit une forte localité sans coût de qualité supplémentaire.
Évaluation Complète : Expériences contrôlées comparant StickyMoE au MoE vanilla, au fine-tuning post-hoc du routeur (ReMoE) et aux bases architecturales (Oracle-MoE) sur des modèles petits et moyens.

4. Résultats Expérimentaux

Les expériences ont été menées sur WikiText-2 avec des modèles MoE petits (~~8,8M de paramètres) et moyens (~~22M de paramètres).

Taux de Changement d'Expert (SR) : StickyMoE réduit le taux de changement d'expert jusqu'à 59 % (ex: de 0,71 à 0,29 sur le modèle moyen).
Taux de Succès du Cache (CHR) : La réduction des changements se traduit par une réduction de 3,92× des échecs de cache (le CHR passant d'environ 0,54 à 0,88) sur le modèle moyen avec un cache de 2 emplacements d'experts.
Perplexité (Qualité) :
- À des valeurs de $\lambda$ modérées (ex: 0,05), StickyMoE améliore la perplexité par rapport à la référence (ex: -4,1 % sur le modèle moyen), suggérant que la perte de cohérence agit comme un régularisateur bénéfique.
- Des valeurs de $\lambda$ agressives introduisent un compromis, mais la frontière de Pareto de StickyMoE domine strictement les autres méthodes.
Échec du Fine-Tuning Post-Hoc : La référence ReMoE (fine-tuning uniquement du routeur sur un modèle pré-entraîné) n'a pas réussi à réduire significativement le taux de changement (<0,5 % de changement), confirmant que la localité ne peut être greffée une fois les représentations fixées.
Analyse par Couche : La perte de cohérence est la plus efficace dans les couches profondes (L1–L3). La couche d'entrée (L0) reste résistante aux améliorations de localité, probablement parce que les plongements de jetons (embeddings) manquent de mélange contextuel.

5. Signification et Revendications

L'article soutient que le routage de la localité temporelle est un objectif d'entraînement légitime qui doit être instillé lors du pré-entraînement plutôt que d'être traité comme une réflexion ultérieure au niveau système.

Changement de Paradigme : Tout comme l'entraînement conscient de la quantification est devenu une norme, les auteurs postulent que l'« entraînement conscient de la localité » devrait être la norme pour les modèles MoE destinés au déploiement sur périphériques.
Résolution de la Cause Racine : StickyMoE s'attaque à la cause racine (décisions de routage incohérentes) plutôt qu'aux symptômes (échecs de cache). En co-entraînant les représentations et les décisions de routage, il évite le décalage de représentation qui frappe les méthodes post-hoc.
Praticité : La méthode est agnostique à l'architecture et ne nécessite aucun changement dans le pipeline d'inférence. Le routeur entraîné produit naturellement des distributions de porte cohérentes, permettant aux stratégies de mise en cache standard (LRU, LFU) de fonctionner de manière nettement plus efficace sans ingénierie supplémentaire.

Les auteurs concluent que bien que les optimisations au niveau système et les ajustements post-hoc restent des compléments utiles, ils ne peuvent remplacer un modèle qui n'a jamais été entraîné pour être respectueux du cache dès le départ. StickyMoE offre une voie pour réaliser l'efficacité mémoire théorique des architectures MoE sur les appareils de bord.