Matching the Training-Time Weight Rule in Inference-Time Expert Deactivation

Summary (Overview)

• This paper systematically compares four expert deactivation implementations (removal, freezing, remaining-expert renormalization, and unnormalized re-routing) on identical disabled sets across five instruction-tuned checkpoints from four model families (OLMoE, DeepSeek-V2, Granite, Qwen3), plus a base OLMoE checkpoint.
• The authors develop a weight-sum model that predicts which implementation causes the most severe performance degradation based on the training-time weight rule: remaining-expert renormalization is most damaging for raw-probability models, while unnormalized re-routing is most damaging for normalized-weight models.
• At 50% deactivation, perplexity ranges from 1.4 to 73 times the intact baseline when the training rule is retained, but 14 to 28,740 times the baseline under a rule mismatch.
• The model extends to inference-time width adjustment: additional renormalization causes increasingly severe perplexity degradation as fewer experts are activated.
• On OLMoE, GSM8K strict-match accuracy falls to zero under all four implementations at 50% deactivation, while multiple-choice task outcomes vary significantly across implementations.

Introduction and Theoretical Foundation

Background

Open-weight mixture-of-experts (MoE) models distribute parameters across expert layers while activating only a small subset of experts per token. Models like OLMoE, DeepSeek-V2, Granite, and Qwen3 follow this architecture. While sparse activation reduces computation per token, the full expert bank still incurs storage, loading, and scheduling costs, motivating expert pruning, offloading, and dynamic routing approaches.

The Core Problem

Deactivating experts requires explicit choices about existing token assignments and their associated weights. The four implementations differ fundamentally:

• Removal (α): Masks disabled experts, re-routes tokens, and recomputes weights using the model's original rule
• Freezing (β): Preserves original routing decisions and surviving weights, removing only disabled experts' weighted contributions
• Remaining-expert renormalization (γ): Leaves vacated slots unfilled but rescales surviving weights to sum to one
• Unnormalized re-routing (δ): Fills vacated slots with replacements, weighting them by raw router probabilities

Training-Time Weight Rules

The training-time weight rule determines how expert weights $g_i$ are computed. For a token representation $h$, the router assigns scores $z_i$ to each of $N$ experts and applies softmax:

$$p_i = \frac{\exp(z_i)}{\sum_{j=1}^{N} \exp(z_j)} \tag{1}$$

The $k$ experts with highest probabilities form the active set $T$. The layer output is:

$$o = h + \sum_{i \in T} g_i y_i \tag{2}$$

Raw-probability models (OLMoE, DeepSeek-V2) use:

$$g_i = p_i, \quad m_T = \sum_{i \in T} p_i < 1 \tag{3}$$

Normalized-weight models (Granite, Qwen3) use:

$$g_i = \frac{p_i}{s_T}, \qquad s_T = \sum_{j \in T} p_j, \qquad m_T = 1 \tag{4}$$

The Weight-Sum Model

For removal, masking disabled experts with total probability $p_D$ gives:

$$q_i = \frac{p_i}{1 - p_D}, \qquad i \notin D \tag{5}$$

Freezing produces the layer output:

$$o_{\beta} = h + \sum_{i \in T \backslash D} g_i y_i \tag{6}$$

with the removed contribution:

$$\Delta_{\beta} = o - o_{\beta} = \sum_{i \in D} g_i y_i \tag{7}$$

Remaining-expert renormalization multiplies surviving contributions by $1/(m_T - m_D)$:

$$o_{\gamma} = h + \frac{1}{m_T - m_D} \sum_{i \in T \backslash D} g_i y_i \tag{8}$$

The weight changes are:

$$\Delta m_{\beta} = m_D, \quad \Delta m_{\gamma} = 1 - m_T \text{ (raw)}, \quad 0 \text{ (renorm)} \tag{9}$$

For unnormalized re-routing in normalized-weight models, the weight ratio for surviving experts is:

$$\frac{g_i^{(\delta)}}{g_i^{(\text{train})}} = s_T, \qquad i \in T \setminus D \tag{10}$$

Key prediction: For a fixed model and disabled set, the largest perplexity increase occurs under γ in raw-probability models and under δ in normalized-weight models.

Methodology

Models Evaluated

Disabled Set Construction

• Calibration corpus: First 2,000 documents from the first C4 shard, packed into 128 sequences of 2,048 tokens
• Selection criteria: Activation-based ranking (primary), selection frequency, and random sampling (seeds 42, 43, 44)
• Deactivation proportions: 12.5%, 25%, 37.5%, and 50%
• For a given model, proportion, and criterion, the disabled set is identical across all four implementations

