Temporally Extended Mixture-of-Experts Models

1. Le Problème : L'Inefficacité de la Commutation dans les MoE Actuels

Les modèles de langage (LLM) modernes utilisent de plus en plus des architectures Mélange d'Experts (MoE). Ces modèles activent un sous-ensemble épars d'experts pour chaque token, permettant de scaler la capacité du modèle sans augmenter le coût de calcul par token. Cependant, une limitation critique émerge lorsque le nombre total d'experts dépasse la mémoire GPU disponible :

• Commutation excessive : Dans les architectures actuelles, l'ensemble des experts actifs change à presque chaque token (taux de commutation > 50-90 %).
• Coûts de latence et de mémoire : Cette "churn" (agitation) constante rend inefficaces les optimisations de mémoire comme le offloading (déchargement vers la RAM) et le prefetching (préchargement). Chaque changement d'expert nécessite un chargement de poids depuis la mémoire hôte, introduisant une latence importante et réduisant le débit.
• Opportunités manquées : L'absence de continuité temporelle empêche des stratégies d'inférence et d'entraînement plus efficaces, car le système ne peut pas prédire quels experts seront nécessaires pour les tokens suivants.

2. Méthodologie : Le Cadre des "Options" et le Critique d'Options

Les auteurs proposent de reformuler le problème de routage des experts comme un problème de prise de décision temporellement étendue, en s'inspirant du cadre des Options en apprentissage par renforcement (RL).

A. Formulation en Semi-MDP

Au lieu de choisir un expert à chaque token (étape primitive), le modèle sélectionne une "option" (un masque d'experts) qui persiste sur plusieurs tokens.

• Option ($\omega$) : Un ensemble de $\hat{k}$ experts autorisés pour une couche donnée.
• Terminaison : Une décision binaire ($d_t$) détermine si l'option actuelle doit être maintenue ou si une nouvelle option doit être chargée.
• Coût de délibération ($\eta$) : Un coût explicite est ajouté à l'objectif pour pénaliser les changements d'option, forçant le modèle à ne basculer que lorsque le gain de qualité attendu justifie le coût de chargement.

B. Architecture du Contrôleur

Pour chaque couche du modèle, un contrôleur léger est ajouté :

1. Entrées : Les états cachés du LLM ($h_t$) et l'option actuelle ($\omega_{t-1}$).
2. Encodage : Un encodeur DeepSets transforme le masque d'experts actuel en une représentation vectorielle invariante par permutation.
3. Têtes de décision :
   • Tête de terminaison ($\beta$) : Prédit la probabilité de changer d'expert.
   • Tête de sélection : Choisit le nouvel ensemble d'experts via une distribution de Plackett-Luce (échantillonnage sans remise).
   • Têtes de valeur ($V_\Omega, Q_\Omega$) : Estiment la valeur de l'état et de l'option pour guider l'apprentissage.

C. Entraînement : Option-Critic avec Distillation

Le modèle est entraîné en utilisant l'algorithme Option-Critic adapté avec un coût de délibération :

• Récompense : Utilisation de la distillation sur politique (on-policy distillation). Le modèle étudiant (avec contrôleur) est entraîné pour imiter la distribution de probabilité d'un modèle enseignant (le MoE de base pré-entraîné, figé). La récompense est basée sur la divergence KL inverse entre les deux distributions.
• Mélange : Pour éviter l'effondrement de la diversité (mode collapse), l'échantillonnage des tokens se fait via un mélange entre la distribution de l'étudiant et celle de l'enseignant.
• Mise à jour : Les gradients sont calculés pour optimiser simultanément la politique intra-option (les poids du LLM et du routeur) et la politique de sélection/terminaison des options.

3. Contributions Clés

1. Nouveau Paradigme de Routage : Introduction du concept de MoE "temporellement étendu", où le routage des experts est traité comme une décision hiérarchique plutôt qu'atomique.
2. Formalisation Mathématique : Cast du chargement dynamique des experts comme un Processus de Décision de Markov Semi-Markovien (s-MDP) avec des coûts de commutation explicites.
3. Preuve de Concept sur Modèle Pré-entraîné : Démonstration qu'un modèle existant (gpt-oss-20b) peut être converti en MoE temporellement étendu avec un entraînement léger (LoRA + contrôleur), sans nécessiter de réentraînement massif.
4. Trade-off Contrôlable : La méthode permet de régler le compromis entre la fréquence de commutation (efficacité mémoire) et la précision du modèle via le paramètre de coût de délibération $\eta$.

4. Résultats Expérimentaux

Les expériences ont été menées sur gpt-oss-20b (32 experts par couche, top-4 actifs) avec des budgets d'experts $\hat{k}$ de 8 et 16.

• Réduction drastique du taux de commutation :
  • Le taux de commutation passe de > 50 % (presque à chaque token) à < 5 % (et même < 1 % selon les configurations).
  • Les experts restent actifs sur de longues séquences de tokens, créant une continuité temporelle forte.
• Préservation de la Performance :
  • Malgré la réduction massive des commutations, le modèle conserve jusqu'à 90 % de la précision du modèle de base sur des benchmarks complexes comme MATH (raisonnement mathématique), MMLU (connaissances générales) et MMMLU (multilingue).
  • Exemple : Sur MATH, avec $\hat{k}=16$ et $\eta=0.02$, le modèle atteint 64,0 % de précision contre 71,5 % pour le modèle de base, tandis que les méthodes de pruning statique (baselines) chutent à ~53 %.
• Comparaison avec les Baselines :
  • Les méthodes de pruning statique (sélection par fréquence, minimisation de la perte de reconstruction, Wanda) souffrent d'une dégradation sévère des performances (souvent < 50 % de précision) car elles ne peuvent pas s'adapter dynamiquement au contexte.
  • Le contrôleur appris surpasse toutes les méthodes statiques, prouvant que la dynamique temporelle est cruciale.
• Stabilité : Le modèle évite le phénomène de "catastrophic repetition" (répétition de tokens) souvent observé dans les MoE contraints, grâce à la distillation sur politique.

5. Signification et Impact

Ce travail ouvre une voie prometteuse pour l'avenir des modèles à très grande échelle :

• Inférence Économe en Mémoire : En réduisant la fréquence des changements d'experts, il devient possible de servir des modèles massifs sur du matériel limité (GPU avec peu de VRAM) en ne gardant en mémoire que les experts actifs pour une longue durée, réduisant ainsi la latence de chargement.
• Apprentissage Continu : L'architecture permet d'ajouter de nouveaux experts à la volée sans augmenter le coût d'inférence par token, facilitant l'adaptation continue à de nouveaux domaines.
• Entraînement par Chunks : La continuité temporelle permet de découper les séquences d'entraînement en blocs où seuls les experts actifs sont chargés, réduisant la mémoire GPU nécessaire à l'entraînement.

En conclusion, les auteurs démontrent que l'intégration de structures temporelles via le cadre des options n'est pas seulement théoriquement élégante, mais constitue une solution pratique et efficace pour surmonter les goulots d'étranglement mémoire des futurs modèles MoE massifs.