adu2021 Optimal Sparsity of MoE Language Models for Reasoning
Core Concept
This work reveals that optimal MoE sparsity differs from traditional dense model scaling. The key insight: memorization and reasoning have opposing sparsity preferences. Memorization improves with more total parameters (dense experts better), while reasoning improves with more active FLOPs (sparse routing better). The paper proposes joint optimization of active FLOPs and tokens-per-parameter (TPP) as the path to compute-optimal reasoning models, revising classical scaling laws.
Architecture Overview
- Two Capability Dimensions: Separate memorization from reasoning evaluation
- Active FLOPs Principle: More compute helps reasoning independent of model size
- TPP Principle: Total tokens per parameter correlates with memorization efficiency
- MoE Configuration Space: Vary experts, top-k routing, and total parameters
- Joint Optimization: Balance both principles for overall performance
Implementation Steps
Stage 1: Establish Evaluation Framework
Separate memorization and reasoning in evaluation.
# Separate memorization from reasoning evaluation
from typing import Dict, List, Tuple
import numpy as np
class CapabilityEvaluator:
"""Evaluate memorization vs reasoning separately"""
def __init__(self):
self.benchmarks = {
"memorization": [
"fact_recall",
"knowledge_qa",
"entity_extraction"
],
"reasoning": [
"math_reasoning",
"logical_deduction",
"reading_comprehension"
]
}
def evaluate_model(
self,
model,
test_datasets: Dict[str, List]
) -> Dict:
"""
Evaluate model on both memorization and reasoning tasks.
"""
results = {
"memorization": {},
"reasoning": {},
"pre_training_loss": 0.0
}
# Memorization benchmarks
for benchmark in self.benchmarks["memorization"]:
dataset = test_datasets.get(benchmark, [])
accuracy = self.evaluate_benchmark(model, dataset)
results["memorization"][benchmark] = accuracy
# Reasoning benchmarks
for benchmark in self.benchmarks["reasoning"]:
dataset = test_datasets.get(benchmark, [])
accuracy = self.evaluate_benchmark(model, dataset)
results["reasoning"][benchmark] = accuracy
# Pre-training loss (proxy for memorization efficiency)
results["pre_training_loss"] = self.get_pretrain_loss(model)
return results
def evaluate_benchmark(self, model, dataset: List) -> float:
"""Evaluate accuracy on benchmark"""
correct = 0
for example in dataset:
prediction = model.generate(example["input"])
if prediction == example["target"]:
correct += 1
return correct / len(dataset) if dataset else 0.0
def get_pretrain_loss(self, model) -> float:
"""Get pre-training loss (approximates memorization capacity)"""
return 0.0 # Placeholder
class CapabilityAnalysis:
"""Analyze memorization vs reasoning trade-offs"""
def __init__(self):
self.results = []
def add_model_result(
self,
model_config: Dict,
eval_result: Dict
):
"""Record model evaluation"""
self.results.append({
"config": model_config,
"evaluation": eval_result
})
def analyze_trade_offs(self):
"""Analyze memorization-reasoning trade-off"""
memo_scores = []
reasoning_scores = []
model_sizes = []
for result in self.results:
config = result["config"]
eval_res = result["evaluation"]
# Aggregate scores
memo_score = np.mean(list(eval_res["memorization"].values()))
reasoning_score = np.mean(list(eval_res["reasoning"].values()))
model_size = config["num_params"]
memo_scores.append(memo_score)
reasoning_scores.append(reasoning_score)
model_sizes.append(model_size)
return {
"memorization": memo_scores,
"reasoning": reasoning_scores,
"model_sizes": model_sizes
}
Stage 2: Design MoE Configuration Space
Create different MoE architectures with varying sparsity.
# MoE sparsity configuration
from dataclasses import dataclass
@dataclass
class MoEConfig:
"""MoE model configuration"""
num_experts: int
expert_size: int # Params per expert
top_k: int # Top-k routing
num_layers: int
hidden_dim: int
class MoESparseFamily:
"""Generate family of MoE models with different sparsity"""
def __init__(self, target_compute_budget: float = 1.0):
self.target_compute = target_compute_budget
self.models = []
def generate_sparse_variants(self) -> List[MoEConfig]:
"""
Generate MoE variants with different sparsity levels.
Vary:
- Number of experts (E)
- Top-k routing
- Expert size
- Keep compute budget constant
"""
variants = []
# Base configuration
base_hidden_dim = 4096
base_num_layers = 32
# Sparse configurations
for num_experts in [8, 16, 32, 64, 128]:
for top_k in [1, 2, 4, 8]:
# Keep total compute roughly constant
# Compute = (hidden_dim^2 * layers * top_k) / num_experts
# Adjust hidden_dim to maintain constant compute
adjusted_dim = int(
base_hidden_dim * np.sqrt(num_experts / (top_k + 1e-6))
)
expert_size = (adjusted_dim * adjusted_dim) // num_experts
config = MoEConfig(
num_experts=num_experts,
expert_size=expert_size,
top_k=top_k,
num_layers=base_num_layers,
hidden_dim=adjusted_dim
)
variants.append(config)
self.models.append(config)
return variants
def compute_statistics(self) -> Dict:
"""Compute statistics for each model"""
stats = []
for config in self.models:
# Total parameters
total_params = (
config.num_experts * config.expert_size + # Expert params
config.num_layers * config.hidden_dim ** 2 # Attention/norm
)
# Active parameters (only top-k experts active)
active_params = (
config.top_k * config.expert_size +
config.num_layers * config.hidden_dim ** 2
)
# Active FLOPs
active_flops = (
config.num_layers * config.hidden_dim * config.top_k * config.expert_size
)
stats.append({
"config": config,
"total_params": total_params,
"active_params": active_params,
"active_flops": active_flops,
"sparsity": 1 - (active_params / total_params) if total_params > 0 else 0
})
return stats
Stage 3: Extract Core Principles
Derive the two scaling principles for MoE optimization.
# Core MoE scaling principles
class MoEPrinciples:
"""Extract active FLOPs and TPP principles"""
@staticmethod
def analyze_active_flops_principle(results: List[Dict]) -> Dict:
"""
Active FLOPs Principle:
'Models with identical training loss but greater active compute
achieve higher reasoning accuracy'
This means: sparse (higher k) > dense (lower k) for reasoning.
"""
# Group by pre-training loss
loss_groups = {}
for result in results:
loss = result["pretrain_loss"]
if loss not in loss_groups:
loss_groups[loss] = []
loss_groups[loss].append(result)
# Within each loss group, compare reasoning accuracy vs active FLOPs
analysis = {}
for loss, models in loss_groups.items():
# Sort by active FLOPs
models.sort(key=lambda x: x["active_flops"])
reasoning_by_flops = [m["reasoning_accuracy"] for m in models]
active_flops = [m["active_flops"] for m in models]
# Check correlation: more FLOPs -> better reasoning
correlation = np.corrcoef(active_flops, reasoning_by_flops)[0, 1]
analysis[loss] = {
"correlation": correlation,
"reasoning_scores": reasoning_by_flops,
"active_flops": active_flops
}
return analysis
@staticmethod
def analyze_tpp_principle(results: List[Dict]) -> Dict:
"""
TPP Principle:
'Memorization improves with increased parameters,
but there exists an optimal TPP for reasoning'
TPP = total_tokens_seen / total_parameters
"""
# Analyze memorization vs total params
total_params_list = [r["total_params"] for r in results]
memo_accuracy_list = [r["memorization_accuracy"] for r in results]
memo_param_corr = np.corrcoef(total_params_list, memo_accuracy_list)[0, 1]
# Analyze reasoning vs TPP
tpp_list = [r["training_tokens"] / r["total_params"] for r in results]
reasoning_list = [r["reasoning_accuracy"] for r in results]
reasoning_tpp_corr = np.corrcoef(tpp_list, reasoning_list)[0, 1]
return {
"memorization_param_correlation": memo_param_corr,
"reasoning_tpp_correlation": reasoning_tpp_corr,
"interpretation": {
"memo": "Higher params -> better memorization",
"reasoning": "Optimal TPP exists; too high or low hurts reasoning"
}
}
Stage 4: Joint Optimization Framework
Optimize both active FLOPs and TPP together.
# Joint MoE optimization
class MoEOptimizer:
"""Find optimal sparsity via joint optimization"""
def __init__(self):
self.pareto_frontier = []
def compute_objective(
self,
reasoning_accuracy: float,
memorization_accuracy: float,
weight_reasoning: float = 0.7
) -> float:
"""
Combined objective balancing reasoning and memorization.
Objective = w * reasoning_acc + (1-w) * memo_acc
"""
return (weight_reasoning * reasoning_accuracy +
(1 - weight_reasoning) * memorization_accuracy)
def find_optimal_config(
self,
results: List[Dict],
compute_budget: float = 1.0
) -> Dict:
"""
Find optimal MoE configuration.
Constraints:
- Maintain compute budget
- Optimize both reasoning and memorization
Returns:
- Recommended config
- Trade-off curve
"""
# Filter by compute budget
valid_results = [
r for r in results
if r.get("active_flops", 1.0) <= compute_budget
]
if not valid_results:
return {}
# Evaluate each configuration
objectives = []
for result in valid_results:
obj = self.compute_objective(
result["reasoning_accuracy"],
result["memorization_accuracy"]
)
objectives.append(obj)
result["objective_score"] = obj
# Find Pareto frontier
self.pareto_frontier = self._compute_pareto_frontier(valid_results)
# Select best overall
best_result = max(valid_results, key=lambda r: r["objective_score"])
return {
"optimal_config": best_result["config"],
"reasoning_accuracy": best_result["reasoning_accuracy"],
"memorization_accuracy": best_result["memorization_accuracy"],
"total_params": best_result["total_params"],
"active_flops": best_result["active_flops"],
"pareto_frontier": self.pareto_frontier
}
def _compute_pareto_frontier(self, results: List[Dict]) -> List[Dict]:
"""Compute Pareto frontier (reasoning vs memorization)"""
frontier = []
for result in results:
# Check if dominated by any other result
dominated = False
for other in results:
if (other["reasoning_accuracy"] >= result["reasoning_accuracy"] and
other["memorization_accuracy"] >= result["memorization_accuracy"] and
(other["reasoning_accuracy"] > result["reasoning_accuracy"] or
other["memorization_accuracy"] > result["memorization_accuracy"])):
dominated = True
break
if not dominated:
frontier.append(result)
return frontier
Stage 5: Evaluate and Validate
Test optimal configurations on benchmarks.
# Comprehensive evaluation
class MoEEvaluator:
"""Full evaluation of MoE sparsity findings"""
def __init__(self):
self.optimizer = MoEOptimizer()
self.evaluator = CapabilityEvaluator()
def run_full_study(
self,
compute_budget: float = 1.0
) -> Dict:
"""
Run complete MoE sparsity study.
1. Generate MoE variants
2. Train each variant
3. Evaluate on memorization and reasoning
4. Optimize for best configuration
"""
# Step 1: Generate variants
family = MoESparseFamily(compute_budget)
configs = family.generate_sparse_variants()
print(f"Generated {len(configs)} MoE configurations")
# Step 2: Train variants (placeholder - in reality this is expensive)
results = []
for config in configs:
# Train model
model = self.train_model(config)
# Evaluate
eval_result = self.evaluator.evaluate_model(
model,
test_datasets=self.get_test_datasets()
)
# Compute statistics
stats = family.compute_statistics()
config_stats = next(
(s for s in stats if s["config"] == config),
{}
)
result = {
"config": config,
"reasoning_accuracy": np.mean(list(eval_result["reasoning"].values())),
"memorization_accuracy": np.mean(list(eval_result["memorization"].values())),
"pretrain_loss": eval_result["pre_training_loss"],
**config_stats
}
results.append(result)
# Step 3: Analyze trade-offs
principles = MoEPrinciples()
flops_analysis = principles.analyze_active_flops_principle(results)
tpp_analysis = principles.analyze_tpp_principle(results)
# Step 4: Optimize
optimal = self.optimizer.find_optimal_config(results, compute_budget)
return {
"total_configs_evaluated": len(results),
"active_flops_principle": flops_analysis,
"tpp_principle": tpp_analysis,
"optimal_configuration": optimal,
"pareto_frontier": self.optimizer.pareto_frontier
}
def train_model(self, config: MoEConfig):
"""Train single MoE model"""
return None # Placeholder
def get_test_datasets(self) -> Dict:
"""Get evaluation datasets"""
return {}
def generate_report(self, study_results: Dict) -> str:
"""Generate findings report"""
report = []
report.append("## MoE Sparsity Study Results\n")
optimal = study_results["optimal_configuration"]
report.append(f"**Optimal Configuration:**")
report.append(f"- Experts: {optimal['optimal_config'].num_experts}")
report.append(f"- Top-K: {optimal['optimal_config'].top_k}")
report.append(f"- Reasoning Accuracy: {optimal['reasoning_accuracy']:.3f}")
report.append(f"- Memorization Accuracy: {optimal['memorization_accuracy']:.3f}\n")
report.append("**Key Findings:**")
report.append("1. Active FLOPs improve reasoning regardless of model size")
report.append("2. Optimal TPP varies by reasoning task complexity")
report.append("3. RL/test-time compute don't change fundamental trade-offs\n")
return "\n".join(report)
Practical Guidance
Recommended MoE Configurations
For Reasoning-Heavy Workloads:
- High sparsity: 32+ experts, top-k 2-4
- Benefit: More active FLOPs per token
- Trade-off: Lower memorization capacity
For Balanced Workloads:
- Medium sparsity: 16 experts, top-k 4-8
- Benefit: Reasonable reasoning + memorization
- Trade-off: Middle-ground performance
For Knowledge-Heavy Tasks:
- Low sparsity: 8 experts, top-k 8+
- Benefit: Better memorization
- Trade-off: Lower reasoning capability
Joint Optimization Strategy
- Define weight between reasoning and memorization (0.7 reasoning is typical)
- Evaluate candidate configurations on both task types
- Identify Pareto frontier
- Select configuration maximizing weighted objective
When to Use This Framework
- Training custom LLMs for specific use cases
- Allocating compute budget between capabilities
- Understanding model trade-offs
- Planning large-scale training runs
When NOT to Use
- Fine-tuning existing models (pre-training decisions already made)
- Deployment-stage optimization (affects training)
- Scenarios without clear reasoning vs memorization distinction
Key Insights
The paper revises compute-optimal scaling laws by showing that more parameters (memorization) and more active FLOPs (reasoning) are orthogonal optimization axes. Classical dense models conflate these: adding parameters increases both FLOPs and capacity. MoE decouples them, enabling expert optimization for reasoning while controlling memorization budget.
Reference
Optimal Sparsity of MoE Language Models for Reasoning. arXiv:2508.18672