Cross-Entropy Guided Routing in Mixture-of-Experts Large Language Models

Yury Nahshan; Nati Daniel; Jacob Goldberger; Yoli Shavit

cs.AISep 29, 2026

Cross-Entropy Guided Routing in Mixture-of-Experts Large Language Models

Authors: Yury Nahshan, Nati Daniel, Jacob Goldberger, Yoli Shavit

Organizations: Bar-Ilan University, Ramat-Gan, Israel · NVIDIA, Israel

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.

Figures & tables

Model	Method	ARC-Challenge	OpenBookQA	SciQ †	MedMCQA
Granite 3.1	CE	$65.33\pm 2.01$	$71.67\pm 1.01$	$88.40\pm 0.53$	$39.60\pm 9.57$
	ACS-ENLL (Ours)	$67.40\pm 0.00$	$72.60\pm 0.53$	$89.27\pm 0.81$	$42.80\pm 0.40$
	ACS-IS (Ours)	$\mathbf{67.87\pm 0.76}$	$\mathbf{72.73\pm 0.95}$	$\mathbf{89.87\pm 0.70}$	$\mathbf{42.87\pm 0.76}$
	CE + affinity tuning	$67.53\pm 1.14$	$\mathbf{74.00\pm 0.72}$	$88.60\pm 0.40$	$41.73\pm 1.22$
	ACS-ENLL + affinity tuning (Ours)	$69.87\pm 1.33$	$73.53\pm 0.90$	$89.20\pm 0.92$	$\mathbf{42.13\pm 1.27}$
	ACS-IS + affinity tuning (Ours)	$\mathbf{70.60\pm 1.80}$	$73.93\pm 0.31$	$\mathbf{89.53\pm 0.42}$	$40.67\pm 1.45$

Table 1: Test answer-choice accuracy (%) for Granite 3.1 3B-A800M and OLMoE-1B-7B-SFT across four MCQA benchmarks. Entries are mean

\pm

sample standard deviation over seeds 42, 43, and 45. Within each backbone, rows form three parameter-matched comparison groups: CE versus ACS with frozen affinity, CE versus ACS with affinity tuning, and Dual Affinity versus TES. All entries use epoch-3 checkpoints except the OLMoE affinity-tuned group, which uses validation-NLL-selected checkpoints. Bold marks the highest mean within each backbone, benchmark, and comparison group. † SciQ is near saturation.

Model	Mechanism	Objective	ARC-C	OpenBookQA	SciQ	MedMCQA
Granite	TES, frozen	IS	$\mathbf{68.27}\,[+1.87]$	$72.47\,[-0.53]$	$\mathbf{89.20}\,[+0.87]$	$43.67\,[+0.00]$
		ENLL	$\mathbf{67.73}\,[+2.20]$	$\mathbf{72.73}\,[+0.47]$	$\mathbf{90.13}\,[+1.60]$	$43.40\,[-0.27]$
Granite	ACS, trainable	IS	$\mathbf{71.20}\,[+0.60]$	$\mathbf{74.53}\,[+0.60]$	$88.80\,[-0.73]$	$\mathbf{41.07}\,[+0.40]$
		ENLL	$\mathbf{70.47}\,[+0.60]$	$\mathbf{75.20}\,[+1.67]$	$89.00\,[-0.20]$	$41.47\,[-0.66]$
OLMoE	ACS, frozen	IS	$\mathbf{63.40}\,[+0.67]$	$69.80\,[-0.13]$	$\mathbf{89.67}\,[+0.80]$	$44.07\,[-0.13]$
		ENLL	$\mathbf{63.93}\,[+1.00]$	$69.20\,[-1.00]$	$\mathbf{90.07}\,[+1.34]$	$44.13\,[-0.40]$

Table 2: Effect of increasing the method-supervision coefficient to

\lambda_{m}=10^{-2}

. Each cell reports epoch-3 test accuracy (%) followed in brackets by its change from the corresponding

\lambda_{m}=10^{-3}

method result, in percentage points. Bold indicates higher mean accuracy than the corresponding

\lambda_{m}=10^{-3}

configuration.

