Summary of "Beyond Routing: Decoupling Expert Dispatch and Aggregation in Sparse Mixture-of-Experts"

Summary (Overview)

• Core contribution: This paper challenges the conventional design in Sparse Mixture-of-Experts (MoE) models where router scores serve dual purposes—selecting which experts to execute (dispatch) and weighting their outputs (aggregation). The authors demonstrate that these two functions can and should be decoupled.

• Key finding: Through carefully controlled counterfactual interventions on OLMoE-1B-7B and DeepSeek-V2-Lite, the authors show that even when expert selection is held fixed, substantial headroom exists in how selected expert outputs are aggregated. The router's top-scored expert is the counterfactual-best choice only 17.2% (OLMoE) and 12.5–16.7% (DeepSeek) of the time.

• Proposed remedy: The authors introduce Fixed-Dispatch Adaptive Aggregation (FDAA), a lightweight 301K-parameter post-compute head that learns token-adaptive aggregation weights under the true language-modeling objective, achieving significant cross-entropy improvements on WikiText-103, C4, and Penn Treebank.

• Cross-architecture validation: The selection–commitment mismatch replicates across two different MoE architectures (Top-8 routing on OLMoE and Top-6 plus shared experts on DeepSeek-V2-Lite), establishing the phenomenon as a general property of pretrained MoE designs.

• Critical negative result: Multiple forward-only credit prediction approaches (counterfactual reranking, gradient-field prediction, task-aware subspace predictors) fail to generalize, motivating direct LM-objective training over representation-based distillation.

Introduction and Theoretical Foundation

Background

Sparse Mixture-of-Experts models scale parameter count without proportionally scaling per-token computation by routing each token to a small subset of experts. Modern language-model MoEs (Mixtral, DeepSeekMoE, OLMoE) use a learned router that simultaneously:

1. Dispatches: selects which experts execute
2. Aggregates: determines how much each selected expert contributes

Key Theoretical Insight

The authors argue these responsibilities need not be identical. A score may be sufficient to rank a useful candidate set while being a poor estimate of each selected expert's downstream causal value. This distinction is operationally significant because changing dispatch is expensive (affects kernel execution, communication patterns, capacity allocation), while changing only aggregation weights preserves the sparse execution pattern exactly.

Mathematical Formulation

For a token $t$ with representation $h_t \in \mathbb{R}^d$, the native MoE output is:

$$m_t = \sum_{j \in S_t} w_{t,j} u_{t,j} \tag{1}$$

where $S_t$ is the selected expert set, $u_{t,j} = E_j(h_t)$ are expert outputs, and $w_{t,j} \geq 0$ are router weights. The total selected router mass is:

$$M_t = \sum_{j \in S_t} w_{t,j} \tag{2}$$

Counterfactual aggregations use alternative weights $w'_{t,j}$ satisfying:

$$\sum_{j \in S_t} w'_{t,j} = M_t \tag{3}$$

This constraint ensures any change in downstream loss is attributable to commitment within the selected set, not to rerouting.

Methodology

Fixed-Dispatch Diagnostic Protocol

The authors employ a strictly controlled intervention framework:

• Keep router-selected Top-K expert IDs fixed
• Compute exactly the same selected experts
• Preserve total selected Top-K router mass
• Only alter relative weights over already-computed expert outputs

Structured Oracle

For OLMoE-1B-7B-0924, the action family includes:

• Native baseline and uniform aggregation
• Temperature transforms (0.5 and 2.0)
• Eight 50% mixtures toward individual expert slots
• Eight vertices placing full mass on one selected expert

FDAA Architecture

The FDAA head predicts residual commitment scores $s_{t,j}$ using:

$$p_{t,j} = \frac{w_{t,j}}{M_t} \tag{5}$$ $$q_{t,j} = \frac{p_{t,j} \exp(s_{t,j})}{\sum_{\ell \in S_t} p_{t,\ell} \exp(s_{t,\ell})}, \qquad w'_{t,j} = M_t q_{t,j} \tag{6}$$

Key properties:

• Zero-initialized final layer ensures native behavior at initialization
• Permutation-equivariant over selected expert slots
• Preserves total selected mass by construction
• 301,313 trainable parameters (negligible vs. 7B checkpoint)

Training Objective

