PaperlayerView exact revision on arXiv

Source-linked AI summary

Dense Mixture-of-Experts as a Reparameterized Wide FFN: A Granularity Sweep at Fixed Compute

Vu Quang Hoang, Nghia Hieu Nguyen

arXiv:2610.02584v1cs.LG

TL;DR

Sparse MoE combines routing with large expert pools and discrete selection, leaving the benefit of dynamic per-token combination independent of those features as a distinct question. The study compares a dense SwiGLU baseline with K complete SwiGLU experts that are always active and combined by a token-wise softmax gate at matched total FFN width. Validation loss is non-monotonic in K: K=2 slightly beats the dense baseline (K=1), while K=4 and K=6 fall increasingly behind and never overtake it.

  • Problem

    Sparse MoE combines routing with large expert pools and discrete selection, leaving the benefit of dynamic per-token combination independent of those features as a distinct question.

    Source: §1 · p. 2

  • Method

    The study compares a dense SwiGLU baseline with K complete SwiGLU experts that are always active and combined by a token-wise softmax gate at matched total FFN width.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    Validation loss is non-monotonic in K: K=2 slightly beats the dense baseline (K=1), while K=4 and K=6 fall increasingly behind and never overtake it.

    Source: §8 · p. 12

  • Takeaways & Limitations

    At fixed width, the tested always-active, softly gated FFNs have a granularity sweet spot, with K=2 performing best among the configurations tested.

    Source: §8 · p. 12

    Each configuration was run once, so the reported validation-loss gaps are single-run figures rather than seed-replicated comparisons.

    Source: §7 · p. 11

Abstract

from arXiv · show

Sparse Mixture-of-Experts (MoE) models combine learned routing with selection from a large expert pool. We isolate the contribution of dynamic expert combination using a dense analogue: $K$ SwiGLU experts, all active for every token and combined by a softmax gate, at fixed total FFN width. With no larger pool or discrete selection, the dense baseline is the $K=1$ case. Validation loss varies non-monotonically with $K$: $K=2$ improves over the baseline by $0.0048$, whereas $K=4$ and $K=6$ worsen it by $0.0053$ and $0.0197$, respectively. Routing generally remains soft and expert usage balanced, except in the first layer of the $K=4$ model, where routing is nearly one-hot. Forcing this gate to uniform increases loss by $2.4$ nats on a diagnostic subset, indicating that its concentrated routing is functionally important. We further show that the architecture is exactly a dense SwiGLU with token-dependent, simplex-constrained group scaling and contains the dense baseline in its function class. These single-run results suggest a granularity sweet spot for always-active, softly gated FFNs at fixed width, with $K=2$ performing best among the configurations tested.

1 Introduction

Sparse MoE uses discrete routing to make large expert pools affordable, but that leaves open whether dynamic expert combination itself improves quality. This work tests always-active, softly gated experts at fixed compute and width, finding a non-monotonic validation-loss pattern across K.

1 Introduction

  • The controlled comparison uses K complete SwiGLU experts that are always active and dynamically weighted, matched to a dense SwiGLU transformer in parameters and FLOPs.The dense baseline is K=1; the expert model has no discrete selection or larger expert pool, isolating dynamic per-token combination.Source: §1 · p. 2
  • The architecture is exactly a wide dense SwiGLU with token-dependent, simplex-constrained group scaling, and the dense baseline is the K=1 case.This makes the K sweep a study of the granularity of the added group-scaling mechanism.Source: §1 · p. 2
  • An early conclusion that the baseline won, based on single-batch logged training loss, did not hold on end-of-annealing validation loss.The reported comparison therefore depends on evaluating validation loss at the end of annealing.Source: §1 · p. 2
  • K=2 finishes below the dense baseline in validation loss, while K=4 and K=6 finish above it by a growing margin.Only K=2’s loss curve crosses the baseline during training; K=4 and K=6 never do and plateau at a persistent deficit.Source: §1 · p. 2
  • A single layer in the K=4 model is a clear exception to generally soft routing and balanced expert usage: its gate is nearly one-hot.Layer-by-layer diagnostics expose a collapse that a layer-averaged summary would miss.Source: §1 · p. 2
  • The authors offer practical guidance for small-scale MoE-inspired research, where implementation details can reverse an apparent result.The comparison is designed to isolate dynamic combination from the systems machinery needed for sparse routing, including capacity limits and all-to-all communication.Source: §1 · p. 2