Objective	Supervised MoE layers	ARC-Challenge	OpenBookQA	SciQ	MedMCQA
ACS-ENLL	Final layer	$62.93\pm 0.64$	$\mathbf{70.20\pm 0.35}$	$88.73\pm 0.23$	$\mathbf{44.53\pm 0.76}$
ACS-ENLL	All 16 layers	$\mathbf{63.13\pm 0.42}$	$69.80\pm 0.40$	$\mathbf{88.80\pm 0.40}$	$44.07\pm 0.81$
ACS-IS	Final layer	$62.73\pm 0.31$	$69.93\pm 0.46$	$88.87\pm 0.12$	$\mathbf{44.20\pm 0.53}$
ACS-IS	All 16 layers	$\mathbf{63.07\pm 0.50}$	$\mathbf{70.13\pm 0.12}$	$\mathbf{89.07\pm 0.31}$	$44.00\pm 0.35$

Table 3: OLMoE epoch-3 test accuracy (%; mean

\pm

sample standard deviation over seeds 42, 43, and 45) for final-layer and all-layer ACS with frozen native affinity and

\lambda_{m}=10^{-3}

. Bold marks the higher mean within each objective and benchmark.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Routing transform	Pointwise $r$	Group $r$	Group $\rho$	Test NLL	Accuracy (%)
$\log\widehat{e}$	0.1586	0.3975	0.4977	0.9723	67.27
$\log(1+\widehat{e})$	0.2486	0.4521	0.5637	0.9673	66.87

Appendix

Table 4: Error-to-routing transformation ablation on ARC-Challenge with Granite and IS supervision. Accuracy, NLL, and pointwise Pearson

r

are epoch-3 means over seeds 42–44. Group-level Pearson

r

and Spearman

\rho

are computed after averaging each layer–token-identity group across the same seeds. Bold marks the better observed value in each column and does not imply statistical significance.

Parameter	Value	Accuracy (%)	Choice NLL
Error-head learning rate	$10^{-3}$	$66.07$	$0.9811$
	$10^{-4}$	$66.60$	$0.9691$
	$10^{-5}$	$\mathbf{66.87}$	$\mathbf{0.9673}$
	$5\times 10^{-6}$	$66.67$	$\mathbf{0.9673}$
Attenuation scale $\gamma$	$0.5$	$65.87$	$0.9729$
	$1$	$\mathbf{66.87}$	$0.9673$

Appendix

Table 5: TES hyperparameter sensitivity on Granite ARC-Challenge with started-log IS supervision. Entries are epoch-3 means over seeds 42–44. The learning-rate panel fixes

\gamma=1

; the attenuation panel fixes the error-head learning rate at

10^{-5}

. Bold marks the best observed value within each panel and metric.

Setting	Granite 3.1 3B-A800M	OLMoE-1B-7B-SFT
Training examples	2,000 (ARC); 4,992 (others)	2,000 (ARC); 5,000 (others)
Validation / test examples	50 / 500	50 / 500
LoRA rank / $\alpha$ / dropout	8 / 8 / 0.05	8 / 8 / 0.05
LoRA targets	Q/K/V and routed-expert projections	Q/K/V/O and routed-expert projections
LoRA learning rate	$10^{-4}$	$2\times 10^{-5}$
Batch / accumulation	8 / 2	8 / 1

Appendix

Table 6: Backbone, optimization, and method configurations for the primary downstream evaluation. ARC denotes ARC-Challenge; “others” denotes OpenBookQA, SciQ, and MedMCQA.

Model	Affinity	Objective	$\lambda_{m}$	ARC-C	OpenBook QA	SciQ	MedMCQA
Granite	Trainable	IS	$10^{-3}$	$70.60$	$73.93$	$\mathbf{89.53}$	$40.67$
			$10^{-2}$	$\mathbf{71.20}$	$\mathbf{74.53}$	$88.80$	$\mathbf{41.07}$
		ENLL	$10^{-3}$	$69.87$	$73.53$	$\mathbf{89.20}$	$\mathbf{42.13}$
			$10^{-2}$	$\mathbf{70.47}$	$\mathbf{75.20}$	$89.00$	$41.47$
OLMoE	Frozen	IS	$10^{-3}$	$62.73$	$\mathbf{69.93}$	$88.87$	$\mathbf{44.20}$
			$10^{-2}$	$\mathbf{63.40}$	$69.80$	$\mathbf{89.67}$	$44.07$