$$L = L_{\mathrm{LM}} + \lambda_{\mathrm{KL}} D_{\mathrm{KL}}(q_t \| p_t) + \lambda_s \|s_t\|_2^2 \tag{7}$$

The regularizers serve as stability terms, not teacher signals. No counterfactual or gradient-utility labels are required.

Gradient Utility Diagnostic

For target output $m_t$ with downstream gradient $g_t = \partial L / \partial m_t$, the first-order predicted improvement for vertex $j$ is:

$$\widehat{G}_{t,j} = -g_t^\top (M_t u_{t,j} - m_t) \tag{4}$$

Experimental Setup

• OLMoE-1B-7B-0924: 16 layers, 64 experts, Top-8 routing, 7B total / ~1B active parameters
• DeepSeek-V2-Lite: 27 layers, 64 routed experts with Top-6 routing, 2 shared experts
• Batch size 1 for all interventions to avoid co-batch numerical coupling
• Strict sanity checks: zero baseline CE mismatch, zero mass error, zero selected-ID mismatch

Empirical Validation / Results

OLMoE Fixed-Dispatch Diagnostic

Table 1: Full-horizon fixed-dispatch diagnostic on OLMoE (n = 128, seed 1337)

Gradient-Based Credit Assignment

• Mean within-Top-8 Spearman: 0.595
• Pairwise accuracy: 0.749
• Exact-best Top1 accuracy: 50.4% (vs. 12.5% for router)
• Selecting gradient-best vertex reduces CE by 0.01073, capturing 53.1% of exact vertex-oracle gain

Negative Results: Post-hoc Predictors

Stage-1 Results on OLMoE

Table 2: Stage-1 locked WikiText evaluation (seed 1337)

Fresh Test, Three Seeds

Table 3: Fresh WikiText-103 test, three optimization seeds

Frozen Confirmatory Evaluation (OLMoE)

Table 4: Three-seed confirmatory evaluation on 1,024 fixed windows per corpus

All three seeds individually significant vs. native, static, and shuffle on all three corpora.

DeepSeek-V2-Lite Replication

Table 5: One-seed DeepSeek-V2-Lite replication

Feature Masking Analysis

Inference-time masking on three WikiText-trained heads:

• Full layer-15 head: -0.1523 ± 0.0031
• Removing expert outputs: only -0.01184 ± 0.00328
• Router + expert-ID only: -0.00377 ± 0.00117
• Removing token hidden state: gain collapses to zero

Theoretical and Practical Implications

Selection Quality ≠ Commitment Quality

The counterfactual results do not imply the router selects a bad expert set—useful alternatives are often already present among the Top-8. The failure lies in mapping selected candidates to commitment weights. This distinction reconciles strong expert specialization with poor within-set router–utility correlation.

Why Direct LM Training Beats Credit Distillation

The gradient audit shows the desired direction is available to backpropagation, but full-horizon credit is not easily recoverable from frozen forward states. Direct LM training avoids this representational bottleneck by allowing the objective that defines downstream utility to reshape the aggregation head's features.

Late-Layer Concentration

The final audited MoE layer accounts for most of the learned gain. This is favorable for deployment (only one post-compute module needed) and raises mechanistic questions about late layers having less opportunity for subsequent blocks to repair misweighted outputs.

Domain Sensitivity

• WikiText-only training harms C4 (+0.01695)
• Mixed training converts this to small but repeatable improvement (-0.00487)
• DeepSeek: significant WikiText/PTB gains, C4 neutral
• Commitment is transferable but not domain-invariant

Conclusion

The paper establishes a clean separation between expert selection and expert commitment in sparse MoE models:

1. Diagnostic phenomenon: Fixed-dispatch headroom persists across two architectures (OLMoE Top-8, DeepSeek Top-6+shared), with router scores providing near-chance guidance for within-set commitment.

2. Credit assignment: Downstream gradients expose useful commitment credit, but forward-only predictors fail to generalize—even with favorable representations and future leakage.

3. Practical remedy: FDAA, a 301K-parameter post-compute head trained directly with the LM objective, achieves significant improvements while preserving sparse execution patterns exactly.

4. Caveats: The magnitude and transferability of learned commitment remain domain- and architecture-dependent. The cross-architecture replication of the phenomenon is stronger than the current replication of the remedy.

Future directions: Additional DeepSeek seeds, latency/throughput systems studies, downstream benchmark evaluation, and investigation of Top-1/Top-2 routing behavior.