💬 NLP

Optimal Sparsity of Mixture-of-Experts Language Models for Reasoning Tasks

This paper demonstrates that for Mixture-of-Experts language models, reasoning performance is optimized not just by minimizing pre-training loss but by jointly balancing active FLOPs and total tokens per parameter, revealing that reasoning tasks are data-hungry and benefit from specific sparsity configurations distinct from memorization tasks.

Original authors: Taishi Nakamura, Satoki Ishikawa, Masaki Kawamura, Takumi Okamoto, Daisuke Nohara, Jun Suzuki, Rio Yokota

Published 2026-03-03

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Taishi Nakamura, Satoki Ishikawa, Masaki Kawamura, Takumi Okamoto, Daisuke Nohara, Jun Suzuki, Rio Yokota

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Specialist" vs. The "Generalist" Dilemma

Imagine you are building a massive library of knowledge to train a super-smart AI. You have a fixed budget for how much time and money you can spend on this project.

In the past, the best way to build this library was to hire one giant, super-broad expert who reads everything. This is a Dense Model.

Recently, the industry switched to a new strategy called Mixture-of-Experts (MoE). Imagine instead of one giant expert, you hire a team of 1,000 tiny specialists. But here's the catch: for every single question the AI gets, it only asks two of those specialists for help. The other 998 stay silent.

Total Parameters: The total number of specialists you hired (the size of the whole team).
Active Parameters: The number of specialists actually doing the work on a specific question (the "Top-k").

The paper asks a simple question: If we have a fixed budget, should we hire a huge team of specialists and only use a few of them (High Sparsity), or should we hire a smaller, denser team where everyone works harder?

The answer depends entirely on what you want the AI to do.

The Two Types of Skills: Memory vs. Reasoning

The researchers discovered that AI has two very different "muscles," and they need different training diets.

1. The "Photographic Memory" Muscle (Memorization)

Tasks: Trivia, facts, reading comprehension (e.g., "Who wrote Hamlet?" or "What is the capital of France?").
The Finding: For these tasks, bigger is better.
The Analogy: Imagine you are training for a trivia night. You want a library with millions of books. It doesn't matter if you only read two pages at a time; as long as the library is huge, you will find the answer.
Result: The more "inactive" specialists you have (high sparsity), the better the AI gets at memorizing facts. It's like having a massive warehouse of information, even if you only pull a few boxes off the shelf.

2. The "Mathematical Logic" Muscle (Reasoning)

Tasks: Solving math problems, coding, complex logic puzzles (e.g., "If a train leaves at 5 PM...").
The Finding: For these tasks, too many specialists is actually bad.
The Analogy: Imagine you are training a team of engineers to build a bridge. If you have 1,000 engineers but only let 2 of them work on the blueprint, those 2 engineers are going to be starved for data. They haven't seen enough examples to learn how to build a bridge. They are "data-hungry."
Result: If you make the team too huge (too many inactive experts), the active experts don't get enough practice. The AI gets worse at reasoning, even though it has "more knowledge" in the warehouse.

The Two Golden Rules Discovered

The paper identifies two main principles that determine the perfect balance:

Rule #1: Active FLOPs (The "Workout" Rule)

What it means: It's not just about how smart the AI could be; it's about how much "brain power" it actually uses to solve a problem.
The Analogy: Two runners might have the same potential (same training loss), but one runs a 5-mile race while the other jogs 1 mile. The one who runs 5 miles (higher Active FLOPs) gets stronger.
Takeaway: For reasoning, you need the AI to actually think harder (use more experts per question), not just have a bigger team sitting on the sidelines.

Rule #2: Tokens Per Parameter (The "Practice Ratio" Rule)

What it means: This is the ratio of "how much data the AI sees" vs. "how many brains it has."
The Analogy:
- Memorization: You want a huge library (many parameters) and a moderate amount of reading.
- Reasoning: You need a sweet spot. The paper found that the AI is happiest when it sees about 20 words of data for every 1 brain cell it has.
- Too many brains, too little data: The brains are bored and under-trained (Reasoning fails).
- Too few brains, too much data: The brains are overwhelmed and can't store it all (Memorization fails).

