Dense vs Sparse Pretraining at Tiny Scale: Active-Parameter vs Total-Parameter Matching

Technical Summary: Dense vs. Sparse Pretraining at Tiny Scale

Problem Statement
Mixture-of-Experts (MoE) transformers decouple total parameter count from per-token active capacity, theoretically allowing models to store larger capacities while activating only a subset of parameters. However, much of the evidence supporting MoE superiority comes from large-scale, distributed settings where routing behavior, communication costs, and systems design are entangled, making it difficult to isolate the modeling contribution. This paper addresses a simpler empirical question: in a tiny, single-GPU-accessible regime, does replacing dense feed-forward networks (FFNs) with sparse routed experts yield a genuine modeling benefit under strictly controlled budgets?

Methodology
The study compares dense and sparse (MoE) transformers within a shared LLaMA-style decoder training stack trained on the TinyStories dataset. The experimental design holds constant the tokenizer, data pipeline, optimizer, schedule, depth, context length, normalization style, and evaluation protocol. The sparse intervention replaces dense FFN blocks with Mixtral-style routed experts.

To ensure fairness, the authors employ two distinct comparison criteria, resizing the dense baselines to match specific budgets:

1. Matched-Active: The dense baseline is resized so its active per-token parameter count matches the MoE's active routed path.
2. Matched-Total: The dense baseline is resized so its total stored parameter count matches the full MoE model.

The primary sparse configuration uses four experts, top-2 routing, dropless dispatch, Switch-style load balancing, and router z-loss. The training objective combines next-token cross-entropy with load-balancing and router z-loss auxiliary terms. The study operates in a sub-25M parameter regime (approx. 14.89M active vs. 21.06M total for the MoE).

Key Contributions

• Disentangling Fairness Criteria: The paper demonstrates that the conclusion of "dense vs. sparse" superiority depends entirely on the fairness metric used. It separates the analysis into matched-active and matched-total comparisons, showing they lead to different rankings.
• Sparse Training Stability at Tiny Scale: The study identifies the specific ingredients required for stable sparse training in small models. Naive top-1 routing collapses immediately; Switch-style balancing is required to prevent collapse, and moving to top-2 routing provides the primary quality gain. Router z-loss is shown to improve routing stability and diagnostics rather than being the primary source of validation loss reduction.
• Routing Diagnostics: The authors provide comprehensive routing diagnostics (busiest-expert fraction, usage variance, entropy, log-z) to confirm that the observed gains are not artifacts of collapsed or pathological routing.

Results
The study presents a three-seed full-data comparison across the three configurations:

• Matched-Active Comparison: The MoE model (approx. 14.77M active params) significantly outperforms the dense active-match baseline (approx. 14.89M active params). The MoE achieves a best validation loss of 1.5788 ± 0.0020, compared to the dense baseline's 1.6545 ± 0.0012. This yields a gap of 0.0758 in favor of the MoE.
• Matched-Total Comparison: The dense model (approx. 21.06M total params) remains slightly stronger than the MoE (approx. 21.06M total params). The dense total-match baseline achieves 1.5608 ± 0.0025, while the MoE achieves 1.5788 ± 0.0020. This yields a gap of 0.0180 in favor of the dense model.
• Training Dynamics: Across training budgets (40.5M to 426.3M tokens), the matched-active advantage for the MoE grows, while the matched-total advantage for the dense model narrows sharply. The MoE consistently sits between the two dense baselines: better than the equal-active dense model but slightly worse than the equal-total dense model.

Significance and Claims
The paper concludes that in this sub-25M parameter regime, conditional computation (MoE) improves validation loss when fairness is defined by active per-token parameters, but it does not surpass dense training when fairness is defined by total stored capacity.

The authors explicitly state that this is a controlled empirical study rather than a new architecture proposal. The significance lies in clarifying that:

1. Sparse training is already useful in tiny-scale regimes if the comparison is based on active capacity.
2. The "MoE advantage" is not a raw total-capacity superiority but a benefit derived from the conditional use of additional stored parameters.
3. The results are specific to the TinyStories regime and do not claim to automatically extend to broader corpora, larger parameter regimes, or downstream transfer settings without further study.

The paper refrains from claiming a wall-clock training win, noting that throughput analyses are secondary context, and avoids making causal claims about dormant capacity mechanisms, focusing instead on the mapping of conclusions to fairness criteria.