PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yury Nahshan, Nati Daniel, Jacob Goldberger, Yoli Shavit

arXiv:2609.37751v1cs.AI

TL;DR

Standard MoE affinities are not directly supervised against token-level error. This paper uses realized next-token loss to supervise expert-error predictions or affinity concentration, and reports gains over corresponding controls across two backbones and four benchmarks. The methods retain the native sparse execution budget and aggregation policy.

  • Problem

    Standard MoE affinities are learned through the language-model objective and routing regularizers, but are not explicitly supervised against token-level error.

    Source: §1 · p. 1, §3.1 · p. 4

  • Method

    TES supervises affinity-weighted expert-error predictions and uses them to attenuate affinity before top-K selection, while ACS directly supervises native affinity concentration without an additional head.

    Source: §1 · pp. 1–2

  • Results

    Across four MCQA benchmarks, TES-IS improves average accuracy by 2.3 points over the parameter-matched Dual Affinity baseline on Granite.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Both mechanisms yield accuracy gains over their corresponding controls while preserving the native sparse execution budget and aggregation policy.

    Source: §1 · p. 2, §5 · p. 8

    The evaluation covers two backbones and four MCQA benchmarks with fixed data splits and three seeds, and does not measure end-to-end latency or TES projection memory overhead.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

The introduction motivates direct token-error supervision because standard MoE affinities are not explicitly aligned with token-level loss. It presents two supervision mechanisms and reports accuracy gains across two backbones and four question-answering benchmarks.

1 Introduction

  • TES predicts expert errors whose affinity-weighted aggregate is supervised against realized token loss, while ACS supervises normalized selected-affinity concentration without an extra head or inference-time transformation.Both use IS divergence, with ENLL explored as a simplified alternative; TES also attenuates affinity before top-K selection using predicted errors.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • Both TES objectives train aggregate error predictions toward the same expected token-error target, but they affect language-model optimization differently.IS leaves the CE update unchanged, whereas ENLL gives relatively greater gradient weight to lower-loss tokens.Source: §1 · pp. 1–2
  • 2.3 points on Granite and 0.5 points on OLMoE: TES-IS improves average accuracy over the parameter-matched Dual Affinity baseline across four benchmarks.ACS-IS improves Granite average accuracy by 2.1 points over CE fine-tuning with frozen native affinity; stronger supervision raises the Granite ARC-Challenge gain to 2.94 points.Source: §1 · p. 2
  • The evaluation compares TES and ACS with standard CE fine-tuning and a parameter-matched routing-head control across two sparse MoE backbones and four MCQA benchmarks.Both methods preserve the native top-K execution budget and architecture-specific aggregation policy.Source: §1 · p. 2

2 Related Work

Related work spans alternative sparse-allocation rules, router supervision, token-difficulty-based computation allocation, uncertainty-aware routing, and predictive-error signals. The paper distinguishes its approach by supervising fixed-budget sparse routing with realized next-token loss.

2 Related Work

  • Expert Choice changes sparse allocation by letting experts select tokens, while ReMoE replaces discontinuous top-K selection with differentiable ReLU routing.The paper frames these as alternatives to learned-affinity top-K routing.Source: §2 · p. 3
  • ERC couples router embeddings to expert capabilities and Expert Divergence uses domain labels, whereas TES and ACS directly supervise routing with realized token loss.TES refines selection using expert-error scores; ACS shapes affinity without an added prediction head.Source: §2 · p. 3
  • Prior token-difficulty methods allocate computation by changing model capacity or experts per token, while this work uses realized next-token loss to supervise sparse routing.Examples include confidence-based expert activation, agreement-derived difficulty labels, and token-dependent expert budgets.Source: §2 · p. 3
  • Uncertainty-aware routers use probabilistic routing signals or adapt capacity and regularization, while this method uses deterministic fixed-budget top-K routing supervised by next-token loss.The passage distinguishes the approach from routers that sample distributions or allocate variable budgets.Source: §2 · p. 3
  • MoGU uses expert-specific predictive error to control mixture weights, but TES predicts routed token loss with native affinities and attenuates experts with larger predicted error.TES scores are not Gaussian likelihoods or calibrated predictive variances.Source: §2 · p. 3

3 Method

The method supplements affinity-based sparse routing with token-error supervision, using both expert-indexed error predictions and direct supervision of affinity concentration.

