Conceptio › Archive › arXiv CS

arXiv CSopen access

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

arXiv:2604.09175v1 [cs.LG] 10 Apr 2026

G ENERALIZATION AND S CALING L AWS FOR M IXTURE - OF -E XPERTS T RANSFORMERS

Mansour Zoubeirou a Mayaki∗ Univ Lyon, Lyon 2, LIRIS UMR 5205 F-69676, Lyon, France mansour.mayaki at liris.cnrs.fr

A BSTRACT 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 ﬁxed 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-speciﬁc 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 certiﬁed by worst-case theory and which must arise from data-dependent routing structure or optimization dynamics.

Introduction

Mixture-of-Experts (MoE) Transformers enable conditional computation: for each token, only a small subset of experts is activated, so the number of active parameters per input is much smaller than the total number of parameters in the model. This yields a practical advantage: compute scales primarily with the number of active experts, while the full expert pool still provides substantial capacity for specialization. Despite striking empirical gains, a theory that cleanly separates the beneﬁts of active capacity from the overheads induced by routing has been missing. In particular, existing scaling laws for dense models track error as a function of total parameters, dataset size, and compute, whereas MoE models require a more reﬁned accounting: (i) approximation should scale with the active parameter budget rather than the total parameter count; (ii) generalization should reﬂect the multiplicity of routing patterns; and (iii) computeoptimal training should couple data and active capacity rather than data and total parameters. In this work, we develop approximation and uniform generalization results for MoE Transformers that make these distinctions explicit. Classical bounds indexed by the ambient dimension D predict unrealistically slow rates for high-dimensional inputs, whereas empirical scaling for Transformers is much faster. A growing body of theory and evidence shows that when data concentrate on a d-dimensional manifold, both approximation and estimation rates and hence the resulting scaling exponents are governed by the intrinsic dimension d rather than D. Recent work (Havrilla & Liao, 2024) develops approximation and statistical theories for Transformers on intrinsically low-dimensional data, deriving data- and model-scaling exponents that depend explicitly on d and validating these predictions on language-modeling benchmarks. Motivated by these results, we adopt d as the geometric control parameter throughout: it aligns the theory more closely with practice, explains faster-than-1/D behavior, and yields MoE scaling predictions that are sensitive to data geometry. ∗

Web page: https://mansmayaki.github.io/

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Our results should be interpreted as a conservative statistical baseline for MoE models rather than as tight predictions of training dynamics. Where empirical behavior is sharper, it must arise from data-dependent routing stability, specialization, or architectural biases not captured by worst-case analysis. Uniform generalization bounds are especially relevant for emerging architectures such as MoE, where empirical regularities are still evolving: they clarify which scaling behaviors are statistically certiﬁed independently of optimization and which require richer, data-dependent explanations. The main theoretical contributions are: • Metric-entropy decomposition for MoE Transformers. We establish a sup-norm covering-number bound in which the complexity of Mixture-of-Experts Transformers decomposes additively into an active-parameter term and a routing combinatorial term, thereby isolating the statistical effect of conditional computation. • Approximation and uniform generalization under intrinsic dimension. Under a d-dimensional manifold model and C β targets, we derive approximation and uniform generalization bounds that separate approximation, estimation, and routing contributions, with rates governed by active capacity. • Neural scaling laws for sparse architectures. We recover dense-network scaling exponents in (d, β) when measured against the active parameter budget, including model, data, and compute-optimal scaling. Routing contributes an additional logarithmic term in the worst-case bound, while specialization can yield stronger empirical gains in practice.

Related Work

Transformers: expressivity, approximation, and generalization. The Transformer architecture (Vaswani et al., 2017) implements content-addressable computation through multi-head self-attention and positionwise MLPs. Formal expressivity results show that (suitably parameterized) Transformers are universal approximators for a broad class of sequence-to-sequence maps (Yun et al., 2020), with positional encodings playing a crucial role in breaking permutation symmetries and enabling length-dependent behaviors. Subsequent analyses reﬁne which architectural ingredients control approximation power number of heads, attention bandwidth, depth vs. width trade-offs, and the effect of prenorm vs. post-norm residual designs often revealing that depth ampliﬁes compositional expressivity while LayerNorm stabilizes Lipschitz constants across layers. From the learning-theoretic perspective, generalization guarantees for deep networks leverage norm-based capacity measures (e.g., spectral-complexity/margin bounds, path norms) (Bartlett et al., 2017; Neyshabur et al., 2015) together with classical complexity tools (Rademacher/Gaussian complexities and covering numbers) (Bartlett & Mendelson, 2002b). Our analysis adopts this toolbox but tailors it to sparse Transformers by (i) conditioning on a ﬁxed routing pattern so that the model reduces to a smooth active subnetwork and (ii) union-bounding over routing patterns to account for combinatorial choices. For approximation, we rely on deep ReLU theory that provides near-optimal rates kf − N k∞ ≲ N −β/d for C β functions on d-dimensional domains (Yarotsky, 2017; Schmidt-Hieber, 2020; Petersen & Voigtlaender, 2018). We lift these Euclidean results to data on manifolds by working in local charts and gluing approximants with partitions of unity, which pairs naturally with the MoE routers k-sparse mixing and underlies our approximation theorem. Mixture-of-Experts (MoE): conditional computation and systems. Sparsely-gated MoE layers (Shazeer et al., 2017) activate only a small subset of experts per token, decoupling total parameters from per-example active compute. This idea enabled large sparse Transformers such as GShard (Lepikhin et al., 2021), Switch Transformers (Fedus et al., 2022), and GLaM (Du et al., 2022); in vision, V-MoE validated similar beneﬁts for ViTs (Riquelme et al., 2021). Practical deployments introduced load-balancing losses and alternative dispatchers (BASE layers, expertchoice routing) to reduce hotspotting and improve stability, and specialized runtimes (DeepSpeed-MoE, Tutel) to scale training/inference (Lewis et al., 2021; Zhou et al., 2022; Rajbhandari et al., 2022; Hwang et al., 2023). Recent systems emphasize expert specialization and the gap between total and active parameters e.g., Mixtral (8×7B, top-2 routing) and DeepSeek-V3 (hundreds of billions total, tens of billions active) (Jiang et al., 2024; DeepSeek-AI, 2024). These trends motivate our theoretical focus on active per-input capacity as the relevant approximation budget and on an explicit routing-combinatorics term that scales like k log(eM/k) in covering-number bounds. In our framework, hard top-k and softmax gating with temperature/load-balancing behave similarly at the level of uniform bounds as long as the effective sparsity is k. Intrinsic-dimension theory for Transformers. Our analysis builds on work showing that Transformer scaling is governed by the intrinsic rather than ambient dimension when data concentrate on low-dimensional structure. In particular, Havrilla & Liao (2024) develop an approximation statistical framework for Transformers on intrinsically low-dimensional data, deriving data/model exponents that depend explicitly on the manifold dimension and validating these predictions empirically. We adopt this viewpoint to interpret MoE: active per-token capacity plays the role of the 2

Generalization and Scaling Laws for Mixture-of-Experts Transformers

effective model-size axis, while routing acts as an additive overhead. Our analysis builds on their dense-Transformer framework and extends it to MoE with routing. Theorem 3.2 mirrors their approximation result but constructs MoE layers with top-k routing and separates active capacity from total parameters. This aligns our MoE scaling laws with intrinsic-dimension theory recovering dense-model exponents measured against active capacity and clarifying when routing can shift constants without changing rates. Scaling laws: dense vs. MoE. Dense language models exhibit empirical power-law scaling of loss with model size, data, and compute (Hestness et al., 2017; Kaplan et al., 2020); the Chinchilla study argues for compute-optimal training by scaling model and data jointly under a ﬁxed budget (Hoffmann et al., 2022). For MoE, recent empirical work proposes MoE-speciﬁc scaling formulations that disentangle total from active parameters and incorporate routing effects. Fine-grained MoE scaling analyzes expert granularity, token/expert budgets, and the efﬁciency frontier as expert shards become smaller and more numerous (Krajewski et al., 2024). Upcycling laws study converting dense checkpoints into MoE while retaining predictable scaling behavior (Liew et al., 2025). Joint denseMoE laws integrate active parameters, dataset size, and expert count to explain memory/throughput trade-offs (Ludziejewski et al., 2025), and uniﬁed efﬁcient-MoE laws relate expert activation ratios, granularity, and compute (Tian et al., 2025). Large-scale case studies (e.g., extending ﬁne-grained MoE beyond 50B parameters) provide further evidence that MoE can match or exceed dense scaling at comparable or lower active compute (Krajewski et al., 2025). Our theory complements these ﬁndings: we recover the dense-network exponents but measure them against the active parameter budget Nact , and we make explicit a routing overhead proportional to k log(eM/k) that shifts constants and determines ﬁnite-sample crossovers (e.g., long sequences or large M k). This yields computeoptimal prescriptions consistent with practice grow data and active capacity together while clarifying when adding experts primarily changes constants rather than exponents.

MoE Model and Approximation Theory

We consider i.i.d. samples (xi , yi )ni=1 drawn from an unknown distribution Q over X × Y. Deﬁnition 3.1 (MoE Transformer). Fix hyperparameters (D, ℓ, demb , m, LT , M, k) and FFN dimensions (LFFN , WFFN ). For an input H (0) ∈ Rℓ×demb , each block j = 1, . . . , LT applies the following residual updates: e (j−1) = H (j−1) + MHAj H (j−1) . H

1. Attention:

e (j−1) , 2. MoE: For each token embedding h̃t ∈ H (j)

= h̃t +

M X

gj,m (h̃t ) Ej,m (h̃t ),

m=1

where {Ej,m }M m=1 are MLP experts of depth LFFN and width WFFN . The routing weights satisfy gj,m (h) ≥ 0,

M X

gj,m (h) = 1,

{m : gj,m (h) 6= 0} ≤ k,

m=1

so that at most k experts are active per token. Deﬁnition 3.1 follows standard sparse MoE architectures (Shazeer et al., 2017; Lepikhin et al., 2021; Fedus et al., 2022). While this framework captures expert sparsity, it is too broad for precise statistical analysis. We therefore analyze a restricted MoE Transformer class in which: (i) submodules have ﬁxed depth and width; (ii) routing is deterministic top-k; (iii) all parameters are uniformly bounded; and (iv) the routing family is expressive enough to implement a k-sparse partition of unity. These conditions retain practical relevance while allowing explicit control of active parameter counts and covering numbers. 3.1

Approximation Bound Under Hölder Smoothness

Theorem 3.2. Let M ⊂ RD be a compact d-dimensional C 1 submanifold, and let f ∈ C β (M) satisfy kf kC β ≤ B, where β > 0. Consider the MoE Transformer class TMoE (Deﬁnition 3.1) with LT blocks, sequence length ℓ, and an MoE FFN sublayer per block with M experts, of which at most k are active per token. Assume each expert is a 3

Generalization and Scaling Laws for Mixture-of-Experts Transformers

ReLU MLP of depth LFFN , width wFFN , and parameter budget Πexp , and let Πattn denote the number of attention / non-expert parameters per block. Assume further that the router can implement a k-sparse partition of unity and that all weights are bounded by κ. Deﬁne the active attention+expert budget Nact := LT Πattn + LT kΠexp . Then, as shown in Appendix F, there exist constants C, c > 0 depending on (B, β, d, M) and polynomially on the architectural quantities (LT , wFFN , LFFN , κ) such that n o −2β/d inf kT − f k2∞ ≤ C · min Nact , M −2β/d . (1) T ∈TMoE

Remark 3.3. The exponent −2β/d is the classical minimax exponent for C β approximation on a d-dimensional manifold. The constants in Theorem 3.2 depend on (B, β, d, M) and polynomially on the architectural parameters. Standard manifold-approximation arguments imply that this dependence is at most polynomial in d for manifolds with bounded regularity; the only exponential-in-d behavior is the unavoidable covering-number dependence that produces the exponent itself. Proof sketch. Appendix F gives a constructive proof of Theorem 3.2. 3.2

MoE Generalization Bound

Theorem 3.4. Let T (ε) be a class of MoE Transformers such that for every f ∈ C β (M) there exists T ∈ T (ε) with kT − f kL2 (Q) ≤ ε. Let T̂n ∈ T (ε) denote an empirical risk minimizer under squared loss based on the i.i.d. sample (xi , yi )ni=1 . Then Nact LT ℓk log(eM/k) EkT̂n − f k2L2 (Q) ≤ ε2 + Õ + . n n Here Õ(g) hides polylogarithmic factors: Õ(g) = O g · polylog n, Nact , LT ℓk, M/k, κ, R, M0

Corollary 3.5 (Generalization bound with MoE approximation rate). Under the same router expressivity assumptions as in Theorem 3.2, substituting −2β/d ε2 Nact gives Nact LT ℓk log(eM/k) −2β/d EkT̂n − f k2L2 (Q) ≲ Nact + + , n n up to polylogarithmic factors. Remark 3.6 (Looseness of the routing union bound). The factor LT ℓk log(eM/k) comes from a worst-case union bound over all top-k routing patterns: log |Π| ≤ LT ℓk log(eM/k). This contribution is conservative and should be interpreted as a worst-case upper bound rather than a tight characterization of practical routing behavior. In practice, router specialization can substantially reduce the effective number of routing patterns, so the bound is expected to be loose but safe. Proof sketch. The proof (Appendix E) conditions on ﬁxed routing patterns, reducing the MoE to deterministic active subnetworks. For each pattern, Lipschitz bounds and a covering argument over the active parameter space yield the dense-style rate, while a union bound over patterns adds the routing term LT ℓk log(eM/k). Lemma 3.7 (MoE covering number). Let X ⊂ RD be compact with kxk∞ ≤ M0 . Consider the Mixture-of-Experts Transformer class TMoE (D, M, k, LT , LFFN , wFFN , demb , m, κ, R) deﬁned as in Deﬁnition 3.1, except that each feed-forward sublayer is an MoE layer with M experts and the router selects hard top-k experts per token per layer. Assume each learned parameter entry satisﬁes kθk∞ ≤ κ and each network output is bounded by kT (x)k∞ ≤ R. 4

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Let Πattn be the number of scalar parameters in the attention / non-expert parts per block, and let Πexp be the number of scalar parameters of a single expert MLP of depth LFFN , width wFFN , and input/output dimension demb . Then for any δ ∈ (0, 1), the covering number of TMoE under the sup-norm satisﬁes C κRM eM 2 0 log N δ, TMoE , k · k∞ ≤ C1 Πattn + LT kΠexp log + C3 LT ℓk log (2) δ k for absolute constants C1 , C2 , C3 > 0 depending polynomially on demb , wFFN , m, LT , LFFN , ℓ. Small polylogarithe notation. mic factors may be hidden in the O(·) Proof sketch. The proof (Appendix H) conditions on a ﬁxed top-k routing pattern, reducing the MoE to a deterministic active subnetwork, and then union-bounds over all routing patterns. For a ﬁxed pattern, a parameter-to-function Lipschitz bound for attention and FFN layers implies that an η-grid of [−κ, κ]Nact induces a δ-cover of the corresponding function subclass, yielding (2) by standard covering-number arguments (Bartlett & Mendelson, 2002b; Vaswani et al., 2017; Ba et al., 2016). Scope and interpretation of the bounds. The theoretical analysis is developed for a structured class of MoE Transformers, including bounded parameters, hard top-k routing, and explicit routing partitions. These assumptions are standard in uniform generalization analysis and are introduced to enable tractable covering-number arguments. The resulting guarantees should be interpreted as worst-case statistical bounds. In particular, the routing complexity term reﬂects a union bound over all possible routing patterns and therefore provides a conservative upper bound rather than a tight characterization of practical MoE behavior.

Neural Scaling Laws

