Optimal Expert-Attention Allocation in Mixture-of-Experts: A Scalable Law for Dynamic Model Design

TL;DR

A scaling-law paper from HKUST(GZ) + Ant Group asking the question prior MoE scaling work has implicitly skipped: given a compute budget, how should per-token FLOPs split between attention and experts? Defining $r$ as the fraction of total per-token FLOPs spent in expert (FFN) layers vs. attention, the authors run GPT-style MoE sweeps across compute scales and sparsity levels and fit an explicit power-law for the optimum $r^*$ as a function of total compute $C$ and sparsity $S$ . The headline empirical finding is that $r^*$ increases with compute — i.e. as FLOPs grow, capacity is better spent on the FFN/expert path than on attention. The result is a closed-form prescription that can plug into existing MoE architecture search, slotting alongside expert-count, granularity, and sparsity as a first-class scaling-law dimension.

Key claims

Prior MoE scaling laws (Ludziejewski 2024; Abnar 2025; Wang 2024) treat the attention-vs-expert FLOPs allocation as fixed and only vary expert count, granularity, or sparsity — leaving the expert–attention trade-off unmodeled [§1, §2].
Define $r$ = fraction of total per-token FLOPs going to expert layers (vs. attention); the optimal $r^*$ is not a constant but follows a predictable scaling behavior in total compute $C$ and sparsity $S$ [§1, §3].
Across controlled experiments at multiple scales and sparsity levels, $r^*$ increases as a power law of total training compute, with the scaling coefficients themselves varying systematically with sparsity [§3, Fig. 5 / Table 3].
This yields an explicit closed-form formula for $r^*(C, S)$ , enabling precise compute-allocation choices given a fixed training budget [Abstract, §3].
Architectural decisions in MoE (expert count, activation sparsity) implicitly determine compute allocation; making $r$ an explicit scaling dimension is the contribution that distinguishes this work from prior MoE scaling laws [§2].
FLOPs accounting used for the sweeps follows the Narayanan et al. (2021) scheme: attention FLOPs are computed per the standard $4 \cdot n_\text{ctx} \cdot d_\text{hidden}$ form, logits are $C_\text{logits} = 2 \cdot n_\text{ctx} \cdot d_\text{hidden} \cdot n_\text{vocab}$ , and backward is forward × a constant factor [Appendix; Eq. 2 / Table 3].
The reported predictions are validated against observed validation loss on held-out MoE configurations [Fig. 5].

Method

The setup is a sweep over GPT-style decoder-only MoE Transformers, varying total training compute $C$ , sparsity $S$ (fraction of experts active per token), and the expert-vs-attention FLOPs ratio $r$ . Because $r$ is itself a function of architecture (hidden dim, head count, expert size, expert count, top-k), the authors hold sparsity and total compute fixed in slices and walk $r$ by adjusting attention dim / expert FFN dim trade-offs. For each $(C, S)$ slice they identify the $r^*$ that minimizes validation loss, then fit a global power-law surface to the collected $(C, S, r^*)$ triples.

The fit machinery follows the standard scaling-laws-by-isoFLOPs recipe: multi-scale runs, parabolic-in-log- $r$ fits at fixed $(C, S)$ to locate $r^*$ , then a separate parametric fit for how $r^*$ evolves with $C$ and $S$ . The FLOPs counting is anchored to the Narayanan et al. 2021 accounting (attention quadratic in sequence length, MLP linear in hidden dim with grouped GEMMs for experts).

Results

The headline qualitative claim — $r^*$ rises with $C$ — is exactly the inversion of the common practitioner default of holding the attention/FFN dim ratio fixed across scales. At large $C$ the compute-optimal split allocates more to experts; at small $C$ , more to attention [§3, Fig. 5].
Scaling coefficients depend on sparsity: denser-MoE configurations (higher fraction of experts active) have different $r^*$ slopes than sparser configs at matched $C$ [§3, Table 3].
Predicted vs. observed validation loss for held-out expert/attention allocations is reported in Figure 5 as validation of the closed-form law.
The paper does not report absolute SOTA numbers — its claim is relative, that within a compute budget, the prescribed $r^*$ reliably beats off-axis choices.

Why it’s interesting

Kyle’s gloss — “as you increase FLOPs, capacity is more efficiently allocated to the MLPs instead of to the attention” — is exactly the direction that’s been bothering everyone running MoE pretraining: as total params grow, the standard 4× FFN/attention ratio drifts off optimum, and there hasn’t been a principled prescription. This paper gives one. It’s directly complementary to Scalable Training of Mixture-of-Experts Models with Megatron Core — the Megatron report names the “parameter-compute mismatch” between attention (TP-friendly) and MoE (EP-friendly) as the systems reason MoE-at-scale is hard, and this paper gives the algorithmic counterpart, prescribing how much compute should land on each side in the first place.

The result also sharpens an open question in Hyperparameter scaling laws: the Fitting paradigm (How to Set the Learning Rate for Large-Scale Pre-training?) extrapolates $\eta^*(N, D)$ from cheap small-scale sweeps but holds architecture fixed. If $r^*$ itself scales with $C$ , then “scale the LR but freeze the architecture” is leaving compute on the table — the architecture-search axis has its own power law and needs to be co-scaled with the optimizer-side one. Adjacent to DynaMoE: Dynamic Token-Level Expert Activation with Layer-Wise Adaptive Capacity for Mixture-of-Experts Neural Networks (which sweeps per-layer expert counts at toy scale) and Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method (which removes the load-balancing tuning knob): all three are “make the MoE design space less arbitrary”, but this paper is the cleanest scaling claim of the three.