3.1 Sparse MoE Routing

  • Native affinity selects and combines experts but receives no explicit token-error supervision, so the method complements it with a separately supervised error signal.Source: §3.1 · p. 4

3.2 Token-Error Head and Affinity Attenuation

  • TES predicts a positive error score per expert, aggregates scores using native affinities, and subtracts a scaled log-error term from affinity before top-K selection.Source: §3.2 · p. 4

3.3 Token-Error Supervision

  • IS and ENLL align aggregate error predictions with realized token loss; across contexts, their shared target is expected loss for a fixed target-token identity, not counterfactual expert losses.Source: §3.3 · pp. 4–5

3.4 Affinity-Concentration Supervision

  • ACS substitutes normalized native affinity for the error prediction; when the IS floor is inactive, its direct gradient favors sharper affinity above the token loss and flatter affinity below it.ACS adds no error head or inference-time routing transformation.Source: §3.4 · p. 5

3.5 Overall Training Objective and Computational Scope

  • Training averages the active supervision objective across supervised tokens and MoE layers alongside the language-model loss and architecture-native router regularizer.For MCQA adaptation, answer-choice CE is primary while next-token CE remains the auxiliary target; TES adds a projection, whereas ACS adds no parameters.Source: §3.5 · pp. 5–6

4 Experiments and Results

The experiments compare ACS and TES with matched baselines across two MoE backbones and four multiple-choice benchmarks, then examine predicted-error association and supervision depth. Results include task-dependent accuracy gains and show that stronger supervision can help several tasks, but not uniformly.

4.1 Experimental Setting

  • The evaluation uses fixed train, validation, and test splits, answer-text candidate scoring, 500-example test sets, three seeds, and primarily final-layer supervision.ACS is compared with CE; TES is compared with a parameter-matched Dual Affinity head, and native expert-execution budgets are preserved.Source: §4.1 · p. 6

4.2 Downstream Accuracy

  • On Granite, ACS-IS improves over frozen-affinity CE on ARC-Challenge, OpenBookQA, and SciQ, while TES-IS improves ARC-Challenge and OpenBookQA on both backbones.ACS-IS also beats affinity-tuned CE on Granite ARC-Challenge and SciQ; TES objectives fall below Dual Affinity on OLMoE SciQ.Source: §4.2 · pp. 6–7, Table 1 · p. 7
  • 68.27%: Granite TES-IS reaches this ARC-Challenge accuracy with λm = 10−2, 2.94 points above Dual Affinity.The larger coefficient does not help every task: Granite TES-IS loses 0.53 points on OpenBookQA, so the study retains a common coefficient.Source: §4.2 · p. 7, Table 2 · p. 7

4.3 Error-Prediction Analysis

  • Test-set predicted and observed token errors correlate at 0.452/0.564 for IS and 0.449/0.562 for ENLL, with NLL rising across eight of nine decile transitions.The standardized plots assess association and out-of-sample shift, not absolute NLL calibration.Source: §4.3 · pp. 7–8, Figure 2 · p. 8

4.4 Ablation Studies

  • Extending ACS to all 16 OLMoE MoE layers changes accuracy by at most 0.46 points, with neither layer scope consistently better.Final-layer supervision remains the simpler competitive default.Source: §4.4 · p. 8, Table 3 · p. 8

5 Conclusion

The paper presents TES and ACS as complementary ways to align sparse MoE routing with token-level loss. Across two backbones and four MCQA benchmarks, both improve on their corresponding controls while preserving the native sparse execution budget and aggregation policy.

5 Conclusion

  • TES predicts expert-error scores to attenuate affinity, while ACS aligns native affinity concentration without an additional head.Both preserve the native sparse execution budget and aggregation policy; stronger supervision raises Granite ARC-Challenge TES-IS by 2.94 points over Dual Affinity.Source: §5 · pp. 8–9
  • The evaluation covers two sparse MoE backbones, four MCQA benchmarks, fixed data splits, and three training seeds; larger-scale models and pretraining-scale optimization remain outside its scope.End-to-end latency and memory overhead from the additional TES projection are not measured.Source: §5 · p. 9

A.1 Scalar Loss Geometry and Token-Identity Prediction Optima

The scalar analysis derives token-identity reference optima for the two error objectives, then distinguishes those references from guarantees about the implemented shared prediction head.

