moe-sparsity-reasoning

name	moe-sparsity-reasoning
title	Optimal Sparsity of MoE Language Models for Reasoning
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2508.18672
keywords	["mixture-of-experts","sparsity","reasoning","memorization","compute-optimal"]
description	Determine optimal MoE sparsity by separating memorization and reasoning trade-offs: active FLOPs improve reasoning while total parameters improve memorization, requiring joint optimization

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

        
         benchmark  .benchmarks[]:
            dataset = test_datasets.get(benchmark, [])
            accuracy = .evaluate_benchmark(model, dataset)
            results[][benchmark] = accuracy

        
        results[] = .get_pretrain_loss(model)

         results

     () -> :
        
        correct = 
         example  dataset:
            prediction = model.generate(example[])
             prediction == example[]:
                correct += 
         correct / (dataset)  dataset  

     () -> :
        
           


 :
    

     ():
        .results = []

     ():
        
        .results.append({
            : model_config,
            : eval_result
        })

     ():
        
        memo_scores = []
        reasoning_scores = []
        model_sizes = []

         result  .results:
            config = result[]
            eval_res = result[]

            
            memo_score = np.mean((eval_res[].values()))
            reasoning_score = np.mean((eval_res[].values()))
            model_size = config[]

            memo_scores.append(memo_score)
            reasoning_scores.append(reasoning_score)
            model_sizes.append(model_size)

         {
            : memo_scores,
            : reasoning_scores,
            : 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, ]:
                
                

                
                adjusted_dim = (
                    base_hidden_dim * np.sqrt(num_experts / (top_k + ))
                )

                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)
                .models.append(config)

         variants

     () -> :
        
        stats = []

         config  .models:
            
            total_params = (
                config.num_experts * config.expert_size +  
                config.num_layers * config.hidden_dim **   
            )

            
            active_params = (
                config.top_k * config.expert_size +
                config.num_layers * config.hidden_dim ** 
            )

            
            active_flops = (
                config.num_layers * config.hidden_dim * config.top_k * config.expert_size
            )

            stats.append({
                : config,
                : total_params,
                : active_params,
                : active_flops,
                :  - (active_params / total_params)  total_params >   
            })

         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
            }

         analysis


     () -> :
        
        
        total_params_list = [r[]  r  results]
        memo_accuracy_list = [r[]  r  results]

        memo_param_corr = np.corrcoef(total_params_list, memo_accuracy_list)[, ]

        
        tpp_list = [r[] / r[]  r  results]
        reasoning_list = [r[]  r  results]

        reasoning_tpp_corr = np.corrcoef(tpp_list, reasoning_list)[, ]

         {
            : memo_param_corr,
            : reasoning_tpp_corr,
            : {
                : ,
                : 
            }
        }

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 = .compute_objective(
                result[],
                result[]
            )
            objectives.append(obj)
            result[] = obj

        
        .pareto_frontier = ._compute_pareto_frontier(valid_results)

        
        best_result = (valid_results, key= r: r[])

         {
            : best_result[],
            : best_result[],
            : best_result[],
            : best_result[],
            : best_result[],
            : .pareto_frontier
        }

     () -> []:
        
        frontier = []

         result  results:
            
            dominated = 
             other  results:
                 (other[] >= result[] 
                    other[] >= result[] 
                    (other[] > result[] 
                     other[] > result[])):
                    dominated = 
                    

              dominated:
                frontier.append(result)

         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,
                : np.mean((eval_result[].values())),
                : np.mean((eval_result[].values())),
                : eval_result[],
                **config_stats
            }

            results.append(result)

        
        principles = MoEPrinciples()
        flops_analysis = principles.analyze_active_flops_principle(results)
        tpp_analysis = principles.analyze_tpp_principle(results)

        
        optimal = .optimizer.find_optimal_config(results, compute_budget)

         {
            : (results),
            : flops_analysis,
            : tpp_analysis,
            : optimal,
            : .optimizer.pareto_frontier
        }

     ():
        
           

     () -> :
        
         {}

     () -> :
        
        report = []
        report.append()

        optimal = study_results[]
        report.append()
        report.append()
        report.append()
        report.append()
        report.append()

        report.append()
        report.append()
        report.append()
        report.append()

         .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

1. Define weight between reasoning and memorization (0.7 reasoning is typical)
2. Evaluate candidate configurations on both task types
3. Identify Pareto frontier
4. 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

• https://arxiv.org/abs/2508.18672