DynaMoE: Dynamic Token-Level Expert Activation with Layer-Wise Adaptive Capacity for Mixture-of-Experts Neural Networks

TL;DR

DynaMoE relaxes two rigid design assumptions in standard MoE — fixed Top-K routing and uniform expert counts per layer — by (a) replacing Top-K with a percentile-threshold gate so the number of active experts per token varies with input complexity, and (b) sweeping six predefined per-layer expert-count schedules (descending, ascending, pyramid-up/down, wave-up/down). On image classification a descending schedule (more experts in early layers, collapsing to 1 in the last) wins by 3–5% over uniform MoE; on language modeling the optimal schedule is scale-dependent (descending for Tiny, ascending for Small, uniform for Medium). All experiments are at toy scale (≤5.6M params, MNIST/CIFAR-10, 1k-sample LM corpus) and the paper explicitly flags this as a feasibility study rather than a scaling result.

Key claims

Standard MoE imposes two rigid constraints — fixed Top-K per token, and uniform expert count per layer — neither of which is theoretically motivated [§1].
DynaMoE replaces Top-K with a percentile threshold on gate values: a token activates every expert whose gate exceeds the τ-th percentile of its own gate distribution, so the active count varies per token between 1 and a layer-dependent maximum [§3.2, Def. 1, Eq. 4-5].
Six layer-wise expert-count schedules are defined as closed-form functions of layer index: descending, ascending, pyramid-up, pyramid-down, wave-up, wave-down [§3.3, Eqs. 8-12].
On MNIST with a 4-layer Small model, descending hits 92.68% accuracy vs 91.35% uniform MoE and 89.42% MLP baseline at matched parameters (~380K), and converges in 6 epochs to 95% of final accuracy vs 7-9 for other schedules [Table 2].
Cross-dataset, descending beats uniform MoE on every image task with the gap widening on harder tasks: +3.26% on MNIST, +4.19% on Fashion-MNIST, +5.47% on CIFAR-10 [Table 4].
On language modeling (1k-sample Recycling-the-Web, GPT-2 tokenization, seq-len 128), the optimal schedule shifts with model size: descending for Tiny (1011.80 PPL), ascending for Small (2308.29, the only DynaMoE variant beating the MLP baseline of 2311.02), uniform for Medium (2383.89 vs MLP’s 2468.16) [Table 5, §6.6.1].
The LM experiments use only 1,000 training samples — perplexities in the 1,000–2,500 range indicate severely underfitted models, and next-token accuracy saturates near unigram baseline (~5%) for Small and Medium, so the paper explicitly frames LM results as a pilot feasibility study, not a competitive benchmark [§6.6 caveat].
The percentile mechanism imposes no capacity factor and no auxiliary load-balancing loss; load imbalance is an acknowledged limitation observed but not catastrophic at the tested scales [§3.2.2].
A “Routing Diversity Gain” theorem shows dynamic routing’s reachable activation-pattern set strictly contains the Top-K set, with the ratio growing combinatorially in expert count; the paper flags this as a routing-pattern count, not a function-space volume [§4.1, Thm 1].
A “Gradient Variance Bound” claim (Thm 2) gives a qualitative characterization of variance reduction under more uniform routing entropy, with the paper explicitly disclaiming that assumptions A1–A3 may not hold in practice [§4.3].
The optimality of descending schedules on image tasks is rationalized via an Information Bottleneck argument: early layers carry high-entropy raw-feature distributions needing diverse experts; deep layers compress to a class-discriminative manifold where a single expert suffices [§7.1.1].

Method

A DynaMoE layer replaces Top-K with a soft, percentile-based selector. For each token the gate logits are computed as usual; the active set is every expert whose gate value exceeds the τ-th percentile of that token’s own gate distribution, with a minimum-activation floor of 1 to prevent dead tokens (Eq. 4, Algorithm 1). The selected experts’ outputs are combined with a temperature-softmaxed weighted sum (Eq. 6), and Gaussian noise is added to the gate values during training for exploration (Eq. 7). There is no auxiliary load-balancing loss and no capacity factor — overflow is handled by simply processing all selected tokens (worst case E_max evaluations per token).

The second axis is per-layer expert count. For a network of L layers with expert-count bounds [E_min, E_max], a schedule is a closed-form function E(l) mapping layer index to expert count. The six families sweep monotonic (ascending/descending), unimodal (pyramid-up/down), and oscillatory (wave-up/down) profiles. The experimental sweep tests which of these — together with E_max and E_min — gives the best accuracy-vs-parameter tradeoff for each dataset and model size. The implementation is MLP-based (no self-attention); Transformer extensions and the attention-entropy probes defined in §5.4 are explicitly deferred to future work.

Results

Headline numbers, all with very small models (85K–5.6M params):

MNIST Small (380K params): Descending 92.68% > Uniform MoE 91.35%

MLP 89.42% > Ascending 90.12%. Descending also fastest to converge [Table 2].
CIFAR-10: Descending 67.85% vs Uniform 65.12% vs MLP 62.38% (+5.47% over MLP) [Table 4].
Expert-count sweep: E_max=8, E_min=1 is the sweet spot — pushing to 16 max gives diminishing returns (92.45%) and raising the floor to 4 hurts (91.24%) because it removes the dynamic-routing flexibility [Table 3].
LM Small: Ascending 2308.29 PPL beats MLP 2311.02; descending 2335.91; uniform 2484.32 [Table 5]. But all configs have ~5% next-token accuracy ≈ unigram baseline — perplexity is the only discriminator and the corpus is 1k samples, so the result is directional at best.
Schedule-task divergence: descending wins on image classification across the board but loses to ascending or uniform on LM at Small and Medium scales — the paper’s takeaway is that there is no universal optimal schedule [§6.6.1].
Training-stability section reports overfitting (final PPL up to 6401.43 for Small Uniform, 4940.24 for Small Descending) — best PPL is always early/middle epoch, so the headline numbers depend on early-stopping the right schedule.

Why it’s interesting

Sam’s framing — “Is this modern day NAS” — is exactly the right read. DynaMoE doesn’t learn an architecture; it sweeps six closed-form per-layer expert-count schedules and reports which one wins per task × scale. That’s a hand-rolled architecture search dressed up as a principled framework, which is why the per-task / per-scale optimal schedule keeps flipping. The genuinely interesting bit is the observation that uniform allocation is empirically suboptimal — something MoEfication-style analyses also flag — but the 1k-sample LM corpus and ≤5.6M-param models mean none of the schedule rankings transfer to anything Luma-relevant. Contrast with Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method, which removes the load-balancing tuning axis at 32B-A5B / 1e22 FLOPs with no schedule sweep — i.e. fixes a real production pain point. And contrast with STEM: Scaling Transformers with Embedding Modules which approaches MoE-style sparsity from the opposite direction (replace the FFN with a static, token-indexed embedding lookup and eliminate routing entirely). The percentile-threshold gate is mechanically nice — it’s the natural soft relaxation of Top-K — but the lack of capacity control and auxiliary balancing means it’s unclear how it behaves once the corpus diversity stops being toy.