A.1 Scalar Loss Geometry and Token-Identity Prediction Optima

  • Occurrence-level IS and ENLL objectives are minimized at the realized target when it is positive; ENLL with a zero target instead tends toward negative infinity as the prediction approaches zero.The occurrence-level ENLL optimum statement assumes a positive target.Source: §A.1 · p. 12
  • For a fixed target-token identity, Equation 10 is a token-identity scalar reference; away from the numerical IS floor, IS and ENLL share the same target.The IS floor ensures a positive optimum, while ENLL requires a positive conditional mean.Source: §A.1 · p. 12
  • The shared contextual head is not guaranteed to match occurrence losses or average loss within each target identity, so Equation 10 is not a calibration guarantee.Its reference applies to grouped evaluation; routing and model updates also change observed losses during joint training.Source: §A.1 · p. 12

A.2 Gradient Allocation and Expert-Score Identifiability

The aggregate readout distributes its direct gradient across executed experts in proportion to their normalized native affinities, but does not uniquely determine their individual scores. For K > 1, local score changes can preserve the aggregate while changing routing and CE.

A.2 Gradient Allocation and Expert-Score Identifiability

  • Each active expert score receives a direct gradient with the aggregate prediction gradient’s sign, scaled in magnitude by its normalized native affinity.Parameter updates also depend on the head’s Jacobian and the optimizer.Source: §A.2 · p. 13
  • For K > 1, one fixed aggregate permits a local (K −1)-dimensional family of active-score perturbations that can change attenuated aggregation weights and CE.This is an occurrence-level readout property, not proof of parameter non-identifiability across the dataset.Source: §A.2 · p. 13
  • Shared-head scores can differ across tokens, but the aggregate target alone does not identify an individual score as that expert’s counterfactual loss.The scores are interpreted as learned error-aware routing signals, not independently calibrated expert losses.Source: §A.2 · p. 13

A.3 Token-Loss Gradient Analysis

The two objectives share an aggregate-prediction gradient toward observed token loss, but differ in how they affect the language-model task-loss gradient. Their matched-prediction conclusions require pointwise equality and do not imply every joint update decreases CE.

A.3 Token-Loss Gradient Analysis

  • Both objectives’ second gradient term trains the aggregate prediction toward observed token loss, while their first terms differ in their contribution to the task-loss gradient.The stated combined-gradient identities concern next-token CE, not the answer-choice CE term used in MCQA experiments.Source: §A.3 · p. 14
  • At pointwise equality between predicted and realized token loss, IS adds no CE gradient for the current error-aware routed model, whereas ENLL scales that model’s CE gradient by an inverse-loss factor.For multiple supervised layers, this conclusion holds when every layer’s prediction matches the token loss.Source: §A.3 · p. 14
  • Away from equality, the IS coefficient on the CE path can be negative, and the prediction-gradient term also contributes, so the analysis does not establish that every joint update decreases CE.Below the IS floor, its gradient through observed loss is zero; the floor boundary requires a subgradient convention.Source: §A.3 · p. 14
  • Along the idealized path where prediction equals realized loss and that loss approaches zero, ENLL is unbounded below; this does not show a finite run reaches that limit.The matched-prediction identities require pointwise equality, not merely a context-averaged optimum.Source: §A.3 · p. 14

A.4 Started-Log Routing Geometry

Started-log attenuation changes routing according to relative predicted errors while retaining sensitivity to their scale against a reference value. Its parameters separately set the operating-scale transition and the overall adjustment strength.

A.4 Started-Log Routing Geometry

  • When one predicted error exceeds another and γ > 0, attenuation lowers that expert’s routing odds relative to the other; equal predicted errors preserve native affinity ratios.At γ = 0, native probabilities are recovered.Source: §A.4 · p. 15
  • Unlike pure-log attenuation, started-log routing can change when all error predictions are rescaled while τ stays fixed, because the errors’ scale relative to τ matters.Pure-log attenuation turns a common multiplicative rescaling into a shared logit shift that cancels under softmax.Source: §A.4 · p. 15
  • Sensitivity to predicted error is bounded by γ/τ near zero, while sensitivity in log-error coordinates increases smoothly toward γ.At low error attenuation is approximately linear; at high error it recovers logarithmic relative-error routing up to a common shift.Source: §A.4 · p. 15
  • The reference scale τ sets the transition between low-error linear and high-error logarithmic regimes, while γ scales the overall routing adjustment.The default configuration sets both τ and γ to one.Source: §A.4 · p. 15

A.5 Error-to-Routing Transformation Ablation