Appendix

Table 7: ACS coefficient sensitivity (accuracy, %) under final-layer supervision. Bold marks the higher mean within each matched pair.

Objective	$10^{-4}$	$10^{-3}$	$10^{-2}$	$10^{-1}$
ACS–IS	$63.27\pm 0.12$	$62.73\pm 0.31$	$\mathbf{63.40\pm 0.40}$	$59.00\pm 0.40$
ACS–ENLL	$62.33\pm 0.46$	$62.93\pm 0.64$	$\mathbf{63.93\pm 0.76}$	$58.33\pm 1.21$

Appendix

Table 8: Complete frozen-affinity OLMoE ARC-Challenge coefficient screen.

Model	Method	ACS and affinity- training scope	ARC-Challenge	OpenBookQA	SciQ	MedMCQA
Granite 3.1	Router CE	Final 1	$\mathbf{67.53\pm 1.14}$	$\mathbf{74.00\pm 0.72}$	$88.60\pm 0.40$	$\mathbf{41.73\pm 1.22}$
		First 8	$65.60\pm 3.02$	$70.93\pm 2.61$	$\mathbf{89.07\pm 0.23}$	$40.87\pm 0.42$
		Last 8	$65.80\pm 1.51$	$71.07\pm 0.61$	$88.07\pm 0.61$	$37.73\pm 0.90$
		All 32	$61.60\pm 4.20$	$61.60\pm 3.40$	$86.27\pm 0.42$	$30.80\pm 3.34$
	ACS-ENLL	Final 1	$\mathbf{69.87\pm 1.33}$	$\mathbf{73.53\pm 0.90}$	$\mathbf{89.20\pm 0.92}$	$\mathbf{42.13\pm 1.27}$
		First 8	$65.00\pm 1.04$	$71.33\pm 1.10$	$88.87\pm 1.53$	$40.53\pm 2.00$

Appendix

Table 9: Answer-choice accuracy (%; mean

\pm

sample standard deviation over seeds 42, 43, and 45) under trainable native-affinity layer scopes. Granite values use the fixed epoch-3 checkpoint and

\lambda_{m}=10^{-3}

. OLMoE ACS values marked ∗ use the checkpoint selected by minimum validation NLL; ACS rows use

\lambda_{m}=10^{-2}

. Bold marks the best evaluated scope within each model, method, and dataset.

Method	Split	$\bar{C}_{\mathcal{D}}$	$N_{\mathrm{eff},\mathcal{D}}$	$r_{\mathcal{D}}$
Router CE	Train	$0.2614\pm 0.0279$	$3.86\pm 0.43$	$-0.152\pm 0.083$
	Test	$0.2589\pm 0.0319$	$3.90\pm 0.50$	$-0.161\pm 0.070$
ACS–IS	Train	$0.8068\pm 0.0395$	$1.24\pm 0.06$	$0.145\pm 0.075$
	Test	$0.8077\pm 0.0397$	$1.24\pm 0.06$	$0.143\pm 0.082$
ACS–ENLL	Train	$0.7813\pm 0.0381$	$1.28\pm 0.06$	$0.169\pm 0.087$
	Test	$0.7830\pm 0.0362$	$1.28\pm 0.06$	$0.158\pm 0.101$

Appendix

Table 10: Final-layer affinity-concentration diagnostics for Granite on ARC-Challenge with the native affinity router trained by task CE alone (Router CE) or jointly with ACS at

\lambda_{m}=10^{-2}

. Entries are epoch-3 means

\pm

sample standard deviations over seeds 42, 43, and 45. Each seed contains 14,679 train-token and 3,550 test-token observations.

Scope	Objective	Depth- scaled	$10\times$ scaled	Change
Last 8	IS	$69.27\pm 0.64$	$57.73\pm 2.32$	$-11.53$
	ENLL	$68.80\pm 1.06$	$56.13\pm 1.22$	$-12.67$
All 32	IS	$67.20\pm 3.47$	$48.60\pm 0.92$	$-18.60$
	ENLL	$69.00\pm 1.11$	$48.13\pm 0.50$	$-20.87$

Appendix

Table 11: ARC-Challenge excessive-scaling control (accuracy, %).