Implementation Details

All four operations are applied directly to released checkpoints with fixed model parameters. Hugging Face Transformers is used for loading, and the EleutherAI Evaluation Harness for evaluation. Implementation checks include reproducing the intact model with an empty disabled set and comparing eager execution with fused kernels.

Evaluation Metrics

• Perplexity: WikiText-2 test perplexity with sliding windows of 1,024 tokens and stride 512
• Multiple-choice: ARC-easy, ARC-challenge, OpenBookQA, HellaSwag, BoolQ, RTE, WinoGrande, and MMLU
• Free-form generation: GSM8K with strict and flexible final-answer matching
• Statistical testing: Exact paired McNemar test with Benjamini-Hochberg correction across 21 paired tests

Empirical Validation / Results

Perplexity at 50% Deactivation

The highest perplexity occurs under the predicted rule mismatch for all five instruction-tuned checkpoints:

Critical isolation result: On identical DeepSeek-V2-Lite disabled sets, β gives perplexity 17.21 vs. 258,621.64 under γ—a difference of more than four orders of magnitude isolating the effect of weight handling alone.

Task Performance on OLMoE at 50% Deactivation

Effect of Disabled Set

The disabled set substantially affects outcomes. On Granite-3B with half the experts frozen, activation-based selection gives perplexity 112.6, while three random sets give 11,428 to 93,759 (101 to 832 times higher).

Width Adjustment Experiments

The perplexity penalty from renormalization increases as the active-expert count decreases, confirming the model's prediction.

Mean Log-Perplexity Increase per Expert Layer

The statistic $\rho = \ln \pi / L_{\text{MoE}}$ (where $\pi$ is the perplexity ratio and $L_{\text{MoE}}$ the number of expert layers) provides a common measure of average degradation:

• DeepSeek-V2-Lite: $\rho = 0.395$ (26 layers)
• OLMoE: $\rho = 0.456$ (16 layers)
• Granite-3B: $\rho = 0.224$ (32 layers)
• Granite-1B: $\rho = 0.254$ (24 layers)
• Qwen3: $\rho = 0.055$ (48 layers)

In OLMoE width experiments, $\rho$ is 0.472 at $k=2$ and 0.469 at $k=4$, within 0.016 of the 50% deactivation value, but decreases to 0.132 at $k=16$.

Theoretical and Practical Implications

Theoretical Contributions

1. Weight-sum model unification: The paper provides a unified mathematical framework explaining how all four deactivation implementations affect the expert-branch scale relative to the residual, using the training-time computation as the reference.

2. Rule-mismatch prediction: The model correctly predicts which implementation causes the most severe degradation for each model class—γ amplifies the expert branch in raw-probability models, while δ attenuates it in normalized-weight models.

3. Extension to width adjustment: The model's predictions extend beyond fixed expert deactivation to inference-time width changes, validated by OLMoE experiments showing renormalization penalty grows at smaller active-expert counts.

Practical Implications

• Implementation choice matters enormously: The same disabled set can produce perplexity differences spanning four orders of magnitude purely from weight-handling choices.
• Training-time rule as reference: Practitioners should use the training-time weight rule as the reference when implementing expert deactivation or dynamic routing.
• Complementary metrics: Perplexity and task accuracy capture different aspects of degradation—configurations with similar answer accuracy can retain very different sequence likelihoods.
• Generation behavior analysis: Even when strict-match accuracy is zero, generated outputs reveal distinct failure patterns: echo/prompt repetition under α, single-token loops under β, and early termination under γ.

Conclusion

This paper establishes that for a fixed model and disabled set, different choices about routing and expert weights produce substantial performance differences. The weight-sum model connects the high-loss implementation to the training-time weight rule: remaining-expert renormalization amplifies the expert branch in raw-probability models, whereas unnormalized re-routing attenuates it in normalized-weight models.

The OLMoE width experiments validate a further prediction: additional renormalization causes more severe degradation when fewer experts are active. The mean log-perplexity increase per expert layer summarizes degradation magnitude, while multiple-choice and free-form evaluations reveal different task outcomes.

Key recommendation: Inference-time changes to expert availability or active-expert counts should use the training-time weight rule as a reference, with the disabled set and target task informing the choice of implementation.

Limitations and Future Directions

• Experiments cover six checkpoints (1B–30B parameters) with rankings based on a fixed C4 calibration corpus; extending scale and calibration distributions would further test predictions
• Free-form generation is evaluated only on GSM8K; other generation tasks would broaden the analysis
• Whether recovery training can mitigate degradation and how implementation differences evolve during such training remains an open question