The "Magic Fix" That Didn't Work

The researchers wondered: "If we make the AI too sparse (too many inactive experts), can we just fix it later?"

They tried two popular modern tricks:

Reinforcement Learning (GRPO): Like a coach yelling "Good job!" when the AI gets a math problem right.
Test-Time Compute: Like giving the AI more time to think or asking it to solve the problem five times and picking the best answer.

The Result: Neither trick fixed the problem.

The Analogy: If you train a chef to cook with only 2 ingredients because you have a kitchen full of 1,000 unused spices, giving them a "cooking bonus" (Reinforcement Learning) or "more time to think" (Test-Time Compute) won't make them a master chef. They were never taught how to use the full kitchen in the first place.

Conclusion: You have to get the team structure right during the initial training. You can't fix a bad architecture later.

Summary for the General Audience

For Trivia/Facts: Go big! Hire a massive team of specialists and only use a few. The more unused experts you have, the more facts the AI can store.
For Math/Logic: Go dense! Don't hire a massive team if you can't let them all work. You need a smaller, tighter team where everyone gets enough practice data to learn how to solve complex problems.
The Lesson: There is no "one size fits all" for AI. If you want a reasoning AI, you can't just throw more money at a sparse model. You have to balance the size of the team with the amount of data they get to practice on.

The paper essentially tells us: Stop assuming bigger is always better. For thinking tasks, a focused, well-fed team beats a massive, underfed one.

1. Problem Statement

While empirical scaling laws have successfully guided the development of dense Large Language Models (LLMs), these laws often fail to account for architectural shifts, specifically the transition to Mixture-of-Experts (MoE) models. MoE models introduce a new dimension of sparsity (routing tokens to a subset of experts) that allows for massive parameter counts without proportional increases in inference cost.

However, it remains unclear how MoE sparsity affects different capability regimes:

Memorization: Tasks requiring knowledge retrieval (e.g., TriviaQA, HellaSwag).
Reasoning: Tasks requiring multi-step logical deduction (e.g., GSM8K, GSM-Plus, coding).

Existing research suggests that increasing sparsity (more experts, fewer active per token) is generally beneficial for fixed compute budgets. This paper challenges that assumption, investigating whether the "optimal" sparsity is universal or if it diverges significantly between memorization and reasoning tasks, and whether post-training techniques (RL) or test-time compute (TTC) can mitigate any negative effects of high sparsity on reasoning.

2. Methodology

The authors conducted a large-scale systematic study using a Mixtral-style architecture (Transformer backbone with RMSNorm, SwiGLU, and rotary embeddings).

Experimental Setup:
- Corpus: A fixed 125B-token pre-training corpus (balanced mix of web text, math, STEM, and code).
- Variables: They swept three architectural hyperparameters:
  1. Model Width ( $d$ ): 512, 1024, 2048.
  2. Number of Experts ( $E$ ): 8 to 256.
  3. Top- $k$ Routing: 2, 4, 8, 16.
- Constraints: Experiments were run under fixed compute budgets (IsoFLOP) and fixed data budgets to isolate the effects of sparsity.
Evaluation Benchmarks:
- Memorization: TriviaQA (Reading Comprehension), HellaSwag (Commonsense).
- Reasoning: GSM8K, GSM-Plus (Math), HumanEval, MBPP (Code).
Post-Training & Test-Time Analysis:
- Applied GRPO (Group Relative Policy Optimization) for RL post-training.
- Applied Test-Time Compute (TTC) via Self-Consistency (repeated sampling) to see if inference-time scaling could recover reasoning performance.
Metrics: Trained loss, validation loss, downstream task loss (cross-entropy on answers), and accuracy (Pass@1 for code, accuracy for others).

3. Key Contributions & Findings

A. Divergence Between Memorization and Reasoning