We derive explicit scaling laws for Mixture-of-Experts transformers that expose the dependence on sample size n, number of experts M , number of active experts per token k, and parameter budgets. Throughout we use the generalization bound from theorem 3.4 in the shorthand form Rroute Nact −2β/d E kT̂n − f k2L2 (Q) ≲ Nact + . + | {z } n} n } | | {z {z approximation estimation routing eM where Rroute := LT ℓk log k . 4.1

(3)

Data Scaling

Let us study how test error decreases with the number of samples n when the architecture is ﬁxed or mildly tuned with n. Ignoring Rroute for the moment (we reintroduce it in §4.4), we balance approximation and estimation by minimizing, over the active budget Nact . Differentiating and setting to 0 gives d n 2β+d d 2β − 2β −1 c ⋆ ⋆ + = 0 =⇒ Nact ⇒ Nact n 2β+d − Nactd d n c ⋆ Substituting Nact back we obtain the data scaling law

2β (4) 2β + d This exponent matches the dense-network exponent under the same intrinsic-dimension assumptions (Havrilla & Liao, 2024); the MoE twist is that Nact is the active budget per input (depending on k), not the total parameter count. E kT̂n − f k2L2 (Q) n−αD ,

4.2

αD =

Model Scaling

We quantify how the error scales with the active parameter budget when the sample size n is ﬁxed. Proposition 4.1 (Model-scaling law in the approximation-dominated regime). For a ﬁxed n, let us suppose that Nact lies in the regime where the estimation and routing terms are dominated by the approximation term, namely there exist absolute constants c1 , c2 > 0 such that Nact Rrt − 2β − 2β ≤ c1 Nactd and ≤ c2 Nactd (5) n n 5

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Then we get 2β (6) d Proposition 4.2 (Characterization of the approximation-dominated regime). Since Nact is the active parameter budget, the ﬁrst condition in (5) is equivalent to −αN E kT̂n − f k2L2 (Q) Nact ,

1+2β/d

Nact ≲ n ⇐⇒ ⋆ that is, below the estimation crossover Nact (n). The routing condition reads −2β/d Rroute ≲ n Nact

⇐⇒

Nact ≲

αN =

⋆ Nact ≲ n 2β+d =: Nact (n),

n LT ℓk log(eM/k)

d 2β

⋆ =: Nact (n, M, k, LT , ℓ).

Therefore, the model-scaling law (6) holds whenever ⋆ ⋆ Nact ≤ min Nact (n), Nact (n, M, k, LT , ℓ) .

(7)

⋆ Remark 4.3 (Crossover and beyond). At the estimation crossover Nact ≈ Nact (n), the approximation and estimation terms balance: −2β/d

Nact

≈

d Nact =⇒ Nact n 2β+d , n

2β

EkT̂n − f k2L2 (Q) n− 2β+d

⋆ For Nact Nact (n) (holding n ﬁxed), the bound is dominated by Nact /n (plus routing), so increasing Nact worsens the bound. Thus (6) describes the subcritical (approximation-limited) regime, analogous to the dense case (Havrilla & Liao, 2024) but with Nact the active parameter budget. Corollary 4.4 (Model-scaling with experts-only approximation). If attention does not contribute to approximation (so Nact = Nexp ), then in the subcritical regime (Eq. (7)) 2β

E kT̂n − f k2L2 (Q) (LT kΠexp )− d 4.3

(8)

Optimizing the Number of Active Experts k

Let A := LT Πexp and absorb the attention budget into constants. With experts dominate approximation, Nact ∝ k, the bound (3) reduces to the function of k: LT ℓ k log(eM/k) Ak E(k) ≲ (Ak)−2β/d + + (9) | {z } n n |{z} | {z } approx

estimation

routing

Ignoring the log ﬁrst. Lets drop log(eM/k) and treat it as a slowly varying constant near the optimizer. Minimize g(k) = (Ak)−2β/d + B̃ n k over k > 0. Differentiating gives (up to architecture constants) 2β

k ⋆ n 2β+d and E(k ⋆ ) n− 2β+d Enforce k ⋆ ≤ M ; otherwise the optimum saturates at k = M . d

(10)

Reintroducing the log. With the routing term present, the derivative of the ﬁrst-order condition gives 2β −2β/d −2β/d−1 1 eM (11) A k = A + LT ℓ log − LT ℓ d n k This matches the stated ﬁrst-order condition whose closed form involves a Lambert-W factor if solved exactly. Since log(eM/k) varies slowly near the optimizer, the solution is d 2β+d n k ⋆ min M, (12) A + LT ℓ log(eM/k ⋆ ) Thus, within the worst-case bound, the dependence of k ⋆ on M enters only logarithmically. This does not preclude stronger empirical gains from larger expert pools arising through specialization effects outside the scope of the uniform analysis. The full proof of these calculations are provided in appendix G.1. We stress that this conclusion holds at the level of the worst-case bound only; the expanded routing ablations in Section 5 show that, for sufﬁciently large M/k, empirical performance can improve substantially due to expert specialization effects outside the scope of the present analysis. 6

Generalization and Scaling Laws for Mixture-of-Experts Transformers

4.4

Routing-dominated vs. Power-law Regime

Comparing the routing term to the power-law term n−2β/(2β+d) yields the threshold 2β+d 2β Rrt d ≳ n− 2β+d ⇐⇒ n ≲ nthr := LT ℓk log eM k n ( EkT̂n − f k2 ≈

LT ℓk log(eM/k) n 2β − 2β+d

n nthr n nthr

(13)

(14)

The bound indicates that routing dominance is mitigated when factors such as LT , ℓ, k, and log(eM/k) are small, a pattern that is consistent with architectural choices observed in prior MoE systems. These strategies align with empirical system-level practices in sparse Transformers: for instance, Switch Transformers (Fedus et al., 2022) and GShard (Lepikhin et al., 2021) emphasized top-1 or top-2 routing for efﬁciency, while later reﬁnements such as BASE Layers (Lewis et al., 2021) and expert-choice routing (Zhou et al., 2022) reduced routing variance and overhead. Large-scale MoE deployments (Riquelme et al., 2021; Du et al., 2022) also highlight the importance of balancing M and k to avoid routing bottlenecks. From the theoretical side, our bound formalizes these intuitions: the routingdominated 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) . While the theory predicts a single smoothness parameter β, empirical estimates obtained from model scaling and data scaling may differ, as these regimes probe different components of the error decomposition (approximation vs estimation and optimization effects). We therefore focus on structural consistency of the scaling exponents rather than exact agreement of inferred β values. 4.5

Sample Complexity for Target Error ε

In the power-law regime, where routing overhead is negligible relative to the approximation and estimation terms, the classical nonparametric rate applies. Speciﬁcally, achieving a target population error ε requires n(ε) ε−

2β+d β

Nact (ε) ε− β

(15)

The ﬁrst relation quantiﬁes the number of samples needed as a function of smoothness β and intrinsic dimension d; the second gives the corresponding effective parameter budget scaling to match the approximation rate. Both coincide with the dense-network theory (Yarotsky, 2017; Schmidt-Hieber, 2020), except that in MoE models Nact is interpreted as the active parameter budget per input rather than the total parameter count. In contrast, in the routing-dominated regime the error ﬂoor is determined by the combinatorial overhead of selecting experts. In this case, sample complexity to achieve accuracy ε grows only linearly in 1/ε: n(ε)

LT ℓk log(eM/k) ε

(16)

This expression highlights that the overhead scales with the number of MoE layers (LT ), sequence length (ℓ), and active experts per layer (k), while depending only logarithmically on the total number of experts M . Thus, if M k, routing dominates and inﬂates sample complexity signiﬁcantly, while M = Θ(k) mitigates this effect. In practice, this explains why architectures such as Switch Transformers (Fedus et al., 2022) or GLaM (Du et al., 2022) adopt small k (top-1 or top-2 routing) and balance M against k to remain in the favorable power-law regime. Our theoretical analysis therefore provides a principled characterization of when MoE gains translate into efﬁcient sample usage and when routing combinatorics instead dictate learning dynamics. 4.6

Compute-optimal Trade-offs

We derive the optimal allocation of active parameters per input Nact and number of training samples n under a ﬁxed training compute budget C. Throughout we use the bound in (3) and we adopt the standard compute model in which per-token (or per-sample) FLOPs scale linearly with the active parameter budget (Hoffmann et al., 2022; Kaplan et al., 2020), so the total training compute satisﬁes C n · Nact . (17) 7

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Proposition 4.5 (Compute-optimal allocation of Nact and n). Assume the compute budget (17). Then the excess error bound is minimized (up to constants) by d

⋆ Nact (C) C 2β+2d

n⋆ (C) =

and

C ⋆ (C) Nact

2β+d

C 2β+2d

(18)

The resulting compute error scaling law is β

E kT̂n⋆ − f k2L2 (Q) C − β+d

(19)

Thus, under a ﬁxed compute budget, both the dataset size and the per-sample active parameters should grow with C according to the exponents in (18). The error decays as a power law in compute with exponent β/(β + d), matching nonparametric theory and providing the MoE analogue of compute scaling laws. See appendix G.3 for proof.

Empirical Scaling and Evaluation on LLMs

Our empirical study is designed to examine the scaling exponents and routing effects suggested by the theoretical analysis under controlled yet practically relevant conditions. We evaluate the framework on a family of causal, decoder-only Transformer language models with Mixture-of-Experts (MoE) feed-forward layers. We consider three public text corpora with differing levels of complexity and estimated intrinsic dimension: TinyStories, a low-dimensional synthetic dataset of children’s stories (Eldan & Li, 2023); WikiText-103, a medium-scale corpus of moderately heterogeneous Wikipedia articles (Merity et al., 2016); and OpenWebText, a large and highly heterogeneous corpus reconstructed from Reddit-linked web pages (Gokaslan & Cohen, 2019). Additional details are provided in Appendix A. Intrinsic dimension d estimates in practice. We estimate the intrinsic dimension d of the data manifold from hidden representations using standard intrinsic-dimension estimators. Concretely, we pass corpus tokens through a ﬁxed pretrained GPT-2 model, collect hidden states at a given layer, and apply the Levina–Bickel k-nearest-neighbor maximumlikelihood estimator, which is known to be more robust than TwoNN in high-dimensional settings (Fukunaga & Olsen, 1971; Levina & Bickel, 2004; Facco et al., 2017). Distances are computed using FAISS, and intrinsic dimension is evaluated on a layer-by-layer basis. To reduce variance, we repeat the procedure over multiple random subsamples and report the median intrinsic dimension. Unless stated otherwise, we use the middle Transformer layers as a representative summary of d, following common practice in prior work. The resulting estimates are stable across layers and subsamples. The estimated intrinsic dimensions are lowest for TinyStories and higher for WikiText-103 and OpenWebText, consistent with their relative heterogeneity; exact values are reported in Table 1, with depth-wise trends reported in Figures 2 and 3 in Appendix A. We emphasize that these estimates are representation-dependent and should be viewed as a coarse proxy for data complexity rather than as a deﬁnitive characterization of the learned MoE representations. Experimental setup. We operate in a small- to medium-scale regime to enable dense sweeps over the capacity grid. Unless noted otherwise, we use sequence lengths ℓ ∈ {128, 256}, m = 4 attention heads, the AdamW optimizer, and a learning rate of 3 × 10−4 . For model scaling, we vary (LT , dﬀ ) over a broad grid and consider multiple MoE conﬁgurations (M, k). For data scaling, we ﬁx a representative architecture and vary the token budget D over the range 5 × 104 –8 × 105 . We examine three empirical regimes corresponding to the theoretical analysis: (i) model scaling, where we estimate the exponent α bN by regressing log L on log Nact ; (ii) data scaling, where we regress log L on log D to extract α bD ; and (iii) routing ablations, where we vary (M, k) to probe the contribution of routing combinatorics and identify the transition from routing-inﬂuenced behavior to a power-law-like regime. Empirical estimates of β obtained from model scaling and data scaling do not coincide exactly, which is expected since these regimes probe different components of the error decomposition. We therefore treat β as an effective smoothness parameter and focus on structural consistency of the exponents rather than exact parameter recovery. Additional discussion is provided in Appendix B. 5.1

Fitting Exponents and Comparison to Theory

In all scaling experiments, we operated in regimes where routing overhead was either negligible relative to the validation loss or explicitly accounted for. When needed, we introduced a small loss ﬂoor c and ﬁt linear models to log(L−c) as a function of the relevant scaling variable (log Nact or log D ), yielding empirical exponents (b αN , α bD ). We compared these exponents to theoretical values (αN , αD ) derived from the intrinsic-dimension framework using an estimated dimension d and a range of smoothness parameters β. Across datasets, the empirical and theoretical exponents are of comparable magnitude and approximately satisfy the compute-optimal consistency relation αD ≈ αN /(1+αN ). 8

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Figure 1: Routing ablation across expert pool size M and active experts k. (a) In the moderate regime (M/k ≤ 8), validation loss increases with the routing complexity term k log(eM/k), consistent with its interpretation as a worstcase routing overhead. (b) Over the full range, performance improves again for sufﬁciently large M/k, suggesting gains from expert specialization not captured by the worst-case analysis. These observations suggest that the intrinsic-dimension-based theory provides a coherent statistical reference for observed model and data-scaling behavior, without implying tight quantitative agreement. Additional sensitivity analyses over β are reported in Appendix A.6. 5.2

Routing ablation.

Figure 1 shows routing ablations over (M, k) under ﬁxed compute, revealing two consistent regimes across scales. In the moderate regime (M/k ≤ 8), validation loss increases approximately monotonically with the routing term k log(eM/k), consistent with the theoretical prediction that routing contributes a statistical overhead. For larger values of M/k, we observe a reversal: increasing the expert pool improves performance despite higher routing complexity. For example, at ﬁxed k = 4, increasing M from 16 to 64 reduces validation loss from 2.73 to 2.37. This specialization-dominated regime is not captured by the worst-case analysis and suggests that practical MoE models beneﬁt from structured routing and expert specialization beyond what uniform bounds predict. Importantly, this two-regime behavior is consistent across different compute budgets, reinforcing the interpretation of the routing term as a conservative upper bound rather than a tight characterization of practical performance.

Discussion and Limitations

Our results 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. When routing complexity remains modest, MoE models admit dense-style scaling behavior when measured against active parameters rather than total model size. Across TinyStories, WikiText-103, and OpenWebText, the observed model- and data-scaling exponents are of comparable magnitude to those suggested by intrinsic-dimension analysis. The expanded routing ablations further reveal two regimes: a moderate regime in which performance degrades with the routing term k log(eM/k), consistent with the worst-case theory, and a specialization-dominated regime for larger M/k in which increasing the expert pool improves performance. The constants in our bounds are conservative, and the regimes studied remain far from asymptotic. From a practical standpoint, three architectural factors appear central to MoE behavior: activation sparsity, expert granularity, and routing diversity. Within worst-case statistical guarantees, expert pools that grow much larger than the number of activated experts are not certiﬁed to improve rates beyond a logarithmic term; however, our routing ablations show that larger expert pools can yield substantial empirical gains through specialization when M/k is sufﬁciently large. Routing entropy and load balance therefore remain useful diagnostics during training. For long sequences, smaller values of k can help limit routing overhead and delay the onset of routing-dominated regimes, although data-dependent effects may substantially alter these trade-offs in practice. 9

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Our analysis focuses on statistical complexity rather than optimization dynamics. In practical MoE systems, mechanisms such as expert specialization, load balancing, and routing regularization can signiﬁcantly inﬂuence behavior, potentially reducing the effective routing complexity and improving performance beyond worst-case predictions. Limitations and future directions. Our analysis relies on several simplifying assumptions, including bounded parameters, hard top-k routing, and squared loss. While these assumptions are standard in uniform convergence analysis, they do not fully capture practical MoE training settings. In particular, the derived bounds are likely conservative. The linear dependence on sequence length ℓ and the number of active experts k arises from worst-case covering arguments and may overestimate the effective complexity of practical models. Moreover, the analysis does not capture approximation gains from expert specialization, which we now observe empirically in regimes where M/k is large. Developing tighter data-dependent analyses, as well as corresponding lower bounds, remains an important direction for future work.

Conclusion

We developed a statistical framework for Mixture-of-Experts (MoE) Transformers that separates active per-token capacity from routing combinatorics. Under standard regularity assumptions, we derive approximation and uniform generalization bounds in which excess risk decomposes into approximation, estimation, and routing terms. When expressed in terms of active parameters, these bounds imply scaling exponents that mirror dense-model behavior, while isolating routing as the distinct MoE-speciﬁc statistical overhead. Within this framework, we characterize the scaling of the optimal activation level up to architectural and logarithmic factors, and show that, under worst-case guarantees, increasing the total number of experts affects the bound only logarithmically through routing complexity. At the same time, our expanded routing ablations reveal a richer empirical picture: in a moderate regime, performance degrades with the routing term k log(eM/k) as predicted by the theory, whereas for sufﬁciently large M/k, larger expert pools improve performance through specialization effects that are not captured by the worst-case analysis. Overall, our work provides a conservative statistical reference point for reasoning about sparse Transformer scaling. It clariﬁes 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.

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Impact Statement This paper presents theoretical and empirical work aimed at improving the scientiﬁc understanding of Mixture-ofExperts Transformers and sparse model scaling. Its primary contribution is methodological rather than directly deployable. A possible positive impact is improved compute efﬁciency in large-scale machine learning through better understanding of active capacity and routing. As with other work on scaling advanced models, these insights could also indirectly support the development of more capable systems, with downstream risks similar to those already associated with large language models, including misuse, bias propagation, and unequal access. We do not believe this paper raises ethical concerns beyond those already well established for research on large-scale machine learning.

References Anthony, M. and Bartlett, P. L. Neural Network Learning: Theoretical Foundations. Cambridge University Press, 1999. Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv:1607.06450, 2016. URL https://arxiv.org/abs/1607.06450. Bartlett, P. L. and Mendelson, S. Rademacher and gaussian complexities: Risk bounds and structural results. JMLR, 3:463–482, 2002a. Bartlett, P. L. and Mendelson, S. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3:463–482, 2002b. Bartlett, P. L., Foster, D. J., and Telgarsky, M. Spectrally-normalized margin bounds for neural networks. In Advances in Neural Information Processing Systems (NeurIPS), volume 30, 2017. Chen, M., Jiang, H., Liao, W., and Zhao, T. Nonparametric regression on low-dimensional manifolds using deep relu networks: function approximation and statistical recovery. Information and Inference: A Journal of the IMA, 11(4): 1203–1253, 03 2022. ISSN 2049-8772. doi:10.1093/imaiai/iaac001. URL https://doi.org/10.1093/imaiai/ iaac001. DeepSeek-AI. Deepseek-v3 technical report. 2024. URL https://arxiv.org/abs/2412.19437. Du, N., Huang, Y., Dai, A. M., Tong, S., Lepikhin, D., Xu, Y., Krikun, M., Zhou, Y., Yu, A. W., Firat, O., Chen, Z., Le, Q. V., Wu, Y., and Chen, Z. Efﬁcient scaling of language models with mixture-of-experts. In International Conference on Machine Learning (ICML), volume 162 of Proceedings of Machine Learning Research, pp. 5547– 5569. PMLR, 2022. URL https://arxiv.org/abs/2112.06905. GLaM. Eldan, R. and Li, Y. Tinystories: How small can language models be and still speak coherent english?, 2023. URL https://arxiv.org/abs/2305.07759. Facco, E., dErrico, M., Rodriguez, A., and Laio, A. Estimating the intrinsic dimension of datasets by a minimal neighborhood information. Scientiﬁc Reports, 7, 2017. URL https://api.semanticscholar.org/CorpusID: 3991422. Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: scaling to trillion parameter models with simple and efﬁcient sparsity. Journal of Machine Learning Research, 23(1), January 2022. ISSN 1532-4435. Fukunaga, K. and Olsen, D. An algorithm for ﬁnding intrinsic dimensionality of data. IEEE Transactions on Computers, C-20(2):176–183, 1971. doi:10.1109/T-C.1971.223208. Gokaslan, A. and Cohen, V. Openwebtext corpus. http://Skylion007.github.io/OpenWebTextCorpus, 2019. Golowich, N., Rakhlin, A., and Shamir, O. Size-independent sample complexity of neural networks. In COLT, 2018. Havrilla, A. and Liao, W. Understanding scaling laws with statistical and approximation theory for transformer neural networks on intrinsically low-dimensional data. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://openreview.net/forum?id=N2wYPMpifA. Hestness, J., Narang, S., Ardalani, N., Diamos, G., Jun, H., Kianinejad, H., Patwary, M. M. A., Yang, Y., and Zhou, Y. Deep learning scaling is predictable, empirically, 2017. URL https://arxiv.org/abs/1712.00409. Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Vinyals, O., Rae, J. W., and Sifre, L. Training compute-optimal large language models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS ’22, Red Hook, NY, USA, 2022. Curran Associates Inc. ISBN 9781713871088. 11

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Hwang, C., Cui, W., Xiong, Y., Yang, Z., Liu, Z., Hu, H., Wang, Z., Salas, R., Jose, J., Ram, P., Chau, J., Cheng, P., Yang, F., Yang, M., and Xiong, Y. Tutel: Adaptive mixture-of-experts at scale, 2023. URL https://arxiv.org/ abs/2206.03382. Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., de las Casas, D., Hanna, E. B., Bressand, F., Lengyel, G., Bour, G., Lample, G., Lavaud, L. R., Saulnier, L., Lachaux, M.-A., Stock, P., Subramanian, S., Yang, S., Antoniak, S., Scao, T. L., Gervet, T., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mixtral of experts, 2024. URL https://arxiv.org/abs/2401.04088. Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models, 2020. URL https://arxiv.org/abs/2001.08361. Krajewski, J., Ludziejewski, J., Adamczewski, K., Píóro, M., Krutul, M., Antoniak, S., Ciebiera, K., Król, K., Odrzygóźdź, T., Sankowski, P., Cygan, M., and Jaszczur, S. Scaling laws for ﬁne-grained mixture of experts. 2024. doi:10.48550/arXiv.2402.07871. URL https://arxiv.org/abs/2402.07871. Krajewski, J., Chochowski, M., and Korzekwa, D. Scaling ﬁne-grained moe beyond 50b parameters: Empirical evaluation and practical insights. 2025. URL https://arxiv.org/abs/2506.02890. Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. {GS}hard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=qrwe7XHTmYb. Levina, E. and Bickel, P. Maximum likelihood estimation of intrinsic dimension. In Saul, L., Weiss, Y., and Bottou, L. (eds.), Advances in Neural Information Processing Systems, volume 17. MIT Press, 2004. URL https://proceedings.neurips.cc/paper_files/paper/2004/file/ 74934548253bcab8490ebd74afed7031-Paper.pdf. Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettlemoyer, L. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning (ICML), volume 139 of Proceedings of Machine Learning Research, pp. 6265–6274. PMLR, 2021. URL https://arxiv.org/abs/2103.16716. Liew, S. P., Kato, T., and Takase, S. Scaling laws for upcycling mixture-of-experts language models. 2025. doi:10.48550/arXiv.2502.03009. URL https://arxiv.org/abs/2502.03009. ICML 2025. Ludziejewski, J., Pióro, M., Krajewski, J., Stefaniak, M., Krutul, M., Małaśnicki, J., Cygan, M., Sankowski, P., Adamczewski, K., Miłoś, P., and Jaszczur, S. Joint moe scaling laws: Mixture of experts can be memory efﬁcient. In Forty-second International Conference on Machine Learning, 2025. URL https://openreview.net/forum? id=70DGIxEHiB. Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models, 2016. Neyshabur, B., Tomioka, R., and Srebro, N. Norm-based capacity control in neural networks. In Conference on Learning Theory (COLT), volume 40 of Proceedings of Machine Learning Research, pp. 1376–1401. PMLR, 2015. Petersen, P. and Voigtlaender, F. Optimal approximation of piecewise smooth functions using deep relu neural networks. Neural Networks, 108:296–330, 2018. ISSN 0893-6080. doi:https://doi.org/10.1016/j.neunet.2018.08.019. URL https://www.sciencedirect.com/science/article/pii/S0893608018302454. Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. Y., Awan, A. A., Rasley, J., and He, Y. Deepspeedmoe: Advancing mixture-of-experts inference and training to power next-generation ai scale. arXiv preprint arXiv:2201.05596, 2022. URL https://arxiv.org/abs/2201.05596. Riquelme, C., Puigcerver, J., Mustafa, B., Neumann, M., Jenatton, R., Pinto, A. S., Keysers, D., and Houlsby, N. Scaling vision with sparse mixture of experts. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pp. 8583–8595, 2021. URL https://arxiv.org/abs/2106.05974. Schmidt-Hieber, J. Nonparametric regression using deep neural networks with relu activation function. The Annals of Statistics, 48(4), August 2020. ISSN 0090-5364. doi:10.1214/19-aos1875. URL http://dx.doi.org/10.1214/ 19-AOS1875. Shalev-Shwartz, S. and Ben-David, S. Understanding Machine Learning: From Theory to Algorithms. Cambridge University Press, 2014. Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017. URL https://openreview.net/forum?id=B1ckMDqlg. Tian, C., Chen, K., Liu, J., Liu, Z., Zhang, Z., and Zhou, J. Towards greater leverage: Scaling laws for efﬁcient mixture-of-experts language models. 2025. doi:10.48550/arXiv.2507.17702. URL https://arxiv.org/abs/ 2507.17702. 12

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998–6008, 2017. URL https: //proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html. Vershynin, R. High-Dimensional Probability. Cambridge University Press, 2018. Yarotsky, D. Error bounds for approximations with deep relu networks, 2017. URL https://arxiv.org/abs/1610. 01145. Yun, C., Bhojanapalli, S., Rawat, A. S., Reddi, S. J., and Kumar, S. Are transformers universal approximators of sequence-to-sequence functions? In International Conference on Learning Representations (ICLR), 2020. arXiv:1912.10077. Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A., Chen, Z., Le, Q., and Laudon, J. Mixture-of-experts with expert choice routing. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pp. 25478–25492, 2022. URL https://arxiv.org/abs/2202.09368.

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Table 1: Intrinsic Dimension (ID) Summary across Models and Datasets. The ID estimates remain remarkably stable across different model sizes for each dataset. For the subsequent scaling analysis (Figure 5 and 4), we chose the median ID values for each dataset. Tinystories Wikitext-103 OpenWebText Model (number of parameters) Mean (CI) Median Mean (CI) Median Mean (CI) Median gpt2 (117 million) 22.9 ± 2.9 23.1 31.6 ± 2.9 32.4 46.7 ± 6.2 49.8 gpt2-medium (345 million) 21.4 ± 1.1 21.9 31.6 ± 1.7 32.9 43.7 ± 2.6 45.0 gpt2-large (774 million) 22.1 ± 0.9 22.4 31.5 ± 1.6 32.1 44.7 ± 2.6 47.7 gpt2-xl (1558 million) 19.7 ± 0.6 19.6 30.3 ± 1.0 31.1 43.0 ± 1.7 43.9 Values Adopted for Scaling Analysis Chosen d (Median ID) d = 23 d = 32 d = 45 Best-Fit β β ≈ 1.0 β ≈ 1.0 β ≈ 1.0 − 1.5

A Empirical Validation on MoE Transformer LLMs This section provides a detailed description of the architectures, hyperparameter grids, and scaling protocols used to ensure the empirical results are consistent with our theoretical predictions. Our goal is to cleanly isolate the roles of (i) effective model capacity, (ii) dataset size, and (iii) routing combinatorics in MoE Transformers. A.1

Datasets.

We consider three publicly available text corpora with varying complexity and intrinsic dimension: Stories / TinyStories is a corpus of short, synthetically generated childrens stories, following (Eldan & Li, 2023). This dataset is known to lie on an intrinsically low-dimensional manifold, and has been used in prior work on theoretical scaling laws for Transformers. WikiText-103 is a curated subset of English Wikipedia articles (Merity et al., 2016), commonly used for language modeling. It provides a medium-scale, moderately heterogeneous domain. OpenWebText is an open recreation of the WebText corpus (Gokaslan & Cohen, 2019), constructed from outgoing Reddit links. It is substantially larger and more heterogeneous than WikiText-103, with higher intrinsic dimension. For each corpus we extract a plain-text training stream (one document per line). For scaling experiments we sub-sample the stream to match the desired token budgets. A.2

Intrinsic Dimension Estimation (Extended Description)

The intrinsic dimension (ID) of a representation manifold quantiﬁes the effective degrees of freedom of the learned features. To estimate the ID of transformer representations, we employ the Levina–Bickel Maximum Likelihood Estimator (MLE) (Levina & Bickel, 2004), which models the local likelihood of k-nearest-neighbor distances on a smooth manifold. Compared to the TwoNN estimator, which relies only on the ratio of the ﬁrst two nearest neighbors, the MLE approach uses the full local neighborhood and exhibits substantially lower variance in high-dimensional settings an essential property when studying large language models. We extract hidden states from each block of a pretrained GPT-2 model (gpt2-medium) by sampling tokens from a large text corpus. For each layer ℓ, we compute all pairwise distances to its k nearest neighbors using FAISS in 32-bit precision, and apply the MLE estimator to the resulting distance matrix Dℓ . The estimation pipeline is repeated for multiple random subsamples to mitigate sampling noise; we report the median MLE estimate for each layer, along with the median absolute deviation (MAD) as a robust measure of uncertainty. Following established practice in intrinsic-dimension analyses of deep networks, we summarize the model-level intrinsic dimension using the middle layers 40–60% depth). Shallow layers typically show inﬂated dimensionality due to local embedding variation, whereas deeper layers tend to collapse. This procedure yields stable and reproducible ID estimates that are consistent with theoretical predictions from neural scaling laws, allowing us to deﬁne a characteristic ID (d) for each corpus that captures its inherent complexity, as summarized in the depth curves (Figure 2) and the scaling law ﬁts (Figure 3). A.3

Model family: MoE decoder-only Transformers

All experiments use a causal decoder-only Transformer with Mixture-of-Experts (MoE) feed-forward layers. Each model consists of token and positional embeddings, followed by LT Transformer blocks. Each block contains: 14

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Table 2: Theoretical Exponents derived from the Chosen d and β. Dataset

Chosen d

Best-Fit β

32 45 45 23

1.0 1.0 1.5 1.0

WikiText-103 OpenWebText OpenWebText TinyStories

αN = 2β d 0.063 0.044 0.067 0.087

In)rin(ic Dimen(ion (MLE) v(. Layer Dep) : OpenWebTex)

Model Size gpt2 (117M para s, A(g ID: 22.9) gpt2- ediu (345M para s, A(g ID: 21.4) gpt2-large (774M para s, A(g ID: 22.1) gpt2-)l (1558M para s, A(g ID: 19.7)

32.5

gp)2-medium (345M param(, Avg ID: 43.7) gp)2-large (774M param(, Avg ID: 44.7) gp)2-,l (1558M param(, Avg ID: 43.0)

30.0

In(rinsic Dimension (d)

Intrinsic Di ension (d)

In)rinsic Dimension (d)

gpt2-medi)m (345M params, Avg ID: 31.6) gp(2-large (774M params, Avg ID: 31.5) gp(2-xl (1558M params, Avg ID: 30.3)

27.5

Model Si-e gp(2 (117M params, Avg ID: 31.6)

50 45

Empirical Exponents bN bD α α 0.060 0.058 0.045 0.043 0.068 0.062 0.085 0.081

In(rinsic Dimension (MLE) vs. Layer Dep(h: Wi iTex(-103

Intrinsic Di ension (MLE) (s. Layer Depth: TinyStories

Model Si.e gp)2 (117M param(, Avg ID: 46.7)

Theoretical Exponents 2β αN αD = 2β+d Consistency 1+α N 0.059 0.059 0.043 0.042 0.062 0.063 0.080 0.080

25.0 22.5 20.0

17.5 15.0

25 0

20 30 Layer Index (Dept )

20 0

(a) OpenWebText (OWT) Dataset

20 30 Layer Index (Depth)

(b) TinyStories Dataset

20 30 Layer Index (Dep(h)

Figure 2: Intrinsic Dimension (ID) Evolution Across Model Depth. The plots show the Maximum Likelihood Estimation (MLE) of the Intrinsic Dimension, d, calculated for the representations at each layer of the four GPT-2 models (varying size). The ID remains stable across different model sizes for a given dataset but varies signiﬁcantly across data corpora. (a) OpenWebText (OWT): Shows the highest intrinsic dimension, reﬂecting its diverse and complex nature. (b) TinyStories: Exhibits the lowest intrinsic dimension, consistent with its synthetic, simple structure. (c) WikiText-103 (WT-103): Has an intermediate intrinsic dimension. • Self-attention: multi-head attention with m heads, model dimension dmodel , causal masking, and standard projection matrices. • MoE feed-forward network: a routing network selecting the top-k experts. Each expert is an MLP with hidden width dﬀ and a GELU nonlinearity. Outputs of the selected experts are combined via a softmax over routing logits. The active parameter budget per input,

Neﬀ = LT Πattn + k Πexp ,

counts only the parameters that contribute to computation for a given token. This is the quantity that appears in our approximation and covering-number bounds, and serves as the scaling variable in model-capacity experiments. A.4

Hyperparameters and Capacity Grids

We operate in the small-to-medium compute regime to allow dense parameter sweeps and low-noise log–log regressions. Unless otherwise stated, we use: • sequence length ℓ ∈ {128, 256}, • m = 4 attention heads, • AdamW with learning rate 3 × 10−4 and weight decay 0.1. Model-scaling grid. To vary the active capacity over 12 orders of magnitude, we sweep (LT , dﬀ ) ∈ {2, 3, 4, 5, 6, 8} × {256, 384, 512, 640, 768, 896, 1024, 1280, 1536}, combined with MoE conﬁgurations M ∈ {4, 8, 16} and k ∈ {1, 2}. We avoid extremely large M or k to remain within GPU memory while still covering a wide range of effective capacities. A.5

Scaling Protocols

We run three classes of experiments that map directly onto the theoretical decomposition of the generalization bound. 15

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Intrinsic Dimension Scaling Law (L(g-L(g): TinySt(ries

Intrin.ic Dimen.ion Sc ling L 2 (Log-Log): OpenWebText

23.0

GPT2

46 45

LARGE

MEDIUM

Me n ID pe- Mode( XL Medi n ID pe- Mode(

43 0.2

0.4

0.6

0.8 1.0 Number of P r meter.

1.2

1.4

1.6 1e9

Average In.rinsi Dimension

32.5

22.5

A.erage Intrinsic Dimension

A1er ge Intrin.ic Dimen.ion

LARGE

22.0 21.5

MEDIUM

21.0 20.5 20.0 19.5

(a) OpenWebText (OWT) Dataset

In.rinsi Dimension S aling La1 (Log-Log): WikiText-103

33.0

GPT2

Mean ID )er M( el XL Me ian ID )er M( el 0.2

0.4

0.6

0.8 1.0 Number of Parameters

1.2

1.4

(b) TinyStories Dataset

32.0 GPT2

MEDIUM

0.2

0.4

LARGE

31.5 31.0 30.5

1.6 1e9

Mea) ID per Model XL Media) ID per Model 0.6

0.8 1.0 Number of Parame.ers

1.2

1.4

1.6 1e9

Figure 3: Intrinsic Dimension Scaling Law: d¯ versus Number of Parameters (N ). The plots show the relationship ¯ (log-log scale) and the Number of Parameters (N ) across the four between the Average Intrinsic Dimension (d) GPT-2 model sizes for each dataset. This relationship is modeled as a power law, d¯ ∼ N β , where the slope of the linear ﬁt determines the ID Scaling Exponent β. The results suggest that d¯ exhibits weak or negligible scaling with model size across all three corpora. (1) Model scaling. For a ﬁxed dataset, we train all models in the sweeping grid under a common token budget Dﬁxed . After convergence, we record the validation loss L and ﬁt a linear regression of log L versus log Neﬀ over the region where routing is negligible. The slope yields the empirical model exponent α bN . (2) Data scaling. We choose a representative “base” architecture (e.g., LT = 4, dmodel = 128, dﬀ = 512, M = 8, k = 2) and vary the number of training tokens D ∈ {5 × 104 , 105 , 2 × 105 , 4 × 105 , 8 × 105 }. For each value of D we adjust the number of optimization steps to keep the total token budget ﬁxed. A regression of log L on log D yields the empirical data exponent α bD . (3) Routing ablations. To isolate the inﬂuence of routing combinatorics, we ﬁx both the base architecture and the token budget, and vary the expert conﬁguration according to k ∈ {1, 2}, M ∈ {2k, 4k, 8k}. For each (M, k) pair we train a model, record its validation loss, and compute the theoretical routing term LT ℓk log(eM/k). Plotting loss against this term reveals the routing-dominated regime and the predicted crossover into the power-law region governed by (Neﬀ , D). The models considered in this part are listed in Tables 3 and 4. A.6

Full Data Scaling and ModelScaling sensitivity

Extended Results: Empirical Comparison to Prior MoE Scaling Laws

We align our theoretical exponents model scaling with the empirical laws reported in EL (Tian et al., 2025) and Joint MoE Scaling Laws (JMSL) (Ludziejewski et al., 2025). In the power-law (routing-negligible) regime, our theory implies two parameter-free identities that can be checked directly: theory αD =

theory αN

theory αC =

, theory

theory αN

(20) theory 1 + αN 2 + αN Accordingly, we test our theory by comparing empirical exponents to these identities, without estimating the intrinsic dimensionality. We align our theoretical exponents model scaling αN , data scaling αD , and compute scaling αC with the empirical laws reported in Joint MoE Scaling Laws (JMSL) (Ludziejewski et al., 2025). They supposed that, for a ﬁxed number of experts E, the validation loss follows a Chinchilla-style surface µ(E)

L(Nact , D | E) = m(E) Nact + n(E) D ν(E) + c, (21) and then introduces an expert-dependent joint law by making both the prefactors and exponents functions of E (ultimately via a smoothed Ê). Concretely, α+γ ln Ê L(Nact , D, Ê) = a Ê δ Nact + b Ê ω D β+ζ ln Ê + c,

(22)

Generalization and Scaling Laws for Mixture-of-Experts Transformers

df f

Nact

dmodel

nheads

2 2 2 2 2

256 384 512 768 1024

3,148,800 3,673,600 4,198,400 5,248,000 6,297,600

512 512 512 512 512

4 4 4 4 4

8 8 8 8 8

2 2 2 2 2

4 4 4 4 4 4

256 384 512 768 1024 1536

6,297,600 7,347,200 8,396,800 10,496,000 12,595,200 16,793,600

512 512 512 512 512 512

4 4 4 4 4 4

8 8 8 8 8 8

2 2 2 2 2 2

6 6

768 1024

15,744,000 18,892,800

512 512

4 4

8 8

2 2

1024

25,190,400

512

Table 3: Model conﬁgurations used in the scaling experiments for each data set. Each row corresponds to a distinct Transformer-MoE model used in the model-scaling sweep. The effective active parameter count Nact varies with the number of layers and FFN width, while dmodel = 512, nheads = 4, M = 8, k = 2, and sequence length ℓ = 256 are held ﬁxed. LT

dﬀ

Nact

dmodel

nheads

4 4 4 4 4

512 512 512 512 512

8,396,800 8,396,800 8,396,800 8,396,800 8,396,800

512 512 512 512 512

4 4 4 4 4

8 8 8 8 8

2 2 2 2 2

5 × 104 1 × 105 2 × 105 4 × 105 8 × 105

Table 4: Model conﬁguration used in the data-scaling experiments. The architecture is held ﬁxed while the training token budget D is varied. Unless stated otherwise, the conﬁguration corresponds to a representative mid-scale Transformer–MoE model with ﬁxed active parameter budget Nact .

which reduces to (21) when E is held ﬁxed, with µ(E) = α + γ ln Ê and ν(E) = β + ζ ln Ê. (Ludziejewski et al., 2025) provides ﬁtted coefﬁcients, per-E reduced exponents, and a compute-optimal analysis under the training-FLOPs proxy F = 6 Nact D. We adopt Their counting conventions for active parameters and FLOPs when ﬁtting/reading off exponents. From their per-expert reduction (their Eq. (5) and Table 4), we read exponents µ(E) and ν(E) and map α̂N (E) := − µ(E),

α̂D (E) := − ν(E).

JMSL reports that as E increases the magnitude of ν(E) increases (improved data efﬁciency) while |µ(E)| slightly decreases. The resulting values and the induced compute exponent are summarized in Table 5. Compute-optimal frontier and induced αC . On compute-optimal slices with the constraint F = 6 Nact D, the separable form (21) implies closed-form optima (JMSL Eq. (9)): opt Nact (F, E) = G(E)

ν(E) µ(E)+ν(E)

F 6

D opt (F, E) = G(E)−1

µ(E) µ(E)+ν(E)

F 6

for a known G(E) depending on m(E), n(E), µ(E), ν(E). Along this frontier, the compute scaling exponent is α̂C (E) =

α̂N (E) α̂D (E) . α̂N (E) + α̂D (E)

(23)

We use JMSL’s F -proxy and the per-E (µ, ν) to report α̂C (E) and to compare the token-to-parameter ratio drift with E (JMSL’s Finding 1). Consistency of β across scaling regimes. Our theoretical framework introduces a smoothness exponent β that appears in both model scaling, through αN = 2β/d, and data scaling, through αD = 2β/(2β +d). When estimating β from empirically ﬁtted scaling slopes α bN and α bD , we ﬁnd that the resulting values βmodel and βdata do not necessarily 17

Generalization and Scaling Laws for Mixture-of-Experts Transformers Da(a Scaling: Varying β, Fi,ed d = 45.0 (ow() 1.0

0.9

log(val_loss)

1.05 1.00 0.95

Da(a Scaling: Varying β, Fi,ed d = 32.0 (w(103)

Gro)nd (r)(h Empirical Fi(: αD̂ = 0.068 β = 0.5, αD = 0.042 β = 1.0, αD = 0.080 β = 1.5, αD = 0.115

0.8

1.00 0.95

0.7

0.90 0.85

Gro)nd (r)(h Empirical Fi(: αD̂ = 0.062 β = 0.5, αD = 0.030 β = 1.0, αD = 0.059 β = 1.5, αD = 0.086

1.05

log(val_loss)

1.10

log(val_loss)

Da(a Scaling: Varying β, Fixed d = 23.0 ((in,s(ories)

Gro)nd (r)(h Empirical Fi(: αD̂ = 0.064 β = 0.5, αD = 0.022 β = 1.0, αD = 0.043 β = 1.5, αD = 0.062

1.15

0.90 0.85 0.80 0.75

0.6

0.80

0.70 13.0

13.5

14.0

14.5

15.0

15.5

log(D(okens)

16.0

16.5

17.0

13.0

13.5

(a) OpenWebText (OWT)

14.0

14.5

15.0

15.5

log(D(okens)

16.0

16.5

17.0

13.0

13.5

(b) TinyStories

14.0

14.5

15.0

15.5

log(D(okens)

16.0

16.5

17.0

Figure 4: Data-scaling results (ﬁxed d, varying β). Empirical estimates of the data-scaling law (loss versus data size D) across datasets, compared with theoretical lines obtained by ﬁxing the dataset-speciﬁc intrinsic dimension d and varying β. The empirical ﬁts are most consistent with β = 1.5 for OpenWebText and β = 1.0 for TinyStories and WikiText-103. Model Scali g: Varyi g β, Fi)ed d = 45.0 (o(t)

1.20

Model Scaling: Varying β, Fixed d = 32.0 ()t103)

1.16

1.00 0.95 0.90

1.14

0.85

1.12

1.10

1.150

Gro(nd tr(th Empirica Fit: αN̂ = 0.022 β = 0.5, αN = 0.043 β = 1.0, αN = 0.087 β = 1.5, αN = 0.130

log()al_loss)

log(val_loss)

1.18

Model Sca ing: Varying β, Fixed d = 23.0 (tinystories)

1.05

Grou d truth Empirical Fit: αN̂ = 0.025 β = 0.5, αN = 0.022 β = 1.0, αN = 0.044 β = 1.5, αN = 0.067

Ground truth E pirical Fit: αN̂ = 0.033 β = 0.5, αN = 0.031 β = 1.0, αN = 0.062 β = 1.5, αN = 0.094

1.125 1.100

log((al_loss)

1.22

1.075 1.050 1.025 1.000

0.80

0.975

1.08

0.950 15.0

15.5

16.0

log(Nact)

16.5

15.0

17.0

(a) OpenWebText (OWT)

15.5

16.0

16.5

log(Nact)

17.0

15.0

(b) TinyStories

15.5

16.0

log(Nact)

16.5

17.0

Figure 5: Model-scaling results (ﬁxed d, varying β). Empirical estimates of the model-scaling law (loss versus active parameter count Nact ) across datasets. For each plot, the dataset-speciﬁc intrinsic dimension d is ﬁxed and theoretical lines are shown for varying β, with theoretical exponent αN = 2β/d. The empirical ﬁts are most consistent with β = 0.5 across all three datasets. coincide exactly (see Figures 4 and 5). Such discrepancies are expected. Model-scaling curves are typically cleaner and more directly dominated by approximation error, whereas data-scaling curves are more sensitive to optimization noise, ﬁnite-sample effects, and, in the MoE setting, routing overhead terms of the form k log(eM/k). Similar differences between model- and data-derived exponents have been reported in prior empirical studies of both dense and MoE language models (Kaplan et al., 2020; Hoffmann et al., 2022; Krajewski et al., 2024; Ludziejewski et al., 2025). While the two estimates of β need not match exactly, they remain broadly consistent in magnitude across datasets and scaling regimes. This qualitative agreement suggests that the smoothness-based exponents used in our analysis provide a reasonable statistical reference point for interpreting the observed scaling behavior. B.1

Our Theory ⇒ Empirical Mapping and Consistency Checks

Theoretical exponents. In the power-law (routing-negligible) regime, our theory characterizes theory αN =

2β , d

theory αD =

2β , 2β + d

theory αC =

β . β+d

This implies two parameter-free identities that can be checked directly against JMSL: theory αD =

theory αN

, theory

1 + αN

theory αC =

theory αN theory 2 + αN

theory theory αN αD theory theory αN + αD

(24)

Generalization and Scaling Laws for Mixture-of-Experts Transformers

(b) Ampliﬁcation factor κ(E)

(a) Data-to-model scaling ratio

Figure 6: Empirical ampliﬁcation patterns across experts. Thus, for each expert count E, a simple test is whether the empirical pair α̂N (E), α̂D (E) approximately satisﬁes α̂D ≈ α̂N /(1 + α̂N ), and whether α̂C (E) computed from (23) aligns with the right-hand identity in (20). Deviations quantify where practice departs from the smooth-manifold asymptotics (ﬁnite-sample regime, tokenization/entropy ﬂoors c, training schedules, or mild routing side-effects). B.2

Per-E identity checks

For each expert count E, we compare the empirical pair α̂N (E), α̂D (E) against the parameter-free identities pred αD (E) =

α̂N (E) , 1 + α̂N (E)

pred αC (E) =

pred α̂N (E) αD (E) α̂N (E) = . pred 2 + α̂N (E) α̂N (E) + αD (E)

Deﬁning residuals pred (E), rD (E) = α̂D (E) − αD

pred rC (E) = α̂C (E) − αC (E),

µ(E)

ν(E)

αN

αD

αC

pred αD

pred αC

rel_errD

rel_errC

1 2 4 8 16 32

-0.1817 -0.1780 -0.1731 -0.1676 -0.1617 -0.1557

-0.1965 -0.2065 -0.2192 -0.2338 -0.2494 -0.2652

0.1817 0.1780 0.1731 0.1676 0.1617 0.1557

0.1965 0.2065 0.2192 0.2338 0.2494 0.2652

0.0944 0.0955 0.0967 0.0976 0.0980 0.0981

0.1537 0.1511 0.1475 0.1435 0.1391 0.1347

0.0832 0.0817 0.0796 0.0773 0.0748 0.0722

0.0427 0.0553 0.0716 0.0902 0.1102 0.1304

0.0111 0.0138 0.0170 0.0203 0.0232 0.0258

0.2779 0.3666 0.4855 0.6287 0.7917 0.9684

0.1335 0.1697 0.2142 0.2625 0.3114 0.3582

Table 5: Per-E identity checks.We set αN (E) = −µ(E) and αD (E) = −ν(E). The compute-optimal exponent is pred pred (E)αD (E) αC (E) = ααNN(E)+α . Predicted exponents use αD = αN /(1 + αN ) and αC = αN /(2 + αN ). Residuals are D (E) pred pred rD = α D − α D and rC = αC − αC ; relative errors are r/pred.

we observe systematic, monotone-positive deviations across (see Table 5). Quantitatively, the mean absolute error is MAE[rD ] = 0.083 and MAE[rC ] = 0.0186, with MAPE[rD ] = 58.7% and MAPE[rC ] = 24.2%. Both gaps grow with E: (rD , rC ) ≈ (0.043, 0.011) at E=1 → (0.130, 0.026) at E=32. Equivalently, Figure 6a show that the ratio α̂D /α̂N increases with E (from ∼ 1.08 at E=1 to ∼ 1.70 at E=32), indicating stronger-than-idealized data-side gains. A parsimonious correction that captures the trend is an E-dependent ampliﬁcation of the identity, α̂D (E) ≈ κ(E)

α̂N (E) , 1 + α̂N (E)

κ(E) > 1,

with empirical multipliers κ(1) ≈ 1.28, κ(2) ≈ 1.37, κ(4) ≈ 1.47, κ(8) ≈ 1.62, κ(16) ≈ 1.79, κ(32) ≈ 1.97. A log-linear form κ(E) = A + B ln E provides a simple, testable ﬁt (see 6b). 19

Generalization and Scaling Laws for Mixture-of-Experts Transformers

(a) Per-E cross-budget ﬁts: D∗ vs. N ∗

(b) Tokens-to-parameters drift across E (per budget)

Figure 7: Cross-budget compute-optimal allocations. (a) Cross-budget slopes for each E. (b) Tokens-to-parameters drift across experts (per budget). ∗ Compute-optimal allocation vs. experts. JMSL reports compute-optimal allocations (Nact (F, E), D ∗ (F, E)), allowing us to test our predicted tokens-to-parameters drift. Our theory (via (20)) imposes a ﬁxed relationship between the model and data exponents, and along the compute frontier (F = 6Nact D) it implies the losscompute exponent ∗ (23). Empirically, JMSL ﬁnds that, at ﬁxed compute F , increasing the number of experts E raises D∗ and lowers Nact (Table 1 and isoFLOPs plots), consistent with a larger effective data exponent at higher E. This pattern appears directly in Fig. 7b: the ratio D∗ /N ∗ increases monotonically with E for every budget C, meaning the compute-optimal allocation shifts toward more tokens per active parameter as expertization grows. Moreover, the slope of D∗ /N ∗ versus E becomes shallower at larger budgets (e.g., the gain from E=1 → 16 is roughly ×4.5 at 1020 but only ×3.0 at 1021 ), suggesting diminishing marginal ampliﬁcation with C.

Theory vs. EL (Tian et al., 2025). Empirically, Tian et al. (2025) corroborate these predictions: their Efﬁciency Leverage metric (EL) grows as a power law in compute and (inverse) activation ratio while expert granularity modulates EL via a log-polynomial term with a stable optimum; their compute-optimal allocations favor more data and smaller active models than dense baselines, and the Ling-mini-beta case study (0.85B active, 17.5B total) matches a 6.1B dense model at > 7× lower compute consistent with our view that active capacity drives gains and routing is a secondary, mostly constant-factor overhead (Tian et al., 2025). Theory vs. JMSL (Ludziejewski et al., 2025). JMSL found that, for a ﬁxed number of experts E, validation loss follows a Chinchilla-style surface and then introduces an expert-dependent joint law by letting both prefactors and exponents be functions of E (via a smoothed Ê). For each E, our theory suggests two simple checks: whether the empirical pair (α̂N (E), α̂D (E)) approximately satisﬁes α̂D ≈ α̂N /(1 + α̂N ), and whether α̂C (E) computed from (23) aligns with the identity in (20). In the JMSL data we observe systematic, monotone-positive deviations: the ratio α̂D /α̂N increases with E (Table 5, Fig. 6a in appendix B), indicating stronger-than-idealized data-side gains, plausibly due to ﬁnite-sample effects, tokenization/entropy ﬂoors c, training schedules, or mild routing side-effects. ∗ Consistent with our compute-frontier geometry, JMSLs compute-optimal allocations (Nact , D ∗ ) show that D∗ /N ∗ increases monotonically with E for every budget C (Fig. 7b in appendix B), i.e., more tokens per active parameter as expertization grows, with a shallower slope at larger budgets, suggesting diminishing marginal ampliﬁcation with compute. Practical implication. At a ﬁxed training budget C, increasing E should allocate relatively fewer active parameters and more tokens (higher D⋆ /N ⋆ ), improving data efﬁciency.

C Additional Preliminaries C.1

Notation and basic deﬁnitions

Sets, measures, and expectations. For a measurable space (X , B) and probability measure Q on (X , B), we write X ∼ Q for a random variable taking values in X , and E[·] for expectation with respect to Q. The symbol P refers to probability. 20

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Covering numbers and metric entropy. Let (F, ρ) be a metric space. An δ-cover of F is a ﬁnite subset {f1 , . . . , fN } ⊂ F such that for every f ∈ F there exists j ∈ [N ] with ρ(f, fj ) ≤ δ. The covering number Vershynin (2018) is N δ, F, ρ := min{ N ∈ N : ∃ δ-cover of F with N elements } The (log) metric entropy is log N δ, F, ρ . We often use ρ(f, g) = kf − gk∞ . Vectors, matrices, tensors. We use bold uppercase (e.g., H) for matrices and bold lowercase (e.g., h) for vectors when convenient. For a vector v ∈ Rd , deﬁne the entrywise ℓ∞ norm kvk∞ := max |vi | 1≤i≤d

For a matrix A ∈ Rdout ×din , we use two norms: kAk∞ := maxi,j |Aij | For an embedding matrix H ∈ Rℓ×demb (sequence length ℓ, embedding size demb ), we adopt kHk∞ :=

max t∈[ℓ], j∈[demb ]

|Ht,j |

Asymptotic notation. For nonnegative quantities a, b, a ≲ b means a ≤ C b for an absolute constant C; a ≃ b e to hide polylogarithmic factors in problem parameters. means a ≲ b and b ≲ a. We write O(·) C.2

Loss and learning rule

Squared regression error. Given a predictor T : X → R, the population (squared) risk under Q is L(T ) := E (T (X) − f (X))2 = kT − f k2L2 (Q) Empirically, for samples {xi }ni=1 , the empirical risk is 2 1X T (xi ) − f (xi ) n i=1 n

Ln (T ) := Empirical risk minimization (ERM). T̂n ∈ T satisfying

An empirical risk minimizer (ERM) over T is any measurable selection T̂n ∈ arg min Ln (T ) T ∈T

C.3

Input-Lipschitz bound (dense)

There exists Lin Rℓ×demb ,

≥ 1, depending polynomially on (ℓ, m, demb , wFFN , LFFN ) and on κ, such that for all H, H ′ ∈

kBdense (θ, H) − Bdense (θ, H ′ )k∞ ≤ Lin kH − H ′ k∞ , with the schematic scaling (hiding absolute constants and activation Lipschitz factors) (dense) LFFN . Lin ≲ 1 + CLN κ ℓ m d2emb + demb wFFN | {z } | {z } (dense)

MHA sensitivity

C.4

FFN sensitivity

Softmax Lipschitz

We work with the temperature-τ softmax στ : Rm → ∆m−1 , eui /τ στ (u)i = Pm u /τ , j j=1 e

τ > 0.

Its Jacobian is

1 Diag(στ (u)) − στ (u)στ (u)⊤ . τ Lemma C.1 (Softmax is (2τ )−1 -Lipschitz in k · k∞ ). For any u, v ∈ Rm , Jτ (u) =

kστ (u) − στ (v)k∞ ≤ 21

1 ku − vk∞ . 2τ

(25)

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Proof. By the mean value theorem, στ (u) − στ (v) =

R1 0

Jτ v + t(u − v) (u − v) dt. Thus

sup kJτ (w)k∞→∞ ku − vk∞ .

kστ (u) − στ (v)k∞ ≤

w∈[u,v]

Row i of Jτ (w) has entries τ1 σi (δij − σj ). Summing absolute values in row i: m 2 X X 1 1 1 σi |δij − σj | = σi (1 − σi ) + σi σj = σi (1 − σi ) ≤ , τ τ τ 2τ j=1 j̸=i

since x(1 − x) ≤ 1/4. Hence kJτ (w)k∞→∞ ≤ 1/(2τ ) uniformly, proving the claim. Corollary C.2 (Row-wise softmax on matrices). Let S, S ′ ∈ Rℓ×ℓ and apply softmaxτ row-wise to obtain A = softmaxτ (S) and A′ = softmaxτ (S ′ ). Then 1 kA − A′ k∞ ≤ kS − S ′ k∞ . 2τ √ In scaled attention, τ = dh , hence kA − A′ k∞ ≤ 2√1d kS − S ′ k∞ . h

Remark C.3 (Row-stochastic contraction). If A is row-stochastic (each row in ∆ℓ−1 ), then kAk∞→∞ = 1 and kAV k∞ ≤ kV k∞ for any V under operator norm. When using entrywise k · k∞ , we will also invoke the safe bound kAV k∞ ≤ ℓkAk∞ kV k∞ (cf. C.4). C.5

Matrix norm identities and product bounds

Lemma C.4 (Basic matrix inequalities). Let A ∈ Rp×q , B ∈ Rq×r , u ∈ Rq , H ∈ Rℓ×q . (a) Operator submultiplicativity: kABk∞→∞ ≤ kAk∞→∞ kBk∞→∞ . (b) Entrywise-to-operator: kAk∞→∞ ≤ q kAk∞ and kAuk∞ ≤ q kAk∞ kuk∞ . (c) Products in entrywise norm: kABk∞ ≤ q kAk∞ kBk∞ and kHBk∞ ≤ q kHk∞ kBk∞ . (d) Row-stochastic contraction: If each row of A sums to 1 and is nonnegative, then kAk∞→∞ = 1 and kABk∞→∞ ≤ kBk∞→∞ . (e) Hadamard product: kA

Bk∞ ≤ kAk∞ kBk∞ .

Proof. Parts (a) and (d) follow from deﬁnitions; (b)–(c) use kAk∞→∞ = maxi norm inequalities; (e) is immediate from entrywise multiplication.

j |Aij | ≤ qkAk∞ and standard

Lemma C.5 (Scaled dot-product attention ingredients). Let H√ ∈ Rℓ×demb , and WQ , WK , WV satisfy kW• k∞ ≤ κ. Set Q = HWQ , K = HWK , V = HWV , scores S = QK ⊤ / dh and A = softmax√dh (S) (row-wise). Then kQk∞ , kKk∞ , kV k∞ ≤ demb κ kHk∞ , p kSk∞ ≤ dh d2emb κ2 kHk2∞ , 1 (row-wise softmax; C.2), kA − A′ k∞ ≤ √ kS − S ′ k∞ 2 dh kAV k∞ ≤ ℓ kAk∞ kV k∞ (entrywise bound; cf. C.3). C.6

Taylor remainder on manifolds

Let M ⊂ RD be a compact C 1 submanifold of intrinsic dimension d. Fix a coordinate chart ϕ : U → V ⊂ Rd that is bi-Lipschitz with constants L, L′ > 0: L−1 kx − yk ≤ kϕ(x) − ϕ(y)k ≤ Lkx − yk for x, y ∈ U and similarly for ϕ−1 on V . Lemma C.6 (Local Taylor remainder for C β functions on M). Let f ∈ C β (M) with β > 0, s = bβc, and α = β − s ∈ [0, 1). Fix a chart (U, ϕ) and a point x0 ∈ U . Let P be the degree-s Taylor polynomial of f ◦ ϕ−1 at z0 = ϕ(x0 ) in local coordinates, mapped back to M by PM := P ◦ ϕ. Then for every x ∈ U , |f (x) − PM (x)| ≤ C(β, d, L, L′ ) kf kC β (U ) kϕ(x) − ϕ(x0 )kβ ≤ C ′ kf kC β (U ) kx − x0 kβ . In particular, on any chart of diameter at most r, supx∈U |f (x) − PM (x)| ≤ C rβ . 22

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Proof. Apply the standard Taylor remainder in Rd to f ◦ ϕ−1 ∈ C β (V ) and use bi-Lipschitz distortion to convert local distance in V to ambient distance in M; constants depend polynomially on L, L′ and on bounds of chart derivatives.

C.7

Partitions of unity with bounded overlap

Lemma C.7 (Partition of unity subordinate to a bounded-overlap cover). Let M ⊂ RD be a compact C 1 manifold and ﬁx r ∈ (0, 1). There exists a ﬁnite cover of M by open sets {Uν }N ν=1 such that: (i) diam(Uν ) ≤ r, (ii) each ∞ x ∈ M belongs to at most s (d) sets (bounded overlap), and (iii) a C partition of unity {φν }N 0 ν=1 subordinate to P {Uν } with 0 ≤ φν ≤ 1, ν φν (x) = 1 for all x, and the Lipschitz bound C(d) for all x ∈ M and all ν. r Consequently, at every x at most s0 (d) of the φν (x) are nonzero. k∇φν (x)k ≤

Proof. Step 1: A bounded-overlap cover by small balls. Work with the metric induced on M by the ambient Euclidean norm on RD . Since M is compact and C 1 , there exists a ﬁnite atlas {(Vj , φj )}Jj=1 such that: • each φj : Vj → Rd is a bi-Lipschitz chart onto its image, and • on each Vj the volume of metric balls is comparable to that of Euclidean d-balls in Rd (with constants depending only on d and the chart regularity). Fix a small radius ρ > 0 proportional to r (we will specify it below), and choose a maximal ρ/2-separated set {xν }N ν=1 ⊂ M , i.e., ρ kxν − xµ k ≥ for all ν 6= µ, 2 and the union of balls B(xν , ρ) (intersection with M ) covers M . Such a set exists by a standard greedy argument: repeatedly add points that are at distance at least ρ/2 from all previously chosen ones until no further point can be added. Deﬁne

Uν := B(xν , ρ) ∩ M By construction, {Uν }N is an open cover of M . Choosing ρ ≤ r/2 ensures diam(Uν ) ≤ r. ν=1 We now show that the overlap multiplicity is bounded. Fix any x ∈ M and consider the index set I(x) := {ν : x ∈ Uν } = {ν : kx − xν k < ρ} For each ν ∈ I(x), the balls B(xν , ρ/4) are pairwise disjoint (by ρ/2-separation) and all lie inside B(x, 2ρ): B(xν , ρ/4) ⊂ B(x, ρ + ρ/4) ⊂ B(x, 2ρ) Using the chart bi-Lipschitz property and volume comparability on M , there exist constants c1 , c2 > 0 depending only on d and the atlas such that for any ball BM (y, t) in M (with respect to the induced metric), c1 td ≤ VolM BM (y, t) ≤ c2 td Hence, the volumes of the disjoint balls {B(xν , ρ/4)}ν∈I(x) satisfy X |I(x)| · c1 (ρ/4)d ≤ VolM B(xν , ρ/4) ≤ VolM B(x, 2ρ) ≤ c2 (2ρ)d ν∈I(x)

Therefore

c2 (2ρ)d = C(d) for some constant C(d), c1 (ρ/4)d depending only on d and the chart regularity. Setting s0 (d) := C(d) yields the bounded-overlap property: each x ∈ M belongs to at most s0 (d) of the sets Uν . |I(x)| ≤

Step 2: Smooth bump functions subordinate to the cover. Let η : RD → [0, 1] be a ﬁxed C ∞ bump function such that: η(z) = 1 if kzk ≤ 1/2, η(z) = 0 if kzk ≥ 1, k∇η(z)k ≤ C0 (d) 23

Generalization and Scaling Laws for Mixture-of-Experts Transformers

For each ν, deﬁne

ψν (x) := η

x − xν ρ

x∈M

Then ψν ∈ C ∞ (M ), ψν (x) ∈ [0, 1], and ρ , ψν (x) = 0 if kx − xν k ≥ ρ 2 Hence supp(ψν ) ⊂ Uν , i.e., each ψν is subordinate to Uν . Moreover, 1 x − xν C0 (d) k∇ψν (x)k = ∇η ≤ ρ ρ ρ ψν (x) = 1

if kx − xν k ≤

Because the balls B(xν , ρ/2) cover M , for every x ∈ M there exists at least one index ν with kx − xν k ≤ ρ/2, hence ψν (x) = 1. Therefore N X S(x) := ψµ (x) ≥ 1 for all x ∈ M. µ=1

On the other hand, for each ﬁxed x, at most s0 (d) of the ψµ (x) are nonzero (because ψµ vanish outside Uµ and the Uµ have overlap at most s0 (d)), and each ψµ (x) ≤ 1, so 1 ≤ S(x) ≤ s0 (d) for all x ∈ M. Step 3: Normalize to obtain a partition of unity and bound gradients. Deﬁne ψν (x) ϕν (x) := , x∈M S(x) Then each ϕν ∈ C ∞ (M ), 0 ≤ ϕν ≤ 1, and N X

1 X ψν (x) = 1 S(x) ν=1 N

ϕν (x) =

ν=1

for all x ∈ M.

Moreover, supp(ϕν ) ⊂ supp(ψν ) ⊂ Uν , so {ϕν } is subordinate to the cover {Uν }. We now bound the gradient. Using the quotient rule, ∇ϕν (x) =

∇ψν (x)S(x) − ψν (x)∇S(x) , S(x)2

∇S(x) =

N X

∇ψµ (x).

µ=1

Thus

k∇ψν (x)k S(x) + |ψν (x)| k∇S(x)k S(x)2 k∇ψν (x)k k∇S(x)k + ≤ S(x) S(x)2 We already know 1 ≤ S(x) ≤ s0 (d); hence 1/S(x), 1/S(x)2 ≤ 1. Also, for each x, at most s0 (d) functions ψµ are nonzero, so X X C0 (d) k∇S(x)k = ∇ψµ (x) ≤ k∇ψµ (x)k ≤ s0 (d) · ρ µ k∇ϕν (x)k ≤

µ:ψµ (x)̸=0

Combining these bounds gives C0 (d) s0 (d)C0 (d) C(d) + ≤ ρ ρ ρ for some constant C(d) depending only on d and the atlas. Choosing ρ proportional to r (e.g., ρ = r/2 as above) yields C(d) k∇ϕν (x)k ≤ r k∇ϕν (x)k ≤

Finally, since at most s0 (d) of the ψν (x) are nonzero at any x, the same holds for the ϕν (x), and the consequently statement follows. Remark C.8 (From partition of unity to k-sparse mixing). SinceP at most s0 (d) functions are nonzero at any x, choosing k ≥ s0 (d) allows a router to implement a k-sparse mixture ν∈S(x) wν (x)Eν (x) with weights wν = φν (or an approximation thereof), as used in the approximation construction. 24

Generalization and Scaling Laws for Mixture-of-Experts Transformers

C.8

Useful norm conversions for attention chains

For attention chains H 7→ Q 7→ K 7→ S 7→ A 7→ O = AV , combining Lemmas C.4,C.5 and C.2 yields the schematic bounds used in the main text: kQ − Q′ k∞ , kK − K ′ k∞ , kV − V ′ k∞ ≤ demb κ kH − H ′ k∞ + demb kW• − W•′ k∞ kH ′ k∞ , p kS − S ′ k∞ ≲ dh d2emb κ kQ − Q′ k∞ + kK − K ′ k∞ , 1 kA − A′ k∞ ≤ √ kS − S ′ k∞ , 2 dh kAV − A′ V ′ k∞ ≤ kA − A′ k∞→∞ kV k∞ + kA′ k∞→∞ kV − V ′ k∞ ≲ kA − A′ k∞ kV k∞ + kV − V ′ k∞ , | {z } | {z } attention change

value change

using kA′ k∞→∞ = 1. These estimates justify the polynomial dependence on (ℓ, m, demb , κ) that appears in the MHA stability bound.

D Full Model Speciﬁcation Deﬁnition D.1 (Dense Transformer Block). We deﬁne the dense residual block based on the canonical Transformer structure Vaswani et al. (2017). Let the input to block j be H ∈ Rℓ×demb with tokenwise ℓ∞ -bound kHk∞ ≤ M0 . The block maps H 7→ Hout via the sequential steps: e = H + MHAψ H , H (26) e + FFNχ H e , Hout = H (27) where θ = {ψ, χ} collects all block parameters. Layer Normalization (LN) Ba et al. (2016).

For each token t ∈ [ℓ],

LNγ,β (H)t = γ

H − µ(Ht )1 pt + β, σ 2 (Ht ) + ϵ

γ, β ∈ Rdemb

with learned scale/shift (γ, β) and small ϵ > 0. We will use the uniform magnitude bounds kγk∞ ≤ κ,

kβk∞ ≤ κ

(28)

Multi-Head Self-Attention (MHA). Let the number of heads be m, and the per-head key/query/value dimension be dh (so dk = dv = dh ) such that m dh = demb . The input to MHA is H. For head h ∈ [m] we have projections (h)

Q(h) = HWQ ,

(h)

K (h) = HWK ,

(h)

V (h) = HWV ,

(h)

WQ , WK , WV

∈ Rdemb ×dh ,

scores S (h) = √1d Q(h) (K (h) )⊤ ∈ Rℓ×ℓ , row-wise A(h) = softmax(S (h) ), head output O(h) = A(h) V (h) ∈ Rℓ×dh , h and (h) WO , WO ∈ R(mdh )×demb MHAψ (H) = Concatm h=1 O We assume the uniform parameter bounds (h)

(h)

kWQ k∞ , kWK k∞ , kWV k∞ , kWO k∞ ≤ κ Positionwise Feed-Forward Network (FFN). most WFFN . For ℓ = 1, . . . , LFFN and token t, (0)

e t, = (H)

(29)

We use an LFFN -layer MLP applied tokenwise. Let widths be at (ℓ)

(ℓ−1)

= σ W (ℓ) zt

+ b(ℓ)

with activations σ (ReLU/GELU), W (1) ∈ RWFFN ×demb , W (ℓ) ∈ RWFFN ×WFFN for 2 ≤ ℓ ≤ LFFN − 1, and e t = zt(LFFN ) , and we assume W (LFFN ) ∈ Rdemb ×WFFN ; biases b(ℓ) are conformal. The FFN output is FFNχ (H) kW (ℓ) k∞ ≤ κ,

kb(ℓ) k∞ ≤ κ, 25

ℓ = 1, . . . , LFFN

(30)

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Output bound and architectural polynomial. We assume the block output is uniformly bounded by kBdense (θ, H)k∞ ≤ R for all admissible θ and kHk∞ ≤ M0 Havrilla & Liao (2024); Chen et al. (2022). Constants in our bounds will be expressed via a ﬁxed architecture polynomial (dense) Parch = poly ℓ, m, demb , wFFN , LFFN Hypothesis classes. We consider a hypothesis class T of predictors. In this paper, T is either the dense Transformer class (D.1) or the MoE class TMoE (D.1). We always assume supT ∈T kT k∞ ≤ R. D.1 MoE Transformer Experts and router. Using Deﬁnition 3.1, each block j has M experts {Ej,i }M i=1 , each a tokenwise MLP Ej,i : (j) d d Remb → Remb of depth LFFN and width wFFN sharing the same architecture. For a token t ∈ [ℓ], let Zt = (j−1) (j) d Ht ∈ Remb be the MoE input. The router in block j computes logits gj,i (Zt ) for i ∈ [M ] (by a bounded parametric map) and selects the hard top-k index set Sj,t (H (j−1) ) ⊂ [M ], |Sj,t (H (j−1) )| = k consisting of the k largest logits. The MoE output at token t is then the k-sparse mixture X X (j) (j) (j) (j) MoE(j) Z (j) t = wj,i (Zt ) Ej,i (Zt ), wj,i (Zt ) ≥ 0, wj,i (Zt ) = 1 (31) i∈Sj,t

i∈Sj,t (H (j−1) )

with router weights wj,i restricted to the selected indices. Readout and parameter set. A ﬁxed linear readout R : Rℓ×demb → R produces the scalar prediction Tθ (x) = R H (LT ) . We deﬁne the MoE Transformer network class as the set of all functions n o TMoE (D, M, k, LT , LFFN , wFFN , demb , m, κ, R) := Tθ : X → R kθk∞ ≤ κ, kTθ k∞ ≤ R

(32)

D.2 Per-block parameter counts (dense baseline). Ignoring biases for clarity and taking the standard choice dh = demb /m: Π(dense) d dh m + (mdh )demb = 4 d2emb , attn = 3 {z } | emb | {z } WQ ,WK ,WV

2 2 + (LFFN − 2) wFFN + wFFN demb = 2 demb wFFN + (LFFN − 2) wFFN , . With biases, add 3 mdh + demb (attention) and wFFN (LFFN − 1) + demb (FFN).

Π(dense) FFN = demb wFFN

Parameter-Lipschitz (baseline). For two parameter sets θ, θ′ with kθ − θ′ k∞ ≤ η and ﬁxed input H, (dense) LFFN kBdense (θ, H) − Bdense (θ′ , H)k∞ ≤ Cblock ℓ m d2emb + demb wFFN κη

(33)

obtained by summing the MHA and FFN parameter-perturbation bounds (each applied after LN) and absorbing LN perturbations into constants. These two baseline inequalities (33)-(25) are the dense counterparts of the MoE block bounds where the FFN sensitivity term gains an additional factor k from the k active experts. D.3 Routing patterns. For each block j ∈ [LT ] and token t ∈ [ℓ], the router chooses a k-subset Sj,t ⊂ [M ]. A routing pattern is the collection π = { Sj,t ⊂ [M ] : |Sj,t | = k, j = 1, . . . , LT , t = 1, . . . , ℓ } Let Π denote the set of all such patterns. Its cardinality is bounded by M L T ℓ e M |Π| ≤ (34) =⇒ log |Π| ≤ LT ℓ k log k k k using M k ≤ (e M/k) . For a ﬁxed π ∈ Π, denote by Tπ ⊂ TMoE the subclass with deterministic routing equal to π. Then TMoE =

[ π∈Π

Tπ

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Proof of theorem 3.4

Assume supT ∈TMoE kT k∞ ≤ R and kf k∞ ≤ R. Let T̂n ∈ arg minT ∈TMoE Ln (T ) be an ERM for squared loss L(T ) = E[(T (X) − f (X))2 ]. Then for any cover scale δ ∈ (0, 1), via the standard biasvariance decomposition described in lemma E.1, we have the following inequality s ! log N δ, TMoE , k · k∞ 2 2 2 E kT̂n − f kL2 (Q) ≤ 3 inf kT − f k∞ + C R + δ (35) T ∈TMoE n for an absolute constant C > 0. Equivalently,

E kT̂n − f k2L2 (Q)

≲

inf

T ∈TMoE

kT − f k2∞

log N δ, TMoE , k · k∞ + n

+ δ

(36)

where ≲ hides universal constants and a factor R2 . Now it remains to bound the two terms of 36. The ﬁrst is bounded in theorem 3.2 as: inf

T ∈TMoE

kT − f k2∞ ≤ ϵ2

(37)

The second terms is bounded using the cover number from lemma 3.7 as follows C κRM eM 2 0 log N δ, TMoE , k · k∞ ≤ C1 Πattn + LT kΠexp log + C3 LT ℓk log δ k Combining (37) and (55) yields to the desired inequality e N −2β/d + N + LT ℓk log(eM/k) E kT̂n − f k2L2 (Q) ≤ O n

(38)

(39) □

Lemma E.1 (ERM bound via symmetrization and covering numbers). Write 1X (T (xi ) − f (xi ))2 . n i=1 n

L(T ) = kT − f k2L2 (Q)

and

Ln (T ) =

|Ln (T ) − Ln (G)| ≤ 4R kT − Gk∞

(40)

(ii) ERM reduction to the ﬁnite cover. Let T ⋆ ∈ arg minT ∈TMoE L(T ) and choose G⋆ ∈ G with kG⋆ − T ⋆ k∞ ≤ δ. By ERM optimality, Ln (T̂n ) ≤ Ln (G⋆ ). Adding and subtracting L and using (40), L(T̂n ) ≤ Ln (G⋆ ) + sup |L(G) − Ln (G)| ≤ L(G⋆ ) + sup |L(G) − Ln (G)| + 4R δ G∈G

G∈G

Also L(G ) ≤ L(T ) + 4R δ. Thus ⋆

⋆

L(T̂n ) ≤

inf

T ∈TMoE

L(T ) + 8R δ + sup |L(G) − Ln (G)|

(41)

G∈G

(iii) Symmetrization over the ﬁnite cover. Deﬁne the bounded loss class F = {x 7→ (G(x) − f (x))2 : G ∈ G} with range in [0, (2R)2 ]. By symmetrization, n i h i X 2 h E sup |L(G) − Ln (G)| ≤ E sup εi (G(xi ) − f (xi ))2 n G∈G G∈G i=1 27

Generalization and Scaling Laws for Mixture-of-Experts Transformers

where (εi )ni=1 are i.i.d. Rademacher variables independent of the sample. By Massarts ﬁnite-class lemma (range bounded by 4R2 ), r r h i log |G| log N (δ, TMoE , k · k∞ ) 2 2 E sup |L(G) − Ln (G)| ≤ C R = CR n n G∈G (iv) Combine and relate L to k · k2∞ . Taking expectations in (41) and using L(T ) = kT − f k2L2 (Q) ≤ kT − f k2∞ yields r log N (δ, TMoE , k · k∞ ) 2 2 2 EkT̂n − f kL2 (Q) ≤ inf kT − f k∞ + C R + C Rδ T ∈TMoE n Absorbing constants and (optionally) replacing R δ by R2 δ (since δ ≤ 1) gives (35). Finally, the elementary inequality √ u ≤ u + 14 for u ≥ 0 (or Youngs ab ≤ a2 /(2λ) + λb2 /2) converts the square-root term into a linear (log N )/n term up to constants, yielding (36).

Proof of theorem 3.2

Deﬁnition F.1 (k-sparse partition-of-unity router). A router realizes a k-sparse partition of unity (under the idealized P router expressivity assumption of Section 3) if there exist nonnegative weights {wν (x)}N with w (x) = 1 for ν ν=1 ν all x ∈ M, such that for each x at most k weights are nonzero, and each wν is supported on a chart Uν of a boundedN of M (cf. Lemma C.7 in App. C). Moreover, the router outputs the top-k indices with weights overlap cover {Uν }ν=1 restricted to these indices. Proof. We give two constructive approximations and then take the minimum rate. Geometric setup (charts and PoU). Fix a resolution r ∈ (0, 1). By Lemma C.7 (App. C), there exists a cover ∞ {Uν }N partition of unity {φν }N ν ) ≤ r, bounded overlap s0 = s0 (d), and a C ν=1 of M with diam(U ν=1 subordinate P to the cover, such that ν φν (x) = 1, 0 ≤ φν ≤ 1, and for each x at most s0 terms are nonzero. Standard volumecomparison yields N ≤ CM r−d . By Lemma C.6 (App. C), for each ν there is a degree s = bβc polynomial Pν (expressed in local coordinates and mapped back to M) such that sup |f (x) − Pν (x)| ≤ C1 B rβ . (42) x∈Uν

Router realization of a k-sparse PoU. Since at most s0 (d) functions φν are nonzero at any x, take k ≥ s0 (d). By assumptionP(Def. F.1), the router computes nonnegative weights wν (x) supported on the same Uν and with k-sparsity, satisfying ν wν (x) = 1. We further assume (without loss) a uniform approximation kwν − φν k∞ ≤ C2 rβ can be achieved by attention-based similarity to chart anchors and a softmax/top-k selection; this is standard with a ﬁxed number of heads and bounded magnitudes (the resulting constants are absorbed into C). We now present two constructions. Construction A (expert-count limited; M regime). Assign one expert Eν to each chart Uν (total experts N ≤ Cr−d ). Within Uν , let expert Eν approximate Pν exactly (ReLU networks implement polynomials with constant depth/width depending on s, d; the parameters stay bounded and independent of r), or to accuracy C3 rβ with a constant per-expert parameter budget Πexp .2 Deﬁne the MoE output T (x) :=

N X

wν (x) Eν (x),

(at most k nonzero summands at each x).

ν=1

We bound the sup error at any x ∈ M as X X |T (x) − f (x)| ≤ wν (x) |Eν (x) − Pν (x)| + wν (x) |Pν (x) − f (x)| ν

≤ max sup |Eν − Pν | + wν (x) max sup |Pν − f | ≤ C4 rβ . ν ν Uν Uν ν | {z } | {z } | {z } ≤C3 r β

(43)

≤C1 Br β

Exact or O(rβ )-accurate polynomial representation by ReLU MLPs with poly(s, d) parameters is classical; see, e.g., Yarotsky (2017), Schmidt-Hieber (2020). 2

Generalization and Scaling Laws for Mixture-of-Experts Transformers

To realize this construction we need M ≥ N , i.e. M ≳ r−d , hence choose r M −1/d , which yields kT − f k∞ ≤ C M −β/d

=⇒

kT − f k2∞ ≤ C M −2β/d .

Since the router uses at most k experts per input, this obeys the MoE constraint. Construction B (active-parameter limited; Nact regime). Here we use few experts (e.g. k = 1 active expert per block) but allocate the entire active budget Nact to the expert(s) and non-expert parts to approximate f globally. By standard ReLU approximation on d-dimensional domains (e.g., Yarotsky (2017); Schmidt-Hieber (2020)), there exists a ReLU network N with N parameters such that sup |N (x) − f (x)| ≤ C5 B N −β/d .

(44)

x∈M

Choose the MoE so that its active subnetwork implements N with N Nact = LT Πattn + LT kΠexp (e.g. k = 1, one expert playing the role of the global MLP; attention acts as an identity/featurizer). This is legitimate because the MoE with ﬁxed routing reduces to a standard feed-forward subnetwork on the active path. Thus we can realize T (x) = N (x) with N Nact and −β/d

kT − f k∞ ≤ C6 Nact

=⇒

−2β/d

kT − f k2∞ ≤ C Nact

Taking the minimum. Constructions A and B are both valid members of TMoE under the parameter bound kθk∞ ≤ κ (after routine rescalings absorbed by LayerNorm and architectural constants). Therefore, −β/d inf kT − f k∞ ≤ C min Nact , M −β/d T ∈TMoE

and squaring gives (1). If M is sufﬁciently large to support the cover at the r achieving the Nact -optimal precision, the M constraint is inactive and (45) follows. inf

−2β/d

T ∈TMoE

kT − f k2∞ ≤ C Nact

(45)

Both regimes yield the rates claimed in (1); taking the better gives the minimum.

Proof of neural scaling laws

G.1

Proof of optimal number of experts

Recall the kdependent objective (experts dominate approximation; constants absorbed) E(k) = (Ak)−2β/d + | {z } approx

A k n |{z} estimation

LT ℓ eM k log , | n {z k }

A := LT Πexp , 1 ≤ k ≤ M.

(46)

routing

Differentiate and set the derivative to zero: E ′ (k) = −

2β −2β/d −2β/d−1 A LT ℓ eM A k + + log − 1 = 0. d n n k

Equivalently,

eM 2β −2β/d −2β/d−1 1 A + LT ℓ log A − LT ℓ . (47) k = d n k This matches the stated ﬁrst-order condition. It has a closed form in terms of the Lambert-W function after algebraic manipulation (since k appears both polynomially and inside log k), but the Lambert-W form is not particularly illuminating. Instead we give (i) a principled ﬁxed-point approximation and (ii) clean sandwich bounds showing that M enters only logarithmically. Fixed-point approximation. Rewrite (47) as k 2β/d+1 =

2β −2β/d n A . d A + LT ℓ [log(eM/k) − 1] 29

(48)

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Treat the slowly varying factor log(eM/k) as (locally) constant to obtain the iterate d " #2β+d 2β −2β/d n (t+1) k := A . (49) d A + LT ℓ [log(eM/k (t) ) − 1] d n 2β+d Initializing with the no-routing optimum k (0) A and taking one step already yields the advertised ﬁxed-point form d " #2β+d n k⋆ ≈ (up to constants), (50) A + LT ℓ log eM k⋆ with the cap k ⋆ ≤ M . Because log(eM/k) varies only between 1 and 1 + log M for k ∈ [1, M ], a couple of iterations of (49) sufﬁce in practice, and the exponent in n is unchanged. Sandwich bounds (log enters only through constants). Note that for 1 ≤ k ≤ M , log(eM/k) ∈ [1, 1 + log M ]. Hence the right-hand side of (47) lies in 1 1 A ≤ RHS ≤ A + LT ℓ log M . n n −2β/d −2β/d−1 2β+d A− 2β+d . Applying this with C = A + Solving the equality 2β k = C d A n for k gives k (n/C) LT ℓ log M (upper RHS) and C = A (lower RHS) yields the sandwich: d d 2β+d n 2β+d n c1 ≤ k ⋆ ≤ c2 , with k ⋆ ≤ M, (51) A + LT ℓ log M A d

for absolute constants c1 , c2 > 0. Thus M affects k ⋆ only through the logarithm, and the rate in n remains nd/(2β+d) . Combining (50) and (51),

d 2β+d n k min M, A + LT ℓ log(eM/k ⋆ ) d and the nexponent remains 2β+d ; M only enters through a logarithm.

⋆

G.2

Effect of the total experts M

For ﬁxed k, increasing M inﬂuences the bound in (3) only through log(eM/k) in Rroute . Therefore, absent specialization effects that improve constants, ∂ k 1 E(k, M ) ≳ · ⇒ larger M with ﬁxed k only raises the bound (logarithmically). ∂M n M A design heuristic is to choose M commensurate with k (e.g., M = O(k)) if one wishes to avoid a large routing overhead; having M k does not help rates in this bound. The details of these calculations are provided in appendix G.2. Proposition G.1 (Monotone (logarithmic) M effect at ﬁxed k). Fix n, k, LT , ℓ, Πattn , Πexp . Deﬁne eM Nact LT ℓk −2β/d E(k, M ) := Nact + + log , 1 ≤ k ≤ M. n n k Then, for any ﬁxed k, E(k, M ) is strictly increasing in M , with LT ℓk 1 ∂ E(k, M ) = · . (52) ∂M n M Consequently, absent additional specialization gains (i.e., changes in the approximation constant that are not captured in (3)), the bound is minimized at the smallest admissible M , namely M = k, and grows only logarithmically with M for ﬁxed k. Proof. The ﬁrst two terms do not depend on M when k is ﬁxed. The routing term is LTnℓk log(eM/k), whose derivative 1 , yielding (52). Monotonicity follows since M 7→ log(eM/k) is strictly increasing for M ≥ k. in M is LTnℓk · M Corollary G.2 (Quantitative overhead from enlarging M at ﬁxed k). For any M2 ≥ M1 ≥ k, M LT ℓk 2 E(k, M2 ) − E(k, M1 ) = log . n M1 In particular, setting M = ρ k with ρ ≥ 1 gives a routing penalty LTnℓk log(eρ) = LTnℓk (1 + log ρ). 30

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Design heuristic. To keep the routing overhead small, choose M commensurate with k, e.g. M = Θ(k), so that log(eM/k) = Θ(1). Having M k increases the bound only through log(eM/k) and does not improve the rate in n or Nact . When is routing negligible? Routing is negligible relative to the estimation term if eM eM LT ℓk Nact Nact log ⇐⇒ log . n k n k LT ℓk

(53)

Since Nact = LT Πattn + LT kΠexp , a sufﬁcient condition is log(eM/k) Πattn /(ℓk) + Πexp /ℓ; in particular, if M = ρk with moderate ρ and ℓ is not enormous, the routing term is dominated by the estimation term. G.3

Proof of proposition 4.5

Under C nNact , we eliminate n = C/Nact and reduce the bound to a single-variable objective −2β/d

Φ(Nact ; C) := A Nact

+ B

2 Nact C

(54)

with positive constants A, B absorbing ﬁxed architectural and loss factors. Differentiating w.r.t. Nact > 0, Φ′ (Nact ) = −

Nact 2β −2β/d−1 A Nact + 2B d C

⋆ The ﬁrst-order condition Φ′ (Nact ) = 0 gives

2B 2β ⋆ 2+2β/d (Nact ) = A C d Hence ⋆ Nact (C) =

β A dB

⋆ (Nact ) 2(1+β/d) =

=⇒ d 2β+2d

β A C d B

C 2β+2d C 2β+2d

⋆ we obtain the second equality in Eq.(18). Substituting which is the ﬁrst equality in Eq. (18). Using n⋆ = C/Nact ⋆ into (54) (the two terms balance at optimum) gives Nact 2β

⋆ ⋆ −2β/d Φ(Nact ; C) (Nact ) C − d · 2β+2d = C − β+d d

which is (19).

Proof of lemma 55 (MoE Covering Number)

Let X ⊂ RD be compact with kxk∞ ≤ M0 . Consider the Mixture-of-Experts transformer class TMoE (D, M, k, LT , LFFN , wFFN , demb , m, κ, R) deﬁned as in Deﬁnition 3.1, except that each feed-forward sublayer is an MoE layer with M experts and the router selects (hard) top-k experts per token per layer. Assume each learned parameter entry satisﬁes kθk∞ ≤ κ and each network output is bounded by kT (x)k∞ ≤ R. Let Πattn be the number of scalar parameters in the attention / non-expert parts per block and let Πexp be the number of scalar parameters of a single expert MLP (depth LFFN , width wFFN , input/output dimension demb ). Then for any δ ∈ (0, 1) the covering number of TMoE under the sup-norm satisﬁes eM C κRM 2 0 + C3 LT ℓk log (55) log N δ, TMoE , k · k∞ ≤ C1 Πattn + LT kΠexp log δ k for absolute constants C1 , C2 , C3 > 0 depending polynomially on demb , wFFN , m, LT , LFFN , ℓ (and we may hide e version). small poly-log factors in the O(·) Proof. The proof follows the structure of the dense-transformer covering-number proof (Lemma 2 in Havrilla & Liao (2024)) with two MoE-speciﬁc modiﬁcations. We divide the proof into four steps: (1) enumerate routing patterns and apply a union bound; (2) bound the sup-norm difference between two networks in terms of the sup-norm difference of their parameters; (3) count active parameters for a ﬁxed pattern and derive a parameter-grid covering bound; (4) combine items to obtain (55). 31

Generalization and Scaling Laws for Mixture-of-Experts Transformers

H.1

decomposition into routing patterns (union bound).

For each MoE layer ℓ and each token position t ∈ {1, . . . , ℓ}, the router chooses a subset Sℓ,t ⊂ [M ] of size |Sℓ,t | = k. A routing pattern π is the collection of these choices for all layers and token positions: π = {Sℓ,t : ℓ = 1, . . . , LT , t = 1, . . . , ℓ}. The number of possible patterns is bounded by |Π| ≤ Using the standard combinatorial bound

M k

M L T ℓ k

≤ (eM/k)k we obtain

log |Π| ≤ LT ℓ · k log

(56)

For a ﬁxed routing pattern π denote by Tπ the subclass of MoE transformers whose router outputs follow π deterministically S (i.e., the same ﬁxed subset of experts is used in each MoE layer and position for all inputs). Since TMoE = π∈Π Tπ , the subadditivity of covering numbers under unions gives N (δ, TMoE , k · k∞ ) ≤

N (δ, Tπ , k · k∞ )

(57)

π∈Π

Log-summax inequality. For nonnegative {aπ }π∈Π one has X aπ ≤ log |Π| + max log aπ log π∈Π

π∈Π

since

π aπ ≤ |Π| · maxπ aπ . Applying this to aπ = N (δ, Tπ , k · k∞ ) in (57) gives

log N (δ, TMoE , k · k∞ ) ≤ log |Π| + max log N (δ, Tπ , k · k∞ ) π∈Π

(58)

Thus it remains to upper-bound log N (δ, Tπ , k · k∞ ) uniformly over π. H.2

sensitivity of the network output to parameter perturbations (Lipschitz in parameter space).

Fix a routing pattern π. Under π the architecture becomes a deterministic (dense) composition of modules where each MoE layer is replaced by the (sparse) subnetwork that contains only the k selected experts for the corresponding tokens/layer. Denote by θ the vector of all scalar parameters that are active under this pattern (embedding + attention + non-expert weights + the selected experts’ weights). Let θ′ be another parameter vector of the same shape and satisfying kθ − θ′ k∞ < η for some η > 0. We will bound sup |Tθ (x) − Tθ′ (x)| x∈X

by a constant times η, with that constant depending polynomially on the architecture hyperparameters. The network is a composition of LT transformer blocks. Let Bj (θj , ·) denote block j (viewed as a map on embedding tensors) parameterized by the block parameters θj . We write the forward pass as H (0) = PE + E(x),

H (j) = Bj (θj , H (j−1) ),

j = 1, . . . , LT ,

and the scalar output Tθ (x) is a ﬁxed linear readout of H (LT ) . Deﬁne the sup-norm on embedding matrices by kHk∞ = maxi,t |Hi,t | (matching the earlier notation). As in the dense case, each module is: (i) Lipschitz in its input (constant Lin polynomial in (demb , m, wFFN , LFFN )) and (ii) Lipschitz in its parameters (constant Cparam polynomial in the same quantities and in ℓ) (Anthony & Bartlett, 1999; Vershynin, 2018; Shalev-Shwartz & Ben-David, 2014). 32

Generalization and Scaling Laws for Mixture-of-Experts Transformers

(A) Multi-head attention sensitivity. Consider a single multi-head attention (MHA) block with m heads, each head parameterized by query/key/value matrices and interaction kernels. Let the MHA parameter vector be ψ. For two MHA parameter sets ψ, ψ ′ with kψ − ψ ′ k∞ ≤ η, and any input embedding H satisfying kHk∞ ≤ M , the goal is to bound the following expression via Lemma H.1 kMHAψ (H) − MHAψ′ (H)k∞ Lemma H.1 (MHA parameter stability). For kψ − ψ ′ k∞ ≤ η and kHk∞ ≤ CM0 , kMHAψ (H) − MHAψ′ (H)k∞ ≤ CMHA ℓ m d2emb κ η. for some architecture-dependent constant CMHA . The dependence ∝ ℓmd2emb arises because each head sums over ℓ tokens and involves inner-products of demb -dimensional projected vectors; the factor κ comes from the bound on parameter magnitudes. (B) FFN (non-expert) sensitivity. For the non-expert FFN parts (the small projection matrices that are outside MoE experts see Deﬁnition 3.1) with depth LFFN and width wFFN , the parameter perturbation bound from the dense proof is given by the following Lemma Lemma H.2 (FFN parameter stability). For a (non-expert) FFN of depth LFFN and width wFFN , LFFN kFFNχ (h) − FFNχ′ (h)k∞ ≤ CFFN demb wFFN κ η.

uniformly over khk∞ ≤ CM0 (cf. (Bartlett & Mendelson, 2002a; Golowich et al., 2018)). For each token input h (and hence on the embedding matrix by taking sup over tokens); the polynomial dependence on wFFN and LFFN reﬂects repeated matrix multiplications. (C) Expert-FFN sensitivity under ﬁxed pattern (key MoE change). Under a ﬁxed routing pattern π, only k experts per token per layer are active, so the FFN stage of the block is effectively a concatenation / parallel application of exactly those k expert MLPs. For a single expert with parameter vector ϕ and sup-norm perturbation kϕ − ϕ′ k∞ ≤ η, the per-token difference is bounded as in (B) by the following lemma Lemma H.3 (Expert stability). For a single expert MLP Eϕ , LFFN κ η. kEϕ (h) − Eϕ′ (h)k∞ ≤ Cexp demb wFFN

Since only k experts are active, the total perturbation coming from expert blocks per token scales linearly in k: X X LFFN Eϕi (h) − Eϕ′i (h) ≤ k Cexp demb wFFN κ η. i∈S

i∈S

∞

(D) Block-level perturbation bound. Combining (A),(B),(C) and the residual connection structure of the transformer block (additive skip-connections and tokenwise FFNs), a single block Bj satisﬁes, for parameter perturbation ≤ η inside the block and any input H with kHk∞ ≤ M yields to Lemma H.4. Lemma H.4 (Block-level stability). With k active experts per block under a ﬁxed routing pattern, LFFN LFFN Bj (θj , H) − Bj (θj′ , H) ∞ ≤ Cblock ℓmd2emb + demb wFFN + kdemb wFFN κη. hence we may write (absorbing polynomial factors into a single constant) kBj (θj , H) − Bj (θj′ , H)k∞ ≤ Cb Parch (ℓ, m, demb , wFFN , LFFN ) κ η, where Parch (·) is a polynomial capturing the block-dependence and includes the linear factor k in the expert term. (E) Induction across blocks. write

Let H (j) and H ′(j) denote the block outputs when using θ and θ′ respectively. We can

kH (j) − H ′(j) k∞ = kBj (θj , H (j−1) ) − Bj (θj′ , H ′(j−1) )k∞ ≤ kBj (θj , H (j−1) ) − Bj (θj , H ′(j−1) )k∞ + kBj (θj , H ′(j−1) ) − Bj (θj′ , H ′(j−1) )k∞ . | {z } {z } | input-Lipschitz term

param-perturbation term

Generalization and Scaling Laws for Mixture-of-Experts Transformers

The input-Lipschitz term is bounded by Lin kH (j−1) − H ′(j−1) k∞ where Lin is the Lipschitz constant of the block as a function of its input (this constant depends polynomially on demb , m, wFFN , LFFN but not on M ). The paramperturbation term is bounded by the block-level bound in the previous paragraph. Thus kH (j) − H ′(j) k∞ ≤ Lin kH (j−1) − H ′(j−1) k∞ + Cb Parch (·) κ η. Iterating this inequality from j = 1 to j = LT and noting that H (0) = H ′(0) (same positional encodings and same input) we obtain kH

(LT )

−H

′(LT )

k∞ ≤

T −1 LX

T LL in Cb Parch (·) κ η. Ltin Cb Parch (·) κ η ≤ Lin − 1 t=0

The previous inequality is proved in details in lemma I.1. Since the decoder/readout is a ﬁxed linear projection, the ﬁnal scalar output difference is bounded by the same order. Thus there exists an overall Lipschitz constant T Lparam = C ′ LL in Parch (ℓ, m, demb , wFFN , LFFN ) κ

such that

sup |Tθ (x) − Tθ′ (x)| ≤ Lparam η. x∈X

Crucially: Parch (·) contains a factor linear in the number of active experts per layer(k) and Lparam does not depend on the total number of experts M for a ﬁxed routing pattern (it depends on M only through how the pattern is chosen which we handle with the union bound). H.3

parameter counting and parameter-space covering (full details).

Fix a routing pattern π. Let θact denote the vector formed by all scalar parameters that are active under π (token embedding/positional terms, attention and non-expert FFN parameters for each block, and the experts that appear in π). Write p := |θact | for the active parameter dimension. (a) Parameter domain and norm. By the magnitude constraint kθk∞ ≤ κ (Assumption in Lemma 3.7), the admissible parameter set is included in the ℓ∞ ball of radius κ: B∞ (κ)p = u ∈ Rp : kuk∞ ≤ κ = [−κ, κ]p . (For ℓ∞ , the ball is the axis-aligned cube.) (b) Constructing an ηgrid in parameter space. For a mesh width η > 0, deﬁne the coordinatewise grid on [−κ, κ] by Gη = − κ + rη : r ∈ Z, −κ ≤ −κ + rη ≤ κ . Along each coordinate, the number of grid points is at most l 2κ m 2κ 3κ N1 ≤ +1 ≤ +2 ≤ (for η ≤ κ; else enlarge the constant). η η η Therefore the full pdimensional grid Gηp has cardinality bounded by 2κ p 3κ p Gηp ≤ +1 ≤ . η η

(59)

Moreover, for every u ∈ [−κ, κ]p , rounding each coordinate of u to its nearest grid value produces ũ ∈ Gηp with ku − ũk∞ ≤ η

(or η/2 if we choose midpoints).

Hence Gηp is an ηnet for the parameter cube in the ℓ∞ norm. (c) Parameter-to-function Lipschitz map. From Step 2, we have the uniform parameter-stability (Lipschitz) inequality sup Tθ (x) − Tθ′ (x) ≤ Lparam kθ − θ′ k∞ ,

x∈X

with

T Lparam = C ′ LL in Parch (ℓ, m, demb , wFFN , LFFN , k) κ, independent of M for ﬁxed routing pattern π (the dependence on k enters via Parch ).

Generalization and Scaling Laws for Mixture-of-Experts Transformers

(d) From parameter cover to function cover. Given δ ∈ (0, 1), choose δ . Lparam

η =

Then any two parameter vectors within ℓ∞ distance η induce network outputs within sup-norm distance at most δ (uniformly over x ∈ X ). Therefore, taking the grid Gηp as in (59) and evaluating the network at each grid parameter gives a δcover of the function class Tπ under k · k∞ : 3κ p 3κL p param N (δ, Tπ , k · k∞ ) ≤ Gηp ≤ = . η δ Taking logarithms gives 3κL 3κL param param log N (δ, Tπ , k · k∞ ) ≤ p · log = |θact | · log . (60) δ δ (e) Bounding the active parameter count p = |θact |. Decompose

|θact | = |θembed | + |θdec | + LT |θattn | + |θﬀn,nonexp | + |θexp,act | .

Under a ﬁxed routing pattern, expert parameter tensors are shared across tokens within a layer, so the active expert count per layer satisﬁes |θexp,act | ≤ k Πexp . Absorbing the (input/output) embedding and decoder constants into C0 and writing Πattn for attention/non-expert parameters per block, we obtain the coarse bound |θact | ≤ C0 + LT Πattn + k Πexp ≲ Πattn + LT kΠexp . Substituting this and the expression of Lparam into (60), and renaming constants, gives C2 κRM0 log N (δ, Tπ , k · k∞ ) ≤ C1 Πattn + LT kΠexp log , δ which is the ﬁxed-pattern covering bound used in Step 4. H.4

□

combine with the routing-pattern union bound.

Plugging the routing count (56) into (58) gives

C κRM eM 2 0 + C1 Πattn + LT kΠexp log log N (δ, TMoE , k · k∞ ) ≤ LT ℓ k log k δ eM C κRM 2 0 ≤ C1 Πattn + LT kΠexp log + C3 LT ℓk log , δ k for some constants C1 , C2 , C3 > 0. This proves the lemma.

Proofs of Stability Lemmas

I.1

Proof of H.1

Fix a sequence embedding H ∈ Rℓ×demb with kHk∞ ≤ M and consider one MHA layer with m heads and parameters (h)

(h)

ψ = {WQ , WK , WV , WO }m h=1 , ′

kψk∞ ≤ κ.

′

Let ψ be another parameter set with kψ − ψ k∞ ≤ η. For head h, deﬁne (h)

Q(h) = HWQ , and the (row-wise) attention map

(h)

K (h) = HWK ,

(h)

V (h) = HWV ,

A(h) = softmax √1d Q(h) (K (h) )⊤ ∈ Rℓ×ℓ . k

The head output is O

(h)

∈R

ℓ×dv

and the MHA output is (h) MHAψ (H) = Concatm WO ∈ Rℓ×demb . h=1 O

We show

kMHAψ (H) − MHAψ′ (H)k∞ ≤ CMHA ℓ m d2emb κ η, for a universal constant CMHA > 0 depending at most polynomially on (demb , dk , dv ) and on the bound M . 35

(61)

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Step 1: linear projections (Q, K, V ). For any matrices X and W , kXW k∞ ≤ din kXk∞ kW k∞ (since each entry (h) (h) (h)′ is a sum of din products). Hence, with kHk∞ ≤ M , kW• k∞ ≤ κ, and kW• − W• k∞ ≤ η, kQ(h) k∞ ≤ demb M κ, kK (h) k∞ ≤ demb M κ, kV (h) k∞ ≤ demb M κ, and

kQ(h) − Q(h)′ k∞ ≤ demb M η,

kK (h) − K (h)′ k∞ ≤ demb M η,

kV (h) − V (h)′ k∞ ≤ demb M η.

Step 2: score matrices and softmax. Let S (h) = √1d Q(h) (K (h) )⊤ ∈ Rℓ×ℓ and similarly S (h)′ . By the product k rule and the matrix ℓ∞ bound kXY ⊤ k∞ ≤ ℓkXk∞ kY k∞ , 1 (h) kS (h) − S (h)′ k∞ ≤ √ kQ − Q(h)′ k∞ kK (h) k∞ + kQ(h)′ k∞ kK (h) − K (h)′ k∞ ℓ ≤ CS ℓ d2emb M 2 κ η, dk √ with CS := 2/ dk . The row-wise softmax σ : Rℓ → Rℓ has Jacobian J(z) = diag(σ(z)) − σ(z)σ(z)⊤ , whose operator norm (for ℓ∞ ) is bounded by a universal constant (e.g., kJ(z)k∞→∞ ≤ 1; tighter bounds give 1/2 or 1/4 but are unnecessary here). Applying the mean value theorem row-wise, kA(h) − A(h)′ k∞ ≤ Csm kS (h) − S (h)′ k∞ ≤ Csm CS ℓ d2emb M 2 κ η, for a universal Csm ∈ [1/4, 1]. Step 3: head outputs (O(h) = A(h) V (h) ). O

(h)

−O

(h)′

Decompose

= (A(h) − A(h)′ )V (h) + A(h)′ (V (h) − V (h)′ ).

Using kXY k∞ ≤ ℓkXk∞ kY k∞ for ℓ×ℓ times ℓ×dv , together with Step 1 bounds on kV (h) k∞ and kV (h) −V (h)′ k∞ , kO(h) − O(h)′ k∞ ≤ ℓ kA(h) − A(h)′ k∞ kV (h) k∞ + ℓ kA(h)′ k∞ kV (h) − V (h)′ k∞ . Since rows of A(h)′ are probability vectors, kA(h)′ k∞ ≤ 1. Hence, kO(h) − O(h)′ k∞ ≤ ℓ Csm CS ℓ d2emb M 2 κ η (demb M κ) + ℓ (demb M η) ≤ CO ℓ d3emb M 3 κ2 η + ℓ demb M η, where CO := Csm CS . We keep the dominant (polynomial) term and absorb the lower-order term into constants. (h) ∈ Rℓ×(mdv ) . Then Step 4: concatenate heads and output projection. Let O = Concatm h=1 O

kO − O′ k∞ ≤ max kO(h) − O(h)′ k∞ ≤ CO ℓ d3emb M 3 κ2 η. h∈[m]

Finally, apply the output projection WO with kWO k∞ ≤ κ and kWO − WO′ k∞ ≤ η: kMHAψ (H) − MHAψ′ (H)k∞ = kOWO − O′ WO′ k∞ ≤ k(O − O′ )WO k∞ + kO′ (WO − WO′ )k∞ ≤ (mdv ) kO − O′ k∞ kWO k∞ + (mdv ) kO′ k∞ η ≤ Cout (mdv ) ℓ d3emb M 3 κ3 η + (mdv ) kO′ k∞ η, ′

and we bound kO k∞ by the same polynomial in (ℓ, demb , M, κ) as above. Absorbing dv and all input/weight bounds into the constant, and noting that the paper tracks only polynomial dependence on dimensions and sequence length (hiding ﬁxed architecture factors into CMHA ), we obtain kMHAψ (H) − MHAψ′ (H)k∞ ≤ CMHA ℓ m d2emb κ η, which is (61). □ I.2

Proof of H.2

Let FFNψ : Rdemb → Rdemb be a depth-LFFN MLP applied tokenwise, with hidden widths at most wFFN , coorFFN dinatewise activation σ (e.g., ReLU/GELU), and parameters ψ = {(W ℓ , bℓ )}L ℓ=1 . Assume the uniform bounds ℓ ℓ ′ ℓ kW k∞ ≤ κ, kb k∞ ≤ κ for all ℓ. Let ψ be another parameter set with kW − W ℓ′ k∞ ≤ η and kbℓ − bℓ′ k∞ ≤ η. Fix an input token h ∈ Rdemb , and denote the layerwise states z 0 = h, z ℓ = σ W ℓ z ℓ−1 + bℓ , z ′ 0 = h, z ′ ℓ = σ W ℓ′ z ′ ℓ−1 + bℓ′ . We prove LFFN kFFNψ (h) − FFNψ′ (h)k∞ = kz LFFN − z ′ LFFN k∞ ≤ CFFN demb wFFN κ η, (62) for a constant CFFN > 0 depending only polynomially on architectural constants and on a uniform bound on khk∞ (absorbed into CFFN ), exactly as done in the paper. 36

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Auxiliary facts. We use (i) kXuk∞ ≤ din kXk∞ kuk∞ for X ∈ Rdout ×din , (ii) kσ(u) − σ(v)k∞ ≤ Lσ ku − vk∞ with Lσ ≤ 1 for ReLU and Lσ = O(1) for GELU, and (iii) kσ(u)k∞ ≤ kuk∞ for ReLU (and similarly for GELU up to a constant). Step 1: bound the forward activations. Let d0 = demb and dℓ ≤ wFFN for 1 ≤ ℓ ≤ LFFN − 1, and dLFFN ≤ demb . Inductively, kz ℓ k∞ ≤ Lσ kW ℓ z ℓ−1 k∞ + kbℓ k∞ ≤ Lσ dℓ−1 κ kz ℓ−1 k∞ + κ . Solving this recursion gives kz ℓ k∞ , kz ′ ℓ k∞ ≤ C0 (dmax κ)ℓ ,

dmax := max{demb , wFFN },

(63) ′

where C0 absorbs Lσ and khk∞ . (Exactly the same bound holds for the primed sequence, since ψ obeys the same magnitude constraints.) Step 2: one-layer perturbation inequality. Fix ℓ ∈ {1, . . . , LFFN } and set uℓ = W ℓ z ℓ−1 + bℓ , u′ ℓ = W ℓ′ z ′ ℓ−1 + bℓ′ . Then kz ℓ − z ′ ℓ k∞ ≤ Lσ kuℓ − u′ ℓ k∞ ≤ Lσ kW ℓ (z ℓ−1 − z ′ ℓ−1 )k∞ + k(W ℓ − W ℓ′ )z ′ ℓ−1 k∞ + kbℓ − bℓ′ k∞ . Using the auxiliary facts and (63), kz ℓ − z ′ ℓ k∞ ≤ Lσ dℓ−1 κ kz ℓ−1 − z ′ ℓ−1 k∞ + dℓ−1 η kz ′ ℓ−1 k∞ + η ≤ αℓ kz ℓ−1 − z ′ ℓ−1 k∞ + βℓ η, where αℓ := Lσ dℓ−1 κ,

(64)

βℓ := Lσ dℓ−1 C0 (dmax κ)ℓ−1 + 1 .

Step 3: iterate the perturbation recursion. Starting with kz 0 − z ′ 0 k∞ = 0 and applying (64) layer by layer, kz LFFN − z ′ LFFN k∞ ≤ η

LX FFN ℓ=1

βℓ

LY FFN

αt .

t=ℓ+1

Since αt ≤ Lσ dmax κ and βℓ ≤ C1 dmax (dmax κ)ℓ−1 + C1 for a constant C1 , kz LFFN − z ′ LFFN k∞ ≤ η · C2

LX FFN

dmax (dmax κ)ℓ−1 (dmax κ)LFFN −ℓ ≤ η · C3 dmax (dmax κ)LFFN −1 .

ℓ=1

Finally, dmax ≤ max{demb , wFFN } ≤ demb + wFFN ≤ 2wFFN (for the hidden layers to be useful we take wFFN ≥ LFFN by (dmax )LFFN ). Absorbing all ﬁxed constants and powers of κ into CFFN , we demb /2; otherwise replace wFFN obtain the clean architectural polynomial LFFN kz LFFN − z ′ LFFN k∞ ≤ CFFN demb wFFN κ η,

□

which is (62). I.3

Proof of H.3

Let Eϕ : Rdemb → Rdemb be a tokenwise MLP (the expert) of depth LFFN with hidden widths at most wFFN , FFN coordinatewise activation σ (ReLU/GELU), and parameters ϕ = {(W ℓ , bℓ )}L ℓ=1 . Assume uniform magnitude bounds ℓ ℓ ′ kW k∞ ≤ κ, kb k∞ ≤ κ for all ℓ. Let ϕ be another parameter set with kW ℓ − W ℓ′ k∞ ≤ η and kbℓ − bℓ′ k∞ ≤ η. For a token h ∈ Rdemb deﬁne the layerwise states z 0 = h,

z ℓ = σ(W ℓ z ℓ−1 + bℓ ),

z ′ 0 = h,

z ′ ℓ = σ(W ℓ′ z ′ ℓ−1 + bℓ′ ),

so that Eϕ (h) = z LFFN and Eϕ′ (h) = z ′ LFFN . Goal. Show

LFFN kEϕ (h) − Eϕ′ (h)k∞ ≤ Cexp demb wFFN κ η,

(65)

with Cexp > 0 depending at most polynomially on (LFFN , kσkLip , khk∞ ) and on ﬁxed architectural constants, in the same sense as the paper. 37

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Auxiliary inequalities. For X ∈ Rdout ×din and u ∈ Rdin , kXuk∞ ≤ din kXk∞ kuk∞ , kσ(u) − σ(v)k∞ ≤ Lσ ku − vk∞ ,

kσ(u)k∞ ≤ Cσ kuk∞ ,

(66) (67)

with Lσ ≤ 1 for ReLU and Lσ = O(1) for GELU (similarly for Cσ ). Step 1: forward activation bound. Let d0 = demb and dℓ ≤ wFFN for 1 ≤ ℓ ≤ LFFN − 1, and dLFFN ≤ demb . Using (66)(67), kz ℓ k∞ ≤ Cσ kW ℓ z ℓ−1 k∞ + kbℓ k∞ ≤ Cσ dℓ−1 κkz ℓ−1 k∞ + κ . Unrolling the recursion yields kz ℓ k∞ , kz ′ ℓ k∞ ≤ C0 (dmax κ)ℓ ,

dmax := max{demb , wFFN },

(68)

for some constant C0 absorbing Cσ , Lσ and khk∞ . Set uℓ = W ℓ z ℓ−1 + bℓ and u′ ℓ = W ℓ′ z ′ ℓ−1 + bℓ′ . Then kz ℓ − z ′ ℓ k∞ ≤ Lσ kuℓ − u′ ℓ k∞ ≤ Lσ kW ℓ (z ℓ−1 − z ′ ℓ−1 )k∞ + k(W ℓ − W ℓ′ )z ′ ℓ−1 k∞ +kbℓ − bℓ′ k∞ . | {z } | {z }

Step 2: one-layer perturbation.

≤dℓ−1 κ ∥z ℓ−1 −z ′ ℓ−1 ∥∞

Using (68),

≤dℓ−1 η ∥z ′ ℓ−1 ∥∞

αℓ := Lσ dℓ−1 κ, βℓ := Lσ dℓ−1 C0 (dmax κ)ℓ−1 + 1 .

kz ℓ − z ′ ℓ k∞ ≤ αℓ kz ℓ−1 − z ′ ℓ−1 k∞ + βℓ η,

(69)

Step 3: iterate the recursion. Starting from kz 0 − z ′ 0 k∞ = 0 and applying (69) layerwise, kz

LFFN

−z

′ LFFN

Since αt ≤ Lσ dmax κ and βℓ ≤ C1 dmax (dmax κ) kz LFFN − z ′ LFFN k∞ ≤ η C2

LX FFN

k∞ ≤ η

ℓ−1

LX FFN ℓ=1

βℓ

LY FFN

αt .

t=ℓ+1

+1 ,

dmax (dmax κ)LFFN −1 ≤ η C3 dmax (dmax κ)LFFN −1 .

ℓ=1

Finally, dmax ≤ max{demb , wFFN } ≤ c0 wFFN for a numerical c0 (or replace by dmax if preferred). Absorbing ﬁxed powers of κ and numerical constants into Cexp , and noting that the input and output layers contribute at most linear factors in demb , we obtain the architectural polynomial LFFN κ η, kEϕ (h) − Eϕ′ (h)k∞ = kz LFFN − z ′ LFFN k∞ ≤ Cexp demb wFFN

□

which is the claimed bound (65). I.4

Network stability & parameter Lipschitzness

Lemma I.1 (Network stability & parameter Lipschitzness). Let Lin ≥ 1 be the block input-Lipschitz constant. Then for j = 1, . . . , LT , + Cb Parch κη, H (j) − H ′(j) ≤ Lin H (j−1) − H ′(j−1) ∞

∞

hence

sup |Tθ (x) − Tθ′ (x)| ≤ Lparam η, x∈X

T Lparam = C ′ LL in Parch (ℓ, m, demb , wFFN , LFFN , k) κ.

Let a (pre-norm) transformer block be

X Bj (θj , H) = H + MHAψ LNγ1 ,β1 (H) + FFNnonexp LNγ2 ,β2 (H) + Eϕi LNγ2 ,β2 (H) , χ i∈Sj

where Sj is the set of k active experts of the MoE sublayer in block j under the ﬁxed routing pattern, and θj = {ψ, χ, {ϕi }i∈Sj , γ1 , β1 , γ2 , β2 }. Let θj′ be another parameter vector with kθj − θj′ k∞ ≤ η, and ﬁx the input H ∈ Rℓ×demb with kHk∞ ≤ CM0 . 38

Generalization and Scaling Laws for Mixture-of-Experts Transformers

Step 1: decompose by residual additivity. Since H is added as a residual with the same value in both blocks, kBj (θj , H) − Bj (θj′ , H)k∞ ≤ MHAψ (LNγ1 ,β1 (H)) − MHAψ′ (LNγ1′ ,β1′ (H)) ∞ + FFNnonexp (LNγ2 ,β2 (H)) − FFNnonexp (LNγ2′ ,β2′ (H)) ∞ χ χ′ X Eϕi (LNγ2 ,β2 (H)) − Eϕ′i (LNγ2′ ,β2′ (H)) ∞ . +

(B.1)

i∈Sj

Step 2: control the LayerNorm parameter perturbations. For ﬁxed H, the (per-token) layer normalization LNγ,β (H) = γ Ĥ + β depends linearly on (γ, β) when H is ﬁxed (the centering/scaling of H is the same for both parameter sets). Hence ′ kLNγ,β (H) − LNγ ′ ,β ′ (H)k∞ ≤ CLN demb kHk∞ k(γ, β) − (γ ′ , β ′ )k∞ ≤ CLN demb κ η.

(B.2)

Both MHA and FFN sublayers are Lipschitz in their inputs (with constants depending polynomially on (ℓ, m, demb , wFFN , LFFN )); thus the contribution of the LayerNorm parameter change to each sublayers output is ′ bounded by a constant times CLN demb κη, which we absorb into the module constants below. Step 3: apply per-module parameter Lipschitz bounds. From the previously established lemmas: (MHA) (non-expert FFN) (single expert)

MHAψ (Z) − MHAψ′ (Z) ∞ ≤ CMHA ℓ m d2emb κ η,

∀Z,

LFFN FFNnonexp (z) − FFNnonexp (z) ∞ ≤ CFFN demb wFFN κ η, χ χ′

Eϕ (z) − Eϕ′ (z) ∞ ≤

LFFN Cexp demb wFFN κ η,

∀z,

∀z.

Under a ﬁxed routing pattern, exactly k experts are active in block j, so the MoE contribution is X LFFN Eϕi (z) − Eϕ′i (z) ∞ ≤ k Cexp demb wFFN κ η.

(B.3)

i∈Sj

Step 4: combine everything. Apply the bounds from Step 3 to (B.1) with Z = LNγ1 ,β1 (H) and z = LNγ2 ,β2 (H); add the (absorbed) LayerNorm terms from (B.2): LFFN LFFN ′′ kBj (θj , H)−Bj (θj′ , H)k∞ ≤ CMHA ℓ m d2emb κ η+CFFN demb wFFN κ η+k Cexp demb wFFN κ η+CLN demb κ η.

Since wFFN ≥ 1 and k, m ≥ 1, the last term is dominated and can be absorbed into the FFN/expert constants. Renaming Cblock to absorb all ﬁxed numerical/polylog factors, LFFN LFFN kBj (θj , H) − Bj (θj′ , H)k∞ ≤ Cblock ℓ m d2emb + demb wFFN + k demb wFFN κ η. □

This is the claimed block-level perturbation bound. I.5

Proof of I.1

Recall the block-level stability (for ﬁxed input H and parameters differing by at most η in ℓ∞ ) kBj (θj , H) − Bj (θj′ , H)k∞ ≤ Cb Parch (ℓ, m, demb , wFFN , LFFN , k) κ η,

(70)

and the input-Lipschitz property of a transformer block kBj (θj , U ) − Bj (θj , V )k∞ ≤ Lin kU − V k∞ ,

(71)

where Parch (·) is a ﬁxed architecture polynomial and Lin ≥ 1 is independent of j. Let ∆j := kH (j) − H ′(j) k∞ ,

H (j) = Bj (θj , H (j−1) ),

H ′(j) = Bj (θj′ , H ′(j−1) ).

Add and subtract Bj (θj , H ′(j−1) ) and use the triangle inequality, (71), and (70): ∆j = kBj (θj , H (j−1) ) − Bj (θj′ , H ′(j−1) )k∞ ≤ kBj (θj , H (j−1) ) − Bj (θj , H ′(j−1) )k∞ + kBj (θj , H ′(j−1) ) − Bj (θj′ , H ′(j−1) )k∞ {z } | | {z } ≤Lin ∆j−1

≤Cb Parch κη

≤ Lin ∆j−1 + Cb Parch (·) κ η.

(72) 39

Generalization and Scaling Laws for Mixture-of-Experts Transformers

This is a linear inhomogeneous recurrence with ∆0 = 0 (since H (0) = H ′(0) share the same input and positional encoding). Solving the recurrence. By induction on j (or by the discrete Grönwall inequality), ∆j ≤ Cb Parch (·) κ η

j−1 X

t Lin .

(73)

t=0

Indeed, the base j = 1 gives ∆1 ≤ Cb Parch κη. Assuming (73) for j − 1, plug into (72): ∆j ≤ Lin · Cb Parch κη

j−2 X

t Lin + Cb Parch κη = Cb Parch κη

t=0

j−1 X

t Lin .

t=0

Geometric series bound. For j = LT , ∆LT = kH (LT ) − H ′(LT ) k∞ ≤

T −1 LX

t Lin Cb Parch (·) κ η.

t=0 L LinT −1

If Lin 6= 1, the sum is Lin −1 , so kH (LT ) − H ′(LT ) k∞ ≤

T LL in − 1 Cb Parch (·) κ η. Lin − 1

PLT −1 t T If 0 ≤ Lin ≤ 1, then t=0 Lin ≤ LT and LL in ≤ 1, so the bound still holds after increasing the prefactor by at most a constant depending on LT . In both cases we may write, for a numerical constant C ′ (absorbing |Lin1−1| or LT ), T kH (LT ) − H ′(LT ) k∞ ≤ C ′ LL in Parch (ℓ, m, demb , wFFN , LFFN , k) κ η.

(74)

From hidden states to the scalar output. Let the ﬁnal readout be Tθ (x) = Rω (H (LT ) ), where Rω is a ﬁxed linear map (e.g., token pooling followed by a linear form) with kωk∞ ≤ κ. Then |Tθ (x)−Tθ′ (x)| = Rω (H (LT ) )−Rω′ (H ′(LT ) ) ≤ kRω k∞→∞ kH (LT ) − H ′(LT ) k∞ + kRω − Rω′ k∞→∞ kH ′(LT ) k∞ . | {z } | {z } via (74)

≤ C ′′ κ η

Using (74) and the uniform forward bound on kH ′(LT ) k∞ (polynomial in the architecture; absorbed into C ′′ ), we obtain ′′ T sup |Tθ (x) − Tθ′ (x)| ≤ C ′ LL in Parch (ℓ, m, demb , wFFN , LFFN , k) κ η + C κ η ≤ Lparam η,

x∈X

with

T Lparam := C ⋆ LL in Parch (ℓ, m, demb , wFFN , LFFN , k) κ, ′ ′′ for a constant C absorbing C , C and ﬁxed readout norms. This is the claimed parameter-stability inequality.

⋆

Record · ID 5989 · SHA-256 966f1ac1936e3581
Conceptio Open Knowledge Archive — every document is proof-bundled with source, license, and retrieval metadata.