Generalization and Scaling Laws for Mixture-of-Experts Transformers

The gist The theory develops a statistical framework for Mixture-of-Experts (MoE) Transformers that separates active per-token capacity from routing combinatorics, deriving approximation and uniform generalization bounds that reveal scaling laws mirroring dense models when measured against active parameters.

Metric-Entropy Decomposition

The main theoretical contributions include establishing a Metric-entropy decomposition for MoE Transformers which decomposes the complexity additively into an active-parameter term and a routing combinatorial term, thereby isolating the statistical effect of conditional computation pg2. This decomposition allows for the isolation of the statistical effect of conditional computation pg3.

Approximation and Uniform Generalization

The paper derives Approximation and uniform generalization under intrinsic dimension bounds, which separate approximation, estimation, and routing contributions, with rates governed by active capacity pg2. Specifically, Theorem 3.2 shows that the error bound is bounded by a term involving the active attention+expert budget Nact and a term involving LT lk log(eM/k) pg4. The generalization bound under this model shows that approximation and estimation trade off as in dense networks once active parameters are accounted for appropriately pg2.

Neural Scaling Laws

The work recovers neural scaling laws for sparse architectures by measuring error against the active parameter budget Nact, showing that the scaling exponents mirror dense-network behavior when measured against active parameters pg4. The data scaling law is derived as E kTˆn − fk 2 L2(Q) n − αD, where αD = 2β / (2β+d) pg4. Furthermore, the model-scaling law shows that under certain conditions, the error scales as E kTˆn − fk 2 L2(Q) N − αN, where αN = 2β / d pg6.

Routing Overhead and Specialization

The analysis explicitly accounts for routing overhead as a term proportional to k log(eM/k) in the worst-case bound pg3. Empirical validation shows that in a moderate regime (M/k ≤ 8), validation loss increases with the routing complexity term k log(eM/k), consistent with worst-case theory pg9. However, for larger values of M/k, an empirical reversal is observed where increasing the expert pool improves performance despite higher routing complexity, suggesting gains from expert specialization not captured by uniform bounds pg9.

Optimal Allocation and Trade-offs

The paper derives Compute-optimal allocation of Nact and n under a fixed compute budget C, showing that the excess error bound is minimized by N⋆act(C) = C d / (2β+2d) pg8. The resulting compute error scaling law is E kTˆn − fk 2 L2(Q) C − β / (β + d) pg8, which provides the MoE analogue of compute scaling laws pg8.

Consistency Checks and Empirical Mapping

The theory is checked against empirical results by comparing theoretical exponents to those derived from intrinsic-dimension analysis, observing systematic, monotone-positive deviations across (see Table 5) pg19. The results suggest that the smoothness-based exponents used in our analysis provide a reasonable statistical reference point for interpreting the observed scaling behavior pg19. The expanded routing ablations reveal two regimes: a moderate regime where performance degrades with the routing term k log(eM/k), and a specialization-dominated regime for larger M/k pg9.

Conclusion and Implications

The findings indicate that active per-token capacity plays a central role in MoE performance, while routing contributes an additional overhead whose effect depends on the regime pg10. The work provides a conservative statistical reference point for reasoning about sparse Transformer scaling, clarifying when MoE models should behave like dense models measured in active capacity and when additional data-dependent structure in routing and expert specialization can lead to behavior beyond uniform guarantees pg10. The analysis suggests that routing entropy and load balance remain useful diagnostics during training, while the theoretical bounds are conservative, reflecting a need for further data-dependent analyses pg13.

How it works

The framework decomposes the error into approximation, estimation, and routing terms where excess risk decomposes into approximation, estimation, and routing terms pg3. The complexity of MoE Transformers is decomposed additively into an active-parameter term and a routing combinatorial term, isolating the statistical effect of conditional computation pg3.

The approximation bound under Hölder smoothness, Theorem 3.2, states that inf T ∈TMoE kT − fk 2 ∞ ≤ C · minn N − 2β/d act, M − 2β/do pg4. This bound separates the approximation rate governed by active capacity from the routing combinatorial term, which contributes an additional logarithmic term LT lk log(eM/k) pg3.