2 Related Work

Prior work studies dense or soft gating, sparse routing, and FFN granularity from different perspectives. The authors position their study as a fixed-width sweep of group count for an always-active, softly gated FFN, analyzed through an exact dense-FFN equivalence.

Classical dense-gated and sparse MoE

  • The original dense-gated MoE evaluates every expert and combines outputs with a differentiable gate; sparse top-k routing later made larger expert pools tractable.GShard and Switch Transformer developed load-balancing and capacity machinery, while ST-MoE introduced router z-loss for stability.Source: §2 · p. 3

Soft routing via token merging

  • Soft MoE combines expert outputs through slots formed as learned weighted averages over all sequence tokens, unlike this paper’s per-token mixing of expert outputs.The cited work notes that whole-sequence normalization is difficult to apply to autoregressive decoders; the studied per-token gate has no cross-token dependency.Source: §2 · p. 3

Soft routing via parameter merging

  • SMEAR merges expert parameters before a single forward pass, whereas this work combines expert outputs at K times a single expert’s FLOPs, matched to a K-times-wider dense baseline.The passage states that parameter merging and output merging are not interchangeable for nonlinear experts.Source: §2 · p. 3

Geometric and theoretical analyses of soft routing

  • A close empirical comparison studies dense, top-k, and fully-soft routing through representational geometry rather than training dynamics or task loss.The passage also notes theory on convergence for jointly trained soft-routed MoE, while relating observed expert redundancy to those analyses is beyond scope.Source: §2 · p. 3

Causality-preserving differentiable routing

  • Differentiable soft top-k relaxes hard per-token routing while retaining sparsity, whereas this design is dense over a small fixed K.The cited soft top-k method also uses a learned per-layer expert budget.Source: §2 · p. 3

Granularity, dense training, and FFN splitting

  • Prior granularity work concerns sparse MoE, while this paper sweeps groups in an always-active, softly gated FFN at fixed total width.Related work also includes dense training with sparse inference, dense router gradients with sparse activation, FFN splitting, and multi-branch FFNs.Source: §2 · pp. 3–4

Positioning

  • Within the searches they ran, the authors found no prior study sweeping K for a softly gated, always-active FFN at fixed total width with the dense baseline as K=1 and an exact wide-dense-FFN equivalence.Prior soft-routing work instead targets capacity gains, single-expert FLOPs, or representational geometry.Source: §2 · p. 4

3 Method

The method replaces the baseline FFN with always-active SwiGLU experts and a token-wise softmax router, matching total hidden width. Algebraically, it is a wide dense SwiGLU with token-dependent group scaling; initialization variants examine the resulting output scale.

3.1 Dense baseline

  • The baseline uses a SwiGLU FFN, and the compared expert configurations retain the same total hidden width while splitting it across K experts.The parameter-count table specifies 12 layers, dmodel=384, and a router addition of at most 2,304 parameters, treated as negligible for matching.Source: §3.1 · p. 4, Table 1 · p. 5

3.2 Dense dynamic mixture of experts

  • The dynamic model replaces the FFN with K complete SwiGLU experts and a linear router; every expert is evaluated for each token, with no top-k selection, capacity factor, or balancing loss.The gate depends only on the current token, so decoding remains causal and incremental.Source: §3.2 · p. 4