The study reveals a fundamental asymmetry in how MoE sparsity impacts different skills:

Memorization Tasks: Performance improves monotonically as sparsity increases (more total parameters, fewer active parameters). Larger total parameter counts with high sparsity yield better results, consistent with "parameter-hungry" behavior.
Reasoning Tasks: Performance follows a non-monotonic (inverted U-shaped) trend.
- Initially, increasing sparsity improves reasoning.
- However, beyond a certain threshold of total parameters (while keeping active parameters fixed), further increasing sparsity degrades reasoning performance, even if the pre-training loss continues to decrease.
- Key Insight: For reasoning, models with identical training loss but higher active FLOPs (larger top- $k$ ) consistently outperform sparser models.

B. The Role of Active FLOPs

The paper establishes that Active FLOPs (compute used during inference/training per token) are a critical determinant of reasoning quality, independent of total parameter count.

Models with larger top- $k$ values (more experts activated per token) achieve higher reasoning accuracy even when total parameters are held constant.
This suggests reasoning requires substantial "active compute" to execute complex logical chains, which high-sparsity configurations (low top- $k$ ) may starve.

C. Total Tokens Per Parameter (TPP)

The authors identify an optimal Total Tokens Per Parameter (TPP) ratio for reasoning tasks:

Memorization: Benefits from low TPP (more parameters relative to data).
Reasoning: Peaks at approximately 20 tokens per parameter.
- If TPP is too low (too many parameters for the data), the model becomes "data-starved" regarding reasoning patterns.
- If TPP is too high (too few parameters), the model is under-parameterized.
This finding refines the Chinchilla scaling law, showing that the optimal TPP is task-dependent, not a universal constant.

D. Ineffectiveness of Post-Training and Test-Time Compute

The study investigates whether advanced techniques can fix the reasoning degradation caused by excessive sparsity:

GRPO (RL Post-training): While GRPO improves overall performance, it does not eliminate the inverted U-shaped trend. The gap between memorization and reasoning performance caused by high sparsity remains.
Test-Time Compute (TTC): Increasing inference compute (e.g., Self-Consistency with 128 samples) improves absolute accuracy but fails to recover the performance drop of overly sparse models.
Conclusion: The optimal sparsity for reasoning must be determined during pre-training; it cannot be fully recovered by post-training or inference-time scaling.

4. Results Visualization

Figure 2 & 3: Show that while training loss decreases monotonically with more experts, GSM8K task loss and accuracy begin to worsen once total parameters exceed a threshold (specifically when active parameters are fixed but total parameters grow).
Figure 5: Demonstrates that for fixed active parameter budgets, increasing sparsity (lower density) helps memorization tasks but eventually hurts reasoning tasks (GSM8K) as the active parameter count grows large.
Figure 7: Illustrates the TPP curve, where reasoning accuracy peaks around TPP $\approx$ 20, whereas memorization accuracy keeps rising as TPP decreases.

5. Significance and Implications

Revising Scaling Laws: The paper argues that the classical view of compute-optimal scaling (focusing solely on total parameters vs. tokens) is incomplete for MoE models. Optimal sparsity must be jointly determined by Active FLOPs and TPP.
Architectural Guidance: For reasoning-heavy applications (math, code, logic), simply adding more experts and increasing sparsity is counterproductive. Instead, a denser configuration (higher top- $k$ or fewer total experts for a given active budget) is often superior in high-compute regimes.
Pre-training Priority: The findings emphasize that the "reasoning capability" of an MoE model is largely baked in during pre-training. Relying on RL or test-time scaling to fix architectural mismatches in sparsity is insufficient.
Open Science: The authors released model checkpoints, code, and logs, providing a reproducible baseline for future MoE scaling research.

In summary, the paper provides a nuanced framework for designing MoE models, asserting that while sparsity is excellent for scaling capacity (memorization), reasoning requires a careful balance of active compute and data density, often favoring less sparse configurations than those optimal for pure knowledge retrieval.