The ablation compares started-log routing with a scale-invariant pure-log control on ARC-Challenge using Granite and IS supervision. Started-log improves test error-prediction correlations, while pure-log has higher answer-choice accuracy in this comparison.

A.5 Error-to-Routing Transformation Ablation

  • Started-log uses the error head’s absolute operating range relative to τ, whereas pure-log applies the same routing adjustment at every absolute scale for a fixed error ratio.Started-log contracts the low-error region and removes pure-log’s singular routing sensitivity near zero.Source: §A.5 · p. 15
  • Test group-level Pearson correlation rises from 0.398 to 0.452 and Spearman correlation from 0.498 to 0.564 with started-log rather than pure-log routing.Pure-log has 0.40 percentage points higher answer-choice accuracy; the comparison supports distinct requirements for error estimation and useful expert selection.Source: Table 4 · p. 15, §A.5 · p. 16
  • The ablation is on ARC-Challenge with Granite and IS supervision; accuracy, NLL, and pointwise correlation are epoch-3 means over seeds 42–44.Group-level correlations average each layer–token-identity group across the same seeds.Source: Table 4 · p. 15, §A.5 · p. 15
  • The observed contrast combines transformation and initialization effects, so it is not a fully isolated causal estimate of the routing transform.The pure-log route starts at be = 1, while started-log uses a zero-parameter softplus initialization.Source: §A.5 · p. 16

A.6 TES Hyperparameter Sensitivity

Across the appendix analyses, token-error ordering is positive but dataset-dependent, ACS outcomes vary with supervision strength and scope, and excessive coefficient scaling reduces accuracy.

A.6 TES Hyperparameter Sensitivity

  • The selected TES defaults—error-head learning rate 10^-5 and γ = 1—give the highest observed accuracy in their sweeps; nearby settings produce similar results.The γ sweep shows that minimizing choice NLL and maximizing answer-choice accuracy need not select the same value.Source: §A.6 · p. 16, Table 5 · p. 17

A.7 Experimental Configuration

  • Primary downstream evaluation uses epoch-3 checkpoints and manifest-based answer-text scoring; custom splits can include official-test examples in training, so results are not official-test evaluations.The secondary robustness analysis selects checkpoints by validation NLL before test evaluation.Source: §A.7 · pp. 16–18

A.8 Cross-Dataset Error-Prediction Analysis

  • Normalized plots show positive token-error ordering across the evaluated datasets, but its strength varies by dataset.MedMCQA IS is weaker than ENLL; the plots assess association and distribution shift, not absolute NLL calibration.Source: §A.8 · p. 18

A.9 Affinity-Concentration Supervision

  • ACS-IS has the highest observed Granite means on ARC-Challenge and SciQ, while the two ACS objectives remain close and neither is uniformly preferred.ACS is headless: token-loss supervision shapes native router distributions without adding error-specific parameters or an inference-time error transformation; its concentration is not a calibrated error estimate.Source: §A.9 · pp. 18, 19, 21

A.10 ACS Supervision-Coefficient Sensitivity

  • The larger ACS coefficient improves five of eight Granite trainable-affinity cells and four of eight cells in each OLMoE affinity setting.The largest increases are 1.67 points for Granite ACS–ENLL on OpenBookQA and 1.34 points for frozen-affinity OLMoE ACS–ENLL on SciQ.Source: §A.10 · p. 21

A.11 Affinity Supervision Depth Ablation

  • Granite ACS accuracy is highest with final-layer supervision on all four datasets, while OLMoE results favor different scopes across datasets.For Granite, all-layer supervision is consistently weakest; the matched Router-CE control also degrades as affinity training broadens.Source: §A.11 · p. 22, Table 9 · p. 22

A.12 Affinity-Concentration Diagnostics

  • ACS raises mean selected-affinity concentration from approximately 0.26 under Router CE to 0.78–0.81, while reciprocal effective count falls from approximately 3.9 to 1.2–1.3.These shifts are descriptive and do not establish that increased concentration causes the accuracy differences; all selected experts still execute.Source: Table 10 · p. 23, §A.12 · p. 23

A.13 Coefficient Scaling Across Supervision Depth

  • All four ARC-Challenge arms improve, and 10 of 12 other-dataset scope–objective arms are unchanged or better when ACS coefficients are scaled with supervision depth.A further tenfold increase sharply reduces accuracy, indicating an intermediate operating range rather than a monotonic benefit.Source: §A.13 · pp. 23–24, Table 12 · p. 24

Related papers

Loading 2609.37751v1…