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

Abstract: Empirical scaling laws have driven the evolution of LLMs, yet their coefficients shift whenever the model architecture or data pipeline changes. Mixture-of-Experts (MoE) models, now standard in state-of-the-art systems, introduce a new sparsity dimension that current dense-model frontiers overlook. We investigate how MoE sparsity influences two distinct capability regimes: memorization skills and reasoning skills. By training MoE families that vary total parameters, active parameters, and top-$k$ routing under fixed compute budgets, we disentangle pre-training loss from downstream accuracy. Our results reveal two principles. First, Active FLOPs: models with identical training loss but greater active compute achieve higher reasoning accuracy. Second, Total tokens per parameter (TPP): memorization tasks improve with more parameters, while reasoning tasks benefit from optimal TPP, indicating that reasoning is data-hungry. Neither reinforcement learning post-training (GRPO) nor increased test-time compute alters these trends. We therefore argue that optimal MoE sparsity must be determined jointly by active FLOPs and TPP, revising the classical picture of compute-optimal scaling. Our model checkpoints, code and logs are open-source at https://github.com/rioyokotalab/optimal-sparsity.

• The paper reveals that increasing model parameters reduces training loss but can worsen reasoning task loss beyond a model size threshold.
• It demonstrates that denser configurations optimize reasoning tasks while memorization tasks benefit from greater sparsity.
• It identifies that maintaining a tokens-per-parameter ratio near 20 is crucial for balancing active FLOPs to support reasoning skills.

Optimal Sparsity of Mixture-of-Experts LLMs for Reasoning Tasks

Introduction

The paper investigates the effects of sparsity in Mixture-of-Experts (MoE) models on LLMs, particularly focusing on two capability regimes: memorization skills and reasoning skills. The study aims to establish a deeper understanding of how active FLOPs and total tokens per parameter (TPP) interplay to optimize MoE LLMs for reasoning tasks.

Influence of Sparsity on Training and Task Performance

Training Loss and Validation Loss:

The study demonstrates that increasing the total number of parameters in an MoE model generally decreases the training and validation loss (Figure 1). However, this trend does not straightforwardly translate into improved downstream task loss for all tasks.

Figure 1: Although training and validation loss decrease as the total number of parameters grows, the task loss on GSM8K can sometimes worsen with larger models.

Task Loss:

Scale-up experiments reveal a dichotomy in task loss behavior across different types of tasks. For datasets like TriviaQA and HellaSwag, larger models consistently reduce task loss as pre-training loss decreases. In contrast, for mathematical reasoning datasets (GSM8K, GSM-Plus), task loss initially declines with decreasing training loss but starts to increase beyond a certain threshold of model size (Figure 2). This U-shaped curve indicates that reasoning task loss does not always benefit from more extensive parameter scales.

Figure 2: For GSM8K and GSM-Plus, once the training loss drops below a certain point, the task loss starts to increase.

Dichotomy of Reasoning and Memorization Tasks

The paper identifies that reasoning skills are notably affected by the level of model sparsity, unlike memorization tasks, which benefit from increased sparsity. Human reasoning tasks such as those in GSM8K show degradation in task performance when model sparsity is increased beyond an optimal point (Figure 3).

Figure 3: Downstream accuracy when scaling total parameters via expert count with width and top-k fixed.

For memorization tasks, increased sparsity consistently enhances performance across varied active parameter budgets, unlike reasoning tasks, where denser models are often more effective beyond certain compute budgets (Figure 4).

Figure 4: At fixed active parameter counts, higher sparsity (lower density) consistently improves performance, but at larger active parameter counts, GSM8K and GSM-Plus shift their optima back toward dense models.

Test-Time Compute and Reinforcement Learning Adjustments

The paper evaluates whether reinforcement learning techniques, such as GRPO, and test-time compute methods can bridge the performance gap observed due to increased sparsity in reasoning tasks. These methods yielded improvements; however, they did not alleviate the inverted U-shaped trade-off characterizing reasoning tasks (Figure 5).

Figure 5: Effect of Test-Time Compute and GRPO on the loss–accuracy trade-off.

Tokens Per Parameter (TPP) Implications

The analysis highlights the TPP's pivotal role in determining task performance. Memorization tasks benefit from a lower TPP, while reasoning tasks demonstrate optimal performance near a TPP of 20 (Figure 6). This finding underscores the importance of balancing parameter count with training tokens to support reasoning capabilities effectively.

Figure 6: Effect of TPP on performance across different tasks.

Conclusion

The study presents an insightful examination of sparsity's dual impact on memorization versus reasoning tasks. Results indicate that maximizing reasoning task performance in MoE models necessitates a careful balance between active FLOPs and TPP. This balance shifts towards dense configurations in high-compute regimes for reasoning tasks, diverging from the scaling laws applicable to purely dense models. The findings prompt a re-evaluation of scaling practices in MoE architectures to optimize reasoning tasks, suggesting a nuanced approach that considers both data availability and computational resources.