3.3 Relation to a wide dense FFN

  • Stacking expert hidden activations yields a dense SwiGLU of width Kde, whose hidden units receive token-dependent scalars shared within each group and constrained to the simplex.At a uniform gate, the function class contains the dense baseline; the only added flexibility over the wide dense FFN is token-dependent group scaling.Source: §3.3 · pp. 4–5

3.4 Initialization and router variants

  • The router is initialized so the gate is approximately uniform, and its simplex weights shrink output variance; the step-0 audit reports fan-in residual ratio ≈1/K and a scaled variant roughly an order of magnitude quieter in standard deviation.The audit uses Gaussian inputs with dmodel=384 and K=3.Source: §3.4 · p. 5, Table 2 · p. 5

4 Experimental Setup

The experiment compares dense SwiGLU and dense dynamic MoE models at fixed total FFN width. Models see the same data in the same order, and comparisons use validation loss.

4 Experimental Setup

  • The models are 12-layer decoders with dmodel=384, six attention heads, RMSNorm, RoPE, tied embeddings, and a 50,257-token GPT-2 BPE vocabulary.Source: §4 · p. 5
  • All FFNs have hidden width 1536; the baseline uses K=1, while expert models use K=2, 4, or 6 with narrower experts.The configurations differ only in the FFN and use one fixed initialization/router variant.Source: §4 · p. 6
  • Training uses 1,343,488,000 tokens from a FineWeb-10B subsample, about 1.03 passes, with identical data order across models.The schedule has 500 warmup steps followed by 20,000 cosine-annealing steps; each optimizer step covers 65,536 tokens.Source: §4 · p. 6
  • Models are compared by validation loss on the full, disjoint FineWeb validation split every 500 steps; each configuration was run once without seed-based replication.Logged training loss is a single-batch measurement and is considered too noisy for cross-run comparison.Source: §4 · p. 6

5 Results

At fixed total FFN width, validation loss is non-monotonic in expert count: K=2 beats the dense baseline, while K=4 and K=6 finish worse. Router behavior is mostly soft and balanced, but K=4 layer 0 is a consequential exception; uniformizing its gate sharply raises loss on a diagnostic subset.

5.1 Dense dynamic MoE vs. dense baseline

  • K=2 finishes 0.0048 below the dense baseline, while K=4 and K=6 finish 0.0053 and 0.0197 above it, respectively.The comparison is at step 20,000; the curves reach distinct plateaus before training ends.Source: Table 3 · p. 6, §5.1 · p. 7

5.2 An early, confounded observation at smaller scale

  • The earlier roughly 0.03-loss advantage at ∼5M parameters is confounded by scaled initialization and router scale, while the later K=2 result without them suggests it was not purely an artifact.The earlier observation opened near a token-to-parameter ratio of roughly 20:1 [6].Source: §5.2 · p. 7

5.3 Reading the granularity sweep

  • With no larger expert pool to provide compensating selection, splitting fixed width into narrower groups costs capacity as K grows; why K=2 improves remains open.The reported ordering is K=2 slightly better than baseline, with K=4 and K=6 progressively worse.Source: §5.3 · pp. 7–8

5.4 Initialization and router variant

  • All reported runs use fan-in expert initialization and no router logit scale; faster early loss under these choices was observed only in uncontrolled, partially trained runs.Those observations were not re-verified as a controlled ablation on completed runs.Source: §5.4 · p. 8

5.5 Router diagnostics

  • Layer-averaged routing is mostly non-collapsed: usage entropy is at least 0.90 and per-token entropy is 0.67–0.72, with mean maximum gate weight above uniform.For K=2, the mean maximum gate weight is 0.762 versus 0.500 under uniform routing.Source: §5.5 · pp. 8–9
  • K=4 layer 0 is the exception: 90.9% of tokens have a maximum gate weight above 0.9 and usage is 93/1/5/1%, unlike the other soft, balanced layers.The final-checkpoint statistics do not reveal when this layer’s collapse emerged during training.Source: §5.5 · p. 9

