Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancing
This addresses the problem of inefficient computation allocation and load balancing in autoregressive language modeling for AI researchers and practitioners, offering a novel method that improves performance with concrete gains.
The paper tackles the limitations of Token-choice Mixture-of-Experts routing by proposing Expert Threshold routing, which uses exponential moving average thresholds to enable dynamic computation allocation and load balancing without auxiliary losses, achieving a 0.067 lower cross-entropy loss and 1.6× fewer tokens for equivalent performance in pretraining experiments.
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.