Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing

2026-03-12

AI Summary

This work addresses the limitations of traditional Token-choice Mixture-of-Experts (TC-MoE) in autoregressive language modeling, particularly its rigid computation allocation and reliance on auxiliary losses to maintain load balancing. The authors propose an Expert Threshold (ET) routing mechanism, wherein each expert maintains an exponential moving average (EMA) threshold derived from the global token distribution, and tokens are independently routed based on whether their routing scores exceed this threshold. This approach achieves fully causal routing without auxiliary losses for the first time, naturally preserving load balance while enabling dynamic computation allocation. Evaluated on a 2.4B-parameter model pretrained on FineWeb-Edu, ET routing reduces cross-entropy loss by 0.067 compared to TC-MoE, equivalent to achieving the same performance with 1.6× fewer training tokens.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Planning, Routing, and Scheduling: Planning with Language Models
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: Trust and reliance of crowd workers and data experts on GenAI
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

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.6$\times$ fewer tokens.

Problem

• Mixture-of-Experts
• autoregressive language modeling
• dynamic computation allocation
• load balancing
• token routing

Innovation

• Expert Threshold Routing
• Dynamic Computation Allocation
• Load Balancing
• Mixture-of-Experts
• Autoregressive Language Modeling