Cross entropy improves routing in large language expert models
Cross-Entropy Guided Routing in Mixture-of-Experts Large Language Models
Artificial Intelligence
Summary
Large language models sometimes use many smaller expert parts to handle different pieces of text. The way the model decides which expert should deal with each word is called routing, but current methods don’t directly consider how well each expert will perform. The authors propose new ways to guide this routing so that the experts chosen better match what will reduce errors. Their methods make the model more accurate on tests without needing more computing power. This could help build smarter language tools that learn to pick the best parts to focus on for each word.
What this means in practice
- •For machine learning engineers: Improve large language model accuracy by better routing tokens to expert subnetworks based on predicted error alignment.
- •For ai service operators: Maintain efficient sparse model execution while increasing question-answering accuracy in deployed language models.
Authors
Yury Nahshan, Nati Daniel, Jacob Goldberger, Yoli Shavit
Abstract
Sparse mixture-of-experts (MoE) large language models scale model capacity by routing each token to a small subset of experts. Their routers are regularized with load balancing terms and learn affinity scores through the language-model objective. However, these objectives do not provide direct alignment between routing affinities and token-level error. We introduce token-error supervision for sparse routing in two forms. The first form predicts an error score per expert. The affinity-weighted aggregate of these scores is aligned to the next-token cross-entropy loss, while the individual scores attenuate affinity before top-$K$ selection. The second directly aligns the router's affinities to the model's objective without requiring an additional head or inference-time modification. Both formulations use the Itakura--Saito divergence or an exponential negative log-likelihood for aligning affinities and token errors. Across two sparse MoE backbones and four multiple-choice question-answering benchmarks, we evaluate both supervision mechanisms. On Granite, our method improves accuracy by approximately 2.3 percentage points on average over a parameter-matched routing baseline. With stronger supervision, the gain on ARC-Challenge reaches 2.94 points. Both mechanisms preserve the native sparse execution budget and aggregation policy. Our code is available in the supplementary materials.