Optimal Sparsity of Mixture-of-Experts Language Models for Reasoning Tasks
Authors: Taishi Nakamura
Satoki Ishikawa
Masaki Kawamura
Takumi Okamoto
Daisuke Nohara
Jun Suzuki
Rio Yokota
Venue: ICLR 2026 Oral (2026)
Keywords: Mixture of Experts
memorization
reasoning
scaling laws
large language models
Primary area: foundation or frontier models, including LLMs
TL;DR: Memorization skills consistently benefit from higher sparsity, while reasoning skills require balancing active FLOPs with total tokens per parameter; the optimal point shifts with the compute budget.

Abstract: Empirical scaling laws have driven the evolution of large language models (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. 
All code, data sources, and logs are released to facilitate reproducibility and future work.

Published as a conference paper at ICLR 2026
OPTIMAL SPARSITY OF MIXTURE-OF-EXPERTS LAN- GUAGE MODELS FOR REASONING TASKS
Taishi Nakamura 1,2, Satoki Ishikawa', Masaki Kawamura', Takumi Okamoto 1.2
Daisuke Nohara' Jun Suzuki Rio Yokota! 1.2 Institute of Science Tokyo, NNII LLMC, TTohoku University, 4RIKEN
(taishi, riwyokota sere. iir, isc.i..ttaac........o
ABSTRACT
Empirical scaling laws have driven the evolution of large language models (LLMs), yet their coefficients shift whenever the model architecture data pipeline changes. offture-of-Experts (MoE) models, now standard in state-of- the-art systems, introduce new sparsity dimension that current cense-model fron- tiers 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, 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 com- pute 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 httttps://githhi http:tttw..co//ooooooooaaaaaaaaa - sparsity
INTRODUCTION
The recent evolution of large language models (LLMs)) has been driven by empirical scaling laws (Hestness et al. (2017) that link training loss model size, dataset size, and compute budget. Kaplan ctt al. (2020) showed that these laws hold across seven orders of magnitude, establishing them as a reliable extrapolation tool for dense Transformers. Subsequent work by [Hofmmann al. (2022) demonstrated that scaling curves can be inverted to choose the compute-optimal combination of parameters and tokens for a fixed budget. Together, these results have made scaling analysis cornerstone of model planning at both academic and industrial labs.
Yet the coefficients of the scaling laws different optimizers or architectures often follow the same loss trajectory but diverge substantially on downstream reasoning benchmarks (Liu et al.) 2023a). Brandfonbrener al. (0255) extend the classic laws with loss-to-loss prediction, showing that the mapping between training and test dis- tributions admits its own power law when the distributions differ substantially. These observations imply that optimal budgets must be re- re-estimated whenever we modify the model or the data pipeline.
universal. Highly expressive models trained under
particularly compelling architectural modification is the Mixture-of offExxerts (MoE) paradigm, offering high capacity at fixed FL.OPs by routing each token through sparse subset of ex- perts (Shazeer et al. [017|Leppikhin et al. 202T Fedus et al. 202T). Modern flagship models, e.g..
Gemini 2.5 Pro (Gemini Team 2025). DeepSeck-V3 (DeepSeek-Al) 2025b), and Qwen3 (Qwen Team 2025) now rely on MoE as facto standard for economical scaling. Abnar et al. (2025) derive a parameters-vs-f -LLOPP frontier and locate optimal sparsity for given compute budget. These findings emphasize that the classical dense-model frontier is an incomplete picture, and one must account for architectural knobs such MoE sparsity and top-k routing.
Evaluating reasoning performance immediately after pre-training overlooks both the benefits of post-training adaptation and the leverage of additional test-time compute. Post-training methods

---

Published as a conference paper at ICLR 2026
such GRPO, which use reinforcement signals to encourage coherent chain- chain-of-thought generation, sharpen a model's reasoning on complex tasks [OpenAl 2024b; DeepSeek-Al] 2025a). Beyond these refinements, models can further improve outputs test time by adopting calibrated decod- ing strategies that mirror how humans pause to reconsider difficult problems. These test-time ap- proaches not only boost routine benchmark performance but, when properly tuned, substantially enhance multi-step mathematical reasoning, demonstrating that adaptive computing at test time is a powerful complement to both model scale and post-training adaptation.
In this paper, we aim to identify how the optimal sparsity of MoE changes between memorization skills (TriviaQA, HellaSwag) and reasoning skills (GSM8K, GSM-Plus) tasks. In this work, we use the term dense models to refer to standard Transformers with single feed-forward network per layer. For MoE models, we define sparsity as sparsity = k following the convention that sparsity measures the fraction of inactive parameters. We train families of MoEs varying not only the total vs. active parameters, but also the number of top-k experts. For each model, we measure the loss on the pre-training data, the task loss the downstream benchmarks, and the accuracy on those benchmarks. This allows us to disentangle the generalization gap between the train vs. test loss, and the gap between loss vs. accuracy. For both memorization and reasoning benchmarks, the train loss decreases monotonically with the total parameters count increases. The task loss and accuracy follow the same monotonic trend as the train loss for memorization benchmarks. In contrast, for reasoning benchmarks, the task loss and accuracy diverge from this monotonic trend as the total parameters increase and training loss decreases. We find that changing the is in top-k routing has negligible effect if the number of active parameters is kept constant. We also consider classic generalization-ga controls by sweeping the learning rate and initialization, showing that their effects align strikingly with the generalization-ga caused by sparsity. This confirms that the gap between the performance on memorization skills and reasoning skills can be induced not only by the sparsity of the MoE, but also by classical hyperparameters like learning rate and initialization. We further investigate whether applying GRPO or additional test-time compute could recover the degraded reasoning ability of sparser models. Our results show that the gap between memorization and reasoning performance caused by increased sparsity remains unchanged even after GRPO and increased test-time compute. This means that finding the optimal sparsity of the MoE during pre- training is crucial for training a reasoning model under a fixed compute budget.
Experts
We characterize these divergences along two key axes: (1) Active FLOPS downstream reasoning performance is not determined by training loss alone. but by the number of active FLOPs at both train and test time; even identical training loss, models with larger top-k consistently outperform smaller ones. (2) Total tokens per parameter (TPP) reasoning ability peaks around 20 tokens per parameter, whereas memorization skills are rraamttrrhungry and improve with lower TPP. Together, these axes define the compute-optimal sparsity for MoE models.
We further demonstrate that neither reinforcement learning post-training (GRPO) nor additional test- time compute eliminates this trade-off, highlighting that pre-training sparsity remains the dominant factor for reasoning ability under fixed budgets.
memorization improves with higher sparsity and more experts whereas reasoning requires careful balance between active FLOPs and mata-per-parameter occasionally favoring denser configurations in high-compute regimes.
Together, our findings refine the scaling laws for MoE LLMs:
We release model checkpoints, code and logs are open-source at hittps:/ggite com/
Fioyokotalab/ hotpialllttmlllllmm aattimmmsspasstt
2
BACKGROUND AND RELATED WORK
2.1 MIXTURE OF EXPERTS
MoE Architecture. Mixture-of-Expert (MoE) networks were introduced by (Jacobs et al. 1991
Jordan & Jacobs 1994) and later brought to large-scale neural language modeling by /Sazzerr et al. (017)). Within the Transformer architecture Vaswani et al., (017). MoE layers have proven espe- cially effective, scaling to hundreds of billions of parameters while maintaining manageable train- ing costs (Lepikhin et al., 202T, Fedus et al. 202T; Du et al. (022). In an MoE layer, a learn- able router assigns each token to sparse subset of experts. Let IRda be token representa-
2

---

Published as a conference paper at ICLR 2026
tion and (FFN(x); the n feed-forward experts. For top-k routing, the router produces scores 8 = x W nouter E IR", and selects the indices K. of the & largest components, then normalizes them: exp(s=) if i K and g(x); 0 otherwise. g(x); - E,ex exp(s;)
The layer output is the weighted sum of the chosen experts: y = --== -i=1(x)): FFN(x);, Modern MoE models typically supplement the token-level cross-entropy loss with two auxiliary terms: a load-balancing loss LLB. which prevents expert collapse (Shazeer et al., 2017). and router-z loss CRZ. which penalizes large router logits for better numerical stability and gradient flow (Zoph et al. 2022 The combined training loss is expressed as LCE + aLLB + fLız. where a and 3 are hyperparameters that control the relative importance of each term in the objective function. This formulation is widely used in recent MoE-based language models and remains unchanged throughout the experiments.
2.2 SCALING LAWS OF LLMs
Existing scaling laws demonstrate power-law relationships between model performance, parameter count, dataset size, and compute budget (Kaplan et al.] 2020 Hoff- mann al. 2022). Scaling laws for MoE models have similarly explored how total parameter count and expert granularity jointly affect scaling behavior (Clark et al., 2022) Ludziejewski et al.) 2024).
Scaling Laws for MoE.
Building on this, Frantar al. (2024) derived sparsity-aware scaling exponents that bridge dense and sparse regimes, while Abnar et al. (2025) empirically charted the optimal trade-offs between total parameters and FLOPs per token in MoE settings. Furthermore, recent analyses indicate that in- creasing sparsity itself can directly improve loss, highlighting sparsity as a key dimension in scaling behavior (Kimi Team 2025).
Task Loss. Since the scaling law for next-token prediction loss does not necessarily align with downstream task loss, it may not be reliable for predicting benchmark performance (Grattafioni et al. 2024 Some work has tried to model downstream accuracy with exponential curve, but accuracy is only predictable when we average over many tasks and carefully choose which ones to include (Gadre et al., 2024). Another line of research instead first quantifies how downstream task loss scales with parameters and data, then converts predicted losses into accuracy estimates, achiev- ing under two points of absolute error for mid-scale models using minimal extra compute (Bhagia et al., (2024). Prior work observes that downstream task loss relates to pre-training loss, where the shifts depend on the minimal achievable losses determined by the intrinsic complexity and distribu- tional mismatch between the pre-training and downstream datasets (Brandfonbrener et al., 2025).
Skills Adding MoE experts tends to improve memorization skills more than reasoning skills, mo- tivating new, generalized scaling frameworks that address scaling laws for reasoning performance (Jelassi et al. 2025). They provide a theoretical explanation for this asymmetry using graph neu- ral networks. Recent work shows that larger top-k can improve compositional generalization in MoE models (Zhao et al., 2025a). though such studies do not observe the non-monotonic effect of sparsity on reasoning performance that identify. Since scaling laws differ across tasks, the opti- mal scaling strategy may also vary; for example, knowledge-basee QA tasks are "eapacity-hungry benefiting more from larger model sizes, whereas code-related tasks are "data-hungry." benefiting more from increased training data (Roberts et al. 2025). Complementing these observations, recent analyses find that models with identical training loss can still exhibit markedly different reasoning performance (GLM-4.5 Team 2025), highlighting that reasoning skill depends not only on loss but also on how compute and data are allocated.
2.3 POST TRAINING AND TEST-TIME COMPUTE (TTC)
Reinforcement Learning (RL) post-training has long been a predominant approach for improving LLMs. Proximal Policy Optimization (PPO) (Schulman et al., 2017) forms the backbone of RLHF pipelines, from the original GPT alignment work (Ouyang et al., (022)) to the GPT-4 family mod- els (OpenAl 2024a). More recently, Group Relative Policy Optimization (GRPO) was introduced as group-relative
a variant of PPO that replaces the value function baseline with a group relative advantage estimator,
thereby improving memory efficiency and stabilizing updates; this approach already powers frontier-
3

---

Published as a conference paper at ICLR 2026
scale systems such DeepSeek-R1, achieving satte-of-the-ar results mathematiical-reasoniig benchmarks (Shao et al. 2024. DeepSeck-At Al, 2025a).
Complementary to these training-time time advances, scaling test-time compute (TTC) offers orthogo- nal approach. TTC denotes accuracy gains obtained without updating model parameters, simply by
2024b Muen- nighoff' et al. 2025b), sampling larger candidate pools (Li et al. 2022 Wang et al. 2023 Brown
allocating more inference resources, e.g., running longer chains of thought (OpenAl
al. 2024| Schaeffer et al., 2025), ) or performing explicit search-and-verify steps (Lightman al,
2024
Shinn et al.) 2024, Snell et al., 2025 Inoue et al. (2025). Among these, self-consistency, re-
peated sampling with majority-vott aggregation, has emerged as a strong TTC baseline (Wang al.,
2023
3
EXPERIMENTS
In this section, empirically demonstrate the scaling of downstream task performance through systematic investigation memorization and reasoning skills in MoE LLMs.
3.1 EXPERIMENTAL. SETUP
We the Mixtral (Jiag et al. 2024 architecture, Transformer backbone with RMSNorm (Zang & Sennrich 2019). SwiGLU activations (Shazeer 2020), and rotary positional embeddings (Su 2024). Each feed-forward block sparsely gated MoE layer, gated by the dropless token- choice top-k routing (Gale et al.] 2023). All models use L = 16 layers, following following/Meenngghool et al. (2025a). We sweep three architectural hyperparameters: (i) the model width d E (512, 1024, 2048J: (ii) the number of experts per layer E E {8, 16, 32, 64, 128, 256); and (iii) the top-k experts per token E {2, 4, 8, 16). Each feed-forward network has a hidden dimension of 2d. When 512 and d 1024, we train every combination of E and k. For d 2048, we limit the search to E < 128 due to computational resource constraints.
We train with AdamW [Loshhilovv & Hutter, 2019) using a peak learning rate of 4 x 10 2k- step linear warm-up followed by cosine decay, and a weight decay of 0.1. Following Xue et al. (2024) and Zoph al. (2022), respectively.
the load-balancing and router 2-losses by 10-2 and 10-3
Hyperparameter Study. To isolate optimization effects, reuse the same 125B-token corpus. For all HP runs, we fix = 16, & = 2, and train two widths, dmodel E (512, 1024), with the same FFN expansion factor 2. We vary (i) LM-head initialization schemes, (ii) peak learning rate, and (iii) AdamW c. Further implementation and environmental details are deferred to Appendix A.3
Pre-training Datasets. We balanced 125B-token mixture consisting of high-quality web text (43B), mathematics corpora (32B), STEM literature and reference (49B), and code (1B). See Appendix A.I for complete statistics.
Evaluation Protocol. We evaluate three capability areas with standard few-shot prompts. Math- ematical Reasoning: GSM8K (Cobbe al. (021) (4-shot) and GSM-Plus (Li et al., 2024) (5-shot CoT). Reading Comprehension: TriviaQA (Jshi et al.] 2017) with 4-shot prompting. Commonsense Reasoning: HellaSwag (Zellers al. 2019), each under 4-shot prompting setup. See Table in Appendix for further details.
3.2 DOWNSTREAM PERFORMANCE DOES NoT NECESSARILY IMPROVE WITH TOTAL PARAMETER SIZE
In this section, we examine how the expert sparsity in MoE models affects the relationship between pre-training loss and downstream performance. We train a series of models with controlled sparsity levels and measure their performance that while increasing the total number of parameters reduces pre- pre-training loss, downstream task loss mathematical reasoning worsens beyond certain model size.
the representative downstream tasks. Our analysis shows

---

Published as a conference paper at ICLR 2026
Train Loss
Validation Loss
HellaSwag
GSMBK
2.3
4.6
2.1
1.4
2.2
4.4
Frall lo5s
Task loss
1.3
4.2
1.8
4.0
1.7
3.8
1.1
1.8
1.6
160
1.7 1.7"."100"
10'
3.6
11""""
"""'""
'100
'66
Total parameters (B)
Total parameters (B)
Total parameters (B)
Total parameters (B)
512k=2
512k=8
d=1024k==
d=1024,k=0
d=2048.k=2
d=2048
512,km4
512,k=16
d-1024km4
d=1024,k-16
d=2048,w44
2-48,
Although training and validation loss decrease the total number of parameters grows, the task loss on GSM8K can sometimes worsen with larger models. Training and valida- tion losses steadily decrease as total active parameters increase. The HellaSwag task loss follows this scaling trend, whereas GSM8K task loss worsens once total parameters exceed threshold. Within each fixed top-k group, moving right on the x-axis corresponds increasing sparsity (be- cause total experts F. increases while k remains fixed), so the right-hand task-loss panels implicitly reflect the same sparsity ordering shown explicitly in Figure
Figure I:
TriviaQA
HellaSwag
GSM8K
GSM-Plus
4.6
2.0
1.4
10
Task loss
140 R.
begii to increase
1.9
1.8
4.0
3.8
1.1
1.6
3.6
1.0
1.6
1.8
2.0
1.6
1.8
1.6
2.0
Final training loss
Final training loss
Final training loss
Final training loss
ds512. km2, A-17OM
de512, km16, Am520M
dmm 1024, kmm,
dm2048, km4,, Am4.A
k=16,
d-512, k- A-320M
1-2424, k-4, A-670M
d-2048, k=2, A-1.56
d-2048, A-1.18
Figure 2: For GSM8K and GSM-Plus, once the training loss drops below certain point, the task loss starts to increase. Results of scaling total parameters by increasing the number of experts, with model width and top-k held constant. For TriviaQA and HellaSwag, the task loss falls monotonically as training loss decreases. By contrast, GSM8K and GSM-Plus show a U-shaped trend: task loss declines with training loss only until threshold, beyond which further reductions in training loss hurt task performance. That threshold moves lower as active parameter count increases, models with more active parameters achieve lower optimal task loss. No such active parameters dependence appears for TriviaQA, HellaSwag.
Task Loss Computation. Following Branndoonbeeeeereeeererrr et al. (2025) and Grattafiori et al. (2024), we compute cross-entropy only over the answer tokens by concatenating prompt with the ground- truth answer. For multiple-choice datasets (c.g.. HellaSwag, TriviaQA) the target sequence is the correct answer string, as [Bhagia et al. (2024). For open-ended mathematics datasets such as GSM8K, and GSM-Plus we likewise compute cross-entropy directly against the ground-truth answer tokens.
Training Loss and Validation Loss. Figure Joureepreee presents the training and validation losses when fix- ing the top-k, MoE layer width constant and increasing only the number of experts (and hence the total parameter count). As the total parameter count grows, both training and validation losses de- crease. Therefore, terms of pre-training loss, increasing total parameters (thereby raising sparsity) reduces pre-training loss, which is consistent with prior work.
Experiments with Task Loss Next, we examine how the downstream task loss responds to in- creases in the total parameter count. Figure 2 shows task loss on several benchmarks as we vary only the number of experts, holding both top-k and each MoE layer widths constant. On TriviaQA and HellaSwag, lower pre-training loss reduces task loss, indicating that larger total parameter mod- els yield better results these datasets. In contrast, for GSM8K and GSM-Plus, further reductions
5

---

Published as a conference paper at ICLR 2026
TriviaQA
HellaSwag
GSM8K
0.50
512.k-4
0.4
512.k=8
Accuracy
0.45
0.3
d=1024,k=
d 512,k=16
d=1024,k=4
0.40
d-1024.k-8
0.2
d=1024.k=16
d=2048,k=2
0.1
k=
--
0.30
k-24
2.0
Final training loss
Final training loss
Final training lass
Figure 3: Downstream accuracy when scaling total parameters via expert count with width and top-k fixed. TriviaQA and HellaSwag exhibit steadily improving accuracy as pre-training loss decreases, whereas GSM8K shows a non non-monotonic trend: further reductions in pre-training loss do not always improve accuracy and can even degrade performance.
Sparsity (1 TopK/ Experts)
Sparsity 0.984 Sparsity 0.969
-Sparsity 0.938
Sparsity 0.500 Sparsity 0.000
TriviaQA
HellaSwag
GSMBK
GSM-Plus
0.00
0.2
0.56
387
0.00
3.03
0.76
0.00
13
21
25
19
42
20 21 22 43 Training loss
Training loss
Training loss
Training loss
Figure 4: Effect of sparsity on performance across different tasks We vary sparsity (1 top- k/Experts) and plot the relationship between pre-training loss and benchmark error rate, including intermediate checkpoints. For TriviaQA and HellaSwag, the error clearly tracks training loss and is largely insensitive to sparsity. In contrast, reasoning skills exhibit strong dependence of error rate on sparsity.
in pre-training loss do not translate into improved task loss; in some cases, the task loss actually worsens. These results suggest that, once top-k and layer width are fixed, an optimal number of experts exists for each task, and adding more beyond that point can harm performance on GSM8K and GSM-Plus.
Dependence on Active Parameter. Can we avoid a decline performance as the total number of experts increases? Figure shows that models with more active parameters begin to overfit at a lower pre-training loss and reach a lower minimum task loss at their optimal expert counts. Consequently, improving results GSM8K and GSM-Plus requires tuning not only the total number of experts but also the top-k size.
Downstream Accuracy. The decline in math-task performance total parameters increase is not limited to task loss; also consistently holds for downstream accuracy (Figure 3}. For Trivi- aQA and HellaSwag, accuracy improves monotonically as training loss decreases. By contrast, on GSM8K, further reductions in pre-training loss do not always translate to higher accuracy. When the number of active parameters is held constant, over-optimizing pre-training loss indeed harm performance. Figure 4 plots benchmark error rate against pre-training loss, including intermediate checkpoints. We observe sparsity dependence for reasoning skills. These results suggest that, for MoE models, downstream accuracy can deviate from the predictions of conventional scaling laws, and these deviations may vary across different tasks, This effect persists even when controlling for optimization hyperparameters such as learning rate (Appendix C.7
3.3 OPTIMAL SPARSITY FOR Iso-FLOP BUDGETS
We next analyze model quality under a constant compute budget, that is, along IsoFLOP contours (Hofmmann et al. 2022 Abnar al. 2025). For fixed per-token FLOP count, vary only the sparsity configuration: the number experts E and the top-k value, while holding the hidden dimension and sequence length.

---

Published as a conference paper at ICLR 2026
TriviaQA
HellaSwag
GSMBK
GSM-Plus
10.5
2.0
16.0
Task loss
0.5
13
Active Parums Pwwnmm.mmtooo (milondi
0.0
8.5
38
1.3
1.6
U112
yA
1028 388
1125
1872
102
MoE Density
MoE Density HellaSwag
MoE Density
MoE Density GSM-Plus
TriviaQA
GSM8K
0.4
0.35
0.20
Accuracy
0.3
0.25
0.2
0.40
0.29
0.30
Actoee Parama
0.10
0.30
0.85
U125
124
1/130
1/66
150
10
1/138
1/16
10
1/28 1na
vA
MoE Density
MoE Density
MoE Density
MoE Density
Figure 5: At fixed active parameter counts, higher sparsity (lower density) consistently im- proves performance, but at larger active parameter counts, GSM8K and GSM-Plus shift their optima back toward dense models. Task loss (top row) and Accuracy (bottom row) against the ratio of active experts l to total experts for a fixed active parameter budget. In the left two tasks (TriviaQA. HellaSwag), increasing sparsity consistently lowers task loss and raises accuracy across all active parameter budgets, in contrast, in the right two tasks (GSM8K, GSM-Plus), once active pa- rameter counts become large, this trend reverses and denser models begin outperform their sparser counterparts. Dashed segments mark the inverse-scaling regime that starts at the black circle; solid segments show the standard scaling region to right.
GSMBK (TTC)
GSM8K (GRPO)
Figure 6: Effeet of Test-Time Compute and GRPO on the loss-accuracy trade-off. Al- though both methods yield performance im- provements that scale with model size, the loss accuracy trade-off on GSM8K remains. Left: Final training loss accuracy under Test- Time Compute (Self-Consistency), Right: Fi- nal training loss accuracy after GRPO post- training.
Accuracy (SC)
0.50
0.
0.45
Accuracy
0.40
0.35
3.30
Final training loss
Final training loss
d=1024.k=0
before GRPO
GRPO
In Figure we show the task-specific optimal sparsity (i.e. 1-TopK/Experts against model perfor- mance under fixed FLOPs budget. Because the dataset size fixed our experiments, increasing the number of active parameters directly increases the training FLOPs. Thus, as the active parame- ter count grows, the plots implicitly trace how model performance changes with increasing FLOOPs For memorization benchmarks, lower density (higher sparsity) consistently yields lower task loss and higher accuracy. This pattern aligns with prior studies showing that, when FLOPPs are fixed to be constant, sparse models outperform denser models QA tasks (Abnar et al.) 2025). By con- trast, on mathematical-reasonii benchmarks, denser models outperform their sparser counterparts. At lower FLOPs, increasing sparsity still reduces loss and improves accuracy; however, once the FLOPs budget grows, denser models begin to perform better, achieving both lower loss and higher accuracy. This shift indicates that reasoning skills admit compute-dependent optimum, rather than monotonically favoring either sparsity or density.
3.4 TOKENS PER PARAMETER
The Chinchilla scaling law (Hofmmann et al. 2022) establishes that, under a fixed compute budget, the optimal trade-off between model parameters and training tokens corresponds approximately
7

---

Published as a conference paper at ICLR 2026
TriviaQA
HellaSwag
GSMBK
GSM-Plus
0.50
0.4
0.20
0.45
0.3
0.3
0.40
0.10
0.35
0.1
10
10-
10
10
10
10
10
10a
TPP
TPP
TPP
TPP
d=512, k=2
d=512, k=B
d==1024 k=2
d=1024, k=8
d=2048,
k=2 * d=2048 k=8
0-512. k-4 d-512.
k-16
d-1024, k-4 # d-1024, k-16
d-2048, k-4
d d-2048, k-16
Figure 7: Effect of TPP on performance across different tasks. For TriviaQA and HellaSwag, performance improves as the number of parameters increases. In contrast, for reasoning skills, performance deteriorates when the number of parameters becomes too large, indicating that there exists an optimal total tokens per paramete ratio for these tasks. Even fixed TPP, models with larger top-k values consistently outperform those with smaller top-k on reasoning tasks.
20 tokens per parameter (TPP) for dense models. More recently, Roberts et al. (2025) refined this view by showing that the optimal TPP ratio is task-dependent: memorization skills benefit from lower TPP (i.e., more parameters), whereas reasoning skills benefit from higher TPP (i.e., more data). These findings highlight that TPP should be interpreted not as a universal constant, but as task-sensitive scaling variable.
In our study, although we varied the number experts while keeping the total FLOPs fixed, this implicitly altered the TPP measured with respect total parameters. As shown in Figure 7 this variation reveals distinct behaviors across task categories. For memorization skills, performance improves monotonically as TPP decreases, consistent with the "parameter- hungry" characterization reported by |Roberts et al. (2025). For reasoning skills, we observe a non-monone monooonoe trend: accu- racy peaks near TPP 20, and degrades when TPP is either too low-when models have too many parameters relative to tokens-or too fiigh-when models have too few parameters relative to tokens.
Furthermore, our experiments reveal that active compute operationalized through the number of top- k experts interacts strongly with TPP. Even at fixed TPP, models with larger top-k values consistently outperform those with smaller top-k on reasoning tasks. This indicates that, MoE models, rea- soning ability depends not only on the Total TPP but also the balance between total and active parameters. In other words, the discussion explicitly consider both total parameter count and the number of activated parameters per token. We further note that depth ablations (Appendix C.6) exhibit the same non-monotonic dependence on TPP.
conpute-optimal scaling in MoE architectures must
3.5 IMPACT OF TTC AND POST-TRAINIE ON DOWNSTREAM PERFORMANCE
Test-Time Compute and RL post-training are standard for boosting reasoning on tasks such as mathematical problem solving. We therefore investigated whether the performance trade-off's ob- served above persist shift when applying (a) Test-Time Compute (TTC) and (b) RL post-training (GRPO). In Teme Compute, we evaluated GSMSK et al. 2021 in a zero-shot setting using Self- Self-Consistency (SC) decoding Wang et al., 2023), generating 2 independent continuations problem and selecting the most frequent answer. In Post- Post-Training, we fine-tuned each model on the GSM8K training dataset using the GRPO algorithm (Shao et al.) 2024). We followed the set- tings of Zhao al. (2025b) including reward function and fixed the learning rate constant across all model configurations. Further details regarding the training setup and hyperparameters are provided in Appendix A.2
As illustrated in Figure neither Time-TTimeao Compute nor GRPO mitigates the GSMSK performance drop that arises when total parameters increase. In other words, although both methods consistently improve overall performance, they do not eliminate the inverted U-shaped relationship between training loss and task accuracy.
8

---

Published as a conference paper at ICLR 2026
TriviaQA
HellaSwag
HumanEval
MBPP
0.0
0.30
Accuracy
0.425
033
Active Parmms www.mmmmmllooo... (miidnd)
0.00
0.25
030
0.375
0.15
025
0.150
0.20
0.10
0.325
020
0.05
0.300
0.15
015
1/12 1/0 1/33 1/39/30 shas
U/18
1022
17238
18A
155,
US
MoE Density
MoE Density
MoE Density
MoE Density
Figure 8: At fixed active parameter counts, higher sparsity (lower density) consistently im- proves performance, but at larger active parameter counts, HumanEval and MBPP shift their optima back toward dense models. Accuracy against the ratio of active experts k to total experts E for fixed active parameter budget. In the left two tasks (TriviaQA, HellaSwag), increasing sparsity consistently raises accuracy across all active parameter budgets, in contrast, in the right two tasks (HumanEval, MBPP), once active parameter counts become large, this trend reverses and denser models begin to outperform their sparser counterparts. Dashed segments mark the inverse-scalin regime that starts at the black circle; solid segments show the standard scaling region to the right.
3.6 CODING TASK ABLATIONS
We evaluate whether the sparsity performance trade offs observed for mathematical reasoning trans- fer to code generation. Unless otherwise noted, we reuse the same architecture and optimization hyperparameters in the experimental setup. Models are trained
125B token corpus com- posed of 95B tokens from Stack-Edu Python (Allal et al. 2025 (high-quality educational Python code trained for four epochs following Muennighoff et al. (2023)) and 30B tokens from DCLM- dedup web text (Zyphra 2024). We assess pass @ accuracy on HumanEval (Chen et al. 2021) and MBPP (Austin al. 2021). See Table 3 in Appendix for further details.
Figure summarizes performance as a function MoE density under matched active compute bud- gets. As active parameters grow, both HumanEval and MBPP exhibit clear shift in their optima toward denser configurations: beyond a task-dependent threshold, further increasing sparsity de- grades pass@1 despite continued improvements in pre-training loss. This echoes our math findings: when compute allows large active capacity, denser MoE layers yield better procedural reasoning for code synthesis, whereas sparser layers are more favorable only in the low-compute regime. Detailed results for the coding tasks are provided in Appendix C.5
DISCUSSION AND LIMITATIONS
All models are trained comparable activated size (Hoffmann et al. 2022). Except for the widest 2048, top- 16 setting, in or near the Chinchilla-optima regime, fall into the overtraining regime. Recent frontier models typically train with fixed token budget (Grrattafiori et al.) 2024 Qwen Team 2025), so holding tokens fixed in our experiments is appropriate. Nevertheless, with larger training corpora, one could train larger models at higher tokens-pers per-parrmeeer, potentially shifting the optimal sparsity toward sparser configurations, even for reasoning skills. We leave this direction for future work.
125B token corpus, which is Chinchilla-optimaal for dense models of
Our study does not attempt to exhaustively explore all architectural and training settings. The design space of MoE models is vast and is infeasible cover every possible combination. Regarding ar- chitectural choices, we adopted the Mixtral architecture (Jiang et al. (024) to ensure comparability with standard dense baselines such as Ilama Touvron et al.) 2023a Grattatiori et al., 2024). The
Mixtral architecture differs from recent state-of-the-art models like Qwen3 MoE (Qwen Team 2025)
primarily the use QK-norm however, since we did not observe training instability, QK- norm was not required in our setting. We also excluded shared experts, as prior work (Muennighoff al.) 2025a) reports mixed or negative results, and our preliminary tests indicated no meaningful perfor- mance changes when active and total FLOPs were matched, adding them would only complicate interpretation. For auxiliary losses such as balaneinn and router z-losses, we followed Muen- nighoff et al. (2025a). Instead, focused representative but systematic sweep width,

---

Published as a conference paper at ICLR 2026
expert count, and top-k routing, which already reveals new regularities that prior work did not cap- ture. particular, we identify an inverted-U relationship between sparsity and reasoning perfor- mance that contradicts the monotonic trends often assumed in scaling analyses. Intuitively, this inverted-U behavior can be explained by two compute-related factors. First, reasoning tasks require substantially higher active compute (i.e., inference FLOOs).. Prior work test-time compute by Snell et al. 2225)) demonstrates that increasing compute directly improves reasoning performance. In MoE models, raising the top-k increases the number of experts contributing each token and thereby increases available inference compute. Second, reasoning requires high training intensity per parameter. Under a fixed training budget, increasing sparsity spreads tokens across more experts, leaving each expert data-starved. While this increased capacity benefits memorization by enlarging model' effective storage, reasoning becomes either data-starved (when TPP too low) or under-paree paaler-prramterizee (when TPP is too high). These intuitive mechanisms align with the empirical TPP patterns analyzed in Section 3.4
5 CONCLUSION
We investigated the optimal sparsity of MoE language models through large-scale exploration of Mixtral-style architectures, varying expert count, top-k routing, and width, and evaluating across pre-training, RL post-training, and test-time compute. Our results reveal two central insights. First, FLOPs matter: downstream reasoning quality is determined not by pre-training loss alone, but by the number of active FLOPs at both train and test time. Larger top-k consistently outper- forms smaller ones even when pre-training loss is matched. Second, TPP matters: the tokens- per-parameter ratio governs task-specific scaling. Memorization tasks parameter-hungr and benefit from sparsity, while reasoning tasks are data-hungry and peak near 20 tokens per parameter, degrading when TPP becomes too low. Together, these findings refine current scaling practice. Spar- sity and more experts improve memorization under fixed budgets, but reasoning requires balancing active FLOPs with TPP. In high-compute regimes, the optimal density depends jointly on compute and data: with limited data, denser MoE layers preserve reasoning, while with abundant data, greater sparsity remains advantageous.
REPRODUCIBILITY STATEMENT
Due to the anonymity requirements and the design of OpenReview, at the time of ICLR submission, the following resources are made available the reviewers. All source codes, including those for MoE training and evaluation are provided in the supplementary material. The training data used in this study is publicly available. The model checkpoints are not shared due to their large file sizes, which makes anonymous sharing infeasible. Similarly, while we plan to release the training logs via wandb, maintaining anonymity remains a challenge, so they are not included at this stage.
The training setup is described in Section 3.1 and Appendix A and the evaluation methodology is provided in Appendix
ACKNOWLEDGEMENTS
This work was supported by the "R&D Hub Aimed at Ensuring Transparency and Reliability of Generative AI Models" project of the Ministry of Education, Culture, Sports, Science and Tech- nology. We used ABCI 3.0 provided by AIST and AIST Solutions with support from "ABCI 3.0 Development Acceleration Use. This work used computational resources TSUBAME4.0 supercom- puter provided by Institute of Science Tokyo through the HPCI System Research Project (Project ID: hp2401 70). The authors would like to thank Yukito Tajima for providing valuable and insightful discussions and feedback during the preparation of this manuscript.
AUTHOR CONTRIBUTIONS
Taishi Nakamura prepared the pretraining datasets, conducted all pre-training experiments and eval- uations (excluding lest-time-compute) and co-designed the overall experimental setup. Satoki
10

---

Published as a conference paper at ICLR 2026
Ishikawa co-designed the experiments and formulated the overall research strategy. Masaki Kawa- mura initiated the post-training and to-mi--empuuee (TTC) experiments. Takumi Okamoto con- ducted the post-training experiments and carried out the Max-Eigen (linear-layer) experiments. Daisuke Nohara conducted TTC experiments. Rio Yokota and Jun Suzuki provided guidance and oversight throughout the project. All authors contributed to manuscript writing and approved the final version.
REFERENCES
Samira Abnar, Harshay Shah, Dan Busbridge, Alaaeldin Mohamed Elnouby Ali, Josh Susskind, and Vimal Thilak. Parameters vs flops: Scaling laws for optimal sparsity for mixture-of-experts language models. International Conference on Machine Learning. 2025.
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Gabriel Martín Blázquez, Guilherme Penedo, Lewis Tunstall, Andrés Marafioti, Hynek Kydlíček, Agustín Piqueres Lajarín, Vaibhav Srivas- tav, et al. Smollm2: When smol goes big data-centric training of small language model. arXiv:2502.02737, 2025.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry. Quoe Le, and Charles Sutton. Program synthesis with large language models. arXiv:2108.07732, 2021.
Akshita Bhagia, Jiacheng Liu, Alexander Wettig, David Heineman, Oyvind Tafjord, Ananya Harsh Jha, Luca Soldaini, Noah A. Smith, Dirk Groeneveld, Pang Wei Koh, et al. Establishing task scaling laws via compute-cfficient model ladders. arXiv:2412.04403. 2024.
David Brandfonbrener, Nikhil Anand, Nikhil Vyas, Eran Malach, and Sham M. Kakade. Loss- to-loss prediction: Scaling laws for all datasets. Transactions on Machine Learning Research, 2025.
Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoe V. Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling. arXiv:2407.21787 2024.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv:2107.03374 2021.
Aidan Clark, Diego De Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoff- mann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et Unified scaling laws for routed language models. In International Conference on Machine Learning. 2022.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv:2110.1411 14168, 2021.
DeepSeek-AL. DeepSeek-R1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv:2501.12948, 2025a.
DeepSeek-AI. DeepSeek-V? V3 technical report.
2025b.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong. Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. GLaM: Efficient scaling of language models with of In International Conference on Machine Learning. 2022.
Runa Eschenhagen, Alexander Immer, Richard Turner, Frank Schneider, and Philipp Hennig. Kronecker- approximate curvature for modern neural network architectures. In Advances in Neural Information Processing Systems, 2023.
William Fedus, Barret Zoph, and Noam M. Shazeer. Switch transformers: Scaling to trillion param- eter models with simple and efficient sparsity. J. Mach. Learn. Res., 2021.
11

---

Published as a conference paper at ICLR 2026
Elias Frantar, Carlos Riquelme Ruiz, Neil Houlsby, Dan Alistarh, and Utku Evci. Scaling laws for sparsely-connectee foundation models. In International Conference on Learning Representations, 2024.
Samir Yitzhak Gadre, Georgios Smymis, Vaishaal Shankar, Suchin Gururangan, Mitchell Worts- man, Rulin Shao, Jean Mercat, Alex Fang, Jeffrey Li, Sedrick Keh, et al. Language models scale reliably with over-training and on downstream tasks. arXiv:2403.085400 2024.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. MegaBlocks: Efficient Sparse Training with Mixture-of-Experts. Proceedings of Machine Learning and Systems, 2023.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h h, et al. The language model evaluation harness, 2024.
Gemini Team. Gemini 2.5:
Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv:2507.06661, 2025.
GLM-4.5 Team. arXiv:2508.06471, 2025.
Glm-4.5: Agentic, reasoning, and coding (arc) foundation models.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The Llama 3 herd of models, arXiv:2407.21783, 2024.
Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, et al. Studying large language model generalization with influence functions. urXiv:2308.03290 2023.
Joel Hestness, Sharan Narang, Newsha Ardalani. Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md. Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically. arXiv:1712.00444 2017.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 1997.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. An empirical analysis of compute-optimal large language model training. In Advances in Neural Information Processing Systems, 2022.
Yuichi Inoue, Kou Misaki, Yuki Imajuku, So Kuroki, Taishi Nakamura, and Takuya Akiba. Wider deeper? scaling Ilm inference-time compute with adaptive branching tree search. MtXiv:2503.04412, 2025.
Robert A Jacobs. Michael Jordan, and Andrew G Barto. Task decomposition through competition in a modular connectionist architecture: The what and where vision tasks. Cognitive science, 1991.
Samy Jelassi, Clara Mohri, David Brandfonbrener, Alex Gu, Nikhil Vyas, Nikhil Anand, David Alvasez-Melis, Yuanzhi Li, Sham M. Kakade, and Eran Malach. Mixture of parrots: Experts improve memorization more than reasoning. In International Conference Learning Represen- tations, 2025.
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv:2401.04088. 2024.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations, 2020.
Michael Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 1994.
12

---

Published as a conference paper at ICLR 2026
Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Anmual Meet- ing of the Association for Computational Linguistics (ACL), 2017.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv:2001 108361, 2020.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference Learning Representations, 2017.
Kimi Team. Kimi k2: Open agentic intelligence. arXiv:2507.205344 2025.
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. In Advances in Neural Information Processing Systems, 2022.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. (Gardhrd: Scaling giant models with condi- tional computation and automatic sharding. In International Conference on Learning Represen- tations, 2021.
Qintong Li, Leyang Cui, Xueliang Zhao, Lingpeng Kong, and Wei Bi. GSM-plus: A comprehensive benchmark for evaluating the robustness of LLMs mathematical problem solvers. In Proceed- ings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), 2024.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. In Advances in Neural Information Processing Systems, 2019.
Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition-level code generation with alphacode. Science, 2022.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In International Conference on Learning Representations, 2024.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better down- stream: Implicit bias matters for language models. In International Conference on Machine Learning, 2023a.
Jiawei Liu, Chunqiu Steven Xia, Yuyao Wang. and LINGMING ZHANG. Is your code generated by chatGPT really correct? rigorous evaluation of large language models for code generation. In Advances in Neural Information Processing Systems, 2023b.
Zechun Liu, Changsheng Zhao, Forrest landola, Chen Lai, Yuandong Tian, Igor Fedorov, Yunyang Xiong, Ernie Chang, Yangyang Shi, Raghuraman Krishnamoorthi, et al. MobileLLM: Optimizing sub-billion parameter language models for on-device cases. In International Conference on Machine Learning, 2024.
Nyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019.
Jan Ludziejewski, Jakub Krajewski, Kamil Adamczewski, Maciej Pióro, Michat Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygó9d5, Piotr Sankowski, et al. Scaling laws for fine-grained mixture of experts. In International Conference Machine Learning, 2024.
James Martens and Roger Grosse. Optimizing neural networks with kronecker- facttordd approximate curvature. In International Conference on Machine Learning, 2015.
13

---

Published as a conference paper at ICLR 2026
Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A Raffel. Scaling data-constrained language models. In Advances in Neural Information Processing Systems, 2023.
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Evan Pete Walsh, Oyvind Tafjord, Nathan Lambert, et al. OLMoE: Open mixture-of-expee language models. In International Conference on Learning Representations, 2025a.
Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hanmaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. s1: Simple test-time scaling. arXiv:2501. 19393,, 2025b.
Team OLMo. 2 OLMo 2 furious. arXiv:2501.00656 2025.
OpenAl. GPT-4 technical report. arXiv:2303.087744 2024a.
OpenAl. Openai ol system card. arXiv:2412.1672) 2024b.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Gray, et al. Training language models to follow instructions with human feedback. In Advances Neurul Information Processing Systems, 2022.
Qwen Team. Qwen3 technical report. ccXiv:2505.09385 2025.
Nicholas Roberts, Niladri Chatterji, Sharan Narang, Mike Lewis, and Dieuwke Hupkes. Compute optimal scaling of skills: Knowledge vs reasoning. arXiv.2503.00661 2025.
Rylan Schaeffer, Joshua Kazdan, John Hughes, Jordan Juravsky, Sara Price, Aengus Lynch, Erik Jones, Robert Kirk, Azalia Mirhoseini, and Sanmi Koyejo. How do large language monkeys get their power (laws)? In International Conference on Machine Learning, 2025.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms.
2017.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song. Xiao Bi. Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, al. DeepSeekMath: Pushing the limits of mathemati- cal reasoning in open language models. arXiv:2402.03300, 20244
Noam Shazeer. Glu variants improve transformer. arXiv:2002.05202 2020.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The paarsely-gated caxture-of-experr layer. In International Conference Learning Representations, 2017.
Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflex- ion: Language agents with verbal reinforcement learning. In Advances in Neural Information Processing Systems, 2024.
Charlie Victor Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling LLM test-time compute optimally can be more effective than scaling parameters for reasoning. In International Confer- ence on Learning Representations, 2025.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 2024.
Ryousei Takano, Shinichiro Takizawa, Yusuke Tanimura, Hidemoto Nakada, and Hirotaka Ogawa. Abci 3.0: Evolution of the leading ai infrastructure in japan. MrXiv:2411.091335 2024.
Gemma Team. Gemma 2: Improving open language models at practical size. arXiv:2408.00118, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. LLaMA: Open and efficient foundation language models. arXiv:2302. 13971, 2023a.
14

---

Published as a conference paper at ICLR 2026
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov. Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv:2307.0928: 099288, 2023b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, t. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems, 2017.
Xuezhi Wang. Jason Wei, Dale Schuurmans, Quoe Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. In International Conference on Learning Representations, 2023.
Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You. OpenMoE: An early effort open mixture-of-experts language models. In International Con- ference on Machine Learning, 2024.
Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, 2021.
Tian Ye, Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu. Physics of language models: Part 2.1, grade-school math and the hidden reasoning process. In International Conference on Learning Representations, 2025.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Biao Zhang and Rico Sennrich. Root Mean Square Layer Normalization. In Advances in Neural Information Processing Systems, 2019.
Jinze Zhao, Peihao Wang, Junjie Yang, Ruisi Cai, Gaowen Liu, Jayanth Srinivasa, Ramana Rao Kompella, Yingbin Liang, and Zhangyang Wang. Sparse mixture-of-experts for compo- sitional generalization: Empirical evidence and theoretical foundations of optimal sparsity. arXiv:2410. 11..411....944 13364,, 2025a.
Rosie Zhao, Alexandru Meterez, Sham Kakade, Cengiz Pehlevan, Samy Jelassi, and Eran Malach. Echo chamber: RI post-training amplifies behaviors learned in pretraining. arXiv:2504.09912, 2025b.
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models. arXiv:2202.08906, 2022.
Zyphra. delm-dedup. http:://mwggiineeeeeo/aaaattttte 2024. Accessed: 2025-05-16.
15

---

OpenAl 2024a)) cannot be attributed any Published as a conference paper at ICLR 2026
Table I: Breakdown of the 125 B-token pre-training corpus.
Source
Type
Tokens
Corpus
Hugging Face GitLab
High Quality Web
DCLM-Deduped
High quality web
33.5B
788.5B
Zypbra/delmn-dddup
Flan decontaninated
High quality web
9.2B
18.5B
allenai/dolmino-n 12x
WebInstructFull
High quality web
14.7M
29.7M
TIGER- CIGER-LabWWeblnstuceuull
Literature Reference
PeS2o
Academic papers
62.9B
allenai/dolma
ArXiv
STEM papers
11.08
22.2B
allenai/dolma
Wikipedia
Encyclopedic
2.3B
4.7B
gitlab www.llmmmjpniiieeeee ae http://ppp......ppppasssssl/llaa jaaaasstttt--------ooo-----
Wikipedia Wikibooks
Encyclopedic
1.9B
3.9B
allenai/dolma
Project Gutenberg
Books
2.7B
5.5B
allenai/olmm
Mathematics
OpenWebMath
Math
6.6B
13.4B
allenai/olma
Algebraic Stack
Math
6.6B
13.3B
allenaidolma
FincMath-44
Math
5.IB
10.3B
HuggingFaaeTBBinnmmea
MathPile commercial subset train split
Math
4.5B
9.2B
GAIR/MathPle Commercial
TiyyGSM M-ND MIND
Synthetic math
3.4B
6.9B
allenai/olmo mix 1124
OpenMathhostruct-2
Synthetic math
2.6B
5.2B
ovadiukKpenMauhhnnnnnttt.
MathCoddr2 Syuthetic
Synthetic Math
2.0B
4.1B
allenai/olmo mix- 124
StackMatiQA
Math
529.00
1070.0M
matbbat/SSackkaatiOAA
NaturalReasoning
General reasoning
506.0M
1022.2M
facebook/nstural reasoning
NuminaMathi-CT train split
CoT reasoning
221.0M
446.4M
Al-MO-NuminaMatt-CCT
OpenMathinstruut- train split
Synthttic math
rvidin/UpenMaahhlssnuuuu
TuluMath
Synthetic math
123.9M
250.4M
Allena/olmo-mix-111
Metamath OWM- Mileeileered
Math
42.3M
85.4M
allenaa/olmo-mix
Orca-Math
Synthetic math
33.SM
67.7M
microsoft/orea-
wernaath-wwr---eeee prordddooolemsss 200k
Dolmino SynthMath
Synthetic math
15.7M
31.7M
allenai/olmo-mix--11 112
GSM8K train split
Math
1.4M
2.8M
allenaiddolmin--ii-1133
GSM8K train split
Math
1.4M
2.8M
openai/gsm8k
CodeSearchNNt Nee-mmttteer
Math
1.IM
2.2M
allenai/ollmino-iix.
Code
StackExchange
CodeText
725.1M
1464.8M
allensa/dolmino-min
Grand total
125.01
973.48
A TRAINING SETUP
A.1 PRE-TRAINING DATASET DETAILS
Table details the pre-training corpus: for each subset, lists the Hugging Face repository, split identifier, and public URI., alongside the original size and the number of subsampled tokens we used (125 tokens in the 99:1 train/validation split, as counted by the IIm-jp tokenizer v3 with 99,487 tokens). Thus, the total token budget is fixed in strict accordance with Kaplan's 's scaling law (Kaplan ct al. 2020), meaning the observed loss increase (and the accompanying puzzling overlitting that mirrors behavior recently reported by (OLMo 2025 change in data volume.
A.2 POST-TRAINING DETAILS
We use GRPo et 2024 with batch size of 1024, train for 15 epochs totaling 105 steps, and truncate prompts and generated sequences to 512 and 1024 tokens respectively. The actor's learning rate fixed at 10-6: 6, the temperature is set to 1.0, the KL-penalty coefficient 10-3 and 5 samples are used per prompt. Optimisation employs Adam with 8 = (0.9, 0.999), e = 10 and weight decay of 10-2 Following Zhao et al. (2025b), we implemented a code-execution-basee evaluator supporting TinyGSM-style and OpenMathInstrruttt outputs. For width of 2048 with 16 or 64 experts, we swept the learning rate (Fig. 9) and subsequently fixed it to 10-6 for all GRPO experiments.
A.3 IMPLEMENTATION & TRAINING ENVIRONMENT
We executed all pre-trainin runs on the ABCI 3.0 supercomputer (Taaano et al. 2024), equipped with NVIDIA H200 GPUs with board-level power capped 500 W per GPU. TTC experiments
16

---

Published as a conference paper at ICLR 2026
Table Detailed composition of the 125 B-token pre-training corpus without GSM8K and its syn- thetic variants (used for the ablation in Section C.3 Token counts and raw corpus sizes are listed for each source, following the same category structure Table T
Source
Type
Tokens
Corpus
Hugging Face GitLab
High Quality Web
DCLM Deduped
High quality web
33.5B
Zyphra/delm dedup
Flan decontaminated
High quality web
9.2B
18.58
allenai/dolmino- mix-
WebInstructFull
High quality web
14.7M
29.7M
TIGER Lab/WeblnstruccFull
STEM Literature Reference
PeS2o
Academic papers
31.1B
62.9B
allenai/dolma
ArXiv
STEM papers
11.0B
22.2B
allenai/dolma
Wikipedia
Encyclopedic
2.3B
4.7B
gitlab pp.iiacc jpdaaaasppdaatttts///li jp-corpus-y v3
Wikipedia Wikibooks
Encyclopedic
1.9B
3.9B
allenai/olma
Project Gutenberg
Books
2.7B
5.5B
allenai/dolma
Mathematies
OpenWehMath
Math
8.2B
13.4B
allenaiddolma
Algebraic Stack
Math
8.1B
13.3B
allenai/dolma
FineMath:4+
Math
HuuginggFaeeBBBiieenalll
MathPile commercial subset train split
Math
5.6B
9.2B
GAIR/MathPilc..om
MathCoder2 Synthetic
Synthetic Math
2.5B
4.1B
allenai/olmo cetaa//oloommii----
StackMathQA
Math
653.9M
1070.0M
matthall/SaciMMattIAAAA
NaturalReasoning
General reasoning
624.7M
1022.2M
facehook/naturil.co reassnnig
NuminaMath CoT train split
reasoning
272.8M
446.4M
MO/NuinaaMMat CoT
TuluMath
Synthetic math
153.0M
250.4M
allenai/olmommix-111
Metamath OWM- OWM-filtere
Math
52.2M
85.4M
allenai/olmo 112
Orca-Math
Synthetic math
414MM
67.7M
microsofioora-matt-
Math
1.MM
2.2M
allenai/ödlmino-mi-1111
Code
Stacknxehange
CodeText
725.1M
1464.8M
allenai/dolmino-niin mix-1124
Grand total
125.0B
961.0B
GSMBK (16 Experts)
GSM8K (32 Experts)
0.6
0.6
Validation Accuracy
0.4
0.3
0.3
0.2
0.2
0.1
0.0
0.0
20
40
80
100
60
100
Step
Step
Ir 1e-6
Ir 2e-6
Ir 5e-6
Ir 1e-5
2e-5
Se-5
Figure 9: Learning-rate sweep for width 2048. We varied the number of experts and swept the
learning rate. For both 16 and 32 experts, 5 x 10-6 produces the most stable training.
were conducted on the TSUBAME 4.0 supercomputer at the Global Scientific Information and Computing Center, Institute of Science Tokyo. They used NVIDIA H100 SXM5 94 GB GPUs (four GPUs per node) and InfiniBand NDR200 interconnects for inter-node communication.
For tre-training, extended the M LMd to add functionality needed for this study, with support for pipeline, tensor, and expert parallelism. Reinforcement learning experiments were implemented using GRPO (Shao al. 2024) top of the veRI framework. Model quality was assessed using Im-evalluation-harees and Laggll.agguaeMMonkyy..
17

---

Published as a conference paper at ICLR 2026
Table 3: Evaluation Benchmark Details
TriviaQA
HellaSwag
GSM-Plus
HumunEval
HumanEval-
Task
MRC EN
Math Reasoning EN
Madh Reasoning
Code Rensoning
Reasoning
Code
Reasoning
Code
Code Reasoning
Langauge
EN
EN
EN
EN
EN
EN
Instances
10,042
Few Few-shot
(0 for TTC)
Metric
Accuracy
Accuracy
Accuracy
CoT Acc.
Pass @1
P
Pass
@
B EVALUATION SETUP
We evaluate our models using the in-evaluation-hammes framework (Gao al. 2024 across four key capability areas. All evaluations employ standard few-shot prompting strategies unless other- wise specified.
We assess logical reasoning capabilities using Mathematical problem-solving is evaluated using GSM8K (Cobbe et al., 202I) with 4-shot prompting and GSM-Plus (Li al. (0024) with 5-shot CoT prompting. We evaluate comprehension abilities using TriviaQA (Joshi et al. 2017) with 4-shot prompting. Common sense reasoning is assessed through HellaSwag (Zellers et 2019) using 4- shot prompting setups. Finally, code reasoning capabilities are benchmarked on HumanEval (Chen) 202T) and HumanEval+ (Liu et al. 2023b) with 0-shot prompting and MBPP (Austin et al. 2021 and MBPP+ (Liu et al. 2023b) with 3-shot prompting, both evaluated using the Pass I met- ric. For Test-Time Compute ( (TTC) experiments specifically. GSMSK evaluation is conducted under zero-shot setting. To accommodate the variety of valid answer formats, extend the strict match patterns provided by the mm-eva uat
cet al.
rness beyond the standard implementation. Our matching criteria accept both the standard GSMSK format (# # # and GSM8K-CoT formats pre- fixed with "The answer is" or "Answer:".
Table 3 provides comprehensive details for all evaluation benchmarks.
C ADDITIONAL EXPERIMENTS
C.1 GRPO
the inverted U- shaped relationship between training loss and task accuracy persists even after applying GRPO. To verify that this phenomenon not due performing GRPO on the GSMSK dataset, additional GRPO experiments on the MATH 500 dataset (Lightman et al., 2024). As illustrated in Figure TO, GRPO on the MATH dataset yields consistent results with those obtained on the GSM8K dataset, confirming that this inverted U-shaped relationship is robust across different GRPO training datasets.
Training on MATH 500 Dataset Following the analysis presented in Section 3.5
conducted
C.2 TEST-TIME COMPUTE
Evaluation Setup We evaluated both (Cobe et al. (021) in purely zero-shot setting using Self-Consistency (SC) decoding Wang et al., 2023), generating 27 independent continuations per problem and selecting the most frequent answer with 128 samples per problem. Specifically, for each prompt we generated up to 1,024 tokens under temperature 0.6 and nucleus sampling (top-p = 0.95), drawing 128 independent continuations and selecting the most frequent answer.
Zero-shot VS Few-shot To set up Test Time Compute appropriately, we investigate how varying the number of prompt shots affected each expert' s behavior (Figure 11). Few shot performance is use zero shot
unstable and dropped significantly for models with small number of experts, inference for Test Time Compute. When few shot chain of thought is used to standardize answer formats, the provided demonstration steps can be internalized as fixed reasoning pattern by the model. As result, the model' inherent inference capabilities may not be fully expressed, and its ability generalize to novel problems could be hindered (Kojima et al. 2022).
18

---

Published as a conference paper at ICLR 2026
GRPO on GSMBK
GRPO on MATH 500
0.45
0.40
Accuracy
0.30
0.30
0.25
0.20
0.20
0.15
0.15
0.10
0.10
1.70
1.75
1.80
1.85
1.90
1.70
1.75
1.80
1.85
1.90
Final training loss
Final training loss
d= 1024,k=2
d= 10=1024,
d=1024,k=8
d=1024,k=16
before GRPO -ffer GRPO
Figure 10: Comparison of GSM8K accuracy for models fine-tuned with GRPO on different training datasets (left: GSM8K, right: MATH 500). Performance decline is consistently observed across different training datasets.
Experts
Experts
64 Experts
128 Experts
256 Experts
Accurary
6.22
0.15
2331
0.22
0.30
3.2%
8200
4.20
0:0
0.3
0.00
3.22
22 Gencration budget
2-
2
budget
budget
Generation budget
Goncration budget
shot
shot
shot
8 shot
Figure 11: GSM8K accuracy of model (d=1024) across different shot counts. Because few shot performance unstable and dropped significantly for models with a small number of experts, zero shot is used for Test-Time Compute.
Temperature Figure 12 shows that the inverted U-shaped performance-decline trend holds across every temperature setting, indicating that sampling temperature does not affect this behavior. This suggests that, although temperature controls inference randomness, the primary drivers of perfor- mance decline inherent to model architecture rather than temperature settings.
Evaluation of Larger Generation Budget We extended the sample size used for Test-Time Com- pute as described in Section 3.5 generating larger set of candidate responses. We then measured the resulting accuracy across different generation budgets to assess how increased sampling influ- ences performance (Figure 13). For an active parameter count of 8 (top-8), the performance decline gradually mitigated, whereas for an active parameter count of (top-2), the decline is instead amplified, resulting in a more pronounced U-shaped trend. Although increasing the sample count further may provide additional insights, remains challenging to identify consistent mitigation pattern across all models.
Increasing Top-k During Inference We compared the performance under TTC for model with hidden dimension of 2048, 128 experts, and top-2 routing by varying the inference time top-k pa- rameter. (Figure 14) Specifically, although doubling top-k sometimes yielded temporary improve- ments Pass @ I, applying TTC ultimately showed that the original top-2 setting maintained the highest performance, suggesting that fundamental performance gain occurs.
19

---

Published as a conference paper at ICLR 2026
T-p-2
Top-4
T-p-8
Top-16
0.30
Accuracy
6.22
0.25
320
0.20
1.5
0.32
6.22
0.15
0.65
loss
loss
Final loss
Final loss
tomp temp
Figure 12: Comparison of performance decline across different temperature settings (pass 1, d=1024). A consistent performance decline is observed regardless of temperature, and overall accu- racy increases as temperature decreases (i.e., approaches greedy).
Top-2
Top-4
Top-8
Top-16
6.4
Accuracy
0.1
0.40
33
0.35
0.2
02
0.7
0.20
1.55
Final Loss
Final Loss
Final Loss
Final Loss
samples
16
32
128
512
1024
Figure 13: Accuracy across generation budgets with increased sample counts. With active parameter count of 8 (top 8), the performance decline gradually alleviated as the budget increases, whereas with active parameter count of 2 (top 2), the decline is amplified, resulting in more pronounced U shaped trend.
C.3 GSMSK OVERFITTING ANALYSIS
To investigate whether our model overfits to GSM8K due to the inclusion of GSM8K training data and synthetic derivatives, we conducted an ablation experiment removing major GSM8K-related datasets from our pre-training corpus as listed in Table
We removed TinGGSM-MIND, both GSM8K train split instances, Dolmino SynthMath, OpenMaahInstruct-Il and OpenMathInstruutt22 which contain either the original GSM8K training data or synthetic problems derived from it.
The results are shown in Figure 15 and T6 We observe that the trends with respect to sparsity on GSM8K remain unchanged, both for Pass and TTC metrics. This indicates that while GSM8K training data and its synthetic derivatives do improve GSM8K scores, they do not alter the underlying performance trends. However, after post-training, we observe some changes in these trends, which leave future work to investigate further.
C.4
GSM8K PROBLEM ANALYSIS
We investigated whether models with varying numbers of experts exhibit differences in their ability to solve specific problems on the GSM8K dataset.
Figure 17 shows the results. We observe that different sparsity levels solve different instances of the problems.
C.5 DETAILED RESULTS FOR CODING TASKS
This appendix provides detailed results for the coding task ablations, mirroring the analyses for mathematical reasoning presented in the main text. Specifically, we detail the non-monotonic lationship between scaling and downstream performance, showing how both task loss (Figure [8] 19 and accuracy (Figure 20) can degrade pre-training loss improves. We then analyze key
20

---

Published as a conference paper at ICLR 2026
GSMBK
MATH
Score
0.2
0.1
0.0
2
2'
2"
2
2
Generation Budget
Generation Budget
Top-k during inference
top top
top
top
Figure 14: Increasing the top-k parameter only at inference time does not improve perfor- mance. Performance comparison under TTC for a Mixture-of-Experts model (hidden dimension 2048, 128 experts, top-2) as the top-k parameter is increased. While doubling k can occasionally improve Pass@1 I, applying TTC ultimately shows that the original top-2 configuration delivers the highest performance.
TriviaQA
HellaSwag
GSM8K
GSM-Plus
1.
2.2
Task loss
4.
Adtiv Params (illlod)
15
2.0
8.5
4.0
1.8
8.0
11. 11.0 115 10 y -t
1012 24
50 s/h y0
1/120 1/0h
200 y
ye
50
(06
yh y y
MoE Density
Density HellaSwag
MoE Density GSM8K
MoE Density
TriviaQA
GSM-Plus
0.50
0.14
24
Accuracy
0.30
0.3
0.08
0.10
Aative Params
0.35
0.04
0.1
0.00
0.00
u128 18A
102
10
va
1/28
102
18
12
1/20
1/2
10
1028
1164
10
MoE Density
MoE Density
MoE Density
MoE Density
Figure 15: Performance versus MoE density after removing GSMSK-related training data. Task loss (top) and accuracy (bottom) plotted against the ratio of active experts k to total experts E for a fixed active parameter budget. While performance on memorization tasks (TriviaQA, Hel- laSwag) improves with sparsity, the trend reverses for math reasoning tasks (GSMSK, GSM-Plus) at larger active parameter counts. Dashed segments mark the inverse-scaling regime.
chitectural factors, including the impact of MoE sparsity (Figure 21 and the optimal Tokens-per. Parameter (TPP) ratio (Figure 23). In addition, we report sparsity accurity trends for the coding benchmarks HumanEval+ and MBPP+ (Figure 22).
C.6 ABLATION ON DEPTH
We conducted additional experiments using 32 layer architecture. Motivated by prior reports suggesting that increased depth can improve performance (Liu et al. 2024 Team 2024 Ye et al. 2025), evaluated whether deeper models exhibit similar trends in setting. For the 32 layer configuration, we observed that the results align with the patterns discussed in the previous section, when analyzed through the lens of TPP, the behavior remains consistent with earlier findings.
21

---

Published as a conference paper at ICLR 2026
GSM8K (Pass@1)
GSMBK (TTC)
GSM8K (GRPO)
0.5
0.25
Accuracy
0.4
0.20
0.3
0.2
GSM8K
-+.Wo GSMBK
0.2
1.7
1.4
Final training loss
Final training loss
Final training loss
Figure 16: GSM8K performance without GSM8K-related training data: Pass@1 @ (left), TTC with
128 budget (center), and after GRPO (right)
500
31 42048
262 02348 K2
Figure 17: Analysis of solvable problems across different numbers of experts on GSM8K. This graph displays the number of problems that were commonly solvable or unsolvable across models with varying numbers of experts.
C.7 INFLUENCE OF OPTIMIZATION HYPERPARAMETER
Thus far, we have demonstrated that the structure of the model, particularly the degree of sparsity, can lead to differences in reasoning performance on downstream tasks, even when the models con- verge to the same training loss. Such differences behavior on unseen data reflects implicit inductive biases rather than mere fit to the training data. Studies neural network generalization have long recognized that not only architectural choices, but also optimization dynamics (i.e., differences in hyperparameter settings, regularization schemes, and optimizer algorithms), play an important role in shaping these inductive biases. Motivated by this insight, we examine the learning- rate scale, which is critical to generalization (Keskar et al., 2017 Li al. [2019] Yang & Hul (021)) Our goal to investigate how these choices influence the model 8 ability to transfer to downstream tasks, beyond what is captured by pre-training loss alone.
similar generalization, in which model's
Figure 25 illustrates our empirical findings, obtained using a MoE architecture with 16 experts. While memorization skills remains largely invariant to these byperparameters, reasoning skills are sensitive to the learning rates: when models converge to the same training loss, trainings with lower learning rates and smaller initialization scales yield superior downstream accuracy. These observations carry an important implication. Studies on generalization large-scale language mod- els should incorporate rigorous reasoning benchmarks rather than relying solely on validation loss curves or standard QA tasks to fully capture the impact of optimization appimization-inducee implicit biases. This enables precise analysis on the generalization of LLMs.
C.8 ABLATION ON TOKEN BUDGET
To validate choice of the 125B token budget used in the main experiments, conducted preliminary ablation study scaling the training duration to IT tokens. We compared the baseline model trained 125B tokens against model trained IT tokens. The training corpus for the
22

---

Published as a conference paper at ICLR 2026
Train Loss
Validation Loss
HellaSwag
HumanEval
1.5
1,5
4.6
Final 1o59
1.4
Task loss
4.4
0.95
1.3
0.90
1.2
1.
0.85
100
"10"
'00
'0'
10
"""" io
'10
'6
Total parameters (B)
Total parameters (B)
Total parameters (B)
Total parameters (B)
512k=2
512k=8
d=1024k==
d=1024
d=2048.k=2
d=2048
512,km4
512,6=16
1024,2m
16
d=2048,k44
2048,
Figure 18: Although training and validation losses generally decrease as the total number of parameters increases, validation loss for the largest models does not fully converge. HellaSwar task loss follows this favorable scaling trend, but HumanEval task loss sometimes worsens once the total number of parameters exceeds certain threshold.
TriviaQA
HellaSwag
HumanEval
MBPP
4.6
1.50
Task loss
begin
8.0
4.2
0.90
7.5
4.0-
7.0
3.8
1.30
1.0
1.2
14
1.0
1.2
1.4
1.0
1.2
1.4
1.0
1.2
1.4
Final training loss
Final training loss
Final training loss
Final training loss
dm512, km2, k-2,A=17000 Am170M
d-512, k-16, A-520M
d-1024, km8, A-1.18
dm 2048,, km4, Am2.38
d-512, A-20M
d-1024, A-470M
d-1024, k-16, A-1.98
d-2048, A-3.98
d=512. k-8.
d=1024, k-4, A-670M
d=2048, k-2, A=1.58
d-2048, k-16, Am7.18
Figure 19: For HumanEval and MBPP, once the training loss drops below certain point, the task loss starts to increase. Results of scaling total parameters by increasing the number of experts, with model width and top-k held constant. For TriviaQA, HellaSwag, and task loss falls monotonically
training loss decreases. By contrast, HumanEval and MBPP show U-shaped trend: task loss declines with training loss only until a threshold, beyond which further reductions in training loss hurt task performance.
1T run was constructed by upsampling the 973.4B token dataset (Table via random sampling to reach the IT target. For this experiment, we used model configuration with width 2048, 16 experts (E 16), and top-2 routing (k: = 2). Hyperparameters followed the settings described in Section 3.1 with the exception of the learning rate schedule, where the cosine decay was stretched to accommodate the IT token duration.
The results are presented in Table 4 Since the extended IT corpus is predominantly composed of general web text, which serves rich source of world knowledge, nemorizationnintensee tasks such as TriviaQA and HellaS wag exhibited significant performance gains. Conversely, reasoning tasks (GSMSK and GSM-Plus) showed comparatively marginal improvements. We attribute this discrepancy to the scarcity of high-quality mathematics and reasoning corpora available in the open source the time of experimentation; simply scaling up the volume of general web text contributes minimally to reasoning capabilities. Therefore, extending the training to IT tokens significantly increases computational costs without yielding proportional gains in reasoning performance, justi- fying our adoption of the 125B token budget to prioritize examining architectural trade-offs.
D THE USE OF LARGE LANGUAGE MODELS
Use of Large Language Models (LLMs). In preparing this paper, we used large language models only for minor editing tasks, such as improving grammar and clarity of exposition. LLMs did not contribute to the research ideas, methodology, experiments, or analysis, and thus played no substantive role the scientific contributions of this work.
23

---

Published as a conference paper at ICLR 2026
TriviaQA
HellaSwag
Evaa
d= 512.k=2
0.3
0.45
512.k=4
0.30
512.k-8
Accuracy
512.k=66
0.2
0.40
0.25
dd1022,k=2
d-1024,k-4
d=1024,k=B
0.1
0.35
0.20
d=1024,k-16
0 k-204,,-2
0 k-04
0.30
k=
d-2048,k-16
1.0
1.2
14
1.0
1.2
1.4
1.0
1.2
14
Final training
Final training loss
Final training loss
Figure 20: Downstream accuracy when scaling total parameters via expert count with width and top-k fixed. TriviaQA and HellaSwag exhibit steadily improving accuracy as pre-training loss decreases, whereas HumanEval shows non-monotonic trend: further reductions in pre-training loss do not always improve accuracy and can even degrade performance.
TriviaQA
HellaSwag
HumanEval
MBPP
4.6
1.50
Task loss
4.4
0.900
Active Parind imllers)
0925
42
1.40
0.875
7.5
4.0
0.800
7.0
3.0
0.825
1728 A0 20.000 10 3 50
2 50
1/12 AAA
J 3
1123 10A
MoE Density
MoE Density
MoE Density
MoE Density
Figure 21: At fixed active parameter counts, higher sparsity (lower density) consistently im- proves performance, but at larger active parameter counts, HumanEval and MBPP shift their optima back toward dense models. Task loss against the ratio of active experts to total experts E for fixed active parameter budget. In the left tasks (TriviaQA, HellaSwag), increasing sparsity consistently lowers task loss across all active parameter budgets, in contrast, in the right two tasks (HumanEval, MBPP), once active parameter counts become large, this trend reverses and denser models begin to outperform their sparser counterparts. Dashed segments mark the inverse -callin regime that starts the black circle; solid segments show the standard scaling region the right.
HvmanEval+
MBPP+
0.275
0.55
0.250
Accuracy
0.225
0.45
222
0.40
0.200
Active Params
0.175
0.30
0.150
0.25
0.125
0.20
1/128
1/64
1/32
1/16
1/8
1/4
1/2
1/128
1/64
1/32
1/16
1/8
1/4
1/2
MoE Density
MoE Density
Figure 22: At fixed active parameter counts, higher sparsity (lower density) consistently im- proves performance, but at larger active parameter counts, HumanEval+ and MBPP+ shift their optima back toward dense models. Accuracy against the ratio of active experts & to total perts E for fixed active parameter budget. Once active parameter counts become large, this trend reverses and denser models begin to outperform their sparser counterparts. Dashed segments mark the inverse-scaling regime that starts at the black circle; solid segments show the standard scaling region the right.
24

---

Published as a conference paper at ICLR 2026
TriviaQA
HellaSwag
HumanEval
MBPP
0.4
0.30
Accuracy
0.2
0.40
0.25
0.3
0.35
0.20
0.1
0.2
0.30
0.15
10
10c
10
10a
10a
10
10:
TPP
TPP
TPP
TPP
k=8
d=1024,
d=2048, d=2048,
d=512,
d=512,
d=1024, d=1024,
d=2048,
Figure 23: Effect of TPP on performance across different tasks. For TriviaQA and HellaSwag, performance improves as the number of parameters increases. In contrast, for ceasoning-intensing tasks such as HumanEval and MBPP, performance deteriorates when the number of parameters becomes too large, indicating that there exists optimal data to parameter ratio for these tasks.
TriviaQA
HellaSwag
GSM8K
GSM-Plus
0.4
0.20
Accuracy
0.45
0.3
0.15
0.2
0.2
0.10
0.35
0.1
0.05
10
10
10
10
10
10
10
102
TPP
TPP
TPP
TPP
d=1024, k=
L=16
d=1024,
k=4, L=16
d=
1024,
k=8,
L=16
d=1024,
k=16, L=16
d d-2048,
2- 1-16
d-2048
k-4,
1-
k-8,
1-16
dd2048, L-16
d-1024,
k-2, 2, L-32
d-1024,
k-4, L-32
1024
k-8, L-32
Figure 24: Effect of model depth TPP-performance trade-offs.
TriviaQA
HellaSwag
GSMBK
GSM-Plus
Max Eigen (Linear Layer)
Curvature Eigen
10*
0.38
10a
Accuracy
0.15
0.34
0.36
0.2
0.10
10a
0.10
10!
0.1
0.05
0.05
10"
0.28
0.26
0.0
0.00
10-1.
2.5
2.3
2.5
Training
Training loss
Training
Training
Training
6.4c-3
3.2e-3
1.66-
Bo.4
40.4
26-4
For reasoning skills, the relationship between training loss and downstream per- formance dependent the choice of optimization hyperparameterss The learning rate also impacts downstream accuracy. For the maximum eigenvalue, we evaluated the maximum eigenvalue eigenvaaluu
Figure 25:
of fisher information matrix under a K-FAC approximation (Martens & Grosse 2015 Eschenhagen et al. 2023). Following (Grosse et al., 2023), we calculate the maximum cigenvalues only for linear layers. We find that higher learning rates lead a lower maximum eigenvalue, which is consistent with existing research indicating that convergence to flatter minima improves generaliza-
tion (Hochreiter & Schmidhuber 1997 Keskar et al. 2017 Jiang et al. 2020).
Table 4: Performance comparison between models trained on 125B tokens versus 1T tokens.
The model configuration is fixed d 2048, E = 16, and = 2.
Token Budget
TriviaQA
HellaSwag
GSM8K
GSM-Plus
125B 1T
0.279
0.426 0.538
0.318
0.188
0.535
0.363
0.214
25


Reviews:
[{"review_id": "nxeoKou1vw", "reviewer": "Reviewer_p7WF", "summary": "The paper studies the optimal sparsity of Mixture-of-Experts models under memorization and reasoning skills by training MoE models with varying total parameters, sparsity, and top-k routing under the fixed budget. Through extensive experiments, the paper concludes that 1. the downstream reasoning quality is decided by both the active FLOPs and pretraining loss and 2. there exist different optimal tokens-per-parameter ratios for memorization and reasoning tasks.", "questions": "Questions are addressed in the Weaknesses section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to understand.\n- The experiments are comprehensive while supporting the major claims of the paper.\n- One of the main findings is surprising, as it shows that higher sparsity only improves performance under memorization instead of reasoning tasks under the iso-FLOP settings.", "weaknesses": "- The paper might need to address more about its intuition and originality from previous works such as [1] and [2], since similar observations regarding the optimal sparsity in MoE models have been made.\n- Theoretical insights are encouraged to explain the experimental findings.\n- The U-shape trend plot for reasoning tasks in Figure 2 is very interesting, and I suggest the authors to verify such finding under more reasoning tasks.\n\n[1] Samira, Abnar, et al. \"Parameters vs FLOPs: Scaling Laws for Optimal Sparsity for Mixture-of-Experts Language Models.\" arXiv:2501.12370 (2025).\n\n[2] Zhao, Jinze, et al. \"Sparse Mixture-of-Experts for Compositional Generalization: Empirical Evidence and Theoretical Foundations of Optimal Sparsity.\" arXiv:2410.13964 (2025)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nSXezajIgt", "reviewer": "Reviewer_pkQZ", "summary": "This paper aims to investigate how MoE sparsity influences two distinct capability regimes: memorization skills and reasoning skills. The work shows how Active FLOP is more important for reasoning, while memorization improves with number of total parameters. Another interesting finding provided in this work is that changing the k in top-k routing has a negligible effect if the number of active parameters is kept constant.", "questions": "1. The number of tokens used seems small to if we are targeting End task performance, specially for MOE models\n2. It would be good to get some ablation for various router choices, though than can be a future work\n3. In Page 9, figure 8, it would be good do the study at k>1 (ideally 8) and E >8\n4. More details about the post training setup is helpful. How many tokens in the post training set?\n5. No details have been provided whether Continuous training is done or learning rate is annealed before evaluating end task", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. It is an important observation that for MoE models, downstream accuracy can deviate from the predictions of conventional scaling laws, and these deviations may vary across different tasks.\n\n2. Exhaustive experimentation is done in reasoning and coding tasks to demonstrate the U shape of tasks performance with the increase of total parameters at a FLOP controlled setting\n\n3. Exhaustive experiments are done to show that post training couldn't improve this.", "weaknesses": "1. The number of tokens used seems small to if we are targeting End task performance, specially for MOE models\n2. It would be good to get some ablation for various router choices, though than can be a future work\n3. In Page 9, figure 8, it would be good do the study at k>1 (ideally 8) and E >8\n4. More details about the post training setup is helpful. How many tokens in the post training set?\n5. No details have been provided whether Continuous training is done or learning rate is annealed before evaluating end task", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GhcntL4HWM", "reviewer": "Reviewer_tKtK", "summary": "The paper empirically studies the effect of sparsity in sparse mixture of experts (MoEs) on downstream tasks. Specifically the paper trains a series of models based on Mixtral architecture and studies downstream performance on tasks that tend to rely on memorization vs tasks that test reasoning capabilities of language models (LMs). The paper studies the relationship between training loss and task loss and provides insights on the effect of active parameter count and tokens per parameter (TPP) on downstream tasks described above. The main findings include active FLOPs may lead to higher reasoning accuracy and that higher tokens per parameter may be preferable for reasoning tasks.", "questions": "- Is there a way to show how downstream performs with sparsity where sparsity is defined in the paper as 1 - (active / total experts)? Sparsity is mentioned in the paper but is not shown explicitly in scaling plots. Please include sparsity value, if possible, with the plots. Only Figure 5 appears to include sparsity (via density term which is its complement).\n\n- The range of accuracy/error rate for GSM8K task appears to be on the lower side? Are these values good enough for readers to draw valid conclusions? A discussion on what is reasonable would be very useful to help the reader.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper studies the effect of sparsity in MoEs on downstream tasks. This area has not been examined in detail so the study in the paper is timely and will likely be of interest to many researchers & practitioners.\n\n- The empirical setup including models, data and downstream tasks are  described clearly in the paper. This gives me confidence that the experiments are reproducible.\n\n- The models considered in the paper are not necessarily compute-optimal. This detail may provide additional insight into how to optimally train and use MoEs at scale.\n\n- The paper proposes tokens per parameter (TPP) as a metric to track in addition to active parameters. This metric provides additional insight into the role of data in training MoEs that work well on downstream reasoning-type tasks. This may be interesting to many readers (and is definitely interesting to this reader)", "weaknesses": "- The paper considers a single architecture inspired by Mixtral family of MoEs in the work. It's understandable why this choice was made (experiment volume) but I do wonder if other architecture choices can change the conclusions made here. If possible, please discuss why Mixtral was chosen as opposed to other choices. \n\n- The fact that memorization depends on total parameter count is known from prior literature. Furthermore, active number of parameters (inference FLOPs) have also observed to improve certain downstream tasks' performance (Abnar et al., 2205). So the claim that these are new contributions is weak.\n\n- While sparsity in mentioned in the paper, the plots (Figure 1) for instance do not show this value but instead show top-K. This makes it hard for the reader to infer the effect of sparsity on empirical observations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sIoWFqOrzq", "reviewer": "Reviewer_fnio", "summary": "This paper explores how the sparsity of MoE models affects their performance on downstream tasks. While prior scaling law works have mostly focused on pretraining loss or efficiency, this work reveals that downstream memorization and reasoning capabilities respond differently to sparsity. The authors conduct very extensive empirical analysis on a range of benchmarks. Through various analyses, they observe interesting points (see Strengths section) that could be further explored and taken into consideration when building or deploying MoEs.", "questions": "In Line 321-323:  \n\n> *At lower FLOPs, increasing sparsity still reduces loss and improves accuracy; however, once the FLOPs budget grows, denser models begin to perform better, achieving both lower loss and higher accuracy.*\n\nHowever, I could not find where this trend is clearly demonstrated. Figure 5 appears to fix the FLOPs budget, so it doesn’t reveal how model performance varies as FLOPs increases.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper presents several insightful findings based on systematic analysis of MoE sparsity across a variety of experiments.\n\n- Active FLOPs has a significant effect on the downstream task performance, not just determined by pretraining loss alone.\n- The study uncovers an important trade-off btw memorization and reasoning with TPP, where memorization skills are parameter-hungry and reasoning skills are data-hungry.\n- Post-training or test-time scaling do not change the memorization-reasoning gap, so the optimal sparsity must be determined pretty much during the pretraining stage.", "weaknesses": "While the observations are novel and well-supported empirically, I would consider it more significant to find any **intuitive or theoretical rationales** behind them. For instance, why does reasoning capability require denser MoEs, while memorization thrives on sparsity? Why doesn’t simply scaling the parameter count work?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)