Routage à Seuil Expert pour la Modélisation du Langage Autoregressif avec Allocation Dynamique de Calcul et Équilibrage de Charge
Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing
March 12, 2026
Auteurs: Hanchi Sun, Yixin Liu, Yonghui Wu, Lichao Sun
cs.AI
Résumé
Le Mixture-of-Experts à choix de token (TC-MoE) achemine chaque token vers un nombre fixe d'experts, limitant l'allocation dynamique des calculs et nécessitant des pertes auxiliaires pour maintenir l'équilibre de charge. Nous proposons un acheminement par seuil d'expert (ET), où chaque expert maintient un seuil de moyenne mobile exponentielle (EMA) estimé à partir de la distribution globale des tokens. Lors de l'entraînement et de l'inférence, chaque token est acheminé indépendamment vers un expert si son score dépasse le seuil de l'expert, permettant une allocation dynamique des calculs tout en atteignant l'équilibre de charge sans pertes auxiliaires. Ce mécanisme entièrement causal élimine la dépendance aux autres tokens du lot, le rendant bien adapté à la modélisation autoregressive du langage. Dans des expériences de pré-entraînement atteignant 2,4 milliards de paramètres sur FineWeb-Edu, ET obtient une perte d'entropie croisée inférieure de 0,067 par rapport au TC-MoE, ce qui équivaut à atteindre les mêmes performances avec 1,6 fois moins de tokens.
English
Token-choice Mixture-of-Experts (TC-MoE) routes each token to a fixed number of experts, limiting dynamic computation allocation and requiring auxiliary losses to maintain load balance. We propose Expert Threshold (ET) routing, where each expert maintains an exponential moving average (EMA) threshold estimated from the global token distribution. At both training and inference, each token is independently routed to an expert if its score exceeds the expert's threshold, enabling dynamic computation allocation while achieving load balance without auxiliary losses. This fully causal mechanism eliminates dependence on other tokens in the batch, making it well-suited for autoregressive language modeling. In pretraining experiments scaling to 2.4B parameters on FineWeb-Edu, ET achieves 0.067 lower cross-entropy loss than TC-MoE, equivalent to reaching the same performance with 1.6times fewer tokens.