The covering number of TMoE under the sup-norm is bounded by a term that includes both the active parameter budget and the routing combinatorial term, specifically log N δ, TMoE, k · k∞ ≤ C1 Πattn + LT kΠexp logC2κRM0 δ + C3LT lk logeM/k" pg5. This bound combines the approximation term and the routing overhead pg3.

The neural scaling laws are derived by minimizing a function g(k) that balances approximation and estimation, leading to an optimal number of active experts k⋆ ≈ nA + LT l log eM/k⋆ d / (2β+d) pg10. This shows that the dependence of k⋆ on M enters only logarithmically pg10. The routing-dominated regime is suppressed precisely when the effective expert activation (k) and the combinatorial explosion in M are controlled, allowing the system to operate in the favorable power-law regime where sample complexity scales as n − 2β / (2β+d) pg13.

How it works

The block-level perturbation bound, Lemma I.1, shows that the total output difference is bounded by Lparam η, where Lparam is a constant depending on Parch(l, m, demb, wFFN, LFFN) κ pg40. This stability bound demonstrates that the perturbation in the network output scales linearly with the parameter perturbation η pg40.

The covering number of TMoE under k · k∞ is bounded by a term that includes LT lk log(eM/k) pg3, which arises from a union bound over all routing patterns pg3. This term is conservative and should be interpreted as a worst-case upper bound rather than a tight characterization of practical routing behavior pg3.

The proof of Theorem 3.2 involves constructing two constructions, Construction A (expert-count limited) and Construction B (active-parameter limited), yielding the minimum rate inf T ∈TMoE kT − fk∞ ≤ C min N − 2β/d act, M−β/d pg32. This demonstrates that the rate is governed by either active capacity or the total number of experts M, with routing entering only logarithmically pg33.

1. Compute efficiency gains through active capacity focus by designing MoE layers where M is commensurate with k (e.g., M = O(k)). This avoids routing overhead inflation and ensures that scaling only affects the approximation term, leading to faster convergence in terms of active parameters rather than total model size.

2. Data-efficient training by optimizing the number of active experts k based on the formula derived from setting the derivative of the error bound to zero (Equation 47). This allows for a precise allocation of Nact and n under a fixed compute budget C, ensuring that both dataset size and active parameter count grow optimally according to Equation 18.

3. Improved generalization by using expert specialization when M is significantly larger than k (M/k is large). The paper shows that in the specialization-dominated regime, increasing the expert pool improves performance despite higher routing complexity, suggesting architectural choices like top-k or expert-choice routing can lead to empirical gains beyond worst-case bounds.

4. Enhanced robustness against parameter perturbations by leveraging the derived block-level stability bound (Equation 70). This bound shows that the total output difference is bounded by a term proportional to k times the active parameter budget, allowing for stable training even when expert weights are perturbed within a small epsilon range.

We develop a theory of generalization and scaling for Mixture-of-Experts (MoE) Transformers that cleanly separates active per-input capacity from routing combinatorics. By conditioning on fixed routing patterns and union-bounding across them, we derive a sup-norm covering-number bound whose metric entropy scales with the active parameter budget and incurs a MoE-specific routing overhead. Combined with a standard ERM analysis for squared loss, this yields a generalization bound under a d-dimensional manifold data model and C β targets, showing that approximation and estimation trade off as in dense networks once active parameters are accounted for appropriately. We further prove a constructive approximation theorem for MoE architectures, showing that, under the approximation construction, error can decrease either by scaling active capacity or by increasing the number of experts, depending on the dominant bottleneck. From these results we derive neural scaling laws for model size, data size, and compute-optimal tradeoffs. Overall, our results provide a transparent statistical reference point for reasoning about MoE scaling, clarifying which behaviors are certified by worst-case theory and which must arise from data-dependent routing structure or optimization dynamics.