TLDR: This research paper investigates the optimal sparsity of Mixture-of-Experts (MoE) language models for different tasks. It finds that while increasing total parameters and sparsity improves memorization, reasoning performance can saturate or even decline. Neither post-training reinforcement learning nor extra test-time compute fully resolves this. The study concludes that optimal sparsity is task-dependent and crucial for reasoning models, with denser MoE layers sometimes outperforming sparser ones for reasoning at higher compute budgets.
Large language models (LLMs) have seen incredible advancements, often guided by “scaling laws” that suggest bigger models and more data lead to better performance. However, a new study delves into a crucial aspect of modern LLMs: Mixture-of-Experts (MoE) models, which introduce a new dimension of “sparsity.” This research explores how this sparsity affects different capabilities of LLMs, specifically memorization and reasoning tasks.
MoE models are a key architectural innovation in today’s most advanced LLMs, allowing them to achieve high capacity while keeping computational costs manageable. Instead of using all parts of the model for every piece of information, MoE models route each input to a small, specialized subset of “experts.” The degree of this specialization is what the researchers refer to as sparsity.
The study, titled “Optimal Sparsity of Mixture-of-Experts Language Models for Reasoning Tasks,” investigated how varying the total number of parameters, the number of active parameters (those actually used for a given input), and the “top-k” routing (how many experts a token is routed to) impacts model performance. They kept the computational budget constant across their experiments, meticulously recording pre-training loss, downstream task loss, and task accuracy.
A significant finding emerged when comparing memorization tasks (like TriviaQA and HellaSwag) with reasoning tasks (like mathematical problem-solving benchmarks GSM8K and GSM-Plus). For memorization, the classic scaling law held true: performance consistently improved as the total number of parameters increased, mirroring the reduction in training loss. Essentially, more parameters meant better memorization.
However, the picture was strikingly different for reasoning tasks. Here, performance didn’t just stop improving; it could actually get worse even as the total number of parameters continued to grow and training loss decreased. This suggests that simply making MoE models larger or sparser isn’t always beneficial for complex reasoning. The researchers also found that changing only the ‘top-k’ value (the number of experts activated per token) had little impact if the total active parameters remained the same.
The study further explored whether common post-training techniques, such as reinforcement learning (specifically GRPO) or providing extra computational resources during testing (Test-Time Compute), could overcome this reasoning deficit in overly sparse models. The answer was no; while these methods generally improved overall performance, they did not eliminate the fundamental trade-off observed between sparsity and reasoning ability. This highlights the critical importance of finding the right sparsity level during the initial pre-training phase for models intended for reasoning tasks.
Another interesting aspect was the “Tokens per Parameter” (TPP) ratio. For memorization tasks, performance improved as TPP decreased (meaning more parameters relative to tokens), aligning with the idea that these tasks are “parameter-hungry.” For reasoning tasks, however, there appeared to be an optimal TPP ratio (around 20), beyond which too many parameters relative to the training data actually hurt reasoning performance. Moreover, models with more active experts (larger top-k values) consistently performed better on reasoning tasks, even when the overall TPP was fixed.
The findings extend to code generation tasks like HumanEval and MBPP, which also showed similar trends to mathematical reasoning. Denser MoE layers were found to be more effective for procedural reasoning in code synthesis when a larger compute budget was available, while sparser layers were better in low-compute scenarios.
Also Read:
- Optimizing Large Language Models on Edge Devices with Smart Expert Management
- GPT-OSS-20B: Unpacking the Deployment Efficiency of OpenAI’s Open-Weight MoE Model
In conclusion, this research challenges the universal applicability of “more experts is better” for MoE models. While increasing sparsity and total parameters can boost memorization, it can hinder reasoning capabilities. The study emphasizes that for reasoning tasks, achieving optimal performance requires a careful balance of total and active parameters, and this optimal balance can shift depending on the available computational budget. This work provides crucial insights for designing and scaling future language models, especially those aimed at complex reasoning. You can find the full paper here: Optimal Sparsity of Mixture-of-Experts Language Models for Reasoning Tasks.
