Robustness of Mixtures of Experts to Feature Noise

TL;DR

The paper addresses why Mixture-of-Experts (MoEs) can outperform dense networks at fixed capacity when inputs exhibit latent modular structure corrupted by feature noise. It develops a tractable linear-model framework with block-diagonal features, compares dense versus sparse (MoE-like) estimators, and shows that activation sparsity filters noise, improving generalization, robustness to input perturbations, convergence speed, and sample efficiency. The routing problem is decoupled and reduced to clustering, enabling a near-perfect router via Quadratic Discriminant Analysis with poly(d, log(1/δ)) samples, while theoretical results show R(Sparse) ≤ R(Dense) under key conditions. Empirical evidence from linear probing on frozen representations and from end-to-end MoE experiments in language models corroborates the theory, highlighting practical benefits for efficient MoE-based inference and MoEfication techniques in large-scale models.

Abstract

Despite their practical success, it remains unclear why Mixture of Experts (MoE) models can outperform dense networks beyond sheer parameter scaling. We study an iso-parameter regime where inputs exhibit latent modular structure but are corrupted by feature noise, a proxy for noisy internal activations. We show that sparse expert activation acts as a noise filter: compared to a dense estimator, MoEs achieve lower generalization error under feature noise, improved robustness to perturbations, and faster convergence speed. Empirical results on synthetic data and real-world language tasks corroborate the theoretical insights, demonstrating consistent robustness and efficiency gains from sparse modular computation.

Figures (9)

• Figure 1: Modular structure in input activations to the up_proj layer within the MLP block of layer-0 of the Llama-2-7B model, revealed using TEAL liu2024training, for uniform magnitude pruning. From left to right, panels show activation percentiles after pruning activations based on 0%, 40%, and 70% uniform activation sparsity thresholds, respectively. Since the activation distributions are skewed due to the presence of outlier activations, we plot activation percentiles.
• Figure 2: We trained both a Dense Baseline and a MoE model from scratch using the MiniMind architecture.The MoE variant employs a "Shared Expert + Routed Experts" architecture. It consists of 4 routed experts and 1 shared expert. Each token selects $K=2$ routed experts. The intermediate dimension of each expert FFN is set to 1024. To ensure a fair comparison based on total parameter count, the intermediate dimension of the dense FFN is set to $5 \times 1024 = 5120$. This guarantees that the total parameters in the FFN layers of the dense model exactly match the sum of parameters of all experts in the MoE model. Despite having fewer active parameters per forward pass ($\sim60\%$ of the dense model), the MoE model exhibits a training loss trajectory that closely tracks and often surpasses the convergence speed of the dense baseline, particularly in the initial phases. This aligns with Theorem \ref{['th_lr_revised']}, suggesting that sparse modularity facilitates efficient optimization.
• Figure 3: The MoE model achieves a lower validation loss faster (i.e., with fewer training samples) than the Dense Baseline. The gap highlights the superior sample efficiency of MoE.
• Figure 4: The actual and fitted curves for the excess risks of the dense estimator and sparse estimator under the setting $d=k=100$.
• Figure 5: Modular structure in input activations to the up_proj layer within the MLP block of (a) layer-0, (b) layer-16, and (c) layer-31 of the Llama-2-7B model, revealed using TEAL for greedy magnitude pruning. From left to right, panels show activation percentiles after pruning activations based on 0%, 40%, and 70% greedy activation sparsity thresholds, respectively.

Theorems & Definitions (7)

• Theorem 4.1: Informal
• Theorem 4.2
• Theorem 4.3
• Theorem 4.4
• Remark 4.5
• Theorem 4.7
• Lemma 1.1: Lemma 5 in li2020benign