Dataset	Scope	Objective	Fixed $10^{-3}$	Depth- scaled	Change
ARC-C	Last 8	IS	$68.47\pm 1.55$	$\mathbf{69.27\pm 0.64}$	$+0.80$
		ENLL	$68.40\pm 1.31$	$\mathbf{68.80\pm 1.06}$	$+0.40$
	All 32	IS	$60.13\pm 10.34$	$\mathbf{67.20\pm 3.47}$	$+7.07$
		ENLL	$63.60\pm 0.92$	$\mathbf{69.00\pm 1.11}$	$+5.40$
OpenBookQA	Last 8	IS	$70.20\pm 0.35$	$\mathbf{70.27\pm 0.12}$	$+0.07$
		ENLL	$69.60\pm 3.67$	$\mathbf{70.20\pm 0.53}$	$+0.60$

Appendix

Table 12: Granite epoch-3 accuracy (%) under fixed and depth-scaled ACS coefficients. Bold marks the higher matched mean.

Explore similar work

When Are Experts Misrouted? Counterfactual Routing Analysis in Mixture-of-Experts Language Models
May 8, 2026Youngsik Yoon, Siwei Wang, Wei Chen +1Mixture-Of-Experts Experts
Department of Computer Science and Engineering, POSTECH, South Korea · Microsoft Research Asia, Beijing, China · Graduate School of Artificial Intelligence, POSTECH, South Korea
Beyond Geometric Complementarity: Coherent Overlap in Sparse Mixture-of-Experts Routing
Jul 30, 2026Huiyuan Tian, Bonan Xu, Shijian LiMixture-Of-Experts Complementarity
College of Computer Science and Technology, Zhejiang University, Hangzhou, China · Department of Aeronautical and Aviation Engineering, The Hong Kong Polytechnic University
TEXAS: Task-Expert-Aware Supervision for Downstream Mixture-of-Experts LLM Adaptation
Jul 31, 2026Guanzhi Deng, Haibo Wang, Kuan Wu +5Token-Level Supervision
City University of Hong Kong, Hong Kong, China · Carnegie Mellon University, Pittsburgh, USA · University of Waterloo, Ontario, Canada +1

May 8, 2026, cs.LGJ/K move · Enter open · S save

When Are Experts Misrouted? Counterfactual Routing Analysis in Mixture-of-Experts Language Models

Youngsik Yoon, Siwei Wang, Wei Chen, Jungseul Ok

Department of Computer Science and Engineering, POSTECH, South Korea · Microsoft Research Asia, Beijing, China · Graduate School of Artificial Intelligence, POSTECH, South Korea

Mixture-Of-Experts Experts Large Language Model Routing Deepseek Loss Function

Open paper AskPDF arXiv

Mixture-of-Experts (MoE) language models route each token to a small subset of experts, but whether the routes selected by a trained top-

k

router are good ones is rarely evaluated directly. Holding the model fixed, we compare each standard route against sampled equal-compute alternatives for the same token and score each by the next-token probability it assigns to the realized token in a verified reasoning trajectory. The result is sharply token-conditional: the standard router is well-aligned with route utility on confident tokens but uninformative on the fragile tokens that drive hard reasoning, where lower-loss equal-compute routes consistently exist inside the frozen model but are not selected. The same pattern holds across Qwen3-30B-A3B, GPT-OSS-20B, DeepSeek-V2-Lite, and OLMoE-1B-7B, and follows structurally from how standard top-

k

training evaluates routing decisions: the language modeling loss scores only the executed route, and load balancing depends only on aggregate routing statistics. A minimal router-only update to the final-layer router, leaving every expert and every other router frozen, is sufficient to shift pass@K on AIME 2024+2025 and HMMT 2025 for both Qwen3-30B-A3B and GPT-OSS-20B, suggesting that at least part of the failure reflects router-reachable misallocation rather than expert capacity alone.

Abstract

Explore similar work

When Are Experts Misrouted? Counterfactual Routing Analysis in Mixture-of-Experts Language Models

Beyond Geometric Complementarity: Coherent Overlap in Sparse Mixture-of-Experts Routing

TEXAS: Task-Expert-Aware Supervision for Downstream Mixture-of-Experts LLM Adaptation