5.5.1 Expert redundancy in the K=4 model

  • Expert-output cosine similarity is highest in K=4 layers 0 and 11, but only layer 0’s router collapses, so the alignment alone does not explain the routing pattern.Their mean similarities are 0.414 and 0.388, respectively; layer 11’s router remains soft and balanced.Source: §5.5.1 · pp. 9–10

5.5.2 Layer 0 is load-bearing, not redundant: a direct ablation

  • Forcing K=4 layer 0’s gate to uniform raises loss by roughly 2.4 nats on the fixed subset, showing that the concentrated routing is load-bearing rather than cheap redundancy.The ablation changes only layer 0’s gate and evaluates both conditions on the same subset.Source: §5.5.2 · p. 10, Table 6 · p. 10

6 Discussion

At fixed width, the K=2 result motivates a possible granularity trade-off, but its mechanism remains a hypothesis. The K=4 model also shows that concentrated routing can be functionally important in one layer.

6 Discussion

  • K=2 beats the dense baseline slightly, while K=4 and K=6 perform worse; the proposed trade-off between reweighting flexibility and narrower groups remains unestablished.The authors predict greater expert redundancy at K=2 and a smaller gain from static K=2 weights, but present these as tests for future work.Source: §6 · pp. 10–11
  • Forcing uniform routing in K=4’s collapsed layer costs roughly 2.4 nats, ruling out the idea that its elevated expert similarity makes the collapse cheap.The authors instead interpret the similarity as position-dependent alignment; why layer 0 collapses while similarly aligned layer 11 does not remains unexplained.Source: §6 · p. 11
  • The model is equivalent to a wide dense FFN with token-dependent group scaling, so observed differences are not due to greater function-class capacity.The design cannot select from a larger expert pool, which the authors note is a proposed source of sparse-MoE benefits.Source: §6 · p. 11
  • The design avoids load-balancing loss, capacity factors, and token dropping while remaining causal and incrementally decodable.These practical properties do not depend on whether K=2’s small performance gain holds up under further scrutiny.Source: §6 · p. 11

7 Limitations

The evidence is limited to small-scale, single-run experiments on one dataset, and the authors caution against treating the observed gaps or proposed mechanism as established.

7 Limitations

  • Results cover one dataset at small scale: 28–29M non-embedding parameters and 48M total parameters.The authors describe the findings as specific to this scale, dataset, and architecture.Source: §7 · p. 11
  • Each configuration was run once without a fixed seed, so the validation-loss gaps are single-run figures rather than seed-replicated comparisons.Stable plateaus are read as evidence of architectural effects, but that interpretation is not established by seed-based replication.Source: §7 · p. 11
  • The study tests only K values {1, 2, 4, 6}, trains no sparse MoE, and has no training-time trajectory for the K=4 collapse.The layer-0 ablation uses a smaller fixed validation subset, and the proposed K=2 trade-off remains a hypothesis.Source: §7 · p. 11

8 Conclusion

At fixed total FFN width, K=2 performs best among the tested expert configurations, while K=4 and K=6 finish behind the dense baseline. The findings are single-run evidence for this setting.

8 Conclusion

  • K=2 slightly beats the dense baseline, while K=4 and K=6 fall increasingly behind and never overtake it; all three gaps stabilize over the final 11,000 of 20,000 steps.The K=2 gap crosses zero partway through training and remains below the baseline thereafter.Source: §8 · p. 12
  • Routing is mostly soft and balanced, but one K=4 layer collapses toward near-hard routing without an explicit mechanism encouraging it.The collapse is a local exception, not a general explanation for the performance pattern.Source: §8 · p. 12
  • The architecture is equivalent to a wide dense FFN with token-dependent group scaling, framing the comparison as a granularity sweep without a larger expert pool.The authors caution that these results come from single runs without seed-based replication.Source: §8 · p. 12

Related papers

Loading 2610.02584v1…